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).
• hasFDerivAt_distillSpec_vec identifies the cancellation-resistant student pullback with the derivative of the teacher-logit loss, with the teacher held fixed.
• hasFDerivAt_crossEntropyLogitsSpec_vec and hasFDerivAt_crossEntropyTargets_vec prove the logits and target-weight derivatives of the weighted loss. Targets need not sum to one.

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_logSoftmaxJvpSpec {n : ℕ} (x dx : TorchLean.Tensor ℝ [n]) : getScalarE (Activation.logSoftmaxJvpSpec 0 (Activation.logSoftmaxSpec 0 x) dx) = logSoftmaxJvp (getScalarE x) (getScalarE dx)

The typed graph's compensated log-softmax JVP agrees with the analytic derivative.

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.

Weighted logits cross-entropy

theorem Proofs.Autograd.crossEntropyLogitsSpec_eq_dot {n : ℕ} (x target : TorchLean.Tensor ℝ [n]) : Spec.crossEntropyLogitsSpec 0 x target = -(dotCLM (getScalarE target)) (getScalarE (Activation.logSoftmaxSpec 0 x))

The overflow-aware weighted loss preserves the real cross-entropy formula. This is a real identity, not bitwise equivalence between the two floating-point evaluation schedules.

theorem Proofs.Autograd.hasFDerivAt_crossEntropyLogitsSpec_vec {n : ℕ} (x : Vec n) (target : TorchLean.Tensor ℝ [n]) : HasFDerivAt (fun (logits : Vec n) => Spec.crossEntropyLogitsSpec 0 (ofFnE logits) target) (dotCLM (getScalarE (Spec.crossEntropyLogitsDerivSpec 0 (ofFnE x) target))) x

The shared logits pullback differentiates the actual weighted loss, including unnormalized and zero target weights. No probability-normalization hypothesis is needed.

theorem Proofs.Autograd.hasFDerivAt_crossEntropyTargets_vec {n : ℕ} (logits : TorchLean.Tensor ℝ [n]) (target : Vec n) : HasFDerivAt (fun (weights : Vec n) => Spec.crossEntropyLogitsSpec 0 logits (ofFnE weights)) (dotCLM (getScalarE (Spec.CrossEntropy.targetGradient 0 logits))) target

Target weights remain differentiable, including at zero. The numerical implementation must not turn a zero-weight forward shortcut into a zero target derivative.

Distillation

theorem Proofs.Autograd.distillDerivSpec_eq_sub_softmax {n : ℕ} (student teacher : TorchLean.Tensor ℝ [n + 1]) : Spec.distillDerivSpec 0 student teacher = (Activation.softmaxVecSpec student).subSpec (Activation.softmaxVecSpec teacher)

The logit-difference gradient is the usual student-minus-teacher probability difference over the reals. Both numerical branches and the residual correction preserve this identity.

theorem Proofs.Autograd.hasFDerivAt_distillSpec_vec {n : ℕ} (student : Vec (n + 1)) (teacher : TorchLean.Tensor ℝ [n + 1]) : HasFDerivAt (fun (x : Vec (n + 1)) => Spec.distillSpec 0 (ofFnE x) teacher) (dotCLM (getScalarE (Spec.distillDerivSpec 0 (ofFnE student) teacher))) student

With teacher logits fixed, distillDerivSpec is the Fréchet derivative of the actual distillation loss on nonempty class vectors. This connects the compensated pullback to the loss, rather than just identifying it with another gradient formula.