First passage attention improves selective memory in language models
First Passage Attention
Computation and Language
Summary
Language models read and remember tokens in a sequence, but they usually can't forget unimportant parts selectively while keeping important older information. The authors propose First Passage Attention (FPA), a new method that can better decide which parts of a sequence to keep or forget by looking at dips in a learned signal rather than adding up effects. Their experiments show FPA can remember key tokens longer and more accurately than existing methods, even on long sequences. This makes it promising for building language models that understand and use context more effectively.
What this means in practice
- •For natural language processing teams: Improve language models to maintain relevant long-range context accurately by incorporating First Passage Attention mechanisms for better token selection and retrieval.
- •For ai service engineers: Build scalable large-scale transformer systems with compressed key-value caches that retain perfect passkey retrieval at longer sequence lengths using First Passage Attention.$Commercial implications: Enables development of more efficient and accurate large transformer-based AI services with longer context windows, valuable to cloud AI providers.
Authors
Yuxuan Liu
Abstract
Augmenting softmax attention with forget gates typically relies on summing decay contributions along the key--query path. Because such accumulation is monotonic, these mechanisms cannot selectively erase intermediate context while preserving older, structurally relevant tokens. We propose First Passage Attention (FPA), which shifts from path summation to path extrema by lowering a key's logit proportionally to the deepest dip of a learned scalar field below that key's level over the intervening span. FPA encompasses the Forgetting Transformer (FoX) as a special case, and when the field and levels follow a shared learned walk, the alive keys with unpenalised logits naturally form a stack. We show that a constant-dimension FPA head realises any subset of keys left alive by subsequent dips, whereas softmax heads with positional biases or token-level forget gates require a head dimension linear in the number of keys. In synthetic experiments, one-layer FPA heads converge to these target sets while comparable softmax and FoX baselines fail on every seed. Scaling to a 110M-parameter language model, FPA matches FoX in perplexity and uniquely achieves perfect passkey retrieval at $8\times$ the training length, even with its key--value cache compressed to about $25\%$.