On the Runway Cascade of Transformers for Language Modeling

TL;DR

The paper addresses failure modes in decoder-only transformers arising from runway cascades, where indirect runway-driven information accumulates beyond what direct-path attention controls. It formalizes runway cascade and derives theoretical bounds on information propagation, demonstrating how misalignment can trigger redundant, noise-like cascades. To mitigate this, it introduces runway-aware rewiring, a parameter-free mechanism that reweights attention based on runway context by computing runway coefficients from a summary token, seamlessly integrating with standard attention. Empirically, the approach yields consistent improvements in general language modeling, information retrieval tasks, and extrapolation capabilities across model sizes and context lengths, underscoring enhanced long-range information routing without added parameters. Overall, the work offers a principled, topology-aware solution to transform the propagation dynamics of causal transformers, with clear practical benefits for robustness and scalability.

Abstract

In decoder-only (causal) transformers, the computation graph created by causal masking routes information through both direct-path attention and indirect paths formed by intermediate tokens. We denote these indirect paths between token pairs as their runways. We argue that certain failure modes of causal transformers as observed by a growing body of recent works are likely exacerbated by a misalignment between these two information propagation modes. We formalize runway cascade as a phenomenon whereby this misalignment results in redundancies and irrelevant information cascading to token representations despite adequately learned attention patterns. As a solution, we propose runway-aware rewiring as a more explicit way of incorporating runway context directly into each token's direct-path attention. This mechanism re-wires the attention pattern for each token based on a summary of its runway landscape, enabling awareness of accumulating representational influences and allowing for more balanced information propagation. Our proposed methodology introduces no additional parameters and can seamlessly be integrated into standard attention mechanism. Empirically, our rewired transformer results in steady improvements in general language modeling as well as noticeably stronger information retrieval and extrapolation abilities compared to standard transformers.

Figures (12)

• Figure 1: Visualizing two types of paths between token $s$ and $d$ along the causal graph: (1) direct path (in dotted red) and (2) indirect paths (in solid blue).
• Figure 2: Validation perplexity on the C4 dataset with (a) varying context window sizes (from 512 to 2048) using a 100 million parameter model and (b) varying model sizes from 50 million to 750 million parameters.
• Figure 3: Passkey retrieval test for standard and our rewired transformers. We use 450 million parameter models trained on general language tasks at context length of 1024 and test their retrieval abilities up to context length of 1536.
• Figure 4: Zero-shot passkey retrieval test for standard and our rewired transformers where passkeys are embedded at varying depths.
• Figure 5: Extrapolation at sequence length 2048 across model sizes for standard and rewired transformers.

Theorems & Definitions (13)

• Definition 3.1: Runway
• Lemma 3.2: Lemma 2 from wu2023demystifying
• Theorem 3.3
• Lemma 3.4: Blindspots for standard attention mechanism
• Theorem 3.5: Runway Cascade
• Theorem 1.1
• proof
• Lemma 1.2: Translation invariance of softmax
• proof
• Lemma 1.3: Blindspots for standard attention mechanism