The Extender: A Log-Structured Transformer

1 Introduction

• Standard Transformer caches grow with model width, depth, and context length, and the aggregate kv-cache can exceed 100 GB before memory-saving optimizations.The residual stream also carries auxiliary layer findings that may compete with other writes and the next-token prediction signal.Source: §1 · p. 1
• Each Extender layer adds a residual update to h and appends a small extension to x; attention kv projections use x, while the FFN and q use h.The fully extended x* contains each layer's x as a prefix, so one shared x*-cache suffices across attention sub-layers.Source: §1 · pp. 1–2
• MHA kv-cache memory is proportional to 2T L d_model, while the Extender x*-cache is proportional to T L d_ε per token.The Extender changes the persistent attention memory requirement by storing one extension per layer and token instead of two model-width vectors per layer and token.Source: §1 · p. 2
• At 198M/436M/920M parameters with d_ε=32, the Extender matches the Reference Transformer on short-context tasks (~+1 on CORE) and leads significantly on RULER tasks.The introduction reports the same short- versus long-context pattern for the tested model sizes.Source: §1 · p. 2
• At inference, layers can materialize kv from x* on demand, amortize that cost across input and output tokens, and release the ephemeral kv-cache between turns.The x*-cache supplies the input for rematerializing each layer's keys and values.Source: §1 · p. 2
• The Extender targets compact attention memory between turns; during turns, memory requirements are unchanged versus a standard Transformer, though existing reduction methods may still be applied.The paper also reports that Grouped Query Attention combines well with the Extender.Source: §1 · p. 2

2 Background

• Attention memory grows with model width, layer depth, and context length, motivating approaches that reduce its footprint.The Extender focuses on memory required between turns and remains compatible with existing approaches during a turn.Source: §2 · p. 2
• GQA reduces attention memory requirements by 4–8× with a modest accuracy loss versus full MHA, while MLA stores a compressed kv pair up to 10× smaller and decompresses it on demand.The Extender changes what attention reads, and GQA can also be applied to it for in-turn memory savings.Source: §2 · p. 2
• YOCO and LCKV reuse caches across layers, CLA and MLKV share a cache among layer groups, and MiniCache merges similar layers' KV states after training.CLA may also be applied to the Extender, but this paper does not evaluate that combination.Source: §2 · p. 3
• Quantization represents features or weights with fewer bits, commonly reducing memory by 2–4×, but the paper does not evaluate it.Applied judiciously, the passage says quantization's accuracy loss can be modest and computational throughput often improves.Source: §2 · p. 3

2.1 Approaches beyond softmax attention

• Alternative attention approaches reduce or eliminate revisiting preceding tokens, whereas the Extender retains softmax attention and changes its kv-projection input to reduce between-turn memory.The passage identifies softmax attention's cost as linear in context length and layer count per token.Source: §2.1 · p. 3

3 The Extender: A Log-Structured Transformer

• Attention kv projections read the log-structured x, while the residual h is read by FFN sub-layers and attention q projections; logits come from h only.Each layer appends its extension εℓ to x and adds its residual update δℓ to h; the final layer emits no εℓ.Source: §3 · p. 3, Figure 1 · p. 4
• Table 1 formalizes the model using a d_model-sized window over xℓ−1; its Last() operator selects the most recently added d_model features.The passage notes that alternatives to this sliding-window design are explored in §3.1.Source: §3 · p. 3
• The FFN produces both a residual update and an extension, while the attention output is added to h after a linear weight.The attention query is not retained for future tokens, so it does not affect attention memory footprint.Source: §3 · p. 4
• The design constrains |s| to d_model to match Transformer attention-layer shape, and layers ℓ>0 read only a subset of xℓ.The passage links the read-window choice to d_ε: small extensions make seeing all extension features more important, while large ones can crowd out earlier features.Source: §3 · p. 4
• With d_ε=32, the 26-layer model has 800 extension features and leaves 864 input-embedding features in s; sizes below 32 or above 64 perform significantly worse.With d_ε=64, the same passage reports 1600 extension features.Source: §3 · p. 5
• For d_ε=64, the k read choice is always Last, while the preferred v choice depends on task; the default instead uses d_ε=32 and ℓ_max=L−2, except |ε0|=64.The passage identifies First, Last, and Random as the v choices considered, with Random a possible compromise at d_ε=64.Source: §3 · p. 5

4 Inference and the x∗-Cache

• At inference, an ephemeral kv-cache amortizes reprojection from x* and can be freed between turns; GQA speeds reprojection by a factor similar to its kv-cache space savings.Repeatedly projecting kv from x* during attention would otherwise incur compute cost.Source: §4 · p. 5

5 Evaluation

• The evaluation uses ClimbMix and ProLong training with CORE and RULER benchmarks, comparing Extender against a Llama-style MHA Reference Transformer and evaluating 920M models at pretraining and ProLong stages.The reference differs from the nanochat design in using Muon for large matrices and SoftCap in attention.Source: §5 · p. 5, §5.2 · p. 8

5.1 Short Context: DCLM CORE Benchmark suite

• Extender and Reference Transformer perform similarly on DCLM CORE with parameter counts within 1% and equal training budgets.One model of each architecture was trained from scratch per size class using Chinchilla token budgets.Source: §5.1 · p. 6

5.2 Long Context: RULER

• Across 1k–32k RULER evaluations, Extender outperforms the Transformer on average, with an advantage on harder multi-key tasks and a Transformer win on variable tracking.The authors did not investigate the suggested task-specific inductive-bias difference further.Source: Figure 3 · p. 7, §5.2 · p. 8
• After ProLong training, the Transformer consistently trails Extender on both NIAH-6 and REST, while differences at 1–2k contexts are marginal.Before ProLong training, Extender is significantly better beyond the 2k training context.Source: §5.2 · p. 8
• At 920M and 64k context, Extender uses 109 MB versus the Reference Transformer’s 11.3 GB of persistent attention memory, a 104× reduction, while matching CORE and somewhat exceeding RULER performance.The memory figures are reported at bf16.Source: §5.2 · p. 8

5.3 Ablation Studies

• Query ablations show that using only h or only x for attention queries significantly underperforms on long-context tasks; concatenating them with a larger query projection does not substantially improve accuracy.The default query input is RMSNorm(x) + RMSNorm(h).Source: §5.3 · p. 9
• Including h in the FFN input lifts performance, but removing it is not catastrophic; h in queries matters more for out-of-domain tasks.The extra RMSNorm ablation gives mixed and likely insignificant results.Source: §5.3 · p. 9

5.4 Limitations

• The tested Extender scales from 199M to 924M parameters, but larger models have not yet been tested because of resource constraints, and token budgets are mostly limited to 20× parameter count.Source: §5.4 · p. 9

• The design targets persistent attention memory, meaning memory required between turns.Source: §6 · p. 9
• 32 attention-memory features per layer match the Reference Transformer on short-context tasks and outperform it on long-context tasks for every tested size class.Source: §6 · p. 9
• The Extender separates the superposition channel h for next-token prediction from the concatenation channel x dedicated to attention memory.Source: §6 · p. 9

• Beyond the 2k training context, Figure 4 shows a clear Extender advantage; it is already visible after the +2.5B ProLong stage in Figure 5.Source: A Additional Long-Context RULER… · p. 13, Figure 4 · p. 13, Figure 5 · p. 13
• Figures 6 and 7 provide per-task RULER results through 32k for 1664-wide, 26-layer models after 64k ProLong mid-training.Source: A Additional Long-Context RULER… · p. 13, Figure 6 · p. 14, Figure 7 · p. 15
• The Table 5 ablations use 920M-class models trained to 9B tokens; most choices differ marginally, while dε=128 is a notable short-context exception.Source: A Additional Long-Context RULER… · p. 16, Table 5 · p. 16
• The default is dε=32 with wide0, while dε=64 with V=Rand is described as a strong choice when fully trained but was not selected.Source: A Additional Long-Context RULER… · p. 16

• The experiment trains 198M Extender students with dε=64 and V=Rand for 4 billion ClimbMix tokens at 2k sequence length, using 924M Extender or 920M Transformer teachers.Source: C Distillation and Extenders · p. 16
• Both distilled students consistently outperform the from-scratch model; the Extender-teacher student is 6 CORE points below its teacher, versus 6.7 points for the Transformer-teacher student.Source: C Distillation and Extenders · pp. 16–17
• On long-context RULER tasks within the 2k training domain, the two students achieve near-identical scores despite the Transformer teacher being weaker.Source: C Distillation and Extenders · pp. 16–17

• The 924M-class GQA4/5 design uses 13 query heads and 3 KV heads in groups of 5, 4, and 4, with 813M parameters after dimensionality reduction.Source: Table 7 · p. 17, D Grouped Query Attention and Extenders · pp. 17–18
• The Extender GQA4/5 model outperforms the Transformer MHA, but no matching Transformer GQA was trained.Source: D Grouped Query Attention and Extenders · p. 18
• On CORE, GQA tracks MHA closely on T-Mono and largely does so on the remaining, noisier tasks.Source: Figure 8 · p. 17, D Grouped Query Attention and Extenders · p. 18
• On RULER, GQA tracks MHA closely on both NIAH-6 and REST; Figure 9 separates these groups across context lengths.Source: D Grouped Query Attention and Extenders · p. 18, Figure 9 · p. 18

E Accuracy across model sizes

• CORE and RULER 1–2k accuracy improve consistently with model size.The passage notes particularly strong consistency on T-mono, where the Transformer also reliably improves with size.Source: E Accuracy across model sizes · p. 19
• Beyond the 2k training context, RULER scores are less predictable but suggest improvement.Source: E Accuracy across model sizes · p. 19
• Table 8 reports the default Extender at each Figure 2 size class: 199M corresponds to 4B tokens, 438M to 9.6B, and 924M to 18B.Source: Table 8 · p. 19
• CORE and T-mono are centered DCLM CORE scores, while RULER is a 13-task mean; italicized lengths exceed the 2k training context.T-mono is the 12-task set in Figure 2.Source: Table 8 · p. 19