Tangential FFN Residuals
Implementation & benchmark of arXiv:2608.02071 — Feed-Forward Steering in Transformer Residual Dynamics
Source paper: Feed-Forward Steering in Transformer Residual Dynamics arXiv:2608.02071 ⓘ · analyzed Aug 31, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Project each FFN residual update onto the tangent space of the current token residual direction before adding it to the stream. This preserves the component that changes representation direction while suppressing norm-only motion, which may reduce residual-norm drift and aggregation-induced representation collapse.
Formulas
Mathematical statement
The paper writes each token residual as x_i=r_i u_i, where r_i=||x_i||>0 is the residual magnitude and u_i=x_i/||x_i|| lies on the unit sphere S^{d-1}. For an ambient update F_i=A_i(U)+Phi_i(u_i), the direction dynamics are dot{u}_i=(1/s_i(t))P_{u_i}[A_i(U)+Phi_i(u_i)], where s_i(t)>0 is the angular-speed scale and P_{u_i}=I-u_i u_i^T is the orthogonal projector onto the tangent space T_{u_i}S^{d-1}. The FFN field decomposes as Phi_i(u_i)=<u_i,Phi_i(u_i)>u_i+P_{u_i}Phi_i(u_i). The first term is radial and changes residual magnitude through dot{r}_i=<u_i,F_i>; the second is tangential and is the only FFN component that directly changes residual direction. The proposed network operation applies the adapted projector P_x=I-xx^T/(||x||^2+epsilon) to the FFN output y, with x the current residual token, y=FFN(Norm(x)), and epsilon>0 for numerical stability. Because the finite-epsilon projector is only approximately orthogonal, measure x^T(P_x y) during experiments and optionally use the exact denominator ||x||^2 when activations are safely nonzero.
Implementation notes
1. Integration point: modify the FFN residual branch in a pre-norm Transformer block, after the FFN produces y=FFN(RMSNorm(x)) and immediately before the residual addition. Treat each x[b,t,:] as one token residual vector. Leave attention unchanged for the first experiment, and use the same residual multiplier alpha as the baseline.
2. Pseudocode:
input x[B,T,D]
h = RMSNorm(x)
y = FFN(h)
denom = sum(x*x, dim=-1, keepdim=True) + eps
radial_coeff = sum(x*y, dim=-1, keepdim=True) / denom
if mode == projected:
y_used = y - radial_coeff * x
elif mode == soft_projected:
y_used = y - beta * radial_coeff * x
else:
y_used = y
x_next = x + alpha * y_used
Use eps=1e-6 in float32, or scale eps relative to the mean squared activation. In the soft version, beta is fixed at 0.25, 0.5, or 1.0; beta=1 is the hard tangent projection. Also test a norm-preserving variant y_used=y_projected*||y||/(||y_projected||+eps), so directional projection is separated from a reduction in update magnitude.
3. Computed directly from the paper's mathematics: the radial coefficient, projected FFN vector, and tangent-space residual. Log the normalized radial fraction rho=|x^T y|/(||x||||y||+eps), the post-projection orthogonality error |x^T y_used|/(||x||||y_used||+eps), residual norm ||x||, and angular displacement arccos(clamp(x^T x_next/(||x||||x_next||),-1,1)). These diagnostics distinguish whether any gain comes from removing radial motion, changing update scale, or regularizing directions.
4. First cheap experiment: train a 6-layer, 8-head small Transformer on WikiText-103 or a 10k-example OpenWebText subset with baseline FFNs, hard projection, beta=0.5 projection, and norm-preserving projection. Match initialization, batch size, optimizer, parameter count, and training FLOPs. Compare validation perplexity and loss versus optimizer steps, residual-norm drift, angular diversity across tokens, gradient variance, and activation overflow frequency. Then apply the variants to GPT-2-small for one short fine-tuning run to test compatibility with pretrained weights. A positive signal is equal-or-better perplexity with lower norm drift or gradient variance and higher directional diversity; a negative result is a persistent loss increase after equal compute or severe degradation when projection is enabled in all layers.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.