On the Runway Cascade of Transformers for Language Modeling

Hunjae Lee; Corey Clark

cs.LGJan 20, 2026

On the Runway Cascade of Transformers for Language Modeling

Authors: Hunjae Lee, Corey Clark

Organizations: Department of Computer Science, Southern Methodist University, Dallas TX USA.

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 & tables

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).

Panel of Figure 2 — 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. Note: lower values indicate better performance.

Panel of Figure 3 — Figure 3: (a) 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. (b) An extension of the same experiment from (a) using longer context (up to 2.5k) on 120 million parameter models (c) Zero-shot passkey retrieval test for standard and our rewired transformers where passkeys are embedded at varying depths. Note: higher values indicate higher accuracy and therefore better performance.

Appendix figures & tables16 assets

Supplementary material from the paper’s appendix.

Appendix

$d=n_{\text{heads}}=n_{\text{layers}}$	$d_{\text{model}}$	Model Parameter Count
8	512	50,927,104
12	768	123,607,296
16	1024	252,890,112
20	1280	457,726,720
24	1536	756,969,984

Appendix

Table 1: Model Parameterization.

Dataset	Std. Transformer	Rewired (Ours)
ARC-Easy	43.81%	42.89%
ARC-Challenge	24.66%	26.11%
HellaSwag	41.75%	42.49%
PIQA	68.61%	68.82%
CommonsSenseQA	19.49%	19.57%

Appendix

Table 2: Reasoning performance of standard and our rewired transformers. Best result from each dataset is marked in blue.

Figure 4: Extrapolation at sequence length 2048 across model sizes for standard and rewired transformers.

Panel of Figure 5 — Figure 5: Analysis of extrapolation performance (validation perplexity). Notably, we find that our rewired models show comparable extrapolation performance to standard transformers with 150 million additional parameters.

Panel of Figure 6 — Figure 6: Extrapolation comparison between standard and rewired transformers at various model sizes: (a) 100 million parameters, (b) 250 million parameters, (c) 450 million parameters, and (d) 750 million parameters.

Figure 7: Dirichlet energy comparison between standard and our rewired transformer.

Sequence Length	Standard Transformer	Rewired (Ours)
1024 (in-context)	4.2967	4.2626
2048	5.5783	5.3423
4096	7.0703	6.8023

Appendix

Table 3: Perplexity comparison between standard and rewired transformers across sequence lengths on PG-19 [ 33 ] . Lower values indicate better performance.

Metric	Standard Transformer	Rewired (Ours)
Training Throughput (tokens/sec)	88955	61039
Inference Time (sec)	22.692	25.203
Inference Throughput (tokens/sec)	9016.20	8117.88
Inference Memory Overhead (MB)	7902.76	9028.30

Appendix

Table 4: Throughput and overhead comparison between standard and our rewired transformer.

$d$	$d_{\text{model}}$	Std. Transformer	Rewired (Bilinear Form)	Parameter $\Delta$
	Total Parameter Count
8	512	50,927,104	50,959,872	+0.06%
12	768	123,607,296	123,656,448	+0.04%
16	1024	252,890,112	252,955,648	+0.03%
20	1280	457,726,720	457,808,640	+0.02%
24	1536	756,969,984	757,068,288	+0.01%

Appendix

Table 5: Model Parameterization and total parameter difference (in percent increase) between the rewired transformer using bilinear form and standard transformer at each model size.

Figure 8: Validation loss comparison at various model sizes between rewired model using dot product (in purple) and bilinear form (in orange). Standard transformer (in green) is included for baseline comparison.

Figure 9: Context window size ablation for rewired transformer with dot product (purple diamond) and bilinear form (orange circle) runway-coefficient calculation. Standard transformer performance (in green square) is also plotted for baseline information.

Figure 10: Extrapolation analysis for our rewired transformers using dot product (in purple diamond) and bilinear form (in orange circle). Standard transformer (in green square) is used for baseline visualization. All models are using 450 million parameters.

Depth	Standard Transformer	Rewired ( $h_{d-1}$ )	Rewired (Global Summary)
$d=12$	3.2547	3.2384	3.2397
$d=24$	2.9053	2.8936	2.8940

Appendix

Table 6: Ablation over rewiring signal choice. Lower values indicate better performance.

Depth	Standard Transformer (ALiBi)	Rewired (ours)
$d=12$	3.2798	3.1945
$d=24$	2.9171	2.8936

Appendix

Table 7: Comparing the effects of uniform, distance-based decay versus our rewiring technique.

Figure 11: Visualization of the strength of edge down-scaling factor $\overline{r}$ in the overall lower-triangular attention map (averaged over all heads). Lower values (which are more purple) indicate stronger down-scaling and therefore more aggressive rewiring.

Panel of Figure 12 — Figure 12: The average down-scaling effect of $\overline{r}$ with respect to each (a) source token position and (b) destination token position.

Transformer Attention

Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.

See all 320 papers

Residual Connections and the Causal Shift: Uncovering a Structural Misalignment in Transformers
Feb 16, 2026Jonathan Lys, Vincent Gripon, Bastien Pasdeloup +4Transformer Residual Learning
IMT Atlantique, Lab-STICC, UMR CNRS 6285, Brest, France · Sony Europe Ltd. Stuttgart Technology Center, EUREC, Germany
Adaptive Computation Depth via Learned Token Routing in Transformers
Apr 18, 2026Ahmed Abdelmuniem Abdalla MohammedEfficient Transformer Inference Transformer
Independent Researcher
RIFT: Token Influence Follows the Path of Least Resistance
Oct 12, 2026Tingan Jin, Shuhang Dong, Haosong LiTransformer Interpretability LLM Interpretability
University of California, Los Angeles Los Angeles, CA, USA · Independent Researcher

Feb 16, 2026J/K move · Enter open · S save

Residual Connections and the Causal Shift: Uncovering a Structural Misalignment in Transformers

Jonathan Lys, Vincent Gripon, Bastien Pasdeloup, Axel Marmoret+3

IMT Atlantique, Lab-STICC, UMR CNRS 6285, Brest, France · Sony Europe Ltd. Stuttgart Technology Center, EUREC, Germany

Transformer Residual Learning Decoder-Only Language Models

Open paper AskPDF arXiv

Large Language Models (LLMs) are trained with next-token prediction, implemented in autoregressive Transformers via causal masking for parallelism. This creates a subtle misalignment: residual connections tie activations to the current token, while supervision targets the next token, potentially propagating mismatched information if the current token is not the most informative for prediction. In this work, we empirically localize this input-output alignment shift in pretrained LLMs, using decoding trajectories over tied embedding spaces and similarity-based metrics. Our experiments reveal that the hidden token representations switch from input alignment to output alignment deep within the network. Motivated by this observation, we propose a lightweight residual-path mitigation based on residual attenuation, implemented either as a fixed-layer intervention or as a learnable gating mechanism. Experiments on multiple benchmarks show that these strategies alleviate the representation misalignment and yield improvements, providing an efficient and general architectural enhancement for autoregressive Transformers.

Abstract

Transformer Attention

Explore similar work

Residual Connections and the Causal Shift: Uncovering a Structural Misalignment in Transformers

Adaptive Computation Depth via Learned Token Routing in Transformers

RIFT: Token Influence Follows the Path of Least Resistance