The Extender: A Log-Structured Transformer

Jakob Eriksson

cs.LGSep 26, 2026

The Extender: A Log-Structured Transformer

Authors: Jakob Eriksson

Organizations: UIC · University of Illinois Chicago

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 $δ_\ell$ which is added to $\mathbf{h}$ , and a much smaller extension $ε_\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|ε_\ell|$ . We find that with $|ε_\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.

Figures & tables

Step	Explanation and Motivation
$\mathbf{x}_{-1}=\mathbf{h}_{-1}=[\operatorname{embed}(\mathrm{token})]$	Both $\mathbf{x}_{-1}$ and $\mathbf{h}_{-1}$ start out as a token embedding, size $d_{model}$ . $embed$ weights tied to $unembed$ below.
$\mathbf{s}_{\ell}=\operatorname{Last}(\mathbf{x}_{\ell-1},d_{model})$	Size $d_{model}$ sliding window of $\mathbf{x}_{\ell\text{-}1}$ . $\mathbf{s}_{0}=\mathbf{x}_{\text{-}1}$ .
$\mathbf{a}_{\ell}=\operatorname{ATTN}(\mathbf{s}_{\ell},\mathbf{h})$	The ATTN $\mathbf{kv}$ projections see only $\mathbf{s_{\ell}}$ , but the $\mathbf{q}$ projection also sees $\mathbf{h}_{\ell}$ .
$\mathbf{y}_{\ell}=\mathbf{a}_{\ell}+\mathbf{h}_{\ell-1}$	$\mathbf{y_{\ell}}$ is the input to the FFN
$\hat{\delta}_{\ell},\epsilon_{\ell}=\operatorname{FFN}(\mathrm{\operatorname{RMSNorm}}(\mathbf{y}_{\ell}))$	The FFN outputs both the residual update $\delta_{\ell}$ , and the extension $\epsilon_{\ell}$ .
$\delta_{\ell}=\hat{\delta}_{\ell}+\mathbf{a}_{\ell}$	$\mathbf{a}_{\ell}$ allows ATTN to write directly to the residual stream, and provides ATTN a clean gradient.

Table 1: A compact description of the Extender. The extended embedding

\mathbf{x}

corresponds to the left part of Figure 1(b) , while the residual stream

\mathbf{h}

appears on the right.

Panel of Figure 1 — Figure 1: Diagrams describing a conventional Transformer (left), and an Extender (right), using the same visual language and annotations. The Extender attention $\mathbf{kv}$ projections see the log-structured extended embedding $\mathbf{x}_{\ell}$ , rather than the residual stream $\mathbf{h}_{\ell}$ .

Size	Width	Depth	Q/K/V	Head
class	$d_{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.

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.

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	–	–	–

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).

Figure 3: Per-task RULER scores for the ClimbMix 18B

\to

+5

B 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.

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 $\mathrm{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 $\mathrm{RMSNorm}([x;h])$	2.7500	20.2	15.6	48.0	37.0	25.9	13.9

Table 4: Extender at

d_{\epsilon}{=}32

(

\sim

200 M,

d{=}1024

L{=}13

\approx

4 B ClimbMix tokens, seed 42). Each later row changes one default choice.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Per-task RULER scores for the ClimbMix 2k pretrain checkpoints, at evaluation lengths 1k / 2k / 4k / 8k / 16k. Extender (

d_{\epsilon}{=}32

) is green, Transformer is red.

Figure 5: Per-task RULER scores for the ClimbMix 18B

\to

+2.5

B 64k ProLong checkpoints, at evaluation lengths 1k / 2k / 4k / 8k / 16k / 32k. All 13 tasks are shown. Avg is the 13-task RULER mean.

Panel of Figure 6 — Figure 6: Per-task RULER scores after 64k ProLong mid-train from the ClimbMix 18B checkpoint ( $+2.5$ B / $+5$ B), at 1k, 2k, and 4k. Extender / Transformer: $d{=}1664$ , $L{=}26$ . Sizes 8k, 16k, and 32k in the next figure.

Panel of Figure 7 — Figure 7: Same models as Figure 6 , at 8k, 16k, and 32k.

Configuration	Val.	T-mono	CORE	RULER
				1k	2k	4k	8k
$d_{\epsilon}{=}32$ wide0, $V{=}\operatorname{Last}$	2.4983	30.5	23.0	60.3	48.1	33.8	14.0
$d_{\epsilon}{=}32$ , $V{=}\operatorname{Last}$	2.4984	30.5	22.9	57.3	44.3	36.9	14.3
$d_{\epsilon}{=}64$ $\ell_{max}{=}13$	2.4781	30.0	22.9	61.9	47.7	33.6	15.1
$d_{\epsilon}{=}64$ , $V{=}\operatorname{Last}$	2.4934	30.8	22.9	63.0	53.9	34.6	15.3
$d_{\epsilon}{=}64$ , $V{=}\mathrm{Rand}$	2.4773	30.4	23.4	58.6	47.7	33.8	13.3

Appendix

Table 5: Window and

\epsilon

choices at 920M pretrain (

\sim

900 M,

d{=}1664

L{=}26

, ClimbMix 2k, 9.6 B tokens). wide0 sets

|\epsilon_{0}|{=}2d_{\epsilon}

\ell_{max}=13

uses

K

and

V

from

\operatorname{Last}(d)

and emits

\epsilon_{\ell}

on only the first 13 layers. RULER is the 13-task mean (%).

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

Appendix

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	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$ B (MHA)	30.8	23.5	70.5	63.7	56.1	48.5	46.6	38.0	22.6
Extender 813M $+2.5$ B (GQA4/5)	31.2	23.3	63.9	57.9	56.3	49.7	45.8	33.8	20.2

Appendix

Table 7: GQA vs. MHA Extender and MHA Transformer RULER is the 13-task mean (%). For all models

d_{model}=1664

. Head dimension 128

\to

13 heads. Thus GQA4/5 is 13 query / 3 KV heads (groups

5,4,4

). Due to the dimensionality reduction of GQA, the parameter count of the GQA4/5 model is 813M rather than 924M.

Figure 8: Centered DCLM CORE subscores for the 18B ClimbMix

+5

B ProLong checkpoints in Table 7 . MHA is the full-attention Extender, GQA is 13 query / 3 KV heads (groups

5,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

B 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

B 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{=}30

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

Appendix

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.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 540 papers

You Do Not Fully Utilize Transformer's Representation Capacity
Feb 13, 2025Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky +2Representation Collapse Memory-Augmented Language Models
1T-Tech · 2Moscow Institute of Physics and Technology · 3HSE University
Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
Oct 13, 2026Sarah Pan, Ali BehrouzLong-Context Language Model Inference Memory-Augmented Language Models
Massachusetts Institute of Technology · Cornell University
Cracks in the Foundation: Seemingly Minor Architectural Choices Impact Long Context Extension
Aug 10, 2026Amanda Bertsch, Luca Soldaini, Matthew R. Gormley +4Transformer Long-Context Language Modeling
Ai2 · Carnegie Mellon University · University of Washington

Feb 13, 2025J/K move · Enter open · S save

You Do Not Fully Utilize Transformer's Representation Capacity

Gleb Gerasimov, Yaroslav Aksenov, Nikita Balagansky, Viacheslav Sinii+1

1T-Tech · 2Moscow Institute of Physics and Technology · 3HSE University

Representation Collapse Memory-Augmented Language Models Self-Attention Cross-Layer Attention Transformer Attention Attention Mechanisms

Open paper AskPDF arXiv

In contrast to RNNs, which compress their history into a single hidden state, Transformers can attend to all past tokens directly. However, standard Transformers rely solely on the hidden state from the previous layer to represent the entire context. We show that this design creates pressure toward representation collapse and can degrade performance. To address this issue, we introduce Layer-Integrated Memory (LIMe), a lightweight extension that leverages existing key-value buffers and learns per-head, per-layer routing weights to integrate representations from previous layers. Across language modeling, synthetic reasoning, and deep architectures, LIMe improves perplexity per FLOP in the studied regimes and yields strong gains on synthetic tasks while preserving higher value-vector entropy and token separability. Finally, learned routing weights reveal systematic reuse of local and long-distance features, showing how LIMe enriches attention-time memory without increasing hidden-state size. Code is available at https://github.com/corl-team/lime.

Abstract

Transformer

Explore similar work

You Do Not Fully Utilize Transformer's Representation Capacity

Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules

Cracks in the Foundation: Seemingly Minor Architectural Choices Impact Long Context Extension