First Passage Attention

Yuxuan Liu

cs.CLOct 15, 2026

First Passage Attention

Authors: Yuxuan Liu

Organizations: Department of Mathematics, University of Bergen, Bergen, Norway

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\%$ .

Figures & tables

Method	$r=10$	$r=14$	$r=18$	durations	walks	Stack
	Alive sets, $d=4$			Varied	Two
FPA	11/15	12/15	13/15	13/15	11/15	14/15
FPA, sum of dips	—	—	—	0/15	—	11/15
Softmax + RoPE	0/15	0/15	0/15	—	—	0/15
Softmax + NoPE	0/15	0/15	0/15	—	—	—
Softmax + ALiBi	0/15	0/15	0/15	—	—	—

Table 1: Seeds out of 15 that succeeded in the one-layer probes. The alive-set columns use the ladder family with head dimension

d=4

, where positional heads realise all alive sets only for

r\leq 6

. The varied durations hold the deepest step of each dip for one or two steps in training, the two walks are the complementary part of the two-walk family (

r=14

d=2

per head), and the stack column uses push and pop sequences. A dash marks a pair that we did not run. With

d=r+2

every positional head and FoX succeeded on all seeds. Selective Attention ∗ lies outside the class of Theorem 3.1 . FPA, First Passage Attention. RoPE, rotary position embedding. NoPE, no positional encoding. ALiBi, attention with linear biases. FoX, Forgetting Transformer. DCP, Dynamic Context Pruning.

	Loss (nats per token) $\downarrow$			Accuracy $\uparrow$
	SlimPajama	WikiText	PG-19	Passkey	Passkey	RULER
Method	test	103	2k–4k	3k–4k	8k–16k	8k
FPA	2.8640	3.1001	3.282	1.00	1.00	0.202
FoX	2.8832	3.1425	3.316	0.73	0.20	0.117
Selective Attention	2.8686	3.1197	3.805	0.11	0.04	0.008
Softmax + RoPE	2.8908	3.1374	5.064	0.00	0.00	0.002

Table 2: Language models with 110M parameters, one seed each. Losses are in nats per token on the SlimPajama test set, on WikiText-103 ( Merity et al., 2017 ) and on PG-19 positions 2048 to 4095 in windows of 4096 tokens. Passkey accuracy averages 50 items at each of two lengths, and RULER averages 13 tasks at length 8192. The models were trained with context 2048. Bold marks the best value in each column. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

Figure 2: Loss far beyond the training length. Mean loss of the four language models on 94 PG-19 test books read in windows of 16,384 tokens, in bins of 256 positions. The models were trained with context 2048 (vertical line). The vertical axis is logarithmic in the loss minus 3.1 nats, which spreads out the lowest losses. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

Figure 3: Passkey accuracy by length of the prompt. Each point averages 50 items, ten passkeys at each of five depths. Lengths up to 4096 come from the preregistered evaluation and the two longer ones from the descriptive one. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

Figure 4: Steps of the walk at brackets in the FPA language model. (a) Increments

\delta_{t}

of layer 1, head 5 at positions 447 to 846 of one test window, with the bracket depth shaded and the mean increments where the depth rises and falls in the whole window. A fixed seed chose the window before the analysis, and the head contrasts the two kinds of bracket most in it. (b) Mean increments at opening and at closing brackets of all 144 heads over 256 test windows. Filled markers show the 56 heads with increments of opposite signs and a familywise significant correlation with the bracket depth. FPA, First Passage Attention.

The same kernel with parts of FPA switched off
	One layer, ms		Model	Decoding at 8k
Method	forward	fwd + bwd	tokens/s	ms/token	bytes/pos.
FPA	0.78	4.31	29,694	22.0	74,880
FoX	0.74	2.87	32,481	20.2	74,304
Softmax (SDPA)	0.87	4.42	26,913	23.9	73,728
Selective Attention	2.00	10.45	20,150	36.7	73,776

Table 3: Cost on one 12 GB partition of a Turing GPU in half precision. One attention layer has 12 heads of dimension 64 at batch 4 and length 2048 (medians of 20 rounds). Model throughput is in training tokens per second of the 110M models in the benchmark. Decoding uses batch size 1 at length 8192, and bytes per position count the keys, values and per-key state of 12 layers in single precision. The lower rows switch off parts of the FPA kernel. SDPA, the scaled dot-product attention of PyTorch, with its memory-efficient kernel on this GPU. FPA, First Passage Attention. FoX, Forgetting Transformer.

Appendix figures & tables38 assets

Supplementary material from the paper’s appendix.

Appendix

Family	Method	Runs	GPU min
Alive sets	FPA †	48	46.3
	FoX †	40	28.9
	Selective Attention †	40	30.2
	Softmax + ALiBi	28	16.6
	Softmax + NoPE	28	16.3
	Softmax + RoPE	24	15.8

Appendix

Table 4: Tuning runs and GPU minutes of the search in the pilots on the two families whose comparisons rest on a declared search. A dagger marks a search that took more than 1.25 times the GPU time of the cheapest compared search of its family. We estimated the GPU minutes of runs that shared the GPU from a calibration of throughput. The configuration of Selective Attention with one mask per head came from later exploratory runs. FPA, First Passage Attention. FoX, Forgetting Transformer. ALiBi, attention with linear biases. NoPE, no positional encoding. RoPE, rotary position embedding. DCP, Dynamic Context Pruning.

Method ( $d=4$ , one layer)	$r=6$	$r=10$	$r=14$	$r=18$
FPA	15/15	11/15	12/15	13/15
Softmax + RoPE	15/15	0/15	0/15	0/15
Softmax + NoPE	14/15	0/15	0/15	0/15
Softmax + ALiBi	14/15	0/15	0/15	0/15
FoX	15/15	0/15	0/15	0/15
Selective Attention, scale 10	15/15	15/15	13/15	15/15

Appendix

Table 5: Every cell of the alive-set family , as seeds out of 15 that ordered the alive keys above the dead keys on at least 90% of the sequences. Selective Attention appears with the selection scale 10 chosen in tuning and with the scale 1 of its paper. The last rows compare the level of the key with the level of the query, whose alive sets are always suffixes. A dash marks a cell that the design did not include. FPA, First Passage Attention. RoPE, rotary position embedding. NoPE, no positional encoding. ALiBi, attention with linear biases. FoX, Forgetting Transformer.

Model	Fixed durations	Varied durations
FPA	9/15	13/15
FPA, sum of dips	15/15	0/15
Selective Attention	15/15	0/15

Appendix

Table 6: Graded depths with fixed and with varied durations , as seeds out of 15 that fitted the training distribution. The sum of dips replaces the minimum of FPA and keeps the same field. FPA, First Passage Attention.

Model	Complementary walks	Identical walks
FPA	11/15	—
Selective Attention, shared mask	0/15	15/15
Selective Attention, one mask per head	13/15	—
DCP	0/15	13/15

Appendix

Table 7: Two walks with two heads of dimension 2 , as seeds out of 15 that ordered the keys of both walks on at least 90% of the sequences. One mask per head is our variant of Selective Attention, in which every head computes its own selection. FPA, First Passage Attention. DCP, Dynamic Context Pruning.

Model	Seeds that succeed
FPA	15/15
FPA without the level of the key ( $h\equiv 0$ )	0/15
FPA with the threshold of the query	0/15
FPA without content retrieval	0/15
FoX	0/15
Stick-breaking attention	0/15

Appendix

Table 8: Retrieval among alive keys , as seeds out of 15 that succeeded. FPA, First Passage Attention. FoX, Forgetting Transformer.

Task	FPA	Softmax + RoPE	FoX	Sel. Att.	Stick-br.	LSTM
Variable assignment	9/15	0/15	15/15	2/15	12/15	15/15
MQRAR	0/15	0/15	0/15	4/15	0/15	0/15
Flip-flop	15/15	0/15	15/15	15/15	0/15	15/15
Copying	15/15	15/15	15/15	15/15	14/15	15/15
Counting	15/15	15/15	15/15	15/15	15/15	15/15
MQAR	0/15	15/15	14/15	12/15	15/15	0/15

Appendix

Table 9: Tasks from earlier work , as seeds out of 15 that succeeded. In variable assignment the softmax column used learned absolute positions in place of RoPE. MQRAR, multi-query repeated associative recall. MQAR, multi-query associative recall. FPA, First Passage Attention. RoPE, rotary position embedding. FoX, Forgetting Transformer. Sel. Att., Selective Attention. Stick-br., stick-breaking attention. LSTM, long short-term memory network.

Figure 5: Training of the four language models. (a) Training loss, averaged over 50 steps of 524,288 tokens, on a vertical axis that is logarithmic in the loss minus 2.75 nats. (b) Validation loss of each baseline minus that of FPA, on 1024 validation windows every 250 steps, so that positive values favour FPA. The triangles mark values beyond the axis, which were

-0.16

for FoX at step 250 and

+0.46

and

+0.25

for softmax attention at steps 250 and 500. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

Figure 6: Loss by position on the PG-19 test books. Mean loss per bin of 64 positions in windows of 2048 and of 4096 tokens, up to four windows per book from 100 and 99 books. The top panels use Unix newlines, as the amended extrapolation test does, and the bottom panels keep the carriage returns of the books. The extrapolation test compared the positions 2048 to 4095 of panel (b). The vertical axes are logarithmic in the loss minus 3.1 nats in the top panels and 4.4 nats in the bottom panels, which spreads out the lowest losses. The models were trained with context 2048. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

FPA minus	FoX	Selective Attention	Softmax + RoPE
PG-19 loss, 0–2k	$-$ 0.029 [ $-$ 0.033, $-$ 0.026]	$-$ 0.014 [ $-$ 0.024, $-$ 0.007]	$-$ 0.038 [ $-$ 0.050, $-$ 0.028]
PG-19 loss, 2k–4k	$-$ 0.031 [ $-$ 0.037, $-$ 0.026]	$-$ 0.516 [ $-$ 0.540, $-$ 0.495]	$-$ 1.847 [ $-$ 1.939, $-$ 1.769]
PG-19 loss, 4k–8k	$-$ 0.020 [ $-$ 0.027, $-$ 0.015]	$-$ 0.300 [ $-$ 0.317, $-$ 0.285]	$-$ 3.666 [ $-$ 3.806, $-$ 3.545]
PG-19 loss, 8k–16k	$+$ 0.026 [ $+$ 0.019, $+$ 0.032]	$-$ 0.528 [ $-$ 0.557, $-$ 0.501]	$-$ 4.258 [ $-$ 4.428, $-$ 4.114]
Passkey 8k	$+$ 0.80 [ $+$ 0.68, $+$ 0.90]	$+$ 1.00 [ $+$ 1.00, $+$ 1.00]	$+$ 1.00 [ $+$ 1.00, $+$ 1.00]
Passkey 16k	$+$ 0.80 [ $+$ 0.68, $+$ 0.90]	$+$ 0.92 [ $+$ 0.84, $+$ 0.98]	$+$ 1.00 [ $+$ 1.00, $+$ 1.00]

Appendix

Table 10: FPA minus each baseline on the descriptive evaluations, with 95% bootstrap intervals in brackets. PG-19 losses come from windows of 16,384 tokens, passkey and NIAH (needle in a haystack) are accuracies on 50 items, and the downstream rows are zero-shot accuracies or log-likelihoods from the evaluation harness of Gao et al. (2024) . Negative losses and positive accuracies favour FPA. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

Figure 7: Copying during training. Greedy accuracy on the second copy of a span of text at gaps of 0, 128 and 320 tokens between the two copies, every 1000 steps and at the last step. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

	RULER, 13 tasks			BABILong, qa1–qa5
Method	4k	8k	16k	0k	4k
FPA	0.208	0.202	0.051	0.082	0.006
FoX	0.218	0.117	0.072	0.078	0.000
Selective Attention	0.005	0.008	0.043	0.080	0.000
Softmax + RoPE	0.002	0.002	0.000	0.088	0.000

Appendix

Table 11: RULER and BABILong , mean accuracy over the 13 tasks of RULER and the five tasks qa1 to qa5 of BABILong, at the context lengths of the columns. All models were trained with context 2048. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding.

	FPA		FoX		Sel. Att.		Softmax
World	recent	older	recent	older	recent	older	recent	older
Scopes in Python code	0.24	0.66	0.01	0.86	0.39	0.77	0.45	0.52
Who saw a hidden object	0.52	0.62	0.47	0.69	0.41	0.67	0.69	0.49
Branches in a version history	0.41	0.60	0.40	0.60	0.27	0.73	0.45	0.56
Stop-loss orders	0.51	0.59	0.79	0.33	0.36	0.41	0.41	0.53
A parenthetical remark	0.55	0.56	0.69	0.30	0.77	0.38	0.77	0.33

Appendix

Table 12: Control items of the zero-shot demonstration , as accuracy over 200 items without inserted filler sentences. The recent items ask for the most recent mention and the older items for an older mention that its content identifies. Chance is 0.5, and bold marks an exact lower 95% bound above chance. A world needed two models with both kinds in bold. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention.

Figure 8: The three preregistered windows of the observational analysis. Each column shows one window with the head that the preregistered rule chose for it, with layers and heads counted from 1. The top panels show the walk

W

of that head, with the bracket depth shaded. The middle and bottom panels show the bias and the attention weight of every key at the last query of the window, with the alive keys marked, among them the query itself. Attention weights below

10^{-7}

are drawn at that value. The walks of all three heads fall almost linearly, so their biases resemble that of FoX. FoX, Forgetting Transformer.

Figure 9: Steps of the walk at brackets in the two other preregistered windows. Increments

\delta_{t}

of layer 1, head 6, the head that the rule chose for both windows, along the whole window, with the bracket depth shaded and darker where brackets nest. The steps are about ten times smaller than in Figure 4 but keep their signs, with means of

+0.41

and

+0.38

where the depth rises and

-0.50

and

-0.51

where it falls.

Figure 10: Alive keys of every head in the final FPA model. Each cell gives the percentage of keys with a bias of at least $-1$ over the queries from position 1024 on in 256 test windows, with layers and heads counted from 1. Dots mark the 19 heads whose alive sets are not suffixes at more than 10% of the queries. FPA, First Passage Attention.

Figure 11: Where keys die in the final FPA model. Keys killed per occurrence of a token, averaged over the 95 local heads, for the 15 tokens with the highest rate among those occurring at least 30 times in 256 test windows. A visible space marks a token that starts with a space, and the dashed line is the mean over all tokens. FPA, First Passage Attention.

	Step 1500	Step 3250	Step 5150
Mean loss over the 12 windows	3.162	2.917	2.821
Mean loss with the bias of every layer removed	6.591	7.335	7.828
Loss increase with alive sets forced into suffixes	$+$ 0.032	$+$ 0.031	$+$ 0.032
windows with a positive increase	12 of 12	12 of 12	12 of 12
Local, middle and global heads	90, 22, 32	90, 20, 34	90, 20, 34
Heads whose alive sets are often not suffixes	21	15	14

Appendix

Table 13: Three checks of the FPA language model , at steps 1500 and 3250 of 5150 and at the final step, on the same 12 validation windows of 2048 tokens. Local, middle and global heads have median alive spans below 64 tokens, from 64 to 1023 tokens and from 1024 on. Alive sets count as often not suffixes when at least 20% of the queries have an alive key older than a dead one. Document heads keep at least 80% of the keys within a document and at most 5% across documents. The copying items differ from those of Figure 7 . FPA, First Passage Attention.

Figure 12: The bias removed at three checkpoints of the FPA model. Loss increase on the first four validation windows when the bias of all layers or of one layer is set to zero, at steps 1500, 3250 and 5150, the last step of training, on a logarithmic axis. FPA, First Passage Attention.

Figure 13: Loss by position at three checkpoints. (a) Loss per bin of 64 positions over the 12 validation windows, with the bias and with the bias of every layer removed. (b) Loss increase when the alive sets of every head are forced into suffixes, the sets that a forget gate can express.

Figure 14: Alive sets forced into suffixes, per validation window. Loss increase on each of the 12 windows at steps 1500, 3250 and 5150, with the means of 0.032, 0.031 and 0.032 nats as dashed lines.

Figure 15: Median alive span of every head in the final FPA model. Spans over the queries 1024 to 2047 of the 12 validation windows, with layers and heads counted from 1. Triangles mark the 21 heads whose span grew or shrank by a factor of two or more since step 1500. FPA, First Passage Attention.

Figure 16: Where keys die at three checkpoints. Keys killed per occurrence of a token, averaged over the 90 local heads of each step, those with a median alive span below 64 tokens. The bars give the final step for its 15 tokens with the highest rate among those occurring at least 30 times in the 12 validation windows. Markers give steps 1500 and 3250 where the token was among the 15 highest at that step. A visible space marks a token that starts with a space, \n is the newline, and the dashed line is the mean over all tokens at the final step.

Figure 17: Attention on earlier documents at three checkpoints. Attention that queries pay to the keys of earlier documents, averaged over the heads of each layer, with the bias and with the bias of every layer removed. The dotted line is the share that weights spread evenly over the visible keys would give.

Figure 18: Retention by token class in the final FPA model. Share of the keys of each class that are still alive 256 positions later, for pairs of key and query in one document, and in the last column for pairs across a document boundary. Content heads lie above the line and the document heads of step 3250 below it, with layers and heads counted from 1. By the final step layer 9, head 5 kept only 76% of the keys within a document and no longer counted as a document head. FPA, First Passage Attention.

Figure 19: The bias of layer 1, head 1 on C++ code in the final FPA model. Weight factor $e^{b_{ts}}$ of every key $s$ for every query $t$ at positions 1100 to 1163 of validation window 0, with factors below 0.02 left white. A key is alive while its factor is at least $e^{-1}$ . Groups of keys stop being alive at once, nine at the end of a statement and three to six at newlines, dots and underscores. FPA, First Passage Attention.

Figure 20: The bias of layer 1, head 10 on a news report in the final FPA model. Weight factor $e^{b_{ts}}$ at positions 1100 to 1163 of validation window 680, with factors below 0.02 left white. Articles, prepositions, punctuation and newlines stop being alive at the next token, while content words such as said and page stay alive to the end of the stretch. FPA, First Passage Attention.

Figure 21: The bias of layer 4, head 9 at a document boundary in the final FPA model. Weight factor $e^{b_{ts}}$ at positions 233 to 296 of validation window 680, with factors below 0.02 left white. At the end-of-document token 36 of the 39 earlier keys of the stretch stop being alive, and at the start-of-document token the others do. No key of the stretch before the boundary stays alive in the next document. FPA, First Passage Attention.

Panel of Figure 22 — Figure 22: Architecture of the FPA language model. (a) Pre-norm decoder with $L=12$ blocks of width 768. (b) Multi-head attention with $H=12$ heads of dimension $d=64$ , in which a fourth linear map gives every head its increments $\delta$ and offsets $g$ and $h$ . (c) One head, in which the bias $B$ with the entries $b_{ts}$ joins the scaled scores before the causal mask. RangeMin returns $\bar{\phi}_{ts}=\min_{s<u\leq t}\phi_{u}$ for every query $t$ and key $s$ , and Sub returns $\bar{\phi}_{ts}-\psi_{s}$ . The parts of FPA are orange, and the one-layer probes use the attention of panel (b) alone on token embeddings. FPA, First Passage Attention. RMSNorm, root mean square normalisation. SwiGLU, gated linear unit with the Swish activation.

Method	Bias $b_{ts}$ added to the logit	Parameters per layer	Positions
FPA	$\beta\min\bigl(0,\min_{s<u\leq t}\phi_{u}-\psi_{s}\bigr)$ with $\phi=W+g$ , $\psi=W+h$	$3H(D+1)+H=27{,}696$	none
FoX	$\sum_{s<u\leq t}\log f_{u}$ with $f_{u}=\sigma(w^{\top}x_{u}+c)$	$H(D+1)=9{,}228$	none
Sel. Att.	$-\sum_{s<u<t}\max\bigl(0,\,q_{u}^{(1)\top}k_{s}^{(1)}/\sqrt{d}\bigr)$	0	RoPE
Softmax	0	0	RoPE

Appendix

Table 14: Bias that each method adds to the logit

q_{t}^{\top}k_{s}/\sqrt{d}

of key

s

at query

t

, with the parameters that it adds to a layer of width

D=768

with

H=12

heads and its positional encoding in the language models. FoX has one gate

f_{u}

per head with a weight vector

w

, a bias

c

and the logistic function

\sigma

. Selective Attention takes

q^{(1)}

and

k^{(1)}

from its first head, shares the bias among all heads and never masks the start token. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention. RoPE, rotary position embedding.

	FPA	FoX	Sel. Att.	Softmax
Parameters	109,862,208	109,640,592	109,529,856	109,529,856
of the attention mechanism	332,352	110,736	0	0
Layers and width	12 and 768
Heads $\times$ head dimension	12 $\times$ 64
MLP	SwiGLU with hidden width 2048
Normalisation	RMSNorm before each sublayer and before the output layer

Appendix

Table 15: The four language models. Parameters count the tied embeddings once. Throughput and time cover the whole run of 5150 steps with its validations and checkpoints on the GPU of Appendix C . For FPA, throughput, time and memory come from an exploratory run of the same configuration with seed 2202 and the kernel of Section 4.3 , outside the preregistration. Table 3 instead gives the throughput that we measured in the benchmark over 12 training steps after 8 steps of warm-up. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention. MLP, multilayer perceptron. SwiGLU, gated linear unit with the Swish activation. RMSNorm, root mean square normalisation. RoPE, rotary position embedding. SDPA, the scaled dot-product attention of PyTorch. GPU, graphics processing unit. GiB, gibibyte.

Setting	Value
Data	SlimPajama, windows of 2049 tokens in one fixed order, each used at most once
Tokenizer	that of Mistral 7B, 32,000 tokens
Context	2048 tokens
Batch	256 windows, 524,288 tokens
Steps	5150, 2.70B tokens
Optimiser	AdamW with $(0.9,0.95)$

Appendix

Table 16: Training of the language models , the same for all four. FPA, First Passage Attention. Sel. Att., Selective Attention. FoX, Forgetting Transformer. ReLU, rectified linear unit.

	One-layer probes	Tasks from earlier work
Optimiser	AdamW with $(0.9,0.98)$ , weight decay 0.01 on every parameter
Batch	64 sequences, drawn afresh at every step
Steps	4000, 16,000 on the graded depths	4000, 8000 for MQAR
Warm-up	200 steps	10% of the steps
Schedule	cosine decay to 0, gradient clipping at 1.0
Precision	single

Appendix

Table 17: Training of the one-layer probes and of the tasks from earlier work. FPA, First Passage Attention. MQAR, multi-query associative recall.

Family	Heads $\times$ $d$	Positions	Final query	Steps	Test set
One-layer probes, width 64
Alive sets	1 $\times$ 4	key offsets, RoPE or ALiBi	MLP of the sequence	4000	all $2^{r-1}$
two and four layers	1 $\times$ 4	RoPE	last token	4000	all $2^{r-1}$
Graded depths	1 $\times$ 8	none	token ?	16,000	2,000 per set
Two walks	2 $\times$ 2	query and key offsets	MLP of the sequence	4000	8,192
Stack	1 $\times$ 8	RoPE	last token	4000	1,000

Appendix

Table 18: Models of the one-layer probes , with one supervised attention layer on token embeddings of width 64 and the head dimension

d

. On the alive sets every key has a learned offset per position and head, FPA, FoX, Selective Attention and softmax attention add RoPE, ALiBi adds its fixed slope, and NoPE and stick-breaking attention use the offsets alone. The positive controls of the alive sets use

d=r+2

, and the field-only FPA of Table 5 uses

d=1

without positional information. The two- and four-layer models put one or three pre-norm blocks with 4 heads and an MLP of width 256 before the supervised layer. On the alive sets and the two walks an MLP of the one-hot sequence with a hidden layer of width 256 gives the final query. FPA, First Passage Attention. FoX, Forgetting Transformer. RoPE, rotary position embedding. ALiBi, attention with linear biases. NoPE, no positional encoding. MLP, multilayer perceptron.

Task	Layers $\times$ width	Heads $\times$ $d$	MLP	Positions	Steps	Test sequences	FPA	Softmax	LSTM
Variable assignment	2 $\times$ 128	2 $\times$ 64	512	learned	4000	2,000	534,672	533,120	393,088
MQRAR	2 $\times$ 128	2 $\times$ 64	512	RoPE	4000	2,000	462,992	461,440	329,856
Flip-flop	1 $\times$ 64	2 $\times$ 32	256	RoPE	4000	2,000	50,632	50,240	33,664
Copying	2 $\times$ 128	2 $\times$ 64	512	RoPE	4000	2,000	399,888	398,336	266,752
Counting	2 $\times$ 128	2 $\times$ 64	512	RoPE	4000	2,000	407,696	406,144	274,560
MQAR	2 $\times$ 128	2 $\times$ 64	512	RoPE	8000	2,000	1,445,904	1,444,352	1,312,768

Appendix

Table 19: Models of the tasks from earlier work , decoders with LayerNorm before each sublayer, an MLP with the Gaussian error linear unit, tied embeddings initialised as

\mathcal{N}(0,0.02^{2})

and the head dimension

d

. The last three columns give the parameters of FPA, of softmax attention with RoPE and of an LSTM of the same depth and width. In variable assignment the attention models use learned absolute positions, elsewhere RoPE, and stick-breaking attention uses none. Variable assignment also has 2000 test sequences with only two values, and flip-flop 10,000 sparse and 2000 dense ones. MLP, multilayer perceptron. MQRAR, multi-query repeated associative recall. MQAR, multi-query associative recall. FPA, First Passage Attention. RoPE, rotary position embedding. LSTM, long short-term memory network.

Method	Second setting	Values	Learning rates	Runs
Alive sets, graded depths and two walks, per method and family
FPA	definition of the minimum	hard, soft with $\tau=1$ , annealed, soft then hard	$\{1,3,10\}\times 10^{-3}$	24
FoX	initial bias of the gates	0, 2, 4, 6	$\{1,3,10\}\times 10^{-3}$	24
Sel. Att.	selection scale $\omega$	1, 1/3, 3, 10	$\{1,3,10\}\times 10^{-3}$	24
DCP	initial interaction bias	2, 0, 4, 8	$\{1,3,10\}\times 10^{-3}$	24
Softmax, stick-breaking	none		12 from $10^{-4}$ to $10^{-1}$	24

Appendix

Table 20: Declared search of the probes and of the tasks from earlier work. Each configuration ran with two tuning seeds, later four, and a grid whose end was selected was extended once by a decade ( Table 4 counts the runs). The 12 learning rates are equally spaced on a logarithmic scale and served softmax attention with RoPE, NoPE or ALiBi and stick-breaking attention on the alive sets. After the pilots the straight-through minimum replaced the four definitions of FPA, and its learning rate was selected again. On retrieval FoX and Selective Attention kept the settings selected on the alive sets and the two walks, on the stack they used the default initialisation of PyTorch and the scale 1, and on the tasks from earlier work the gate bias 0 and the scale 1. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention. DCP, Dynamic Context Pruning. RoPE, rotary position embedding. NoPE, no positional encoding. ALiBi, attention with linear biases.

	Alive sets	Graded depths	Two walks	Stack	Retrieval
FPA	0.03	0.003	0.001	0.1	0.01
FPA, sum of dips		0.003		0.1
FoX	0.03 (6)			0.003	0.03 (6)
Sel. Att.	0.03 (10)	0.03 (1)	0.03 (1/3)	0.001 (1)	0.03 (1/3)
DCP			0.01 (0)
Softmax with RoPE	0.0534			0.001

Appendix

Table 21: Learning rates of the reported runs , with the second setting in brackets, namely the initial gate bias of FoX, the selection scale of Selective Attention and the initial bias of the interaction scores of DCP. The alive sets used one configuration for every

r

and, for the competitors, for their positive controls. Selective Attention with the scale 1 of its paper used

10^{-3}

on the alive sets, and with one mask per head the scale 3 and

3\times 10^{-2}

on the two walks. The two- and four-layer models of FoX and of softmax attention used

3\times 10^{-3}

and

10^{-2}

. FPA used

3\times 10^{-3}

with the field only and

10^{-2}

with the level of the query on the alive sets. On retrieval it used

3\times 10^{-3}

without the level of the key,

3\times 10^{-2}

with the level of the query and

10^{-3}

without content, and on MQAR

10^{-3}

without a bias in its first layer. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention. DCP, Dynamic Context Pruning. RoPE, rotary position embedding. NoPE, no positional encoding. ALiBi, attention with linear biases. MQRAR, multi-query repeated associative recall. MQAR, multi-query associative recall. LSTM, long short-term memory network.

	Forward, ms				Forward and backward, ms
Shape	FPA	FoX	SDPA	Sel. Att.	FPA	FoX	SDPA	Sel. Att.
16 $\times$ 512	0.37	0.34	0.27	0.67	1.76	1.13	1.64	3.54
8 $\times$ 1024	0.51	0.47	0.46	1.12	2.61	1.68	2.50	5.79
4 $\times$ 2048	0.78	0.74	0.87	2.00	4.31	2.87	4.42	10.45
2 $\times$ 4096	1.33	1.27	1.68	3.75	7.76	5.47	8.14	19.50
1 $\times$ 8192	2.39	2.33	3.30	7.29	15.19	10.45	15.79	37.37

Appendix

Table 22: One attention layer with 12 heads of dimension 64 at batch

\times

length, in milliseconds (medians of 20 rounds). The ratios in Section 4.3 are medians over the rounds of the ratio within each round. SDPA, the scaled dot-product attention of PyTorch, with its memory-efficient kernel on this GPU. FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention.

		Length 2048		Length 8192
Method	Rule	loss change	kept	loss change	kept
FPA	bias $\leq-5$	$+$ 0.0063	29%	$+$ 0.0054	23%
	bias $\leq-10$	$+$ 0.0004	34%	$+$ 0.0005	26%
	bias $\leq-20$	$+$ 0.0001	41%	0.0000	30%
FoX	bias $\leq-5$	$+$ 0.1836	13%	$+$ 0.2675	6%
	bias $\leq-10$	$+$ 0.0195	15%	$+$ 0.0227	7%

Appendix

Table 23: Eviction of the keys whose bias fell below

-\theta

(FPA and FoX) or whose mask rose above

\theta

(Selective Attention), with the loss change in nats per token and the share of keys kept on SlimPajama test windows. Softmax attention used a sliding window of the mean size that FPA keeps. Block skipping is the pruning of Lin et al. (2025b) . FPA, First Passage Attention. FoX, Forgetting Transformer. Sel. Att., Selective Attention.

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 189 papers

Why Attend to Everything? Focus is the Key
Mar 12, 2026Hengshuai Yao, Xing Chen, Ahmed Murtadha +8Attention Mechanisms Efficient Attention
Sapient Intelligence · University of Alberta · Tsinghua University
Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers
Jun 20, 2026Xin GaoTransformer Inference KV Caching
Department of Mathematics and Statistics York University Toronto, M3J 1P3
Express Language Modeling
Jun 9, 2026Albert Gong, Annabelle Michael Carrell, Raaz Dwivedi +1Long-Context Language Model Inference Attention Mechanisms
Cornell Tech · University of Cambridge · Microsoft Research

Mar 12, 2026J/K move · Enter open · S save

Why Attend to Everything? Focus is the Key

Hengshuai Yao, Xing Chen, Ahmed Murtadha, Jin Li+7

Sapient Intelligence · University of Alberta · Tsinghua University

Attention Mechanisms Efficient Attention Attention Approximation Sparse Attention Long-Context Modeling

Open paper AskPDF arXiv

Standard attention scales quadratically with sequence length. Efficient attention methods reduce this O(n^2) cost, but when retrofitted into pretrained models, they often degrade perplexity, downstream accuracy, or both. We introduce Focus, a method that learns which token pairs matter. Focus adds a small set of learnable centroids--as few as 148K parameters per layer--that act as gates: only token pairs belonging to the same centroid group attend to each other over long ranges. Focus is composable: it can be added to any pretrained model by training only the centroids while keeping all original weights frozen. Experiments show that composing Focus onto pretrained models yields zero degradation on downstream benchmarks across model sizes from 124M to 70B parameters and five attention architectures. Surprisingly, sparse Focus attention outperforms full attention at 124M scale (30.3 vs. 31.4 perplexity) and matches full attention when trained from scratch at 7B scale (13.82 vs. 13.89). Focus is also fast: top-k group membership gives a 2x speedup with better quality than the original pretrained model. Using our FlashAttention decomposition, Focus achieves an 8.6x speedup at 1M tokens without custom kernels.

Abstract

Long-Context Language Modeling

Explore similar work

Why Attend to Everything? Focus is the Key

Keyless Attention: Value-Space Routing and Value-Only Caching for Efficient Transformers

Express Language Modeling