Paper deep dive
The Extender: A Log-Structured Transformer
Jakob Eriksson
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
Abstract:We introduce the Extender, a log-structured variant of the standard Transformer architecture. In a standard Transformer, each layer communicates with subsequent layers exclusively via the residual $\mathbf{h}$, a superposition channel. The Extender adds a concatenation channel $\mathbf{x}$: each layer $\ell$ emits both a residual update $\delta_\ell$ which is added to $\mathbf{h}$, and a much smaller extension $\epsilon_\ell$ which is appended to $\mathbf{x}$. While both the FFN and $\mathbf{q}$ see $\mathbf{h}$, the attention $\mathbf{kv}$ projections take only $\mathbf{x}$ as input. As a result, the fully extended $\mathbf{x}$ contains the complete input for the $\mathbf{kv}$ projections of all layers, reducing the persistent attention memory footprint from $2Ld_{model}$ to $\sum|\epsilon_\ell|$. We find that with $|\epsilon_\ell|=32$, the Extender matches Transformer accuracy on short-context (CORE) tasks at 199M-924M parameters, and exceeds Transformer accuracy on long-context (RULER) workloads, again at 924M parameters. For our 1664-wide, 924M model, the Extender's persistent attention memory footprint is $104\times$ smaller than MHA. The memory savings grow with model width.
Tags
Links
- Source: https://arxiv.org/abs/2609.32759v1
- Canonical: https://arxiv.org/abs/2609.32759v1
Trouble viewing inline? Open PDF directly →
Full Text
54,541 characters extracted from source content.
Expand or collapse full text
The Extender: A Log-Structured Transformer Jakob Eriksson Affiliation: University of Illinois Chicago Email: jakob@uic.edu July 21, 2026 Abstract We introduce the Extender, a log-structured variant of the standard Transformer architecture. In a standard Transformer, each layer communicates with subsequent layers exclusively via the residual h, a superposition channel. The Extender adds a concatenation channel x: each layer ℓ emits both a residual update δℓ _ which is added to h, and a much smaller extension ϵℓ _ which is appended to x. While both the FFN and q see h, the attention kv projections take only x as input. As a result, the fully extended x contains the complete input for the kv projections of all layers, reducing the persistent attention memory footprint from 2Ldmodel2Ld_model to ∑|ϵℓ|Σ| _ |. We find that with |ϵℓ|=32| _ |=32, the Extender matches Transformer accuracy on short-context (CORE) tasks at 199M-924M parameters, and exceeds Transformer accuracy on long-context (RULER) workloads, again at 924M parameters. For our 1664-wide, 924M model, the Extender’s persistent attention memory footprint is 104×104× smaller than MHA. The memory savings grow with model width. 1 Introduction The residual stream is fundamental to the standard Transformer model: it is the sole input to every layer; the sole output of each layer is an update to the residual stream. Arguably, its most important function is that of mitigating vanishing gradients in deep models: each individual layer sees a direct gradient from the final loss computation, enabling stable optimization. In other words, the loss directly teaches each layer how to predict the next token. However, this design is not without drawbacks. While each layer may in practice only make small changes, the residual update from each layer is full width, making each feature of the input to the subsequent attention layer unique. During inference, this requires each attention layer to maintain its own, full-width residual (h) or key-value (kv) cache to avoid a very costly recomputation. The memory footprint of these caches grows linearly with model width dmodeld_model, depth L and context length T. For a large language model, the aggregate kv-cache alone can exceed 100 GB before dimensionality reduction, layer reuse, quantization and other space-saving optimizations. Moreover, beyond predicting the next token, each layer ℓ may also need to communicate its own findings, say ϕℓ _ to subsequent layers. Here, ϕℓ _ may not look anything like the next token; for example, it may encode a distance from another token, or a repetition counter, yet this too must be communicated through the residual, potentially in contention with other writes, and against the teachings of the loss. The Extender is a log-structured variation on the classic Transformer architecture that separates the next-token prediction stream from the attention input. Similar to the Transformer, the Extender is made up of a stack of Extender layers, each consisting of several individual sub-layers, but dominated by softmax attention and FFN. However, instead of giving the attention sub-layers full visibility into the output of the previous layer, the Extender provides two distinct channels: h The conventional residual, used primarily for predicting the next token. Each layer contributes a full-width residual update δℓ _ , such that ℓ=ℓ−1+wℓδℓh_ =h_ -1+w_ _ , where wℓw_ is a scalar weight. ℓh_ is part of the input to the FFN sub-layer of the next block, as well as to attention (q only, not kv), and the final logit projection. −1=[embed(t)]h_-1=[embed(t)]. x Each layer ℓ contributes information to subsequent attention kv and q projections via a small, read-only feature extension ϵℓ _ to a log-structured token embedding. The extended embedding after layer ℓ is ℓ=[ℓ−1;RMSNorm(ϵℓ)]x_ =[x_ -1;RMSNorm( _ )] where -1=[embed(t)]x_-1=[embed(t)]. We use ∗x_ to refer to the fully extended embedding after the token has traversed all L layers. As the ∗ in the notation implies, the append-only nature of x in the Extender architecture means that ∗x_ contains every xℓ,ℓ∈[0..L-1]x_ , ∈[0..L-1] as a prefix. Therefore, while every layer sees a different ℓx_ , a single shared ∗x_ -cache suffices across all attention sub-layers. Thus while the Multi-Head Attention (MHA) Vaswani et al. (2017) Transformer requires kv-cache memory proportional to 2TLdmodel2TLd_ model (two dmodeld_model-wide vectors per layer, per token), the Extender instead requires ∗x_ -cache memory TLdϵTLd_ε (one dϵd_ε-wide extension per layer, per token), reducing the persistent attention memory footprint at long context lengths by up to 2dmodeldϵ× 2d_modeld_ε× vs. MHA. The model sizes we tested, 198M/436M/920M, with dϵ=32d_ε=32, the Extender matches our Reference Transformer on short-context tasks (∼+1 +1 on CORE), and leads significantly on long-context (RULER) tasks. To compute softmax attention without a kv-cache, ∗x_ -cache based attention would have to repeatedly re-materialize kv from ∗x_ , potentially resulting in significant computational overhead. To resolve this, the Extender features an ephemeral kv-cache at inference time. The ephemeral kv-cache of a given layer may be materialized on-demand from the ∗x_ -cache by multiplication with the relevant attention WkW_k and WvW_v matrices. Amortized over multiple input and output tokens, this kv materialization cost is small. Thus the ephemeral kv-cache may be released between turns, as needed. Because of the large attention memory footprint of MHA, it is rarely used in its pure form in large models. Instead, models that retain softmax attention use techniques such as dimensionality reduction Shazeer (2019); Ainslie et al. (2023); DeepSeek-AI (2024), layer reuse Liu et al. (2024a); Sun et al. (2024); Wu and Tu (2024); Brandon et al. (2024); Zuhri et al. (2024) and quantization Xiao et al. (2023); Frantar et al. (2023); Liu et al. (2024b) to limit the footprint. The Extender instead changes the architecture, separating the per-layer next-token prediction δℓ _ from the per-token features ϵℓ _ that persist in attention memory. The result is a compact attention memory footprint between turns, while achieving similar accuracy on short-context and better accuracy on long-context tasks. During turns, memory requirements are unchanged vs a standard Transformer. However, existing dimensionality reduction, layer reuse and quantization methods may still be applied to attention during turns, if desired. Our experiments with GQA (See Appendix D) suggest that Grouped Query Attention combines well with the Extender. The remainder of the paper is organized as follows. §2 briefly introduces past work on reducing attention memory footprint. §3 presents the Extender in more detail, and §4 discusses inference considerations. §5 compares Extender accuracy vs. a Reference Transformer on short-context (DCLM CORE) and long-context (RULER) tasks, and §6 concludes. Additional results and analysis, including GQA, knowledge distillation and scaling results are available in the Appendix. 2 Background Because the attention memory footprint of a standard softmax multi-head attention Transformer grows linearly with model width dmodeld_model, layer depth L and context length T, a variety of memory-reduction approaches have been proposed. The Extender reduces persistent attention memory: the memory required between turns, while remaining compatible with existing approaches during a turn. That said, because a method that applies during a turn applies equally between turns, we briefly review the related work below. Dimensionality Reduction With grouped query attention (GQA) Ainslie et al. (2023), an attention layer with H query heads may maintain only H/4H/4 or H/8H/8 keys and values per token. This in turn reduces attention memory requirements by 4–8×, at a modest accuracy loss vs. full MHA. Multi-head Latent Attention (MLA) DeepSeek-AI (2024) instead stores a compressed version of the kv pair, as much as 10× smaller. This is decompressed on-demand before attention. As a result, MLA does not need to reduce the number of kv heads, but may suffer compression losses instead. The Extender instead changes what the attention layer reads. However, GQA can be applied to the Extender to reduce memory footprint during turns (see Appendix D). Layer Reuse YOCO and LCKV move toward one shared cache Sun et al. (2024); Wu and Tu (2024): YOCO reuses a single global KV cache in its cross-decoder. LCKV has most layers use the same first layer KV cache. Cross-Layer Attention (CLA) and MLKV instead share one cache among a group of layers Brandon et al. (2024); Zuhri et al. (2024). MiniCache instead merges similar layers’ KV states after training Liu et al. (2024a). The Extender significantly reuses attention input data, but does not change the attention mechanism. If desired, CLA may be applied to the Extender to save attention memory during turns. We do not evaluate this here. Quantization With quantization Xiao et al. (2023); Frantar et al. (2023); Liu et al. (2024b), individual features or weights are represented using fewer bits, leading to both reduced memory footprint and often improved computational throughput. Applied judiciously, the accuracy loss from quantization can be modest, and savings on the order of 2–4× are not uncommon. We do not evaluate quantization in this paper. 2.1 Approaches beyond softmax attention Softmax attention incurs a cost that is linear in the context length and number of layers, per token. Several alternative approaches Katharopoulos et al. (2020); Sun et al. (2023); Gu and Dao (2023); Peng et al. (2023); Peng et al. (2024); Goldstein et al. (2024) have been proposed which reduce or eliminate the need to revisit each preceding token during attention. The Extender preserves softmax attention, but changes the input to the attention kv projections, to achieve a reduction in memory footprint between turns. Step Explanation and Motivation −1=−1=[embed(token)]x_-1=h_-1=[embed(token)] Both −1x_-1 and −1h_-1 start out as a token embedding, size dmodeld_model. embedembed weights tied to unembedunembed below. ℓ=Last(ℓ−1,dmodel)s_ =Last(x_ -1,d_model) Size dmodeld_model sliding window of ℓ-1x_ -1. 0=-1s_0=x_-1. ℓ=ATTN(ℓ,)a_ =ATTN(s_ ,h) The ATTN kv projections see only ℓs_ , but the q projection also sees ℓh_ . ℓ=ℓ+ℓ−1y_ =a_ +h_ -1 ℓy_ is the input to the FFN δ^ℓ,ϵℓ=FFN(RMSNorm(ℓ)) δ_ , _ =FFN( RMSNorm(y_ )) The FFN outputs both the residual update δℓ _ , and the extension ϵℓ _ . δℓ=δ^ℓ+ℓ _ = δ_ +a_ ℓa_ allows ATTN to write directly to the residual stream, and provides ATTN a clean gradient. ℓ=[ℓ−1;RMSNorm(ϵℓ)]x_ =[x_ -1;RMSNorm( _ )] Every layer appends its own read-only ϵℓ _ to x. ℓ=ℓ−1+wℓδℓh_ =h_ -1+w_ _ Weighted residual stream update. The learned layer weight wℓw_ tends to give more weight to later layers. logits=unembed(RMSNorm(L−1))logits=unembed(RMSNorm(h_L-1)) Logits from the final hidden state. unembedunembed weights tied to embedembed above. Table 1: A compact description of the Extender. The extended embedding x corresponds to the left part of Figure 1(b), while the residual stream h appears on the right. 3 The Extender: A Log-Structured Transformer (a) 2-layer Reference Transformer: a conventional Llama-style Touvron et al. (2023) Transformer architecture used for comparative evaluation purposes. Each attention layer caches unique KV-pairs derived from the input to the attention layer. (b) 3-Layer Extender. The extended embedding x grows through concatenation with the ϵℓ _ generated by each layer. There are two layer outputs: the extension ϵℓ _ , and δℓ _ . Hexagons indicate RMSNorm. Figure 1: Diagrams describing a conventional Transformer (left), and an Extender (right), using the same visual language and annotations. The Extender attention kv projections see the log-structured extended embedding ℓx_ , rather than the residual stream ℓh_ . Figure 1(b) shows a high-level view of the Extender architecture. Compared to the Transformer in 1(a), the biggest difference lies in the separation of the extended embedding x (left) seen by the attention kv projections, and the residual stream h (right) which is seen only by the FFN sub-layers and the attention q projections. Key to the Extender are the extensions ϵℓ _ and residual updates δℓ _ produced by each layer. The extended embedding ℓx_ is log-structured, updated by concatenating ϵℓ _ to the end, rather than the addition used to update h by δℓ _ . After x and h have passed every layer, logits are produced from h only. Thus the final layer produces no ϵℓ _ . Table 1 describes the model more formally, for ℓ∈0..L-1 ∈ 0..L-1. Here, unembedunembed and embedembed have tied weights. ℓs_ is a size dmodeld_model window over ℓ−1x_ -1. Shown here is a sliding window using a Last()Last() operator, which takes a slice of the most recently added dmodeld_model features out of the extended embedding. In §3.1 we explore alternatives to this design choice. Note that while every ϵε persists between turns, the dmodeld_model-wide initial embedding is fixed and does not have a persistent memory footprint beyond the existing embed table and a token identifier. Beyond the extended embedding x, a few aspects of the Extender design bear explicit mention. Dual FFN outputs the FFN has two outputs: δℓ _ and ϵℓ _ , thus the output of the FFN is dmodel+dϵd_model+d_ε wide. The hidden dimension of the FFN is governed by the width of its input. Attention Query the Attention query projection takes RMSNorm(ℓ)+RMSNorm(ℓ−1)RMSNorm(s_ )+RMSNorm(h_ -1) as input. This does not affect the attention memory footprint, since q is not retained for future tokens to attend to. FFN skip The attention output is added after a linear weight to h. This mirrors the design of a Llama-style transformer, and allows the Attention layer to learn directly from the loss. 3.1 Choosing dϵd_ε and Designing the kv Windows In this paper, we use ||=dmodel|s|=d_model. This makes the Extender attention layer shape identical to that of the Transformer, making direct comparison easier. Given this constraint, we must choose dϵd_ε, or more generally |ϵℓ|| _ |. Here, we decompose that into setting dϵd_ε and ℓmax _max, where ℓmax _max is the final layer that emits an ϵε. By default, ℓmax=L−2 _max=L-2. Moreover, since for ℓ>0 >0, |ℓ|≥dmodel|x_ |≥ d_model each layer ℓ>0 >0 can only read a subset of ℓx_ . These two design decisions are closely tied: with a small dϵd_ε, it becomes more important that all the features of ϵε are seen by subsequent layers. Similarly, with a large dϵd_ε, late layer writes may crowd out the initial embedding and early layer writes, making late layers unable to see critical features. For example our 26-layer model has dmodel=1664d_model=1664 (widthheight=64 widthheight=64). With dϵ=32d_ε=32, ∑|ϵℓ|=800Σ| _ |=800 features, leaving room for 864 input embedding features in 25s_25. with dϵ=64d_ε=64, ∑|ϵℓ|=1600Σ| _ |=1600 features. Empirically, sizes larger than 64 and smaller than 32 perform significantly worse, and are not evaluated here. With dϵ=64d_ε=64, the choice of read window is particularly important. Here, a sliding window leaves the final layer with only 64 features of the original embedding. Three choices we considered were FirstFirst, LastLast, and RandomRandom, for both the k and the v projection. For k, we found that LastLast is always the better choice. For v however, the choice is less obvious: many important tasks involve an identity mapping. Thus v must be able to recall enough of the input embedding to produce the correct logit. For these, LastLast would be a poor choice when dϵ=64d_ε=64 For many other tasks, FirstFirst would instead be a poor choice. For dϵ=64d_ε=64 and ℓmax=L−2 _max=L-2, =Lastk=Last, =Randomv=Random may be the better compromise. Another possibility is changing ℓmax _max, preventing later layers from writing to x. For example, setting ℓmax=L/2 _max=L/2 cuts ∑|ϵℓ|Σ| _ | by half, leaving room for the embedding. We evaluate a few of these choices in Appendix B. Our default model uses dϵ=32d_ε=32, ℓmax=L−2 _max=L-2, except |ϵ0|=64| _0|=64. 4 Inference and the ∗x_ -Cache By replacing the kv-cache with an ∗x_ -cache and computing kv from ∗x_ during attention, the Extender could achieve the memory savings claimed above both during and between turns. However, this would incur the compute cost of repeatedly reprojecting kv from ∗x_ . To address this, the Extender uses an ephemeral kv-cache, which may be freed between turns as desired. This effectively amortizes the cost of reprojecting the kv-cache from the ∗x_ -cache. It is also worth noting that because the reprojection cost grows linearly with the total kv dimension, using GQA for attention dimensionality reduction kv-cache speeds up the reprojection time by a factor similar to the GQA space savings. 5 Evaluation Below, we present results from Extender and Transformer language models trained on the ClimbMix Diao et al. (2025) and ProLong Gao et al. (2025) datasets, using the short-context DCLM CORE Li et al. (2024) and long-context RULER Hsieh et al. (2024) benchmark suites. Reference Transformer In order to better understand the strengths and weaknesses of the Extender, we compare it to a Reference Transformer. The Reference Transformer is a Llama-style MHA transformer with two small changes, following Andrej Karpathy’s nanochat Karpathy (2025) design: the Muon optimizer Jordan et al. (2024) instead of AdamW for the large matrices, and SoftCap Gemma Team (2024) in the attention layer (included with Flash Attention-2 Dao (2024)). The Reference Transformer is a decoder-only Llama-style Transformer with a standard pre-norm residual stream. Token embeddings of dimension dmodeld_model are processed by L identical blocks, each of the form ←+Attn(RMSNorm()),←+FFN(RMSNorm()), +Attn(RMSNorm(h)),~~~~~h +FFN(RMSNorm(h)), (1) followed by a final RMSNormRMSNorm before unembedding. RMSNormRMSNorm uses ε=10−5 =10^-5 and a learnable scale. The Extender mirrors the Reference Transformer as closely as possible in terms of shape, activation functions, optimizer, training schedule, weights, norms etc. Below are several common characteristics. By default, we use common multi-head causal attention, with Rotary Position Embedding (RoPE, default θ=106θ=10^6) applied to queries and keys, using the standard half-dimension split. GQA is addressed briefly in Appendix D Attention logits are soft-capped as s⋅tanh(z/s)s· (z/s) with default s=50s=50, using a FlashAttention-2 kernel. The feed-forward network is SwiGLU Shazeer (2020), FFN(h)=W2(SiLU(W1)⊙W3),FFN(h)=W_2 (SiLU(W_1h) W_3h ), with hidden width ⌊8d/3⌋ 8d/3 rounded up to a multiple of 3232, and all linear layers bias-free. The default tokenizer is nanochat’s byte-level BPE, vocabulary 32k. Documents are prefixed with a beginning-of-sequence token, and we use document masking for best-fit packed training sequences. Input embeddings and the output projection share a single weight matrix (tied embeddings). Parameters are initialized as: embedding rows (0,1/d)N(0,1/ d); multi-dimensional weights Uniform(−1/fan_in,1/fan_in)Uniform(-1/ fan\_in,1/ fan\_in); RMSNormRMSNorm scales are initialized to 1. Models are trained using a WSD training schedule Hägele et al. (2024). We use the Muon optimizer Jordan et al. (2024) for the large weight matrices, and the AdamW optimizer Loshchilov and Hutter (2019) for the embedding weights and other one-dimensional tensors, such as norms. Effective batch for ClimbMix 2k was 256 sequences, scaled peak learning rates were: Muon 2.1×10−22.1× 10^-2, AdamW 1.41×10−21.41× 10^-2 on embeddings, and 8.5×10−48.5× 10^-4 on other 1-D parameters. For ProLong 64k, effective batch was 4 sequences, scaled peak LR: Muon 8.8×10−48.8× 10^-4 and AdamW 3.5×10−53.5× 10^-5 on both AdamW groups. For efficiency, GEMMs run in bfloat16; RoPE Su et al. (2024) frequencies and RMSNorm Zhang and Sennrich (2019) statistics stay in fp32. The design intentionally omits QK-norm, value embeddings, attention residuals, and output logit softcap, all of which could be adapted to the Extender, but would not help characterize the difference between the two architectures. The model parameter counts vary slightly between our Extender and Reference Transformer models due to architecture differences, but model width dmodeld_model and depth L are identical, and as follows. We’ve made an effort to use the actual sizes in the tables, plots and analysis, but will use the size class name when discussing multiple models of the same size class but different exact sizes. Size Width Depth Q/K/V Head class dmodeld_model L heads dim 198M 1024 13 8 128 436M 1280 20 10 128 920M 1664 26 13 128 Table 2: Model shape for each size class. The 436M and 920M models have width-to-depth ratio of 64. The 198M model departs slightly from this ratio, to retain 128-feature heads. The size classes are named after the Transformer size. The Extender versions have approximately 0.5% more parameters. We apply an auxiliary regularization cost to each token to prevent runaway residual writes: 1L∑ℓrelu(RMS(δℓ)−τ)2 1L _ relu(RMS( _ )-τ)^2 with τ=8τ=8 and coefficient 3×10−33×10^-3. After this, later layers still receive more weight (wℓw_ is often 10×10× that of the earliest layers), but features of δℓ _ are in a healthy range. 5.1 Short Context: DCLM CORE Benchmark suite Figure 2: Centered DCLM Core v1 scores. One Transformer (hatched) and one Extender per size class, plus Llama-3.2-1B. 198M/436M/920M trained to Chinchilla token budgets. T-Mono: tasks where the Transformer score increased monotonically with parameter count under Chinchilla. We trained one Extender and one Reference Transformer of each size class, from scratch, using Chinchilla Hoffmann et al. (2022) token budgets of ≈ 20 tokens per model parameter, or approximately 4B, 9B and 18B tokens for the 198M, 436M and 920M parameter models respectively. No fine-tuning or other task-specific training was applied. Figure 2 shows the performance of these models on the DCLM CORE v1 Li et al. (2024) benchmark suite, using the nanochat Karpathy (2025) CORE harness. We group the 22 CORE benchmarks into categories T-Mono/other. Here, T-Mono is the subset of tasks on which the Transformer score improved monotonically with parameter count under Chinchilla-matched training. Intuitively, monotonic improvement with size suggests a more reliable benchmark in this size range. For reference, we also include Llama-3.2-1B scores. This model was pretrained on up to 9 trillion tokens Meta (2024), explaining the large performance gap on some tasks. On other tasks, such as arc_easy, Llama’s 500× more training does not improve performance, suggesting the 1B model size is the primary limitation rather than the training or the architecture. Overall, we find that the Extender and Transformer models perform similarly on DCLM CORE tasks, given parameter counts within 1% of each other, and equal training budgets. Table 3: Mean accuracy (%) after each mid-train stage, for Extender and Transformer (920M unless noted). SCORE is NIAH-6 (six needle tasks: S1–S3, MK-1, MV, MQ), REST (the other seven RULER tasks: MK-2/3, VT, CWE, FWE, QA-1/2), or RULER (13-task mean). Model Stage SCORE 1k 2k 4k 8k 16k 32k 64k T-mono CORE Extender ClimbMix 18B NIAH-6 91.4 82.8 57.6 37.2 13.3 4.0 0.7 32.7 23.7 REST 31.2 22.5 11.6 5.0 4.1 3.7 2.9 RULER 59.0 50.3 32.9 19.9 8.4 3.8 1.9 Transformer ClimbMix 18B NIAH-6 90.0 84.4 17.5 0.0 – – – 32.3 22.5 REST 33.1 24.1 5.9 0.0 – – – RULER 59.4 51.9 11.2 0.0 – – – Extender ProLong 64k +2.5+2.5B NIAH-6 92.8 87.6 85.1 77.5 78.6 73.2 43.3 30.8 23.5 REST 51.4 43.1 31.1 23.7 19.2 7.9 4.8 RULER 70.5 63.7 56.1 48.5 46.6 38.0 22.6 Transformer ProLong 64k +2.5+2.5B NIAH-6 91.0 83.6 75.4 55.7 48.9 42.9 21.0 29.9 21.8 REST 38.5 18.6 11.1 8.2 7.6 5.0 3.0 RULER 62.7 48.6 40.8 30.1 26.7 22.5 11.3 Extender prolong 64k +5+5B NIAH-6 90.7 84.3 83.3 75.0 78.6 74.7 59.3 31.5 22.8 REST 49.9 40.0 41.4 28.3 28.6 17.3 9.9 RULER 68.7 60.4 60.8 49.9 51.7 43.8 32.7 Transformer ProLong 64k +5B NIAH-6 87.8 84.3 84.0 71.0 65.3 64.6 28.5 30.8 21.8 REST 42.2 29.7 20.4 14.4 17.2 9.9 7.1 RULER 63.2 54.9 49.8 40.5 39.4 35.1 17.0 5.2 Long Context: RULER Figure 3: Per-task RULER scores for the ClimbMix 18B → +5+5B 64k ProLong checkpoints, at evaluation lengths 1k / 2k / 4k / 8k / 16k / 32k. Extender is green, Transformer is red; darker bars are longer contexts. All 13 tasks are shown. Avg is the 13-task RULER mean. In this section, we evaluate Reference Transformer and Extender performance on long-context tasks using the RULER Hsieh et al. (2024) benchmark suite. The purpose of the evaluation is an apples-to-apples comparison between Extender and Transformer architectures, not to maximize RULER scores. Models were not fine-tuned for instruction-following or RULER. As a result, these absolute scores are lower than overtrained or fine-tuned public models. Here, we focus on the 920M parameter models, and evaluate them at the pretrain stage, after 2.5B tokens of 64k ProLong data Gao et al. (2025), and after 5B tokens of 64k ProLong respectively, using the recommended 60/40 mix of long vs. short context documents. For 2.5B ProLong training, we kept θ=106θ=10^6, and resumed training from the pre-decay pretrain checkpoint, then trained with ProLong tokens finished by a brief decay. For 5B ProLong training, we similarly resumed training from the pre-decay checkpoint of the 2.5B training run. The resulting aggregate token count for the 5B ProLong models is 21.3B tokens (some ClimbMix tokens lost here due to pre-decay resume). Figure 3 shows how the Extender and the Reference Transformer performed on RULER for context lengths between 1k and 32k tokens. Results are based on n=30n=30 samples per task, with greedy decoding. On average, we find that the Extender outperforms the Reference Transformer on these long-context tasks. Two interesting details stand out: the Extender has a large advantage on the harder multi-key tasks, while the Transformer wins convincingly at the variable tracking task. This suggests a possible difference in inductive bias on these three tasks. However, we did not investigate this further. We define a subset NIAH-6 as the six easier tasks: single-key needle-in-a-haystack variants 1–3 (S1–S3), multivalue (MV), multiquery (MQ), and multikey variant 1 (MK1). We use the label REST to refer to the remaining group of tasks (MK2–3, VT, CWE, FWE, QA1 and QA2). Table 3 reports NIAH-6, REST, and the 13-task RULER mean at each evaluation length, plus T-mono and centered CORE (%) for the ClimbMix 2k base and the 64k ProLong mid-train stages. After the ClimbMix pretrain, and before ProLong training, the Extender is significantly better than the Transformer at context lengths beyond the 2k training context. For 1-2k contexts, any differences are marginal. After ProLong training, the Reference Transformer improves significantly on longer contexts over its pretrain equivalent. However, Transformer performance on long-context tasks consistently trails the Extender, on both the easier NIAH-6 and the harder REST tasks. For the 920M models shown, the persistent attention memory footprint at 64k context is approximately 54.5 million features (109 MB at bf16) for the Extender, vs. 5.67 billion features (11.3 GB at bf16) for the (MHA) Reference Transformer, a factor 104×104× reduction. Despite this, we find the Extender to be interchangeable with the Transformer on short context DCLM CORE tasks, and somewhat ahead on long-context RULER tasks. 5.3 Ablation Studies Table 4: Extender at dϵ=32d_ε=32 (∼ 200 M, d=1024d=1024, L=13L=13, ≈ 4 B ClimbMix tokens, seed 42). Each later row changes one default choice. Configuration Val. T-mono CORE RULER 1k 2k 4k 8k default Extender 2.7547 20.1 14.6 50.1 39.3 25.7 13.5 Q from RMSNorm(h)RMSNorm(h) 2.7555 20.4 15.7 47.0 41.8 12.6 1.2 Q from h 2.7617 19.5 15.1 44.7 41.5 22.1 12.4 Q from x 2.7695 19.2 15.4 44.9 38.5 18.1 8.0 Q from RMSNorm([x;h])RMSNorm([x;h]) 2.7500 20.2 15.6 48.0 37.0 25.9 13.9 Q from [x;h][x;h] 2.7586 20.0 14.3 48.2 39.4 23.1 10.6 FFN input: RMSNorm(ℓ)RMSNorm( a_ ) 2.8398 17.8 14.5 42.4 35.3 25.3 10.2 FFN input: RMSNorm(ℓ+RMSNorm())RMSNorm( a_ +RMSNorm(h)) 2.7727 19.6 14.8 46.0 39.2 33.0 15.4 To better understand how each aspect of the Extender contributes to its performance, we run a series of ablation studies. For these, we train matched ∼ 200 M Extenders (dmodel=1024d_model=1024, L=13L=13, dϵ=32d_ε=32) on ClimbMix for ≈ 4 B tokens with a shared WSD recipe, varying one architectural choice at a time relative to the default Extender residual graph. Table 4 reports final validation loss, DCLM CORE (T-mono mean and overall CORE, centered) and RULER, for a subset of the ablation tests performed. Attention Query Input In contrast to the Transformer, Extender has a choice of what input to take for the attention query projection matrix: x, h, or some combination of the two, with or without RMSNorm applied to the input. The default setting is Q from RMSNorm()+RMSNorm()RMSNorm(x)+RMSNorm(h). In the table, [x;h][x;h] denotes concatenation rather than addition of inputs. Many of the differences are marginal, however, two significant results stand out: Q from either h or x alone significantly underperforms on long-context tasks. Note that a concatenated input, with a WqW_q that is twice as large, did not substantially improve accuracy. Residual Input into FFN As seen in Table 1, the default configuration passes RMSNorm(ℓ+ℓ−1)RMSNorm( a_ +h_ -1) as input to the FFN. Here, we test two ablations: eliminating h as input, and applying an additional norm on h: RMSNorm(ℓ+RMSNorm(ℓ−))RMSNorm( a_ +RMSNorm(h_ -1)). From these ablations, it is clear that including h in the FFN input lifts performance. However, compared to a conventional Transformer, the loss from dropping h is not catastrophic: much of the communication between layers is carried by x rather than relying fully on h. For out-of-domain tasks, having h in Q is much more important than having h in the FFN input. The extra RMSNorm ablation reports mixed and likely insignificant results. 5.4 Limitations While the Extender with dϵ=32d_ε=32 scales well from 199M to 924M parameters, we have not yet been able to test it at much larger model sizes due to resource constraints. Similarly, our token budgets are largely limited to 20× the parameter count. 6 Conclusion In contrast with prior work on reducing the attention memory footprint, the Extender focuses on the persistent memory footprint: memory required between turns. The Extender separates the superposition channel h focused on next-token prediction, from the concatenation channel x dedicated to attention memory. For every size class we tested, 32 features of attention memory per layer was enough to match the Reference Transformer on short-context tasks, and outperform it on long-context tasks. References Ainslie et al. (2023) J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai GQA: training generalized multi-query transformer models from multi-head checkpoints. In Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing (EMNLP), p. 4895–4901. Cited by: §1, §2. Battaglia et al. (2018) P. W. Battaglia, J. B. Hamrick, V. Bapst, A. Sanchez-Gonzalez, V. Zambaldi, M. Malinowski, A. Tacchetti, D. Raposo, A. Santoro, R. Faulkner, C. Gulcehre, F. Song, A. Ballard, J. Gilmer, G. Dahl, A. Vaswani, K. Allen, C. Nash, V. Langston, C. Dyer, N. Heess, D. Wierstra, P. Kohli, M. Botvinick, O. Vinyals, Y. Li, and R. Pascanu Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261. Cited by: Appendix C. Brandon et al. (2024) W. Brandon, M. Mishra, A. Nrusimha, R. Panda, and J. Ragan-Kelley Reducing transformer key-value cache size with cross-layer attention. In Advances in Neural Information Processing Systems, Note: NeurIPS 2024 Cited by: §1, §2. Dao (2024) T. Dao FlashAttention-2: faster attention with better parallelism and work partitioning. In International Conference on Learning Representations (ICLR), Cited by: §5. DeepSeek-AI (2024) DeepSeek-AI DeepSeek-V2: a strong, economical, and efficient mixture-of-experts language model. arXiv preprint arXiv:2405.04434. Cited by: §1, §2. Diao et al. (2025) S. Diao, Y. Yang, Y. Fu, X. Dong, D. Su, M. Kliegl, Z. Chen, P. Belcak, Y. Suhara, H. Yin, M. Patwary, Y. Lin, J. Kautz, and P. Molchanov Nemotron-CLIMB: CLustering-based iterative data mixture bootstrapping for language model pre-training. In Advances in Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track, External Links: 2504.13161 Cited by: §5. Frantar et al. (2023) E. Frantar, S. Ashkboos, T. Hoefler, and D. Alistarh GPTQ: accurate post-training quantization for generative pre-trained transformers. In International Conference on Learning Representations (ICLR), Cited by: §1, §2. Gao et al. (2025) T. Gao, A. Wettig, H. Yen, and D. Chen How to train long-context language models (effectively). In Proceedings of the 63rd Annual Meeting of the Association for Computational Linguistics (ACL), External Links: 2410.02660 Cited by: §5.2, §5. Gemma Team (2024) Gemma Team Gemma 2: improving open language models at a practical size. arXiv preprint arXiv:2408.00118. Cited by: §5. Goldstein et al. (2024) D. Goldstein, F. Obeid, E. Alcaide, G. Song, and E. Cheah GoldFinch: high performance RWKV/transformer hybrid with linear pre-fill and extreme KV-cache compression. arXiv preprint arXiv:2407.12077. Cited by: §2.1. Goyal and Bengio (2022) A. Goyal and Y. Bengio Inductive biases for deep learning of higher-level cognition. Proceedings of the Royal Society A 478 (2266), p. 20210068. Cited by: Appendix C. Gu and Dao (2023) A. Gu and T. Dao Mamba: linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752. Cited by: §2.1. Gu et al. (2024) Y. Gu, L. Dong, F. Wei, and M. Huang MiniLLM: knowledge distillation of large language models. In International Conference on Learning Representations (ICLR), Cited by: Appendix C. Hägele et al. (2024) A. Hägele, E. Bakouch, A. Kosson, L. B. Allal, L. von Werra, and M. Jaggi Scaling laws and compute-optimal training beyond fixed training durations. In Advances in Neural Information Processing Systems, Note: NeurIPS 2024. Warmup-stable-decay schedule External Links: 2405.18392 Cited by: §5. Hinton et al. (2015) G. Hinton, O. Vinyals, and J. Dean Distilling the knowledge in a neural network. arXiv preprint arXiv:1503.02531. Note: NIPS 2014 Deep Learning Workshop Cited by: Appendix C. Hoffmann et al. (2022) J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre Training compute-optimal large language models. arXiv preprint arXiv:2203.15556. Cited by: §5.1. Hsieh et al. (2024) C. Hsieh, S. Sun, S. Kriman, S. Acharya, D. Rekesh, F. Jia, Y. Zhang, and B. Ginsburg RULER: what’s the real context size of your long-context language models?. In First Conference on Language Modeling (COLM), External Links: 2404.06654 Cited by: §5.2, §5. Jordan et al. (2024) K. Jordan, Y. Jin, V. Boza, J. You, F. Cesista, L. Newhouse, and J. Bernstein Muon: an optimizer for hidden layers in neural networks. Note: Blog posthttps://kellerjordan.github.io/posts/muon/ Cited by: §5, §5. Karpathy (2025) A. Karpathy Nanochat. Note: GitHub repositoryhttps://github.com/karpathy/nanochat Cited by: §5, §5.1. Katharopoulos et al. (2020) A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret Transformers are RNNs: fast autoregressive transformers with linear attention. In Proceedings of the 37th International Conference on Machine Learning (ICML), p. 5156–5165. Cited by: §2.1. Li et al. (2024) J. Li, A. Fang, G. Smyrnis, M. Ivgi, M. Jordan, S. Gadre, H. Bansal, E. Guha, S. Keh, K. Arora, S. Garg, R. Xin, N. Muennighoff, R. Heckel, J. Mercat, M. Chen, S. Gururangan, M. Wortsman, A. Albalak, Y. Bitton, M. Nezhurina, A. Abbas, C. Hsieh, D. Ghosh, J. Gardner, M. Kilian, H. Zhang, R. Shao, S. Pratt, S. Sanyal, G. Ilharco, G. Daras, K. Marathe, A. Gokaslan, J. Zhang, K. Chandu, T. Nguyen, I. Vasiljevic, S. Kakade, S. Song, S. Sanghavi, F. Faghri, S. Oh, L. Zettlemoyer, K. Lo, A. El-Nouby, H. Pouransari, A. Toshev, S. Wang, D. Groeneveld, L. Soldaini, P. W. Koh, J. Jitsev, T. Kollar, A. G. Dimakis, Y. Carmon, A. Dave, L. Schmidt, and V. Shankar DataComp-LM: in search of the next generation of training sets for language models. In Advances in Neural Information Processing Systems (NeurIPS) Datasets and Benchmarks Track, External Links: 2406.11794 Cited by: §5.1, §5. Liu et al. (2024a) A. Liu, J. Liu, Z. Pan, Y. He, G. Haffari, and B. Zhuang MiniCache: KV cache compression in depth dimension for large language models. In Advances in Neural Information Processing Systems, Note: NeurIPS 2024 Cited by: §1, §2. Liu et al. (2024b) Z. Liu, J. Yuan, H. Jin, S. Zhong, Z. Xu, V. Braverman, B. Chen, and X. Hu KIVI: a tuning-free asymmetric 2bit quantization for KV cache. In Proceedings of the 41st International Conference on Machine Learning (ICML), p. 32332–32344. Cited by: §1, §2. Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter Decoupled weight decay regularization. In International Conference on Learning Representations (ICLR), Cited by: §5. Meta (2024) Meta Llama 3.2 model card. Note: Hugging Facehttps://huggingface.co/meta-llama/Llama-3.2-1B Cited by: §5.1. Peng et al. (2023) B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, L. Derczynski, X. Du, M. Grella, K. K. GV, X. He, H. Hou, P. Kazienko, J. Kocon, J. Kong, B. Koptyra, H. Lau, J. Lin, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, J. S. Wind, S. Wozniak, Z. Zhang, Q. Zhou, J. Zhu, and R. Zhu RWKV: reinventing RNNs for the transformer era. In Findings of the Association for Computational Linguistics: EMNLP 2023, p. 14048–14077. Cited by: §2.1. Peng et al. (2024) B. Peng, D. Goldstein, Q. Anthony, A. Albalak, E. Alcaide, S. Biderman, E. Cheah, X. Du, T. Ferdinan, H. Hou, P. Kazienko, K. K. GV, J. Kocoń, B. Koptyra, S. Krishna, R. M. Jr., J. Lin, N. Muennighoff, F. Obeid, A. Saito, G. Song, H. Tu, C. Wirawan, S. Woźniak, R. Zhang, B. Zhao, Q. Zhao, P. Zhou, J. Zhu, and R. Zhu Eagle and finch: RWKV with matrix-valued states and dynamic recurrence. In First Conference on Language Modeling (COLM), External Links: 2404.05892 Cited by: §2.1. Sanh et al. (2019) V. Sanh, L. Debut, J. Chaumond, and T. Wolf DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter. arXiv preprint arXiv:1910.01108. Note: NeurIPS 2019 EMC2 Workshop Cited by: Appendix C. Shazeer (2019) N. Shazeer Fast transformer decoding: one write-head is all you need. arXiv preprint arXiv:1911.02150. Cited by: §1. Shazeer (2020) N. Shazeer GLU variants improve transformer. arXiv preprint arXiv:2002.05202. Note: Introduces SwiGLU Cited by: §5. Su et al. (2024) J. Su, Y. Lu, S. Pan, A. Murtadha, B. Wen, and Y. Liu RoFormer: enhanced transformer with rotary position embedding. Neurocomputing 568, p. 127063. Note: arXiv:2104.09864 Cited by: §5. Sun et al. (2023) Y. Sun, L. Dong, S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei Retentive network: a successor to transformer for large language models. arXiv preprint arXiv:2307.08621. Cited by: §2.1. Sun et al. (2024) Y. Sun, L. Dong, Y. Zhu, S. Huang, W. Wang, S. Ma, Q. Zhang, J. Wang, and F. Wei You only cache once: decoder-decoder architectures for language models. In Advances in Neural Information Processing Systems, Note: NeurIPS 2024 Cited by: §1, §2. Touvron et al. (2023) H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample LLaMA: open and efficient foundation language models. arXiv preprint arXiv:2302.13971. Cited by: 1(a). Vaswani et al. (2017) A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin Attention is all you need. In Advances in Neural Information Processing Systems, p. 5998–6008. Note: NIPS 2017 Cited by: §1. Wu and Tu (2024) H. Wu and K. Tu Layer-condensed KV cache for efficient inference of large language models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (ACL), Note: ACL 2024 Cited by: §1, §2. Xiao et al. (2023) G. Xiao, J. Lin, M. Seznec, H. Wu, J. Demouth, and S. Han SmoothQuant: accurate and efficient post-training quantization for large language models. In Proceedings of the 40th International Conference on Machine Learning (ICML), p. 38087–38099. Cited by: §1, §2. Zhang and Sennrich (2019) B. Zhang and R. Sennrich Root mean square layer normalization. In Advances in Neural Information Processing Systems, Note: NeurIPS 2019 Cited by: §5. Zuhri et al. (2024) Z. M. K. Zuhri, M. F. Adilazuarda, A. Purwarianti, and A. F. Aji MLKV: multi-layer key-value heads for memory efficient transformer decoding. arXiv preprint arXiv:2406.09297. Cited by: §1, §2. Appendix A Additional Long-Context RULER Evaluation Details Figure 4: Per-task RULER scores for the ClimbMix 2k pretrain checkpoints, at evaluation lengths 1k / 2k / 4k / 8k / 16k. Extender (dϵ=32d_ε=32) is green, Transformer is red. Figure 5: Per-task RULER scores for the ClimbMix 18B → +2.5+2.5B 64k ProLong checkpoints, at evaluation lengths 1k / 2k / 4k / 8k / 16k / 32k. All 13 tasks are shown. Avg is the 13-task RULER mean. Figures 4 and 5 offer additional training stage RULER performance details. In Figure 4, the Extender has a clear advantage at context lengths beyond the 2k training context. Figure 5 is the +2.5+2.5B stage of Figure 3, with the Extender advantage already clearly visible. Figures 6(a)–7(c) show detailed per-task RULER results per context length, up to 32k. (a) Context length 1k. (b) Context length 2k. (c) Context length 4k. Figure 6: Per-task RULER scores after 64k ProLong mid-train from the ClimbMix 18B checkpoint (+2.5+2.5B / +5+5B), at 1k, 2k, and 4k. Extender / Transformer: d=1664d=1664, L=26L=26. Sizes 8k, 16k, and 32k in the next figure. (a) Context length 8k. (b) Context length 16k. (c) Context length 32k. Figure 7: Same models as Figure 6, at 8k, 16k, and 32k. Appendix B ϵℓ _ size and V window design Table 5: Window and ϵε choices at 920M pretrain (∼ 900 M, d=1664d=1664, L=26L=26, ClimbMix 2k, 9.6 B tokens). wide0 sets |ϵ0|=2dϵ| _0|=2d_ε. ℓmax=13 _max=13 uses K and V from Last(d)Last(d) and emits ϵℓ _ on only the first 13 layers. RULER is the 13-task mean (%). Configuration Val. T-mono CORE RULER 1k 2k 4k 8k dϵ=32d_ε=32 wide0, V=LastV=Last 2.4983 30.5 23.0 60.3 48.1 33.8 14.0 dϵ=32d_ε=32, V=LastV=Last 2.4984 30.5 22.9 57.3 44.3 36.9 14.3 dϵ=64d_ε=64 ℓmax=13 _max=13 2.4781 30.0 22.9 61.9 47.7 33.6 15.1 dϵ=64d_ε=64, V=LastV=Last 2.4934 30.8 22.9 63.0 53.9 34.6 15.3 dϵ=64d_ε=64, V=RandV=Rand 2.4773 30.4 23.4 58.6 47.7 33.8 13.3 dϵ=128d_ε=128, V=RandV=Rand 2.4842 29.2 20.6 61.4 52.0 32.8 14.7 Transformer 2.4983 32.3 23.5 58.5 46.2 19.7 0.8 Table 5 evaluates dϵd_ε vs. window design for V on our 920M model class. Due to resource constraints, these ablation models are trained to 9B tokens instead of the full 18B. In our experience, relative 9B token model performance is a good indicator of relative 18B model performance. We evaluate six different choices. Here, wide0 doubles the width of the first layer: by iteratively zeroing out ϵℓ _ during evaluations, we have anecdotally found that ϵ0 _0 is critical to Extender performance on most tasks, while others are task dependent. dϵ=32d_ε=32 with wide0 is the default Extender. However, we find the differences between most models are marginal in these ablations, dϵ=128d_ε=128 being a notable exception for short contexts. Anecdotally, dϵ=64d_ε=64 V=Rand is a strong choice when fully trained. We did not choose it as the default model due to the cleaner and more memory and compute-efficient design of the dϵ=32,V=Lastd_ε=32,V=Last options. Appendix C Distillation and Extenders Table 6: Logit distillation into a 198M Extender. T-mono and CORE are centered DCLM CORE (%); T-mono is the 12-task set in Figure 2. RULER is the 13-task mean (%) at 1k–8k. Model Tokens T-mono CORE RULER 1k 2k 4k 8k 199M Extender, from scratch 4B 19.9 15.9 46.8 38.5 24.4 12.0 199M Extender student, Extender teacher 4B 24.1 19.1 51.7 44.0 29.6 14.8 199M Extender student, Transformer teacher 4B 23.8 17.2 51.9 44.0 27.8 16.0 924M Extender (teacher) 14.5B 33.2 25.1 70.7 64.1 62.5 53.5 920M Transformer (teacher) 14.5B 32.1 23.9 68.3 58.1 54.7 50.1 Knowledge distillation Hinton et al. (2015); Sanh et al. (2019); Gu et al. (2024) trains a student using a larger teacher model, often achieving higher token efficiency than training from scratch. To better understand differences in inductive bias Goyal and Bengio (2022); Battaglia et al. (2018) between Transformer and Extender models, we train three 198M class extenders: one from scratch, one with a 924M Extender teacher, and one with a 920M Transformer teacher. Here, all three 198M Extenders use dϵ=64d_ε=64 and v selects features from x using the RandomRandom policy. All three are trained to 4 Billion ClimbMix tokens with a 2k token sequence length. Both teachers were trained with 9B ClimbMix + 5B ProLong tokens - earlier models than the 18B reported in the main body of the paper. The Extender teacher is somewhat stronger than the Transformer teacher, on both CORE and RULER scores. Table 6 shows the results of this experiment. The two student models both consistently outperform the from-scratch model, showing that Extenders train well under a distillation objective. The student led by an Extender teacher achieved a CORE score 6 points lower than its teacher, while the Transformer teacher transferred slightly less knowledge, with the student scoring 6.7 points below the teacher. On long-context RULER tasks, the two students achieved near-identical scores within the 2k training domain despite the Transformer teacher being weaker. These results indicate that under logit distillation, an Extender student learns approximately as well from an Extender teacher as from a Transformer teacher. Appendix D Grouped Query Attention and Extenders Table 7: GQA vs. MHA Extender and MHA Transformer RULER is the 13-task mean (%). For all models dmodel=1664d_model=1664. Head dimension 128 → 13 heads. Thus GQA4/5 is 13 query / 3 KV heads (groups 5,4,45,4,4). Due to the dimensionality reduction of GQA, the parameter count of the GQA4/5 model is 813M rather than 924M. Model T-mono CORE RULER 1k 2k 4k 8k 16k 32k 64k Extender 924M (MHA) 32.7 23.7 59.0 50.3 32.9 19.9 8.4 3.8 1.9 Extender 813M (GQA4/5) 33.1 22.9 61.2 52.4 36.3 18.4 6.6 2.1 1.2 Transformer 920M (MHA) 32.3 22.5 59.4 51.9 11.2 0.0 – – – Extender 924M +2.5+2.5B (MHA) 30.8 23.5 70.5 63.7 56.1 48.5 46.6 38.0 22.6 Extender 813M +2.5+2.5B (GQA4/5) 31.2 23.3 63.9 57.9 56.3 49.7 45.8 33.8 20.2 Transformer 920M +2.5+2.5B (MHA) 29.9 21.8 62.7 48.6 40.8 30.1 26.7 22.5 11.3 Extender 924M +5+5B (MHA) 31.5 22.8 68.7 60.4 60.8 49.9 51.7 43.8 32.7 Extender 813M +5+5B (GQA4/5) 31.6 22.5 67.3 61.2 58.4 45.2 40.2 37.7 26.3 Transformer 920M +5+5B (MHA) 30.8 21.8 63.2 54.9 49.8 40.5 39.4 35.1 17.0 Figure 8: Centered DCLM CORE subscores for the 18B ClimbMix +5+5B ProLong checkpoints in Table 7. MHA is the full-attention Extender, GQA is 13 query / 3 KV heads (groups 5,4,45,4,4), and TF is the document-mask Transformer (hatched). T-mono is the 12-task set from Figure 2; CORE is the 22-task mean. Figure 9: RULER, NIAH-6, and REST versus context length for the +5+5B checkpoints in Table 7. NIAH-6 is S1–S3, MK1, MV, and MQ; REST is the other seven tasks (MK2–3, VT, CWE, FWE, QA1–2). lengths 1k–64k. Figure 10: Per-task RULER scores for the 920M 18B ClimbMix +5+5B ProLong MHA and GQA checkpoints, at 1k / 2k / 4k / 8k / 16k / 32k / 64k. All 13 tasks are shown. Avg is the 13-task RULER mean, n=30n=30. Table 7 shows the performance of GQA on our 924M model, vs. our regular MHA Extender and MHA Transformer, at the 18B ClimbMix pretrain, +2.5B ProLong and +5B ProLong checkpoints. Because our 920M class model has dmodel=1664d_model=1664, there are 13 heads each 128 features wide. To accomodate this, we use GQA4/5: two groups of 4 queries, one group of 5 queries. Due to the GQA dimensionality reduction, the actual parameter count of the GQA4/5 model is 813M rather than 924M. We find that GQA interacts well with Extender, and the Extender GQA4/5 model outperforms the Transformer MHA. Due to resource constraints, we did not train a matching Transformer GQA. Figure 8 provides per-task performance on the CORE benchmark. GQA tracks MHA performance closely on the more stable T-Mono tasks. On the remaining tasks results are a little noisier, but GQA is largely tracking MHA there as well. Figure 9 shows the RULER performance divided into 13-task mean, NIAH-6, and REST. GQA is overall tracking MHA closely on both the easier NIAH-6 tasks and REST. Finally, Figure 10 shows individual task performance vs. context length. Appendix E Accuracy across model sizes Table 8 reports centered CORE and RULER for the default Extender for our three size classes, for models trained on ClimbMix only, to a Chinchilla token budget. Accuracy on both CORE and RULER 1–2k improves consistently with model size, particularly on T-mono, where the Transformer is also reliably improving with model size. At context lengths outside of the 2k training context length, RULER scores are less predictable, but suggest an improvement as well. Table 8: Default Extender at each size class in Figure 2. 199M is 4B tokens, 438M is 9.6B, and 924M is 18B. T-mono and CORE are centered DCLM CORE (%); T-mono is the 12-task set in Figure 2. RULER is the 13-task mean (%). Italics are lengths longer than the 2k training context. “–”: not evaluated. Model T-mono CORE RULER 1k 2k 4k 8k 199M 20.4 15.7 49.3 41.0 28.0 16.9 438M 26.0 19.5 57.3 45.1 33.4 15.6 924M 32.7 23.7 59.0 50.3 32.9 19.9