First Passage Attention Replaces the Sum of Decays With an Extremum: Passkey Retrieval at 8× Length

Forget gates in attention almost always work like a penalty jar: the farther a key is from the query and the more tokens lie between them, the more its logit drops. The accumulation is monotonic, so the model cannot erase a chunk of the middle while keeping old, structurally important tokens. The First Passage Attention (FPA) paper proposes different math. Instead of summing decays along the path, it takes an extremum. On a 110M-parameter model, this yields perfect passkey retrieval at 8 times the training length, even with the KV-cache compressed to roughly 25%. Perplexity is on par with Forgetting Transformer (FoX).

How the mechanism works

FPA has a learnable scalar field and levels attached to keys. A key's logit is reduced in proportion to the deepest dip of the field below that key's level on the segment between the key and the query. The name comes from random-walk theory, where "first passage" means the moment a trajectory first crosses a given boundary.

The main consequence concerns the order of events. If the field dropped deep after the key, the key is penalized. If the dip was shallow, or the key sits below it, the key is left untouched. A sum cannot tell these cases apart; an extremum can. The authors call FoX a special case of FPA, which means the new scheme is strictly broader than the old one.

If the field and the levels follow the same learnable walk, the "live" keys with unpenalized logits form a stack. A new dip displaces everything above it, while everything deeper remains accessible. This resembles the structure of brackets or nested scopes: close a block and forget its contents, but the outer context stays intact.

What the experiments showed

The theoretical result is stated in terms of head dimension. A constant-dimension FPA head can implement any subset of keys that remain alive after successive dips. Softmax heads with positional biases, and FoX with token-level gates, need a dimension linear in the number of keys to do the same. This is a gain in expressiveness at a fixed parameter budget.

On synthetic tasks, single-layer FPA heads converge to the target key sets. Comparable softmax and FoX baselines failed on every seed. The authors then scaled the method to a language model:

• model size: 110M parameters;
• perplexity: on par with FoX, with no drop relative to the closest analog;
• passkey retrieval: perfect result at 8 times the training length, the only one among the compared methods;
• KV-cache: compression to roughly 25% without losing that result;
• synthetic tasks: FPA converges to the target sets on single-layer heads, while the baselines fail on every seed.

Practical takeaways

The combination of length extrapolation and a compressed cache hits the biggest cost item of long context: KV memory. If a stack of live keys makes it safe to discard the remaining tokens, the cache can be trimmed by a rule the model learned itself rather than by a heuristic like "keep the last N." For agent inference with long sessions and nested tasks, this structure looks like a natural fit.

There are three reasons for caution. First, 110M parameters is far from production scale, and behavior at billions of parameters has not been tested. Second, passkey retrieval remains a simple test: it shows the ability to find a single needle, not the quality of reasoning over a long document. Third, the abstract does not make clear how FPA works with optimized attention kernels or what it costs in training speed.

A team working on architectures with forgetting should read the paper in full and reproduce it on its own baseline. The first check is simple: compare FPA with FoX on tasks with nested structure and at lengths several times the training length. If the result holds on models with billions of parameters, extremums instead of sums could become a standard part of forget gate design.