Documentation

LeanMlir.Proofs.Architectures.Softmax

Softmax and the softmax–cross-entropy gradient #

softmax and crossEntropy are defined in MLP.lean; this file differentiates them: softmax_diff, the Jacobian pdiv_softmax, the global softmax_has_vjp, and softmaxCE_grad — ∂(crossEntropy ∘ softmax)/∂logits = softmax − onehot, the cotangent every classifier's train step starts from. Attention (rowSoftmax) and the small-net IR build on it.

source
theorem Proofs.softmax_diff (c : ℕ) :
Differentiable ℝ (softmax c)

Differentiability of softmax c: each coordinate is exp(z k) · (Σ_j exp(z j))⁻¹, and the denominator is positive.

The softmax Jacobian #

For p = softmax(z) with p_j = exp(z_j) / sum_k exp(z_k), the quotient rule gives:

dp_j/dz_i = p_j * (delta_{ij} - p_i)

This is the famous "diag minus outer product" form:

J = diag(p) - p * p^T

Dense (every output depends on every input), but rank-1 correction to a diagonal — which means the VJP has a closed-form collapse, just like BatchNorm did.

source
theorem Proofs.pdiv_softmax (c : ℕ) (z : Vec c) (i j : Fin c) :
pdiv (softmax c) z i j = softmax c z j * ((if i = j then 1 else 0) - softmax c z i)

Partial derivative of softmax (quotient rule on the exponentials).

d(softmax(z))_j/dz_i = softmax(z)_j * (delta_{ij} - softmax(z)_i)

Proved (was an axiom). The j-th coord of softmax c z is Real.exp (z j) / S with S := Σ_k Real.exp (z k) > 0, so the j-th output coord function z' ↦ exp(z' j) * (Σ_k exp(z' k))⁻¹ has HasFDerivAt derivative built from HasFDerivAt.exp, HasFDerivAt.fun_sum, (hasDerivAt_inv ·).comp_hasFDerivAt, and HasFDerivAt.mul. Evaluating that CLM at basisVec i and collapsing Σ_k exp(z k) · δ_{ki} = exp(z i) gives the formula.

source
noncomputable def Proofs.softmax_has_vjp (c : ℕ) :
HasVJP (softmax c)

Softmax VJP — the closed-form collapse.

back(z, dy)_i = p_i * (dy_i - <p, dy>)

where p = softmax(z) and <p, dy> = sum_j p_j * dy_j is one scalar.

Read this carefully. The naive VJP would be: dz_i = sum_j J_{ji} * dy_j = sum_j (p_j * (delta_{ij} - p_i)) * dy_j

That's O(c) per entry, O(c^2) total. But expanding: dz_i = p_i * dy_i - p_i * sum_j p_j * dy_j = p_i * (dy_i - <p, dy>)

The rank-1 correction lets you precompute one scalar (<p, dy>) and apply it to every entry. Total work: O(c). Same optimization pattern as BN (one reduction + a broadcast) and max-pool (one comparison + a select).

Interpretation. Softmax outputs a probability distribution. Its backward subtracts the "weighted average of the incoming gradient under that distribution" from each entry, then scales by the entry's probability. Entries with low probability get small gradients (because the softmax flattened them in the forward); entries with high probability get gradients proportional to how much they deviate from the weighted-average cotangent.

This is the one place where "softmax means softly select one thing" maps directly to "softmax backward selectively amplifies the gradient for the winning class."

Equations
  • One or more equations did not get rendered due to their size.
Instances For
    source
    theorem Proofs.crossEntropy_differentiable (c : ℕ) (label : Fin c) :
    Differentiable ℝ fun (z : Vec c) => crossEntropy c z label

    crossEntropy is differentiable in the logits. softmax > 0 lets Real.log (hence crossEntropy = -log(softmax · label)) inherit smoothness.

    source
    theorem Proofs.softmaxCE_grad (c : ℕ) (logits : Vec c) (label j : Fin c) :
    pdiv (fun (z : Vec c) (x : Fin 1) => crossEntropy c z label) logits j 0 = softmax c logits j - oneHot c label j

    Softmax cross-entropy scalar gradient — proved (was an axiom in MLP.lean; relocated here to use pdiv_softmax).

    ∂(-log softmax(z)[label])/∂z_j = softmax(z)_j - onehot(label)_j

    Stated using pdiv on a Vec 1-valued wrapper (cross-entropy is naturally scalar, but pdiv is defined for Vec → Vec; we just take the only output index). Proof: pdiv_eq_fderiv_coord extracts the only coord, then HasFDerivAt.log (with softmax z label > 0) composed with softmax_diff gives the derivative of the inner Real.log. Negating and evaluating at basisVec j reduces via pdiv_softmax to the expected formula.