🤖 AI Summary
This work proposes a dual-model coupled recurrent Transformer architecture to enhance the modeling of long-range dependencies in sliding window attention while preserving training parallelism. The framework consists of a prefill model Q that employs full-history attention and a decoder P that uses only sliding window attention augmented with recurrent key-value memory. The two models are aligned during training via a memory consistency loss, yet only the lightweight decoder P is required at inference time. By integrating full-history context with localized recurrent memory and enabling parameter sharing between the models, the approach unifies efficient parallel training with strong long-range modeling capabilities while reducing memory overhead. Experimental results demonstrate that the proposed method significantly outperforms both pure sliding window and latent recurrent Transformer baselines in terms of validation loss and downstream pretraining tasks, with most performance gains retained even under parameter sharing.
📝 Abstract
We introduce \ours{}, a recurrent Transformer architecture with fixed-size memory that generalizes sliding-window attention while remaining parallelizable during training. \ours{} consists of two coupled models: a prefiller $Q$, which leverages full attention\footnote{In practice, we use interleaved full and sliding-window attention for $Q$, as this yields stronger performance. The essential requirement is that $Q$ be more expressive than $P$, with access to the full history.} to produce memory targets $m'_t$, and a decoder $P$, which uses only sliding-window attention and recurrent K/V injection to produce decoder memories $m_t$ for next-token prediction. We train \ours{} with a memory consistency loss that aligns $m_t$ with $m'_t$, allowing inference to use $P$ alone. Empirically, \ours{} improves validation loss and downstream pretraining benchmarks over sliding-window and latent recurrent transformer baselines. Moreover, sharing parameters between $P$ and $Q$ reduces parameter memory while preserving most of the gains.