Softmax

Fréchet-derivative facts for axis softmax on Euclidean vectors.

This is the analytic ($\mathbb R$) ingredient used to justify attention-style row softmax nodes (Vec n → Vec n) in the tape/DAG autograd proofs.

References

• Baydin et al., Automatic Differentiation in Machine Learning: a Survey (JMLR 2018).
• The Matrix Cookbook (softmax Jacobian identities / vector calculus conventions).
• PyTorch docs for naming/behavior alignment (not used for theorems): https://pytorch.org/docs/stable/generated/torch.nn.functional.softmax.html

noncomputable def Proofs.Autograd.sumExp {n : ℕ} (x : Vec n) : ℝ

sumExp x = ∑ᵢ exp(xᵢ).

This is the normalizing denominator in softmax.

theorem Proofs.Autograd.sumExp_pos {n : ℕ} (x : Vec n.succ) : 0 < sumExp x

sumExp x is strictly positive when the index type is nonempty.

theorem Proofs.Autograd.sumExp_ne_zero {n : ℕ} (x : Vec n.succ) : sumExp x ≠ 0

Convenience corollary: sumExp x ≠ 0 (for n = succ _).

theorem Proofs.Autograd.hasFDerivAt_exp_apply {n : ℕ} (x : Vec n) (j : Fin n) : HasFDerivAt (fun (x : Vec n) => Real.exp (x.ofLp j)) (Real.exp (x.ofLp j) • evalCLM j) x

The coordinate exponential x ↦ exp (x j) has derivative dx ↦ exp (x j) * dx j.

theorem Proofs.Autograd.hasFDerivAt_sumExp {n : ℕ} (x : Vec n) : HasFDerivAt sumExp (∑ j : Fin n, Real.exp (x.ofLp j) • evalCLM j) x

sumExp has derivative dx ↦ ∑ j, exp (x j) * dx j.

noncomputable def Proofs.Autograd.softmaxVec {n : ℕ} : Vec n → Vec n

Softmax on Euclidean vectors.

For n = succ _: softmaxVec x i = exp(xᵢ) / sumExp x.

The n = 0 branch is the identity on the trivial space.

noncomputable def Proofs.Autograd.dotCLM {n : ℕ} (y : Vec n) : Vec n →L[ℝ] ℝ

The dot-product functional x ↦ ∑ᵢ yᵢ * xᵢ, packaged as a continuous linear map.

This is used to express the softmax Jacobian in a clean coordinate-free way.

@[simp]

theorem Proofs.Autograd.dotCLM_apply {n : ℕ} (y x : Vec n) : (dotCLM y) x = ∑ j : Fin n, y.ofLp j * x.ofLp j

The dot functional evaluates to the expected sum of products.

noncomputable def Proofs.Autograd.softmaxDerivCoord {n : ℕ} (x : Vec n) (i : Fin n) : Vec n →L[ℝ] ℝ

The ith output coordinate of the softmax derivative at x, as a continuous linear map.

If y = softmaxVec x, then this is the linear functional: dx ↦ yᵢ * dxᵢ - yᵢ * ⟪y, dx⟫.

noncomputable def Proofs.Autograd.softmaxDerivCLM {n : ℕ} (x : Vec n) : Vec n →L[ℝ] Vec n

The full Fréchet derivative of softmaxVec at x, packaged as a CLM Vec n →L Vec n.

noncomputable def Proofs.Autograd.softmaxJvp {n : ℕ} : Vec n → Vec n → Vec n

Closed-form JVP (directional derivative) for softmax.

If y = softmaxVec x and s = ⟪y, dx⟫, then (softmaxJvp x dx)ᵢ = yᵢ * (dxᵢ - s).

theorem Proofs.Autograd.softmaxJvp_eq_deriv {n : ℕ} (x dx : Vec n) : softmaxJvp x dx = (softmaxDerivCLM x) dx

The closed-form JVP softmaxJvp agrees with the CLM derivative softmaxDerivCLM.

theorem Proofs.Autograd.inner_softmaxJvp_comm {n : ℕ} (x dx δ : Vec n) : inner ℝ (softmaxJvp x dx) δ = inner ℝ dx (softmaxJvp x δ)

Self-adjointness identity for the softmax Jacobian in this inner-product encoding.

This lemma is used to show the VJP can be expressed by reusing the JVP formula.

theorem Proofs.Autograd.hasFDerivAt_softmaxVec {n : ℕ} (x : Vec n) : HasFDerivAt softmaxVec (softmaxDerivCLM x) x

Softmax is Fréchet-differentiable everywhere, with derivative softmaxDerivCLM.