On the Runway Cascade of Transformers for Language Modeling

TL;DR

This paper introduces runway-aware rewiring in decoder-only transformers to address information propagation issues, leading to improved language modeling, retrieval, and extrapolation without adding extra parameters.

Contribution

It formalizes the runway cascade phenomenon and proposes a parameter-free rewiring method to enhance information flow in causal transformers.

Findings

Improved language modeling performance

Enhanced information retrieval capabilities

Better extrapolation abilities

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.