Softmax and log-softmax: spec definitions versus analytic presentation

NN.Proofs.Autograd.FDeriv.Softmax and NN.Proofs.Autograd.FDeriv.LogSoftmax prove Fréchet-derivative facts about softmaxVec and logSoftmaxVec, which are stated directly on Euclidean vectors. The specification layer instead defines Activation.softmaxVecSpec and Activation.logSoftmaxVecSpec on tensors, using the numerically stable max-shifted form.

This file closes that gap over ℝ:

• getScalarE_softmaxVecSpec and getScalarE_logSoftmaxVecSpec identify the spec kernels with the analytic definitions after vectorization;
• hasFDerivAt_softmaxSpec_vec and hasFDerivAt_logSoftmaxSpec_vec transfer differentiability to Activation.softmaxSpec 0 and Activation.logSoftmaxSpec 0;
• softmaxFDerivCorrect and logSoftmaxFDerivCorrect package Spec.softmaxOp 0 and Spec.logSoftmaxOp 0 as OpSpecFDerivCorrect, so their backward rules are proved to be the adjoint of the true derivative (softmaxBackwardSpec_eq_vjp, logSoftmaxBackwardSpec_eq_vjp).

Softmax

theorem Proofs.Autograd.getScalarE_softmaxVecSpec {n : ℕ} (t : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.softmaxVecSpec t) = softmaxVec (getScalarE t)

The spec softmax kernel is the analytic softmaxVec after vectorization.

theorem Proofs.Autograd.softmaxVecSpec_ofFnE {n : ℕ} (x : Vec n) : Activation.softmaxVecSpec (ofFnE x) = ofFnE (softmaxVec x)

softmaxVecSpec on a vector built from Euclidean coordinates.

theorem Proofs.Autograd.softmaxSpec_zero_forwardVec {n : ℕ} : (fun (xV : Vec n) => getScalarE (Activation.softmaxSpec 0 (ofFnE xV))) = softmaxVec

The vectorized forward map of Spec.softmaxOp 0 on vectors is softmaxVec.

theorem Proofs.Autograd.hasFDerivAt_softmaxSpec_vec {n : ℕ} (xV : Vec n) : HasFDerivAt (fun (xV : Vec n) => getScalarE (Activation.softmaxSpec 0 (ofFnE xV))) (softmaxDerivCLM xV) xV

Axis-0 spec softmax on vectors is Fréchet-differentiable with derivative softmaxDerivCLM.

theorem Proofs.Autograd.getScalarE_softmaxBackwardSpec {n : ℕ} (x dY : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.softmaxBackwardSpec 0 x dY) = softmaxJvp (getScalarE x) (getScalarE dY)

The spec softmax VJP is the analytic softmaxJvp (the softmax Jacobian is symmetric).

noncomputable def Proofs.Autograd.softmaxCorrect (n : ℕ) : OpSpecCorrect (Spec.Shape.dim n Spec.Shape.scalar) (Spec.Shape.dim n Spec.Shape.scalar)

Spec.softmaxOp 0 on vectors with its analytic JVP and the VJP/JVP adjointness law.

noncomputable def Proofs.Autograd.softmaxFDerivCorrect (n : ℕ) : OpSpecFDerivCorrect n n

Spec.softmaxOp 0 on vectors is analytically correct: its JVP is the Fréchet derivative.

theorem Proofs.Autograd.softmaxBackwardSpec_eq_vjp {n : ℕ} (x δ : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.softmaxBackwardSpec 0 x δ) = (vjp (fun (xV : Vec n) => getScalarE (Activation.softmaxSpec 0 (ofFnE xV))) (getScalarE x)) (getScalarE δ)

The spec softmax backward is the vector-Jacobian product of the spec softmax forward.

Log-softmax

theorem Proofs.Autograd.getScalarE_logSoftmaxVecSpec {n : ℕ} (t : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.logSoftmaxVecSpec t) = logSoftmaxVec (getScalarE t)

The spec log-softmax kernel is the analytic logSoftmaxVec after vectorization.

theorem Proofs.Autograd.logSoftmaxSpec_zero_forwardVec {n : ℕ} : (fun (xV : Vec n) => getScalarE (Activation.logSoftmaxSpec 0 (ofFnE xV))) = logSoftmaxVec

The vectorized forward map of Spec.logSoftmaxOp 0 on vectors is logSoftmaxVec.

theorem Proofs.Autograd.hasFDerivAt_logSoftmaxSpec_vec {n : ℕ} (xV : Vec n) : HasFDerivAt (fun (xV : Vec n) => getScalarE (Activation.logSoftmaxSpec 0 (ofFnE xV))) (logSoftmaxDerivCLM xV) xV

Axis-0 spec log-softmax on vectors is Fréchet-differentiable with derivative logSoftmaxDerivCLM.

theorem Proofs.Autograd.getScalar_expSpec_logSoftmaxVecSpec {n : ℕ} (x : TorchLean.Tensor ℝ [n.succ]) (i : Fin n.succ) : (Activation.logSoftmaxVecSpec x).expSpec.getScalar i = (softmaxVec (getScalarE x)).ofLp i

Exponentiating the spec log-softmax output recovers the analytic softmax coordinates.

theorem Proofs.Autograd.getScalarE_logSoftmaxBackwardSpec {n : ℕ} (x dY : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.logSoftmaxBackwardSpec 0 (Activation.logSoftmaxSpec 0 x) dY) = logSoftmaxVjp (getScalarE x) (getScalarE dY)

The spec log-softmax VJP, evaluated on the recomputed forward output, is the analytic logSoftmaxVjp.

noncomputable def Proofs.Autograd.logSoftmaxCorrect (n : ℕ) : OpSpecCorrect (Spec.Shape.dim n Spec.Shape.scalar) (Spec.Shape.dim n Spec.Shape.scalar)

Spec.logSoftmaxOp 0 on vectors with its analytic JVP and the VJP/JVP adjointness law.

noncomputable def Proofs.Autograd.logSoftmaxFDerivCorrect (n : ℕ) : OpSpecFDerivCorrect n n

Spec.logSoftmaxOp 0 on vectors is analytically correct: its JVP is the Fréchet derivative.

theorem Proofs.Autograd.logSoftmaxBackwardSpec_eq_vjp {n : ℕ} (x δ : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.logSoftmaxBackwardSpec 0 (Activation.logSoftmaxSpec 0 x) δ) = (vjp (fun (xV : Vec n) => getScalarE (Activation.logSoftmaxSpec 0 (ofFnE xV))) (getScalarE x)) (getScalarE δ)

The spec log-softmax backward is the vector-Jacobian product of the spec log-softmax forward.