🤖 machine learning

Feed-Forward Steering in Transformer Residual Dynamics

This paper extends attention-only dynamical theories of Transformers by modeling feed-forward networks as local steering fields, demonstrating that their tangential components are essential for residual motion and model performance while identifying commutator defects as a key metric for determining the feasibility of parallelizing attention and FFN blocks.

Original authors: Timur Mudarisov, Mikhail Burtsev, Radu State

Published 2026-08-04

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Timur Mudarisov, Mikhail Burtsev, Radu State

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are watching a massive, invisible dance party inside a computer brain. This isn't a party with music and flashing lights, but with words and numbers. In the world of artificial intelligence, specifically the "Transformers" that power chatbots and search engines, there is a constant flow of information called a "residual stream." Think of this stream as a river carrying tiny boats (tokens) representing words. For a long time, scientists believed the only thing steering these boats was a mechanism called "Attention." You can picture Attention as a giant, invisible magnet that pulls boats toward each other if they are talking about similar things. If a boat says "king," the magnet might pull it closer to a boat saying "queen." This pulling force is called "aggregation," and it's why words that belong together start to clump up in the same direction.

However, there's another part of the dance floor: the Feed-Forward Network, or FFN. If Attention is the magnet pulling boats together, the FFN is like a local DJ or a personal coach standing right next to each boat, nudging it in a specific direction based on what that specific boat is saying. For years, researchers thought the FFN was just a sidekick, maybe just adjusting the speed of the boats or fixing their volume. But a new paper asks a big question: Is the FFN just a volume knob, or is it actually the steering wheel that decides where the boats go? Understanding this matters because if we know exactly how these boats move, we can make AI faster, smarter, and easier to fix when it makes mistakes.

The Paper's Big Discovery: The DJ is the Captain

This paper, titled "Feed-Forward Steering in Transformer Residual Dynamics," suggests that the old view was missing a crucial piece of the puzzle. The authors propose a new way to look at the dance floor: Attention is the "aggregation" that pulls everyone toward a shared group, but the FFN is a "steering field" that actively pushes individual boats in new directions. They treat the movement of these word-boats as a physics problem, where the boats are particles moving on a sphere (like the surface of a globe).

The researchers found that the FFN doesn't just push the boats forward or backward (which would just change their speed or size). Instead, the most important part of the FFN's job is to push them sideways, or "tangentially." Imagine a boat on a globe; pushing it toward the North Pole changes its size relative to the center, but pushing it sideways changes its actual direction on the map. The paper shows that this sideways push is what actually changes the meaning and path of the word.

What They Did and What They Found

To prove this, the team ran a series of experiments on several famous AI models, including GPT-2, Pythia, Mistral, and Llama-3. They treated the AI like a science lab where they could turn parts on and off or change how they worked.

First, they checked if Attention alone could explain where the boats went. They found a huge "angular-alignment deficit." In plain English, if they only let the Attention magnet work, the boats ended up pointing in the wrong direction compared to where the real AI sent them. This deficit grew larger in bigger, newer models, going from about 0.41 in GPT-2 to 0.63 in Llama-3-8B. This suggests that without the FFN's steering, the AI would be lost.

Next, they performed a "surgery" on the FFN. They split the FFN's push into two parts: the "radial" part (pushing forward/backward) and the "tangential" part (pushing sideways).

When they kept only the radial part, the AI's performance collapsed completely, just as if they had removed the FFN entirely.
When they kept only the tangential part, the AI kept working almost perfectly, with only a tiny drop in quality.

This was a smoking gun. It proved that the FFN's main job is to steer the direction of the words, not just adjust their size.

They also looked at what happens when the Attention magnet pulls too hard, trying to squish all the boats into one tiny clump (a problem called "rank collapse"). They found that the FFN acts like a counter-force. While Attention tries to squeeze the boats together, the FFN pushes them apart, keeping the group diverse and preventing them from all becoming the same word. It's like a DJ who notices everyone is crowding the center and starts playing a beat that makes people spread out again.

A Practical Trick: Speeding Up the Dance

The most exciting part of the paper is a practical trick they discovered. In a standard AI, the boats first get pulled by the Attention magnet, and then the FFN coach gives them a nudge. This happens one after the other (sequentially). The researchers asked: "What if we let the magnet and the coach work at the same time (in parallel)?"

Usually, you can't do this because the coach needs to see where the magnet pulled the boat first. But the team measured a "defect score" to see how much the order mattered for each layer of the AI. They found that for some layers, the order didn't matter much. By identifying these "low-defect" layers, they could run the Attention and FFN parts at the same time.

In GPT-2-large, they managed to speed up the process by about 1.24 times with almost no loss in quality (only +0.02 loss).
In Mistral, they got a 1.10 times speedup with a tiny loss of +0.009.

However, if they tried to do this on layers where the order did matter (high-defect layers), the AI's performance crashed. This suggests that while we can't just make the whole AI run in parallel instantly, we can speed up specific parts of it by knowing exactly how the steering works.

What This Means

The paper doesn't claim to have solved everything about AI, but it changes how we see the engine. It suggests that the FFN isn't just a helper; it's a directional steering field that shapes the geometry of the AI's thoughts. It keeps the AI from getting stuck in a single direction and allows it to navigate complex ideas. By understanding that the FFN is the "steering wheel" and not just the "gas pedal," researchers might be able to build faster, more efficient AI models that don't lose their way. The findings are based on careful measurements and simulations across multiple models, suggesting that this "aggregation-steering" view is a robust way to understand how these digital brains move.

Technical Summary: Feed-Forward Steering in Transformer Residual Dynamics

Problem Statement
Existing continuous theories of Transformer dynamics primarily model the residual stream as a system of particles aggregating on a sphere, driven solely by the attention mechanism. While this "attention-only" framework explains phenomena like rank collapse and clustering, it fails to account for the role of the Feed-Forward Network (FFN). Concurrently, interpretability research identifies FFNs as storage sites for key-value memories and effective targets for editing, yet these views remain disconnected from the dynamical systems perspective. The central problem addressed is how to integrate the FFN into a unified dynamical model of residual stream evolution to explain its geometric impact on token directions and to determine the feasibility of approximating sequential blocks with parallel ones.

Methodology
The authors extend the continuous attention-dynamics framework by incorporating the FFN as a local steering field acting on token states. The model treats the residual stream $x_i = r_i u_i$ (where $u_i$ is the direction on the unit sphere and $r_i$ is the magnitude) as evolving according to:
$\dot{u}_i = \frac{1}{s_i(t)} P_{u_i} \left[ A_i(U) + \Phi_i(u_i) \right]$
Here, $A_i(U)$ represents the non-local aggregation field from attention, and $\Phi_i(u_i)$ is the local steering field provided by the FFN. $P_{u_i}$ is the orthogonal projection onto the tangent space, ensuring that only the tangential component of the vector field alters the direction, while the radial component affects magnitude (and thus angular speed).

The study employs four primary experimental protocols across GPT-2, Pythia, Mistral, and Llama models:

Angular-Deficit Diagnostic: Measuring the alignment between the attention-only update and the realized block update to quantify the FFN's contribution to angular motion.
Radial/Tangential Ablation: Decomposing the FFN output into radial ( $\langle u, \Phi(u) \rangle u$ ) and tangential ( $P_u \Phi(u)$ ) components to isolate which drives direction changes.
Diversity under Aggregation: Analyzing spherical diversity statistics before and after attention and FFN layers to observe how FFN counteracts attention-induced concentration.
Sequential-to-Parallel Defect: Estimating a layer-wise sensitivity score ( $C_\ell$ ) to determine if sequential blocks ( $x + A(x) + \Phi(x+A(x))$ ) can be approximated by parallel surrogates ( $x + A(x) + \Phi(x)$ ) without significant performance loss.

Additional diagnostics include layerwise FFN gate calibration, critical-point geometry analysis (identifying attractors, saddles, and repellers), and prefix-sensitivity tests.

Key Contributions

Aggregation-Steering Framework: The paper formulates Transformer residual dynamics as an aggregation-steering flow on a sphere, where attention provides non-local aggregation and the FFN provides local steering.
Identification of Tangential Steering: The authors mathematically and empirically demonstrate that the tangential component of the FFN field is necessary for motion in residual-direction space. The radial component primarily rescales magnitude, while the tangential component dictates direction.
Critical Direction Geometry: The study identifies that critical residual directions correspond to nonlinear projective equilibria. Contrary to the intuition that FFNs create stable attractors, the analysis reveals that FFN steering actively reorganizes the local geometry toward saddle-like and rotational transport modes, preserving diversity.
Sequential-to-Parallel Defect Metric: A direct defect metric is derived to characterize when a sequential attention-FFN block can be accurately approximated by a parallel, additive map. This defect is governed by the directional derivative term $D_\Phi A$ , distinct from the Lie bracket alone.

Results

Angular Misalignment: The attention branch alone exhibits a substantial angular-alignment deficit relative to the realized block update. This deficit increases with model scale, ranging from 0.41 in GPT-2 to 0.63 in Llama-3-8B, indicating that the FFN is not merely a secondary correction but a primary driver of angular updates.
Ablation Findings: Retaining only the tangential FFN component preserves most model quality (loss increases of 0.07–1.72), whereas retaining only the radial component causes performance to collapse (loss increases of 7.19–12.96), comparable to removing the FFN entirely.
Diversity Preservation: Attention typically contracts residual directions (reducing diversity), while the FFN often restores or reshapes this diversity. The tangential FFN component injects saddle-like and rotational transport modes that prevent excessive concentration under aggregation pressure.
Parallelization Feasibility: Layers with small sequential-to-parallel defects (low sensitivity scores) can be approximately parallelized with modest loss increases. For example, parallelizing 18 low-defect layers in GPT-2-large yields an estimated 1.24× speedup with only +0.02 loss, whereas parallelizing high-defect layers degrades performance rapidly.
Critical Geometry: Component ablation of the critical-direction field shows that FFN-only candidates are predominantly saddle-like (mixed-sign spectra). The FFN does not simply create stable attractors but restructures the phase portrait to maintain high-dimensional representation geometry.

Significance and Claims
The paper claims to support the interpretation of FFN layers as directional steering fields that shape Transformer residual geometry. It argues that the FFN is essential for governing the feasibility of block-level interventions and for maintaining representation diversity against the collapsing pressure of attention.

The authors position their findings as a bridge between dynamical systems theory and interpretability, suggesting that the "steering" view explains why FFNs are effective editing sites and why they are necessary for the model's geometric stability. The practical significance lies in the identification of layers that can be safely parallelized to reduce inference latency without significant performance degradation.

The paper remains modest regarding its scope, noting that the theoretical mean-field results describe full attention, while causal decoder-only models require prefix-indexed dynamics. The critical-point analyses are presented as local diagnostics rather than global descriptions of the full residual space. Furthermore, the practical parallelization results are framed as estimated submodule-time speedups rather than claims of optimized production kernel performance. The authors explicitly state that their findings are diagnostic and do not claim universal guarantees for anti-collapse or perfect parallelization across all architectures.

Technical Summary: Feed-Forward Steering in Transformer Residual Dynamics

More like this