First Passage Attention

First Passage Attention. By Yuxuan Liu, from the University of Bergen. Published on the fifteenth of October. Its code repository is public but slight: three source files.

One way to let softmax attention forget is to add forget gates, and these typically work by summing decay contributions along the path from a key to a query. That accumulation is monotonic. It rules out a particular kind of selectivity: such a mechanism cannot erase intermediate context while preserving older tokens that are still structurally relevant.

First Passage Attention, or FPA, replaces path summation with path extrema. A key's logit is lowered in proportion to the deepest dip of a learned scalar field below that key's level, over the span that lies between the key and the query. What counts is not how much decay piled up along the way, but how low the field ever went.

Two structural facts follow. FPA contains the Forgetting Transformer, FoX, as a special case. And when the field and the levels follow a shared learned walk, the keys left alive, with unpenalised logits, naturally form a stack. On expressiveness, the author shows that one FPA head of constant dimension can realise any subset of keys left alive by subsequent dips. Softmax heads with positional biases, or with token-level forget gates, need a head dimension that grows linearly with 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. Scaled up to a language model of one hundred and ten million parameters, FPA matches FoX in perplexity. It is also the only model compared to achieve perfect passkey retrieval at eight times the training length, and it does so even with its key-value cache compressed to about twenty-five percent.

Moving from a sum to an extreme is a small change of formula with a large change of reach. A sum can only grow with distance, so a gate built on it forgets in order of age. A deepest dip depends on what happened in between, which is what lets older, relevant tokens stay alive while intermediate context is dropped. Passkey retrieval far past the training length, under a compressed cache, is where that shows up in practice.