Source-linked AI summary
RoPE attention is an exact forward-pass gradient step with softmax intact
Julie Huang, Maggie Chlon, Leon Chlon
TL;DR
The paper asks how to express full RoPE–softmax attention as an exact gradient step in original projection-weight coordinates, including biases. It constructs a query-conditioned effective matrix using the exponential divided difference; the representation is exact for every query and finite projection weights, while frozen-matrix reuse has substantial reported error against its constant-anchor baseline.
Problem
The paper addresses how to express full RoPE–softmax attention as an exact gradient step in projection-weight coordinates, including biases.
Source: §1 · p. 1Method
It uses the exponential divided difference to construct a query-conditioned effective matrix and per-token decomposition for the RoPE–softmax readout.
Source: §1 · p. 1, §2.2 · p. 2Results
The representation is exact for every query and finite projection weights; frozen reuse has rRMSE 2.608 against the constant-anchor output for anchor 0.
Source: §2.2 · p. 2, §5.2 · p. 5Takeaways & Limitations
The cache supplies outer-product directions while the query selects positive coefficients, yielding rank(ΔM_i) ≤ d in this representation.
Source: §2.2 · p. 2For a fixed finite bank and query position, a nonconstant head has no globally exact affine query readout; changing position or bank can also change features and mean.
Source: §4.2 · p. 3, §4.3 · p. 4
Abstract
from arXiv · showhide
We derive an exact gradient-step representation of the RoPE-softmax forward pass. For every deterministic RoPE-softmax attention head with arbitrary affine projection weights, we construct a query-dependent effective matrix $\Delta M_i$ satisfying $y_i = \mu_i + u_i^\top \Delta M_i$, where $\mu_i$ is the uniform mean of the attended values and $u_i$ is the augmented query input. The construction applies the classical exponential divided difference $\rho = \phi_1$ to retain the softmax exactly. Its positive coefficients give a unit gradient-step representation on a query-conditioned quadratic objective. The same function connects the RoPE generator to exact positional finite differences. We derive a tokenwise formula for the error of reusing one query's matrix and prove that a nonconstant finite-cache head cannot admit a globally exact affine query readout. Reconstruction checks and frozen-reuse calibration on one pretrained Qwen2.5-0.5B layer verify the representation and quantify the correction required when one query's matrix is reused.
1 Contribution and relation to existing constructions
The paper derives an exact gradient-step representation of full RoPE–softmax attention in original projection-weight coordinates, including biases. It distinguishes this construction from linear-attention gradient views and connects its matrix function to positional derivatives and finite differences.
1 Contribution and relation to existing constructions
- The exponential divided difference yields an exact query-conditioned matrix and per-token decomposition for full RoPE–softmax attention, including projection biases.The construction also has a corresponding quadratic gradient-step representation.Source: §1 · p. 1
- The same matrix-function identity connects the construction to positional derivatives and exact finite differences.A boundedness proposition also characterises when a globally shared affine readout is obstructed.Source: §1 · p. 1
- The factorisation uses the classical φ1 function from exponential integrators to retain the softmax readout.Softmax attention is described as a normalised kernel average, with RoPE entering the scores through relative rotation.Source: §1 · p. 1
- Unlike linear-attention gradient constructions, its coefficients retain the exponential and its normalisation for arbitrary projection weights.The paper contrasts this with prior linear-attention gradient-step and fast-weight constructions.Source: §1 · p. 1
2 Exact effective weights with softmax intact
For finite attended banks and arbitrary affine projections, the paper represents RoPE–softmax attention exactly as a uniform value mean plus an augmented query acting on a query-conditioned matrix. The exponential divided difference preserves softmax, while the construction also yields a unit gradient step and a low-score linear-attention limit.
2.1 RoPE, affine projections, and the attended bank
- The setup uses affine query, key, and value projections with RoPE, a finite nonempty permitted-key set, and its uniform value mean.Causal masking selects the permitted keys.Source: §2.1 · pp. 1–2
2.2 The exponential divided difference
- The positive exponential divided-difference coefficients preserve the exact exponential response; as scores approach zero, they approach 1/N_i.In that limit, the readout becomes the centred outer-product form underlying linear-attention gradient constructions.Source: §2.2 · p. 2
- Theorem 1 gives an exact effective-weight representation for every query and all finite projection weights and inputs.The readout adds the query-dependent matrix term to the uniform mean of the permitted values.Source: §2.2 · p. 2, §2.1 · p. 2
- The cache supplies the outer-product directions, while the query selects their positive coefficients, giving rank(ΔM_i) ≤ d in this representation.The factorisation is expressed through the matrix product ΔM_i = L T_i.Source: §2.2 · p. 2
- For each fixed query, the effective matrix is a unit gradient step on a query-conditioned quadratic objective.The objective holds the query, features, centred values, and coefficients fixed.Source: §2.2 · p. 2
3 The same function links RoPE derivatives and finite differences
The same entire function links the RoPE generator to both infinitesimal positional response and exact finite differences. This connection applies to frozen rotated features, including zero-frequency blocks.
3 The same function links RoPE derivatives and finite differences
- For frozen content rotated as z(p) = R(p)z_0, A gives the infinitesimal response and Aρ(δA) gives the exact finite-difference quotient.The power-series definition covers zero-frequency blocks, and the same function converts scalar softmax and matrix rotary increments.Source: §3 · p. 3
4 Residual coordinates, gauge, and exact reuse error
The residual write is an exact query-dependent readout, and score shifts generate a family of matrices for the same output. A nonconstant finite-bank head has no globally exact affine query readout, while freezing one query’s matrix omits a tokenwise correction.
4.1 Residual write and score-shift family
- Projected head writes are summed for multi-head attention, and a pre-normalised head uses the normalised input in its exact residual write.The readout’s differential also includes query derivatives of the coefficients.Source: §4.1 · p. 3
- Subtracting the same query-dependent score shift from every attended score preserves softmax and yields a family of exact matrices for the same output map.Key shifts form a subfamily of dimension at most d at a fixed position.Source: §4.1 · p. 3
4.2 Global affine representability
- For a finite key/value bank at a fixed query position, a nonconstant softmax head cannot have one fixed affine readout over all query coordinates.The output is bounded because it is a convex combination of fixed values, whereas a nonzero affine matrix would be unbounded along some direction.Source: §4.2 · p. 3
4.3 The exact term omitted by freezing
- Freezing an anchor query’s matrix omits an exact projected correction resolved token by token; changing position or bank can also change the shared features and mean.The fixed-position condition keeps features shared; changing position changes L_i, while changing the bank can change both features and mean.Source: §4.3 · p. 4
5 Pretrained reconstruction and frozen-reuse calibration
The pretrained-layer checks reconstruct prefix-only and full causal attention readouts, then calibrate the residual from reusing a frozen query matrix. Reconstruction errors are very small, while frozen reuse has a substantial error relative to the constant-anchor output.
5.1 Protocol and reconstruction precision
- The checks use 16 article prefixes of 511 tokens, all 14 heads, and 32 alternative one-token continuations, including each continuation’s own key/value in the full causal readout.Projection biases and pre-normalised head inputs are included.Source: §5.1 · p. 4
- Float64 combined prefix and full reconstruction MSEs are approximately 9.8 × 10^-30 and 9.6 × 10^-30; against the native FP32 combined write, MSE is 7.01 × 10^-14 and maximum absolute difference is 1.4 × 10^-5.The float64 figures compare algebraically equivalent evaluations; the native comparison measures agreement at model execution precision.Source: §5.1 · p. 4
5.2 Frozen reuse as calibration of the omitted term
- The frozen-reuse calibration reports pooled MSE for the combined projected prefix write, comparing the observed-next-token anchor with averages over 32 anchors within each of 16 contexts.The matrices use the original-score gauge.Source: Table 1 · p. 5
- Frozen reuse has rRMSE 2.608 for anchor 0 and 2.795 across all anchors, while the target RMS is 0.670545 and 0.670824, respectively.The anchor-0 RMSE ratio against the cache mean is 1.615; per-head ratios are relative to each head’s constant-anchor output.Source: §5.2 · pp. 4–5
- Adding the tokenwise correction omitted by the frozen matrix recovers the exact attention output for each query.The frozen-matrix row measures the MSE of that omitted correction.Source: §5.2 · pp. 4–5