Softmax analysis properties #
This module proves theorem-level facts about TorchLean's spec-level softmax operators. The
definitions themselves live in NN.Spec.Layers.Activation; this file belongs under
NN.Proofs.Analysis because it imports real-analysis and finite-sum proof machinery to establish
properties of those definitions.
Current theorem surface:
getScalar_le_maxVecSpecandexists_getScalar_eq_maxVecSpec: the shared stable-shift maximum is an attained upper bound;softmax_shift_nonposandexists_softmax_shift_eq_zero: shifted logits are nonpositive and one is exactly zero;softmax_shift_exp_le_oneandsoftmax_shift_denom_bounds: exponentials cannot overflow and their denominator lies between1and the axis length;softmax_vec_spec_normalized: exposes the positive normalized weights used by the stable max-shifted implementation;softmax_vec_spec_pos: every coordinate is strictly positive;softmax_vec_spec_mem_unitInterval: every coordinate lies in[0,1];sum_spec_softmax_vec_spec: a nonempty vector softmax sums to1;sum_spec_softmax_backward_spec: the concrete stable softmax VJP has coordinate sum zero;abs_getScalar_softmax_backward_spec_le_two_mul: bounded upstream coordinates give a dimension-independent coordinate bound for that VJP;sum_spec_softmax_spec_row: axis-1matrix softmax has rows summing to1when the key dimension is nonempty;getScalar_softmaxVecSpec_eq_exp_divandgetScalar_logSoftmaxVecSpec_eq_sub_log: overℝthe stable kernels agree coordinatewise with the textbook formulas, which is what connects them to the analyticsoftmaxVecandlogSoftmaxVecinNN.Proofs.Autograd.FDeriv.SoftmaxSpec;softmaxSpec_zero_vec,softmaxSpec_one_matrix, andget_softmaxSpec_one: the axis-parametric operators reduce to the vector kernels on vectors and matrix rows.
We intentionally state these over ℝ: positivity of exp and division by a positive denominator
are the mathematical facts that make the probabilistic interpretation precise.
Stable max shift #
Every coordinate is bounded above by the exact maximum used by softmax and log-softmax.
The maximum used by stable softmax is attained by an input coordinate.
Every logit shifted by the implementation's maximum is nonpositive.
At least one max-shifted logit is exactly zero.
Exponentiating a max-shifted real logit produces a value at most one. This is the central overflow-prevention fact behind stable softmax.
Max-shifted real exponentials remain strictly positive.
One max-shifted exponential is exactly one, because the maximum is attained.
The stable softmax denominator lies in [1,n] for a nonempty vector of length n.
The lower bound rules out division by zero. The upper bound follows because every shifted
exponential is at most one. Together with softmax_shift_exp_le_one, this makes overflow
prevention an explicit theorem of the max-shifted implementation rather than an empirical claim.
Normalized coordinates #
The stable vector softmax has positive weights normalized by their sum.
This lemma exposes exactly one reusable algebraic description of the implementation. The weights
are the max-shifted exponentials computed by softmaxVecSpec; subsequent proofs of positivity,
range, and normalization do not unfold the implementation again.
Coordinate equation for the concrete stable vector softmax.
This is the small unfolding lemma that downstream algebraic proofs should use. It exposes the max-shifted numerator and its tensor sum while hiding the implementation chosen for tensor reduction.
Coordinate formula for the max-shifted exponentials.
The stable softmax denominator is the plain exponential sum times the shift factor.
Over ℝ, the stable max-shifted softmax agrees coordinatewise with the textbook formula
exp(xᵢ) / ∑ⱼ exp(xⱼ).
This is the lemma that lets analytic developments work with the unshifted formula while the spec keeps its overflow-safe implementation.
Over ℝ, the stable log-softmax agrees coordinatewise with xᵢ - log ∑ⱼ exp(xⱼ).
Axis softmax on vectors and matrices #
The axis-parametric operators move the selected axis to the innermost position and back. For a
vector (axis 0) and for the key axis of a matrix (axis 1) no permutation is needed, so the
operators reduce definitionally to the vector kernels. These lemmas record that reduction so that
downstream files do not depend on how the permutation bookkeeping is implemented.
Axis-0 softmax on a vector is the vector kernel.
Axis-0 softmax backward on a vector is the vector kernel.
Axis-0 log-softmax on a vector is the vector kernel.
Axis-0 log-softmax backward on a vector is the vector kernel.
Axis-1 softmax on a matrix is the row-wise innermost kernel.
Row i of an axis-1 matrix softmax is the vector softmax of row i.
Probability-simplex properties #
Every coordinate of a nonempty real softmax vector is strictly positive.
softmaxVecSpec produces a vector whose entries sum to 1 over ℝ.
Every coordinate of a nonempty real softmax vector lies in the closed unit interval.
Backward conservation #
The concrete stable softmax backward is tangent to the probability simplex.
Activation.softmaxBackwardSpec 0 is the vector VJP used by the spec and tape layers. Its
coordinate sum is zero because the stable forward weights sum to one. This statement is about the
actual tensor definition, not the separate analytic EuclideanSpace presentation of the same
derivative.
Coordinatewise bound for the concrete stable softmax VJP.
If every upstream coordinate has magnitude at most G, each input-gradient coordinate has
magnitude at most 2G. The estimate does not grow with the axis length because the softmax output
is a nonnegative vector of total mass one.
Axis-1 softmax on matrices is rowwise, so each row sums to 1.
This is the attention-shaped theorem: for score matrices, the key axis is the last/vector axis, and softmax is applied independently to every query row.