cs.LGMay 18, 2026

InfoFlow: A Framework for Multi-Layer Transformer Analysis

Authors: Penghao YuHaotian JiangZeyu BaoQianxiao Li

Organizations: Department of Mathematics National University of Singapore · Institute for Functional Intelligent Materials National University of Singapore · National University of Singapore · Department of Mathematics Institute for Functional Intelligent Materials National University of Singapore

Abstract

While the approximation properties of single-layer Transformer architectures have been studied in recent works, a rigorous theoretical understanding of the multi-layer setting remains limited. In this work, we establish that multi-layer Transformers possess fundamentally different approximation capabilities from single-layer ones: for certain retrieval tasks, any single-layer Transformer requires least Ω(εk)Ω(\varepsilon^{-k}) parameters to achieve precision ε\varepsilon, where kk grows linearly with sequence length TT, whereas a two-layer Transformer with a single head per layer achieves the same approximation precision with at most O(ε1)O (\varepsilon^{-1}) parameters. To understand this separation, we identify two structural mechanisms underlying multi-layer approximation. Specifically, softmax attention can only efficiently retrieve the token attaining the maximum attention score, incurring exponential-in-length parameter cost for kk-th largest retrieval with k2k \geq 2. Moreover, the parameter cost of decoding coupled information scales with the size of the retrieved token set. Motivated by these findings, we propose InfoFlow, a framework for multi-layer Transformers. The framework tracks an information set of accessible input positions at each token and layer, assigning an explicit approximation rate to each mode of information propagation. This abstraction recovers known approximation bounds, remains consistent with experimental observations on trained networks, and yields concrete predictions in settings where direct theoretical analysis is currently intractable. Our results provide a principled framework for reasoning about the approximation efficiency of multi-layer Transformers.

Explore similar work

Oct 4, 2025cs.LG

Performance-Efficiency Tradeoffs in Transformers: An Approximation Theory Perspective

Transformers have achieved remarkable successes across a wide range of applications, yet the theoretical foundation of their model efficiency remains underexplored. In this work, we investigate how the model parameters -- mainly attention heads and head dimensions -- should be allocated across layers to balance expressivity and efficiency. We first provide mathematical analysis on the role of early layers in information extraction from an approximation perspective, with a theoretical characterization on the trade-off between the number of heads and head dimension under a fixed parameter budget. In addition, we uncover and prove the \emph{saturation} behavior of softmax activations: Continuously increasing head dimensions can lead to diminishing returns in learning errors, particularly for long sequences. Supported by both theory and experiments, this saturation pattern suggests that later layers can operate more efficiently with reduced parameters. Combining these insights, we propose principled strategies for allocating attention heads and dimensions across Transformers' layers, shedding light on theoretically-grounded model efficiency of Transformer-based architectures.
Ruoxi Yu, Haotian Jiang, Jingpu Cheng +3
Apr 23, 2026cs.LG

The Recurrent Transformer: Greater Effective Depth and Efficient Decoding

Transformers process tokens in parallel but are temporally shallow: at position tt, each layer attends to key-value pairs computed based on the previous layer, yielding a depth capped by the number of layers. Recurrent models offer unbounded temporal depth but suffer from optimization instability and historically underutilize modern accelerators. We introduce the Recurrent Transformer, a simple architectural change where each layer attends to key-value pairs computed off its own activations, yielding layerwise recurrent memory while preserving standard autoregressive decoding cost. We show that the architecture can emulate both (i) a conventional Transformer and (ii) token-to-token recurrent updates under mild assumptions, while avoiding optimization instability. Naively, prefill/training appears bandwidth-bound with effective arithmetic intensity near 11 because keys and values are revealed sequentially; we give an exact tiling-based algorithm that preserves the mathematical computation while reducing HBM traffic from Θ(N2)Θ(N^2) to Θ(NlogN)Θ(N\log N), increasing effective arithmetic intensity to Θ(N/logN)Θ(N/\log N) for sequence length NN. On 150M and 300M parameter C4 pretraining, Recurrent Transformers improve cross-entropy over a parameter-matched Transformer baseline and achieve the improvement with fewer layers (fixed parameters), suggesting that recurrence can trade depth for width, thus reducing KV cache memory footprint and inference latency.
Costin-Andrei Oncescu, Depen Morwani, Samy Jelassi +3
Aug 4, 2026cs.CL

LoopMTP: A looped transformer guided by latent multi-token prediction

Looped transformers have emerged as a parameter-efficient alternative to scaling depth for strong reasoning. By reusing one stack of layers across TT iterations, they attain the effective depth and reasoning capabilities of larger models at a fixed parameter count. Yet existing approaches suffer from latent overthinking and undifferentiated computation, largely because intermediate representations receive no guidance across loops. Multi-token prediction (MTP) supplies exactly the dense, forward-looking supervision the loop is missing. We propose \textsc{LoopMTP}, which links the two through a structural correspondence in latent space: a model that loops TT times can anticipate TT future tokens. \textsc{LoopMTP} realizes this by softly aligning the hidden state of loop tt with the embedding of the token tt steps ahead, while a lightweight gate preserves useful information across iterations. \textsc{LoopMTP} improves average accuracy by up to 8.1% (relative) over the non-looped baseline, with training remaining stable for up to 15 loops.
Behzad Shomali, Markus Frey, David Berghaus +2