Source-linked AI summary
First Passage Attention
Yuxuan Liu
TL;DR
Softmax attention mechanisms that accumulate decay along the key–query path can struggle to erase newer context while preserving older, structurally relevant keys. First Passage Attention (FPA) instead penalizes keys by the deepest dip of a learned field below their levels. FPA matches FoX in perplexity and uniquely achieves perfect passkey retrieval at extended training lengths, including with a compressed key–value cache.
Problem
Path-accumulating attention mechanisms cannot selectively erase a newer key while preserving older keys, and pairwise accumulation also grows with intervening steps and incurs computation overhead.
Source: §1 · p. 1Method
FPA lowers each key’s logit according to the deepest later dip of a learned scalar field below that key’s level, and its shared-walk form evolves alive keys as a stack.
Source: §1 · p. 1Results
FPA matches FoX in perplexity and uniquely achieves 100% passkey retrieval at 4× and 8× training length, even with about one-quarter of its key–value cache.
Source: §1 · p. 2Takeaways & Limitations
A constant-dimension FPA head realizes alive sets induced by later dips, whereas softmax heads with positional biases or token-level forget gates require dimension linear in the number of keys.
Source: §1 · p. 1The language-model comparison used 110M-parameter models with one seed each, so it may change at larger scale.
Source: §6 · p. 10
Abstract
from arXiv · showhide
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\%$.
1 Introduction
First Passage Attention (FPA) replaces path-summed attention decay with a penalty based on the deepest dip of a learned field, allowing later events to erase newer keys while preserving older ones. The paper analyzes this expressivity and reports synthetic-probe and language-model results.
1 Introduction
- Standard attention scores earlier tokens using key–query matching, often with relative positional encodings, but intervening tokens do not directly enter that score.The motivation is that events such as backspaces, closing brackets, and reassignments can invalidate newer information while leaving older information intact.Source: §1 · p. 1
- Path-summed or multiplied mechanisms face structural limits: token-level forget gates attenuate keys monotonically with age, while pairwise accumulation grows with intervening steps and incurs pairwise computation overhead.Thus, a later event cannot erase a recent token through FoX gates without also suppressing older ones.Source: §1 · p. 1
- FPA lowers a key’s logit according to the deepest dip of a learned scalar field below that key’s level, regardless of how long the dip lasts.With a shared learned walk for the field and levels, a later drop can remove newer keys while preserving older ones.Source: §1 · p. 1
- A constant-dimension FPA head realizes alive sets induced by later dips, whereas softmax heads with positional biases or token-level forget gates require dimension linear in the number of keys.The paper also proves that range extrema provide duration-invariant suppression and that head-specific fields overcome the dimension bottleneck of layer-shared masks.Source: §1 · pp. 1–2
- In the typed-word example, FPA’s stack form lets a later backspace remove the newer key while retaining older keys.The shared walk’s alive keys evolve as a stack.Source: §1 · pp. 1–2
- FPA matches FoX in perplexity and uniquely reaches 100% passkey retrieval at 4× and 8× training length, even with about one-quarter of its key–value cache.In preregistered one-layer probes, FPA learns the target families on most seeds, while the covered baselines fail on every seed.Source: §1 · p. 2
2 First Passage Attention
FPA adds a first-passage bias to attention logits using the minimum of a learned field after each key. Its parameterisation includes several existing biases, and its stack form yields permanent key death and stack-structured alive sets.
2 First Passage Attention
- FPA adds a first-passage bias to attention logits, based on the range minimum of the field after each key and a learned positive strength.The attention output is a softmax-weighted sum of values using logits that combine the usual query–key term with this bias.Source: §2 · p. 2
- A key is alive when the later field minimum stays at or above its level; otherwise its logit penalty is proportional to the depth of the deepest dip.As the bias strength tends to infinity, a dead key’s weight tends to zero.Source: §2 · p. 2
- Each head derives a walk increment and two offsets from each input position, with the walk shifting later comparisons and the offsets setting field and key levels.A layer of width D with H heads adds 3H(D + 1) + H parameters, or 332 thousand in the 110M language model.Source: §2 · p. 2
- With a shared walk as field and level, alive-key levels do not decrease with writing time, keys killed at a step form a suffix, and dead keys stay dead.The broader parameterisation also includes FoX, a bottleneck bias, and learned expiry; FoX corresponds to a stack-form walk that never rises, while a rising walk can preserve an old key as it removes newer ones.Source: §2 · p. 3
- Training uses a straight-through estimator because the outer minimum can close early and stop gradients to the field; evaluation retains the hard range-minimum bias.The backward pass substitutes a soft-minimum gradient with τ = 1.Source: §2 · p. 3
- Blocked computation uses three field-minimum summaries per query–key block pair, while autoregressive decoding updates each bias row in O(t) operations and stores two scalars per cached key and head.The fused Triton kernel computes attention in blocks with an online softmax.Source: §2 · p. 3
3 Theory
The theory compares which key sets attention heads can represent, how dip duration affects suppression, and whether heads need separate fields. Its bounds and constructions apply within specified head, precision, and mask settings.
3.1 Alive sets chosen by later dips
- On the ladder family, positional-bias heads need d ≥ r−2 and prefix-fixed keys need dp ≥ r−2, while one FPA head realizes every alive set with O(log r)-bit scalars.The FPA construction uses constants β and M of order log r at ε = 1/(2r + 2).Source: §3.1 · p. 4
3 Theory
- These bounds do not cover biases where later tokens act on keys separately, and they constrain one head’s weights rather than deeper models or a whole layer’s function.Without a precision bound, no dimension lower bound follows for the ordering condition when keys read their prefixes.Source: §3.1 · p. 4
3.2 Depth, not duration
- For variable-duration dips, one FPA head exactly fits the stated targets and its bias is unchanged by holding a step longer, while token-only additive logits cannot fit them exactly.The target weights are proportional to e^−κηj, with ηj = 0 when a key has no dip.Source: §3.2 · p. 5
3.3 One field per head
- With a shared mask, two heads require 2d ≥ r−1 to realize every complementary-walk alive set at the stated tolerance; one walk per FPA head realizes both sets with one-dimensional dot products.For r = 14 and d = 2, the shared-mask bound covers 1,093 of 8,192 sequences.Source: §3.3 · p. 5
4 Experiments
The experiments test one-layer expressivity, language-model performance, and computational cost. FPA succeeds on varied synthetic targets and matches FoX on language-model loss, while retaining long-context passkey retrieval with a compressed cache.
4.1 One-layer probes
- FPA succeeded on 11–13 seeds for alive sets, while dimension-4 positional and token-gated heads failed at r = 10, 14, and 18; their d = r + 2 controls succeeded on all seeds.FPA succeeded on 13 seeds for variable-duration dips, whereas Selective Attention and the sum-of-dips variant succeeded on none; both competitors fit fixed-duration cases on all 15 seeds.Source: §4.1 · p. 6, Table 1 · p. 6
- FPA succeeded on 11 seeds for complementary two-walk tasks and 14 for push-pop sequences, while shared-mask Selective Attention and DCP failed on the two-walk task and one-layer FoX, softmax, and Selective Attention failed on push-pop.For name-occurrence retrieval among alive keys, FPA succeeded on all 15 seeds; variants omitting the key level, content retrieval, or using the query level succeeded on none.Source: §4.1 · p. 6
- On six earlier-work tasks, FPA succeeded on every seed for flip-flop language modelling, copying, and counting, but failed on every seed for multi-query associative recall, where softmax succeeded on all.Removing the first-layer bias restored FPA’s success on all 15 seeds for multi-query associative recall; the authors diagnose early suppression of each value’s attention to its own key.Source: §4.1 · p. 6
4.2 Language models
- FPA retrieved every passkey at lengths 8192 and 16,384, where FoX retrieved 20%; its loss was 0.02–0.03 nats below FoX through position 8192 but 0.026 above it from 8192 to 16,384.On PG-19 positions 2048–4095, FPA loss was 0.034 nats lower than FoX; at lengths 3072 and 4096 it retrieved all 100 passkeys against FoX’s 73.Source: §4.2 · p. 7, Figure 3 · p. 8
4.3 Cost
- At 24–27% of its cache, FPA still retrieved every passkey at 8192 and 16,384, while threshold eviction cost 0.00044 nats at length 2048 and 0.00047 at 8192.FPA trained at 0.914 times FoX throughput; at batch size 1 and length 8192, decoding took 22.0 ms per token versus 20.2 ms for FoX and 23.9 ms for softmax attention.Source: §4.3 · p. 9, Table 3 · p. 9
5 Related Work
Prior work addresses path-dependent logit biases, range-based selection, and eviction through several mechanisms. FPA is distinguished by scalar path extrema, learned levels per key, and a walk that can rise as well as fall.
Path-dependent biases
- FoX sums log forget gates and is the stack form of FPA with a never-rising walk, so its visible sets are suffixes; Selective Attention and DCP can realize FPA’s alive sets in the hard limit but accumulate pairwise scores along the path [Lin et al., 2025a; Leviathan et al., 2025; Anagnostidis et al., 2023].FPA’s bias depends on scalars at individual positions and costs less to compute than Selective Attention; unlike its suppression, FPA suppression does not depend on dip duration [Proposition 3.4].Source: §5 · p. 9
Range extremes and stacks
- StructFormer, PRPN, and syntax-guided local attention compare range extrema with query-linked levels, while FPA learns its walk without structural supervision [Shen et al., 2021; Shen et al., 2018; Hou et al., 2022].FPA includes learned expiry as the case of a clock field [Sukhbaatar et al., 2021].Source: §5 · p. 9
Expressivity and eviction
- FPA’s one-layer separation is specific to alive sets created by later dips; bounded Dyck languages need linear size in one layer, while a few layers suffice [Bhattamishra et al., 2024; Yao et al., 2021; Hayakawa and Sato, 2025].FPA’s bias also supports threshold-based key eviction, dropping exactly dead keys in the hard limit [Xiao et al., 2024; Lin et al., 2025b].Source: §5 · p. 10
6 Limitations and Conclusion
The section gives proof properties and constructions for FPA, including stack behavior, blocked computation, and a dimensional lower bound for positional heads.
6 Limitations and Conclusion
- The language-model comparison is limited to 110M parameters, one seed, 2.7B training tokens, and one 12 GB Turing GPU partition.The authors note that the comparison may change at larger scale.Source: §6 · p. 10
A Proofs
- The proofs assume finite logits and exact arithmetic, with the field and key level fixed once their positions are written.Source: A Proofs · p. 15
A.1 The ladder family
- In the ladder family, key j < r is alive at the final query exactly when its following block is U D; key r is always alive.Each block-and-key pair has one dead position, leaving r −1 dead positions and 2r + 1 alive positions.Source: §A.1 · p. 15
A.2 Proof of Proposition 2.2
- The queue-reflection quantity L_ts expresses the deepest path dip, while a nonincreasing walk makes the FPA bias equal the accumulated log forget gates.The same proposition also relates the bias to learned minima and an expiry rule.Source: §A.2 · p. 15
A.3 Proof of Proposition 2.3
- FPA’s death rule is permanent, and when field and levels share a walk, each step kills a suffix of alive keys so the set evolves as a stack.Source: §A.3 · p. 15
A.4 Proof of Proposition 2.4
- Blocked computation derives each query/key-block bias from suffix, intervening-block, and prefix minima, with running minima updating cached-key biases during decoding.The passage describes O(B) operations for within-block suffix and prefix minima and O(T/B) block minima for a sequence of length T.Source: §A.4 · p. 16
A.5 Proof of Theorem 3.1
- For the ladder family, fixed positional heads need dimension d ≥ r −2 to realize every alive-set pattern.The proof uses pattern ordering of logits and a sign-pattern argument; token-only path terms do not remove this bound.Source: §A.5 · p. 16
A.6 Proof of Theorem 3.2
The appendix proofs establish bounds on which alive-key patterns softmax-style heads can represent and develop FPA constructions for dip-based selection. Synthetic probes test these claims across alive sets, variable-duration dips, two walks, stack operations, and retrieval among alive keys.
A.6 Proof of Theorem 3.2; A.10 Proof of Proposition 3.6
- The fixed-key proof gives dp ≥ r−2 only when query values are bounded; the two-head argument bounds shared-mask sign patterns and covers Selective Attention and DCP.The two-head result applies to masks that may depend arbitrarily on the sequence, with the pair of heads fixed in advance.Source: §A.6 · p. 17, §A.10 · p. 18
A.7 Proof of Proposition 3.3
- In the stack construction, alive keys have zero bias while dead keys receive bias at most −β/2, and the dot product keeps other positions below them.The construction gives a tolerance bound using the fact that the last key is always alive.Source: §A.7 · p. 17
A.8 Proof of Proposition 3.4; A.9 Proof of Proposition 3.5
- The A.8 argument contrasts maximum-depth bias, unchanged by inserting shallower dips, with summed dips; A.9 shows varying a dip’s duration can defeat exact realization in the tested class.A separate construction in A.9 realizes the target by setting each key’s bias to −κη_s.Source: §A.8 · p. 17, §A.9 · pp. 17–18
B.1 Common setup
- Several graded-depth configurations and Selective Attention with one mask per head came from exploratory runs rather than the declared search, and four selected learning rates remained at grid ends.The passage identifies FPA and its sum-of-dips variant on the stack family, stick-breaking retrieval, and softmax with RoPE on variable assignment as the four end-of-grid cases.Source: §B.1 · p. 19
B.2 Alive sets
- At r=10 and r=14, a dimension-1 field-only FPA head learned the alive-set family on 12 and 10 seeds, while four bounded-class competitors failed every seed above r=6.The comparison is reported for the alive-set probes; the positive controls succeeded on every seed.Source: Table 5 · p. 20, §B.2 · p. 20
B.3 Graded depths and varied durations
- With dip bottoms held 0 or 16 steps, ordered sequences fell by at most 0.1 on 14 FPA seeds and 10 sum-of-dips seeds; only FPA met the preregistered line of 12 seeds.FPA and the sum-of-dips variant fitted the separate walk-level family on all 15 seeds.Source: §B.3 · p. 21, Table 6 · p. 21
B.4 Two walks; B.5 Stack
- On two-walk tests, no shared-mask run exceeded 2,380 ordered sequences or 1,093 meeting 1/30; one-layer FPA succeeded on 14 stack seeds, versus none for FoX, RoPE softmax, or Selective Attention.The stack comparison is qualified: the authors had no proof that the one-layer competitors could not learn that family.Source: §B.4 · p. 21, §B.5 · p. 21, Table 7 · p. 21
B.6 Retrieval among alive keys
- The retrieval task requires both key-level selection and content retrieval among queried-name occurrences; the best tested ablation reached 0.4905, below the 0.6 success line.The task design excludes solutions based only on position or level, and the ablation evidence supports the stated need among the tested variants.Source: §B.6 · p. 22, Table 8 · p. 22
B.7 Tasks from earlier work
The paper evaluates smaller-scale versions of six published tasks and reports exploratory checks on multi-query associative recall (MQAR). In those checks, FPA’s first layer suppressed attention from each value to its own key, while language-model previous-token heads appeared in later layers.
B.7 Tasks from earlier work
- The smaller-scale task suite includes variable assignment, repeated and ordinary associative recall, flip-flop language modelling, copying, and counting.The models were two-layer decoders except for the one-layer flip-flop model, with an LSTM serving as a recurrent reference.Source: §B.7 · p. 22
- Table 9 reports successful seeds out of 15 across the earlier-work tasks, with task-specific model and positional-encoding settings described in its caption.The caption identifies FPA, FoX, Selective Attention, stick-breaking attention, and LSTM among the compared methods.Source: Table 9 · p. 22, §B.7 · p. 22
- In exploratory MQAR runs, two-layer FPA models suppressed attention from each value to its own key; three- and four-layer models failed on all 4 seeds, while removing the first-layer bias yielded success on both tested seeds.The authors’ exploratory inspection found previous-token heads in the second and third layers of the 12-layer FPA language model.Source: §B.7 · p. 23
C.1 Models, data and training
FPA had the lowest validation loss of the four models at every evaluation from step 500 onward, despite starting more slowly. The section describes the model and evaluation setup and reports amendments to the preregistered tests.
C.1 Models, data and training
- The four 12-layer models trained on 2.70B SlimPajama tokens; evaluation included 20M-token test sets, PG-19 extrapolation, and passkey retrieval.The models had about 110M parameters and were trained with context 2048; tests used resampling or exact tests as specified.Source: §C.1 · p. 23, §C.2 · p. 23
- 0.020 nats below FoX, 0.006 below Selective Attention, and 0.027 below softmax attention was FPA’s final validation-loss advantage.From step 500 onward, FPA had the lowest validation loss at every evaluation.Source: §C.1 · p. 23, Figure 5 · p. 24
C.3 Amendments
- The PG-19 extrapolation test was amended to use Unix newlines; under both newline conventions, FPA had the lowest loss beyond training length.With carriage returns, the test would also have supported FPA, with a difference of −0.250 nats.Source: §C.3 · pp. 23–24
C.4 Descriptive evaluations
On descriptive evaluations, FPA had lower SlimPajama test loss than FoX at every position bucket and retrieved every passkey across the reported lengths, while copying accuracy was similar across models.
C.4 Descriptive evaluations
- FPA had 0.008 to 0.022 nats lower SlimPajama test loss than FoX at every position bucket.FPA retrieved every passkey at all six lengths from 1024 to 16,384 tokens; FoX retrieval fell from 1.00 at 2048 tokens to 0.58 at 4096 and 0.20 beyond.Source: §C.4 · p. 24
- At 4096 tokens, FoX retrieved 0.58 of passkeys, while FPA retrieved every passkey through 16,384 tokens.Selective Attention retrieved 0.80 at 2048 tokens and almost none beyond 3072; softmax retrieved about half even within its training length.Source: §C.4 · p. 24
- Across models, copying accuracy ranged from 0.91 to 0.93 at zero gap, 0.83 to 0.85 at 128 tokens, and 0.68 to 0.71 at 320 tokens.Softmax attention had the highest copying accuracy at every gap.Source: §C.4 · p. 24
C.5 A demonstration without a usable task
The preregistered zero-shot demonstration yielded no usable world under its control-item rule, so the authors drew no conclusion from it.
C.5 A demonstration without a usable task
- No world met the preregistered requirement that at least two models answer both control-item types above chance, so no test was run and no conclusion was drawn.Several models answered one control type well but the other below chance.Source: §C.5 · p. 24
- The PG-19 position-loss figure compares 2048- and 4096-token windows, with the extrapolation test focused on positions 2048 to 4095.The models were trained with context 2048; panels also distinguish Unix newlines from book carriage returns.Source: Figure 6 · p. 25
- The training-copying figure tracks greedy accuracy on the second copy of a text span at gaps of 0, 128, and 320 tokens.The descriptive-evaluation tables cover PG-19, passkey and needle-in-a-haystack accuracy, downstream evaluations, RULER, and BABILong.Source: Figure 7 · p. 25, Table 10 · p. 26, Table 11 · p. 26
C.6 Observational analysis
The observational analysis found diverse alive-key patterns and associations between some attention walks and bracket structure, but the authors treat the correlations as descriptive rather than evidence of bracket tracking.
C.6 Observational analysis
- The median head marked 96.5% of keys dead and 0.3% alive; 19 heads had non-suffix alive sets in more than 10% of cases.In two first-layer heads, almost all alive sets were suffixes, with most attention on alive keys older than the newest dead key.Source: §C.6 · p. 26
- The detrended bracket-depth correlations were small and descriptive, not evidence that the model tracks brackets.The median absolute correlation was 0.032 and the maximum was 0.248; 82 heads passed the familywise test, 81 with positive correlation.Source: §C.6 · p. 26
- In window 268, the only alive key at the final query was the query itself, which received 0.90 of attention.Figure 8 displays three preregistered windows with the selected head’s walk, bracket depth, bias, and attention weights.Source: §C.6 · p. 26, Figure 8 · p. 27
- In window 268, the selected head stepped up 4.70 on average where bracket depth rose and down 5.51 where it fell.Across 256 windows, 56 heads combined opposite-sign opening and closing bracket increments with a familywise-significant correlation.Source: §C.6 · pp. 26, 28
- Across heads, the median share of deaths at document-boundary tokens was 6.6 times their share of tokens.The analysis classified 95 heads as local; on exploratory checks, the threshold agreed with a median alive span below 64 tokens for 138 of 144 heads.Source: §C.6 · p. 28
C.7 Exploratory checks at three checkpoints
Exploratory checks tracked the final FPA model and two earlier checkpoints on the same 12 validation windows. They examine how removing or restricting the bias affects loss, and how learned alive-key patterns vary across heads and tokens.
C.7 Exploratory checks at three checkpoints
- Removing all FPA biases raised loss by 3.43, 4.42, and 5.01 nats; forcing alive sets into suffixes raised loss by 0.032, 0.031, and 0.032 nats on average.On the first four windows, the first layer's share of the loss increase grew from 1.49 to 2.67 and 3.48 nats.Source: §C.7 · p. 28, Figure 12 · p. 30, Figure 14 · p. 31
- 90 of 144 heads had median alive spans below 64 positions, and local heads killed about one key per token occurrence at each checkpoint.The mean rates were 0.99, 0.98, and 0.98; the newline and the token The without a preceding space were among the highest-rate tokens throughout.Source: §C.7 · p. 28, Figure 15 · p. 31, Figure 16 · p. 32
- With bias, earlier-document attention stayed between 0.002 and 0.10 across layers, while inspected content heads retained capitalised words and continuations but no digits.Document heads kept nearly every key within a document and none across a document boundary.Source: §C.7 · pp. 28, 32, Figure 17 · p. 32, Figure 18 · p. 33
- Examples show code heads killing key groups at statement ends and punctuation, news heads dropping function words while retaining content words, and a document head clearing prior-document keys.The ablations show reliance of the trained FPA model on these structures, not that a FoX-trained model could not learn the same function another way.Source: §C.7 · p. 32, Figure 19 · p. 34, Figure 20 · p. 35
D.1 Architecture
- FPA adds a bias to every head before softmax, computed from a per-head field, a range minimum, and the depth of its dip; it uses no other positional information.The forward pass uses a hard minimum and the backward pass the gradient of a soft minimum with τ = 1.Source: §D.1 · p. 37, Figure 22 · p. 37
D.2 Language models
- The four language models used the same data order, schedule, and FoX 125M-recipe hyperparameters, while differing in attention, positional encoding, kernel, and micro-batch split.The micro-batch split changes the gradient only through rounding.Source: §D.2 · p. 37, Table 14 · p. 38, Table 15 · p. 38
D.3 Software and compute
- The preregistered probes and tasks comprised 1,815 runs across 121 settings and took 73.5 hours; training the four language models took 110.5 hours.Source: §D.3 · p. 39
D.4 One-layer probes and tasks from earlier work
This section documents the models, searches, and training configurations used for one-layer probes and tasks from earlier work. The supplied passages describe experimental setups rather than their outcomes.
D.4 One-layer probes and tasks from earlier work
- One-layer probes used a supervised attention layer on width-64 token embeddings; the field-only FPA configuration used head dimension 1 without positional information.The alive-set positive controls used d = r + 2, while other probe variants included FoX, Selective Attention, softmax, NoPE, and stick-breaking attention.Source: Table 18 · p. 39, §D.4 · p. 39, Table 17 · p. 39
- Tasks from earlier work used decoder models with LayerNorm, tied embeddings, and task-specific positional encodings; variable assignment and flip-flop included additional test-set variants.The table also compares parameter counts for FPA, softmax attention with RoPE, and an LSTM of matching depth and width.Source: Table 19 · p. 40
- The declared search used tuning seeds and extended a selected grid endpoint by a decade; after pilots, FPA switched to the straight-through minimum and its learning rate was selected again.The configurations were selected by the stated pilot rules and frozen before reported runs.Source: §D.4 · p. 39, Table 20 · p. 40
- Reported learning rates varied by task and FPA configuration, including 3 × 10^-3 for field-only alive-set probes and 10^-3 for MQAR without a first-layer bias.Source: Table 21 · p. 41
E Kernel and cost
The appendix describes a fused Triton implementation, benchmark setup, and key eviction measurements. It also reports perfect FPA passkey retrieval beyond training length under all tested thresholds with a reduced cache.
E Kernel and cost
- The fused Triton kernel computes range minima in 64-position blocks without a T × T matrix, using deterministic gradient calculations with no atomic additions.The evaluated language model used an earlier kernel version; exploratory tests found comparable error bounds for the current version.Source: E Kernel and cost · p. 42
- Against a double-precision reference, the earlier kernel's largest reported relative errors were 5.1 × 10^-4 on GPU and 3.3 × 10^-3 for strength gradients at tied field values.The current version had the same largest errors in exploratory checks.Source: E Kernel and cost · p. 42
- The benchmark timed methods in the same process across five batch-by-length shapes; the FoX bias inside the FPA kernel took 1.05 times as long as the FoX kernel.The authors attribute FPA's additional time in Table 3 to the operations FPA adds; decoding time changed little from prompts of 1,024 to 16,384 tokens at batch size 1.Source: E Kernel and cost · p. 42, Table 22 · p. 42
- FPA retrieved every passkey at 8,192 and 16,384 tokens under every threshold, using 24% to 27% of its cache.Sliding-window softmax retrieved at most 20%.Source: E Kernel and cost · pp. 42–43
Related papers
This paper cites
- Attention Is All You Need2017
- RoFormer: Enhanced Transformer with Rotary Position Embedding2021
- Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge2018
- Pointer Sentinel Mixture Models2016
- Efficient Streaming Language Models with Attention Sinks2023
- Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation2021