Push, Pop, Parallelize: Stack-Augmented Linear Attention via the Delta Rule

Linear attention architectures based on the Delta rule, such as DeltaNet and RWKV-7, combine Transformer-level performance with RNN-like efficiency and provably solve regular language tasks. However, their fixed-size states struggle to capture the recursive, hierarchical structures intrinsic to natural languages. To bridge this gap, we introduce DeltaStack, which augments DeltaNet's associative memory with a lightweight, differentiable stack. Unlike prior approaches that rely on sequential recurrence, DeltaStack formulates stack operations as linear delta-rule updates, enabling a hardware-aware implementation fully parallelizable over sequence length. Theoretically, we prove DeltaStack extends DeltaNet's expressivity to model both regular and hierarchical languages. Empirically, DeltaStack outperforms DeltaNet and Stack-Attention on formal language benchmarks and consistently surpasses DeltaNet baselines in language modeling perplexity and zero-shot performance across scales up to 760M parameters. Our code is publicly available at https://github.com/teeann/DeltaStack.

Current AI language models face a difficult tradeoff: they are either highly capable but computationally expensive to train, or they are fast but struggle to understand complex, layered structures like mathematics and computer code. We created a new architecture called DeltaStack to bridge this gap. Think of it like giving a highly efficient AI a "stack of sticky notes." When the model encounters a nested problem, it can "push" its current context onto the stack, process the inner detail, and then "pop" the note back off to seamlessly resume its previous thought. Crucially, we designed this memory stack mathematically so that the AI can still process massive amounts of text simultaneously during training, maintaining extreme efficiency. This allows the model to accurately track and solve complex algorithmic tasks that require understanding nested logic, all without sacrificing speed. By combining the rapid training efficiency of streamlined models with the deep reasoning capabilities required for complex language, our work paves the way for smarter and more resource-efficient LLMs.