Technical Summary: InfoFlow: A Framework for Multi-Layer Transformer Analysis

Problem Statement

While the approximation properties of single-layer Transformer architectures have been rigorously studied, a theoretical understanding of the multi-layer setting remains limited. Existing analytical techniques for single layers do not extend directly to multi-layer settings due to the complex interplay between attention layers and feed-forward networks across depth. Specifically, it is unclear how the parameter cost of $\epsilon$-approximation scales with architectural parameters (sequence length $T$, depth $L$, embedding dimension $E$, and head count $h$) and which complexity measures of the target function determine this scaling. Current theory lacks a principled framework to explain why multi-layer Transformers exhibit fundamentally different approximation capabilities compared to single-layer ones, particularly regarding retrieval tasks.

Methodology: The InfoFlow Framework

To address the intractability of direct theoretical analysis in multi-layer settings, the authors propose InfoFlow, a coarse-grained framework inspired by effective theories in physical sciences (e.g., thermodynamics). Instead of tracking continuous hidden vector representations, InfoFlow abstracts the computation by tracking discrete information sets.

Core Components

1. Information Sets ($I(t, l)$): For each token position $t$ at layer $l$, the framework defines an information set $I(t, l) \subseteq \{1, \dots, T+1\}$ representing the collection of input positions whose information remains accessible from the token state.
2. Structural Mechanisms: The framework is grounded in three identified modes of information propagation, derived from theoretical bounds and experimental observations:
   • Max-Position Retrieval: Softmax attention efficiently retrieves only the token attaining the maximum attention score. Retrieving the $k$-th largest token ($k \ge 2$) incurs a parameter cost growing exponentially with sequence length $T$.
   • Global Information Aggregation: A single head can aggregate information from the entire sequence, but the parameter cost for decoding this aggregated state scales exponentially with the number of input tokens ($O(\epsilon^{-T/E})$).
   • Specific Position Aggregation: With positional encoding, attention can select a fixed set of positions independent of token content.
3. Update Rules: The framework defines how $I(t, l)$ updates to $I(t, l+1)$ based on the chosen propagation mode (e.g., union of argmax positions and residual connections).
4. Parameter Cost Law: The framework associates a parameter cost with each update rule. Crucially, decoding information from a set of size $T_0$ requires a parameter count of $O(\epsilon^{1 - T_0 d / E})$, where $d$ is the input dimension and $E$ is the embedding dimension.
5. Number of Comparison: To quantify task complexity, the authors define the "Number of Comparison" using a "Tree of Comparison" structure. This measures how many tuples of tokens must be compared to compute a target output, characterized by a "Dimension of Comparison" ($\beta_1$) and "Order of Comparison" ($\beta'$).

Key Contributions

The paper makes three primary theoretical and methodological contributions:

1. Depth Separation Theorem (Theorem 1): The authors establish that single-layer and multi-layer Transformers possess fundamentally different approximation capabilities.

   • For certain retrieval tasks (e.g., finding the minimum of pairwise dot products), a single-layer Transformer requires a parameter count of $\Omega(\epsilon^{-k(T)})$, where $k(T)$ grows linearly with sequence length $T$.
   • In contrast, a two-layer Transformer with a single head per layer achieves the same precision $\epsilon$ with a parameter count of only $O(\epsilon^{-1})$, independent of $T$.
   • This separation arises because multi-layer models can decompose global minimization into successive local minimizations, whereas single-layer models must resolve all comparisons simultaneously.

2. Structural Limitations of Attention (Theorems 2 & 3):

   • Theorem 2: Formalizes the limitation that softmax attention can only efficiently retrieve the maximum score token. Retrieving the $k$-th largest token ($k \ge 2$) requires $\Omega(\epsilon^{-(T-k-1)/(n+1)})$ parameters, confirming the exponential cost of non-max retrieval.
   • Theorem 3: Quantifies the cost of decoding aggregated information, showing that recovering a set of $T_0$ tokens from an embedding of dimension $E$ requires parameters scaling as $O(\epsilon^{1 - T_0 d / E})$.

3. The InfoFlow Framework: A unified abstraction that combines information set updates with the derived parameter cost laws. This framework allows for the estimation of approximation rates for multi-layer Transformers prior to training by comparing the "Number of Comparison" of the target against the capacity of the InfoFlow model.

Results and Validation

The authors validate InfoFlow both theoretically and experimentally:

• Theoretical Consistency: The framework recovers known approximation bounds for single-layer Transformers (e.g., results from Yu et al. [2026] regarding head count and universality) and is consistent with the established depth separation result.
• Intrinsic Dimension Phenomenon (Experimental):
  • Prediction: For a target requiring $D$ distinct pairwise comparisons (intrinsic dimension $D$), a two-layer Transformer requires at least $D$ heads in the first layer to approximate the target efficiently.
  • Validation: Experiments on a permutation-invariant target with $D \in \{2, \dots, 6\}$ show a sharp phase transition. Configurations with $(h_1, h_2) = (D, D)$ achieve validation NMSE $\approx 10^{-5}$, while configurations with fewer heads (e.g., $D-1$) result in NMSE $\approx 10^{-2}$ to $10^{-3}$, regardless of sequence length. This confirms that insufficient head count leads to a failure to track necessary information sets.
• Approximation Hardness of Higher-Order Tasks (Experimental):
  • Prediction: Targets requiring simultaneous comparison of three or more tokens (Order of Comparison $\beta' > 2$), such as the "triangle-center" problem ($\min \|x(t_1) + x(t_2) + x(t_3)\|^2$), cannot be approximated efficiently by any fixed-size Transformer as $T$ grows. The parameter cost grows super-polynomially with $1/\epsilon$.
  • Validation: Training two-layer Transformers on the triangle-center task for varying $T$ shows that NMSE rapidly approaches 1.0 (random guessing) as $T$ increases, regardless of model size (heads or embedding dimension). This confirms the theoretical prediction that fixed architectures cannot handle the cubic growth in comparison complexity.

Significance and Claims

The paper claims that InfoFlow provides a principled framework for reasoning about the approximation efficiency of multi-layer Transformers. Its significance lies in:

1. Bridging Theory and Practice: It offers a tractable abstraction that remains consistent with experimental observations on trained networks, explaining phenomena like the "intrinsic dimension" requirement for head counts.
2. Predictive Power: It yields concrete predictions for settings where direct theoretical analysis is currently intractable, specifically predicting the failure of fixed-size models on higher-order retrieval tasks.
3. Architectural Insight: It clarifies that the depth of a Transformer is not merely a stacking of layers but a mechanism that fundamentally alters the complexity of information retrieval, allowing for the decomposition of global tasks into local steps.

The authors note limitations, including that the cost law currently abstracts away the rank of attention weight matrices and that the framework is primarily validated on retrieval-type targets where the active index set does not grow with $T$. Future work is suggested to extend these tools to more general targets and practical Transformer variants.