A research paper from arXiv presents a novel method for measuring exactly how linear each block in a transformer feed-forward network (FFN) actually is. The study, titled "How Linear Is a Transformer Feed-Forward Block? Per-Block Linear Recoverability Is Learned, Not Architectural," challenges the common assumption that FFN blocks serve uniformly as nonlinear stores of computation. Instead, the authors find that linearity is a learned property of individual trained blocks and varies dramatically across layers.
Measuring Per-Block Linearity
The authors treat each FFN block as a position-wise input-to-output map and decompose it into an exact least-squares linear approximation plus a residual. The key metric introduced is R^2_lin, which measures the held-out variance explained by the closed-form linear map. This provides an optimiser-free, per-block assessment of linearity. The study applies this method to three transformer models: GPT-2, Pythia-160m, and llama-160m, spanning all twelve blocks of each.
Heterogeneous and Learned Behavior
The results show that R^2_lin is highly heterogeneous and non-monotone with depth, ranging from near-linear (>0.99) to strongly nonlinear (<0.3) between adjacent blocks. Critically, this variation is not set by the activation function. For example, both GPT-2 and Pythia-160m use GELU activations and have the same width, yet their R^2_lin profiles are sharply different. The authors conclude that recoverability is a learned property of individual trained blocks, not an architectural one.
| Model | R^2_lin Range (across 12 blocks) | Activation Function |
|---|---|---|
| GPT-2 | <0.3 to >0.99 (non-monotone) | GELU |
| Pythia-160m | <0.3 to >0.99 (different profile) | GELU |
| llama-160m | <0.3 to >0.99 (heterogeneous) | (assumed similar) |
The paper also explores the nature of the residual—the part not captured by the linear approximation. A low-rank bilinear probe of the residual recovers only a few additional points of R^2, with gain uncorrelated with residual nonlinearity. This indicates that the unrecovered computation is not a single position-wise product but likely higher-order or distributed structure.
Implications for Model Compression
The R^2_lin measurement serves as a targeted compression signal. Recoverable blocks (those with high R^2_lin) admit large single-layer replacements. For instance, in GPT-2's early FFN, a replacement with 8x fewer parameters results in only +0.77 perplexity increase. Conversely, low-recoverability blocks flag where such compression would be unsafe. This insight could enable more efficient deployment of transformer models in resource-constrained environments.
The study also exposes a methodological pitfall: trained linear baselines can badly under-converge on ill-conditioned transformer activations. The authors address this by reporting the exact closed-form least-squares ceiling throughout, ensuring reliable comparisons.
For enterprise technology decision-makers, this research provides a principled way to identify which layers of a transformer model can be replaced with simpler linear operations without significant performance loss. Such compression could reduce computational costs and memory footprint, making AI more feasible for edge and real-time applications in supply chain and logistics, such as document processing or demand forecasting.
