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:

• toVec_le_maxVecSpec and exists_toVec_eq_maxVecSpec: the shared stable-shift maximum is an attained upper bound;
• softmax_shift_nonpos and exists_softmax_shift_eq_zero: shifted logits are nonpositive and one is exactly zero;
• softmax_shift_exp_le_one and softmax_shift_denom_bounds: exponentials cannot overflow and their denominator lies between 1 and 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 to 1;
• sum_spec_softmax_backward_spec: the concrete stable softmax VJP has coordinate sum zero;
• abs_toVec_softmax_backward_spec_le_two_mul: bounded upstream coordinates give a dimension-independent coordinate bound for that VJP;
• sum_spec_softmax_spec_row: matrix softmax is rowwise, so each nonempty row sums to 1;
• sum_spec_softmax_spec_row_of_ne_zero: the same row theorem with nonemptiness supplied as nK ≠ 0.

We intentionally state these over ℝ: positivity of exp and division by a positive denominator are the mathematical facts that make the probabilistic interpretation precise.

Scalar helpers

softmaxVecSpec is written over tensors, so even one coordinate has type Tensor ℝ .scalar. Local helper definitions expose scalar coordinates to the proof without adding public API.

Stable max shift

theorem Proofs.toVec_le_maxVecSpec {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : TensorAlgebra.toVec t i ≤ (Activation.maxVecSpec t).toScalar

Every coordinate is bounded above by the exact maximum used by softmax and log-softmax.

theorem Proofs.exists_toVec_eq_maxVecSpec {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : ∃ (i : Fin n.succ), TensorAlgebra.toVec t i = (Activation.maxVecSpec t).toScalar

The maximum used by stable softmax is attained by an input coordinate.

theorem Proofs.softmax_shift_nonpos {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : TensorAlgebra.toVec (t.subSpec (Spec.replicate (Activation.maxVecSpec t))) i ≤ 0

Every logit shifted by the implementation's maximum is nonpositive.

theorem Proofs.exists_softmax_shift_eq_zero {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : ∃ (i : Fin n.succ), TensorAlgebra.toVec (t.subSpec (Spec.replicate (Activation.maxVecSpec t))) i = 0

At least one max-shifted logit is exactly zero.

theorem Proofs.softmax_shift_exp_le_one {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : TensorAlgebra.toVec (Activation.maxShiftedExpVecSpec t) i ≤ 1

Exponentiating a max-shifted real logit produces a value at most one. This is the central overflow-prevention fact behind stable softmax.

theorem Proofs.softmax_shift_exp_pos {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : 0 < TensorAlgebra.toVec (Activation.maxShiftedExpVecSpec t) i

Max-shifted real exponentials remain strictly positive.

theorem Proofs.exists_softmax_shift_exp_eq_one {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : ∃ (i : Fin n.succ), TensorAlgebra.toVec (Activation.maxShiftedExpVecSpec t) i = 1

One max-shifted exponential is exactly one, because the maximum is attained.

theorem Proofs.softmax_shift_denom_bounds {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : 1 ≤ (Activation.maxShiftedExpVecSpec t).sumSpec ∧ (Activation.maxShiftedExpVecSpec t).sumSpec ≤ ↑n.succ

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

theorem Proofs.softmax_vec_spec_normalized {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : ∃ (weights : Fin n.succ → ℝ), (∀ (i : Fin n.succ), 0 < weights i) ∧ ∀ (i : Fin n.succ), TensorAlgebra.toVec (Activation.softmaxVecSpec t) i = weights i / ∑ j : Fin n.succ, weights j

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.

theorem Proofs.toVec_softmaxVecSpec {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : TensorAlgebra.toVec (Activation.softmaxVecSpec t) i = TensorAlgebra.toVec (Activation.maxShiftedExpVecSpec t) i / (Activation.maxShiftedExpVecSpec t).sumSpec

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.

Probability-simplex properties

theorem Proofs.softmax_vec_spec_pos {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : 0 < TensorAlgebra.toVec (Activation.softmaxVecSpec t) i

Every coordinate of a nonempty real softmax vector is strictly positive.

theorem Proofs.sum_spec_softmax_vec_spec {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : (Activation.softmaxVecSpec t).sumSpec = 1

softmaxVecSpec produces a vector whose entries sum to 1 over ℝ.

theorem Proofs.softmax_vec_spec_mem_unitInterval {n : ℕ} (t : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (i : Fin n.succ) : TensorAlgebra.toVec (Activation.softmaxVecSpec t) i ∈ Set.Icc 0 1

Every coordinate of a nonempty real softmax vector lies in the closed unit interval.

Backward conservation

theorem Proofs.sum_spec_softmax_backward_spec {n : ℕ} (x dY : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) : (Activation.softmaxBackwardSpec x dY).sumSpec = 0

The concrete stable softmax backward is tangent to the probability simplex.

Activation.softmaxBackwardSpec is the 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.

theorem Proofs.abs_toVec_softmax_backward_spec_le_two_mul {n : ℕ} (x dY : Spec.Tensor ℝ (Spec.Shape.dim n.succ Spec.Shape.scalar)) (G : ℝ) (hdY : ∀ (i : Fin n.succ), |TensorAlgebra.toVec dY i| ≤ G) (i : Fin n.succ) : |TensorAlgebra.toVec (Activation.softmaxBackwardSpec x dY) i| ≤ 2 * G

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.

softmaxSpec 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.

theorem Proofs.sum_spec_softmax_spec_row {nQ nK : ℕ} (maskedScores : Spec.Tensor ℝ (Spec.Shape.dim nQ (Spec.Shape.dim nK.succ Spec.Shape.scalar))) (i : Fin nQ) : (Spec.get (Activation.softmaxSpec maskedScores) i).sumSpec = 1

Convenience row-sum theorem when the key dimension is written as an arbitrary nK plus a proof nK ≠ 0.

Many model statements quantify over a natural key length nK; this wrapper converts that style into the Nat.succ _ shape required by sum_spec_softmax_spec_row.

theorem Proofs.sum_spec_softmax_spec_row_of_ne_zero {nQ nK : ℕ} (hK : nK ≠ 0) (scores : Spec.Tensor ℝ (Spec.Shape.dim nQ (Spec.Shape.dim nK Spec.Shape.scalar))) (i : Fin nQ) : (Spec.get (Activation.softmaxSpec scores) i).sumSpec = 1