Axis softmax #
This file collects the mathematical facts needed by numerical bounds for the coupled, vector-valued softmax used in attention. It is separate from the NF scalar logistic helper: axis softmax has a shared denominator and a dense Jacobian, while logistic acts independently on each tensor entry.
The stable spec implementation is Activation.softmaxVecSpec. Proofs.Analysis.Softmax proves
that its entries are positive, sum to one, and lie in [0,1]. The analytic derivative is
Proofs.Autograd.softmaxJvp; its Jacobian is self-adjoint, so the same formula implements the VJP.
The rounded-real bounds follow max subtraction, exponential, sequential summation, and division.
They require denominator-error margins supplied as proofs; this noncomputable file does not
implement an executable certificate checker or verify a native softmax kernel. The analytic
conservation law says that the exact real JVP has coordinate sum zero, not that a rounded VJP
satisfies that identity exactly.
References:
- A. Griewank and A. Walther, Evaluating Derivatives, 2nd ed., 2008, for forward/reverse differentiation of coupled maps.
- A. A. Baydin et al., "Automatic Differentiation in Machine Learning: a Survey," JMLR 2018.
- PyTorch
torch.nn.functional.softmaxdocumentation for the runtime axis convention.
Maximum and max shift #
Forgetting the NF format commutes with the nonempty-vector maximum exactly.
NF.max only selects one operand; it performs no arithmetic and therefore introduces no rounding
error. The proof uses one fold homomorphism rather than repeating a coordinate induction in every
stable normalization operator.
The maximum of a rounded vector approximates the real maximum with the same infinity-norm budget as the vector itself. No additional ULP term appears because maximum is a selection.
Rounded stable softmax #
Error after subtracting the rounded maximum from every logit.
Instances For
Error after exponentiating the max-shifted logits.
Instances For
Error in the sequentially rounded denominator reduction.
Instances For
Per-coordinate output budget for stable softmax.
The exact denominator is at least one. The forward theorem additionally requires a proof of
denominatorErrorBound eps xR < 1, which keeps the rounded denominator positive. This real-valued
bound is not itself an executable check.
Instances For
Infinity-norm forward-error budget for stable vector softmax.
Instances For
The max-shifted NF implementation approximates real vector softmax.
The theorem follows the rounded-real stages: maximum, subtraction, exponential, sequential sum,
and division. In addition to the input approximation, it requires a proof that the denominator
error remains below the exact denominator's lower bound 1.
Row-wise error tensor for axis-1 softmax on a matrix.
Instances For
Global infinity-norm budget for row-wise axis-1 softmax.
Instances For
Matrix-level stable softmax theorem, obtained by applying the vector theorem independently to each row.
The denominator premise is proved separately for each row. The output uses one global infinity-norm budget, as required by the matrix multiplication and graph composition contracts.
Exact hard-masked softmax #
Stable softmax numerators for a row whose allowed maximum is already known.
Blocked coordinates are set to literal zero after exponentiation. This expression is equivalent to
the some rowMax branch of Spec.hardMaskedSoftmaxVecSpec and never introduces a finite masking
sentinel.
Instances For
The staged numerator computation equals the entrywise expression in the public spec. This is an equality of tensor expressions, not a claim about a fused native kernel.
Error after subtracting an approximate allowed-row maximum from every score.
Instances For
Error in the hard-masked numerator vector, independent of the mask because applying it adds no rounding error.
Instances For
Error in the sequentially rounded sum of the allowed numerators.
Instances For
Final per-coordinate budget for a nonempty hard-masked softmax row.
Instances For
Rounded-real error theorem for the nonempty branch of hard-masked vector softmax.
hmaxS and hmaxR identify the selected allowed-row maxima. The theorem does not trust those
values blindly: hmax must relate them numerically, and hdenomLower supplies the exact positive
lower bound used by division. For the canonical selected maximum this lower bound is 1, because
one allowed shifted score is zero and contributes exp 0 = 1.
If both semantics find no allowed key, hard-masked softmax agrees exactly on the zero row.
Proof-bearing evidence for nonempty hard-masked softmax rows.
The data fields record real-valued maxima and margins; the proposition fields relate them to the exact and rounded rows. The structure is indexed by both score matrices, the mask, and the score error budget, so its proofs cannot be reused for unrelated inputs without establishing a bridge. It does not construct or decide those proofs automatically.
Exact allowed-row maximum of each row.
- rowMaxR : Fin m → FloatLib.Floats.Formats.Flocq.NF β fexp rnd
Rounded allowed-row maximum of each row.
Error budget relating the two maxima of each row.
Positive lower bound on the exact denominator of each row.
The rounded maximum approximates the exact one within
epsMax.- specMax (i : Fin m) : Spec.hardMaskedMax? (Spec.get scoresS i) (Spec.get mask i) = some (self.rowMaxS i)
rowMaxSis the maximum selected by the exact hard-masked softmax. - runtimeMax (i : Fin m) : Spec.hardMaskedMax? (Spec.get scoresR i) (Spec.get mask i) = some (self.rowMaxR i)
rowMaxRis the maximum selected by the rounded hard-masked softmax. - denominatorLower (i : Fin m) : self.eta i ≤ (hardMaskedNumerators (Spec.get scoresS i) (Spec.get mask i) (self.rowMaxS i)).sumSpec
The exact denominator of each row is at least
eta. - denominatorMargin (i : Fin m) : hardMaskedDenominatorError epsScores (self.epsMax i) (Spec.get scoresR i) (Spec.get mask i) (self.rowMaxR i) < self.eta i
The rounded denominator error of each row stays strictly below
eta.
Instances For
Rowwise hard-masked softmax bounds with one independently certified maximum per row.
Instances For
Matrix-level hard-masked softmax when every row has at least one allowed coordinate.
The selected maxima and denominator premises remain row-local. This matches causal attention, where
row i always admits key i, and avoids replacing all rows by the worst intermediate scale before
the final infinity norm is taken.
Rounded softmax backward #
Error in the dY * softmax(x) product used by the softmax VJP.
Instances For
Error in the rounded dot product sum (dY * softmax(x)).
Instances For
Error after subtracting the replicated rounded softmax dot product from dY.
Instances For
End-to-end infinity-norm budget for the rounded softmax VJP.
Supplying center includes the subtraction of that upstream coordinate before the weighted
reduction. Selecting a coordinate introduces no rounding, but subtracting it does; both upstream
values carry the original error budget. Omitting it describes the uncentered formula used for
arbitrary already-supplied weights. Attention opts into the centered formula.
Instances For
Rounded VJP theorem for any already-certified softmax weight vector.
This formulation is shared by ordinary and hard-masked softmax. In the masked case blocked weights are exactly zero, so the common formula also returns exactly zero gradient at blocked logits; no separate finite-sentinel derivative rule is required.
Numerical bound for the centered helper used by attention's JVP and VJP.
The exact row must be normalized or identically zero. The rounded weights select the anchor; the exact weights need not select the same maximum. The budget includes both copies of the upstream-coordinate error and the actual rounded subtraction before the weighted reduction. No native backend schedule is inferred from this rounded-real statement.
Error theorem for the rounded-real softmax VJP, including its upstream centering.
This is the training counterpart of approxTensor_softmaxVecSpec. The rounded argmax selects
the upstream anchor; it need not agree with the exact argmax. Coordinate selection adds no
rounding, while the first subtraction is charged to the budget before applying the existing
weighted-reduction theorem. This does not establish a native backend's arithmetic schedule.
The analytic softmax on a nonempty vector sums to one.
A softmax JVP is tangent to the probability simplex: its coordinates sum to zero.
Coordinatewise VJP/JVP bound in the infinity norm.
If every upstream coordinate has magnitude at most G, then every softmax input gradient has
magnitude at most 2G. The estimate is dimension-free because softmax weights are nonnegative and
sum to one. It is intentionally conservative; tighter certificates may retain the factor
2 * y_i * (1 - y_i) for each coordinate.