Axis softmax #
This file collects the mathematical facts needed by numerical bounds for the coupled, vector-valued softmax used in attention. It is deliberately separate from the older 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 theorem below adds the conservation law needed for backward error analysis: every softmax JVP
has coordinate sum zero.
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 checker must additionally establish
denominatorErrorBound eps xR < 1; this prevents the rounded denominator from crossing zero and
turns the division condition into an explicit, checkable certificate obligation.
Instances For
Infinity-norm forward-error budget for stable vector softmax.
Instances For
The max-shifted NF implementation approximates real vector softmax.
Unlike a blanket continuity statement, the theorem follows the executable stages: maximum,
subtraction, exponential, sequential sum, and division. The sole side condition is the numerical
certificate check that the denominator error remains below its proved real 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 obligation remains row-specific: a certificate may accept well-conditioned rows without replacing them by a single pessimistic analytic assumption. The output uses one global infinity-norm budget because that is the contract consumed by matrix multiplication and graph composition.
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 fused expression used by the public spec.
Error after subtracting an approximate allowed-row maximum from every score.
Instances For
Error in the hard-masked numerator vector; applying the mask 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
Numerical certificate 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.
Checkable evidence for nonempty hard-masked softmax rows.
The data fields record the maxima and margins used by execution; the proposition fields establish that they describe the exact and rounded rows. Keeping this evidence together prevents a caller from accidentally pairing a denominator check with a different score matrix or mask.
- rowMaxR : Fin m → FloatLib.Floats.Formats.Flocq.NF β fexp rnd
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 checks 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.
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.
Forward-error theorem for the executable softmax VJP.
This is the training counterpart of approxTensor_softmaxVecSpec. It follows the implementation's
factorization y * (dY - sum (dY * y)); the proof never materializes a dense Jacobian and reuses
the same rounded multiplication, reduction, replication, and subtraction contracts as ordinary
model execution.
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.