InfoFlow: A Framework for Multi-Layer Transformer Analysis

TL;DR

This paper introduces InfoFlow, a framework that analyzes the approximation capabilities of multi-layer Transformers, revealing their efficiency advantages over single-layer models for certain retrieval tasks.

Contribution

The work provides a theoretical framework for understanding multi-layer Transformer approximation properties and introduces InfoFlow to analyze information propagation and efficiency.

Findings

Multi-layer Transformers require fewer parameters than single-layer ones for certain tasks.

Softmax attention efficiently retrieves only the maximum scoring token, leading to exponential costs for other retrievals.

InfoFlow accurately predicts approximation bounds and aligns with experimental observations.

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 $\Omega (\varepsilon^{-k})$ parameters to achieve precision $\varepsilon$, where $k$ grows linearly with sequence length $T$, whereas a two-layer Transformer with a single head per layer achieves the same approximation precision with at most $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 $k$-th largest retrieval with $k \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.