面向循环 Transformer 的因果潜在工作区

arXiv cs.CL 论文

摘要

该论文提出因果潜在工作区 Transformer(causal latent workspace Transformer),通过在共享深度的循环块之间加入可读写的前缀专属潜在槽位,提升了循环 Transformer 的语言建模困惑度,并在三个字节级语料上的多数设置中优于五个基线对照。

arXiv:2610.22796v1 Announce Type: new Abstract: Reusing a Transformer block across depth increases computation without replicating its parameters, but successive passes have no dedicated state through which positions can exchange compact summaries of their causal prefix. We introduce the causal latent workspace Transformer: an anchored, four-round shared-depth architecture in which each position writes gated vectors into prefix-specific latent slots and retrieves a content-addressed summary before the next shared block application. This path complements causal self-attention while the original input embedding remains accessible at every round. Across three byte-level corpora, three seeds, and two matched parameter budgets, the proposed model achieves the best mean held-out perplexity in five of six settings against five Transformer controls. At 101.71M parameters it leads on WikiText-103 raw, enwik8, and TinyStories, reaching perplexities 2.7820, 2.7843, and 1.6823; these improve upon the strongest non-workspace control on each corpus by 1.06%, 0.73%, and 1.01%. All nine seed-labelled comparisons at this budget favor the workspace. At 0.918M parameters it leads on WikiText-103 and TinyStories, whereas a conditioned loop is stronger on enwik8. Architectural ablations show that adding the workspace improves upon the anchored loop in all six settings and support refreshing its causal summaries at each round. Together, these results identify a causal, content-addressable prefix workspace as an effective structural augmentation of looped Transformers.
查看原文
查看缓存全文

缓存时间: 2026/10/19 07:14

# Causal Latent Workspaces for Looped Transformers
Source: [https://arxiv.org/html/2610.22796](https://arxiv.org/html/2610.22796)
October 2026

###### Abstract

Reusing a Transformer block across depth increases computation without replicating its parameters, but successive passes have no dedicated state through which positions can exchange compact summaries of their causal prefix\. We introduce the*causal latent workspace Transformer*: an anchored, four\-round shared\-depth architecture in which each position writes gated vectors into prefix\-specific latent slots and retrieves a content\-addressed summary before the next shared block application\. This path complements causal self\-attention while the original input embedding remains accessible at every round\. Across three byte\-level corpora, three seeds, and two matched parameter budgets, the proposed model achieves the best mean held\-out perplexity in five of six settings against five Transformer controls\. At 101\.71M parameters it leads on WikiText\-103 raw, enwik8, and TinyStories, reaching perplexities 2\.7820, 2\.7843, and 1\.6823; these improve upon the strongest non\-workspace control on each corpus by 1\.06%, 0\.73%, and 1\.01%\. All nine seed\-labelled comparisons at this budget favor the workspace\. At 0\.918M parameters it leads on WikiText\-103 and TinyStories, whereas a conditioned loop is stronger on enwik8\. Architectural ablations show that adding the workspace improves upon the anchored loop in all six settings and support refreshing its causal summaries at each round\. Together, these results identify a causal, content\-addressable prefix workspace as an effective structural augmentation of looped Transformers\.

*K*eywordsTransformer⋅\\cdotrecurrent Transformer⋅\\cdotparameter sharing⋅\\cdotneural memory⋅\\cdotbyte\-level language modeling⋅\\cdotcausal attention

## 1Introduction

Self\-attention provides the Transformer with a direct, parallel mechanism for integrating information across a sequence\[[1](https://arxiv.org/html/2610.22796#bib.bib1)\]\. Its effectiveness also raises a continuing architectural question: how should this information be organized as representations are refined through depth? In a conventional stack, independently parameterized layers successively transform the hidden sequence\. A looped Transformer instead reuses a shared block, allocating more computation to refinement without replicating the block’s parameters\[[2](https://arxiv.org/html/2610.22796#bib.bib2),[5](https://arxiv.org/html/2610.22796#bib.bib5)\]\. This makes the transition between rounds a central design choice\. A standard shared loop carries forward the hidden sequence, but provides no separate, addressable representation of the causal prefix\.

We introduce the*causal latent workspace Transformer*, which augments shared\-depth refinement with a learned write/read pathway\. At each round, gated writes accumulate information from the causal prefix into a small set of vector slots\. A query derived from the current token state retrieves a weighted summary from these slots before the next shared block application\. The original input embedding remains available through an anchor, allowing refinement to combine the input, the evolving hidden sequence, and a content\-addressed prefix summary\. Causal self\-attention and the workspace consequently provide complementary routes: the former operates on individual token states, while the latter operates on learned prefix summaries\. Each prefix has its own workspace, reconstructed from the current round and reset at the next input block\.

Our experiments examine whether this additional information pathway improves language modeling beyond shared depth, input anchoring, and more expressive recurrent transitions\. We compare the proposed model with five Transformer controls at approximately 0\.918M and 101\.71M parameters, using the same training\-target budget within each size\. Across WikiText\-103 raw, enwik8, and TinyStories, with three seeds per condition, the workspace achieves the lowest mean held\-out perplexity in five of six corpus–size settings\. It improves upon the anchored loop in all six settings\. At 101\.71M parameters it also outperforms the strongest alternative on every corpus and for every corresponding seed label\. At the smaller budget it leads on WikiText\-103 and TinyStories, while the conditioned\-dense control leads on enwik8\. These comparisons establish the value of a causal workspace under matched parameter and training\-target budgets\.

The contribution is a causal write/read mechanism that makes compact prefix summaries available throughout recurrent refinement, together with an empirical study of its effect relative to alternative recurrent architectures\. The following sections place the design in the context of recurrent depth and neural memory, derive the architecture, and illustrate its computation with a fully specified example\. We then describe the matched evaluation and use architectural ablations to distinguish the benefit of adding a workspace from the choice of how its state is updated across rounds\.

## 2Related work

Recurrent depth provides a starting point for organizing repeated computation\. Universal Transformers share a transition over depth\[[2](https://arxiv.org/html/2610.22796#bib.bib2)\], while adaptive computation time varies the number of applications according to the input\[[3](https://arxiv.org/html/2610.22796#bib.bib3)\]\. Block\-Recurrent Transformers combine attention with recurrent block\-level state\[[4](https://arxiv.org/html/2610.22796#bib.bib4)\]\. More recent latent\-recurrence and Mixture\-of\-Recursions architectures study shared computation and selective depth allocation\[[5](https://arxiv.org/html/2610.22796#bib.bib5),[6](https://arxiv.org/html/2610.22796#bib.bib6)\]\. Together, these approaches establish repeated refinement as an alternative to an independently parameterized stack\. The present work addresses the transition within that refinement: we keep the number of shared\-block applications fixed and introduce a separate pathway for constructing and retrieving causal prefix summaries\.

Neural memory provides the second foundation of the design\. Transformer\-XL carries representations across segments\[[7](https://arxiv.org/html/2610.22796#bib.bib7)\], Compressive Transformers compress older states\[[8](https://arxiv.org/html/2610.22796#bib.bib8)\], and Memorizing Transformers retrieve past key–value pairs\[[9](https://arxiv.org/html/2610.22796#bib.bib9)\]\. Differentiable memory systems such as Neural Turing Machines and Differentiable Neural Computers separate a controller from an addressable store\[[10](https://arxiv.org/html/2610.22796#bib.bib10),[11](https://arxiv.org/html/2610.22796#bib.bib11)\]\. These approaches extend accessible history or maintain a store as a sequence is processed\. The causal workspace instead organizes information within the current context during depth refinement\. It reconstructs its summaries from the current token states at every round, allowing the write/read pathway to respond to changes in those states\. This function is distinct from a decoder key–value cache, which reuses previously computed attention projections during autoregressive decoding\.

Latent bottlenecks offer a way to implement such a compact information pathway\. The shared global workspace of Goyal et al\. coordinates specialized modules\[[12](https://arxiv.org/html/2610.22796#bib.bib12)\]; Slot Attention learns reusable latent slots\[[13](https://arxiv.org/html/2610.22796#bib.bib13)\]; and the Perceiver exchanges information through a compact latent array\[[14](https://arxiv.org/html/2610.22796#bib.bib14)\]\. Applying this principle to causal language modeling requires a separate summary at each sequence position\. A slot available at positioniican include writes only from positionsj≤ij\\leq i\. Our gated prefix accumulations satisfy this constraint while sharing slot keys and projection parameters across positions and rounds\. The resulting workspace is a collection of causal summaries rather than a single global summary of the entire sequence\.

This separation of shared parameters from transient input\-dependent state also connects the model to fast\-weight networks\[[15](https://arxiv.org/html/2610.22796#bib.bib15),[16](https://arxiv.org/html/2610.22796#bib.bib16)\]\. In our architecture, the transient state is a gated vector accumulation; the learned projection matrices remain fixed during a forward pass\. Other sequence architectures pursue different interventions: Linear Transformers and Performers reformulate attention\[[17](https://arxiv.org/html/2610.22796#bib.bib17),[18](https://arxiv.org/html/2610.22796#bib.bib18)\], Longformer and Big Bird restrict interaction patterns\[[19](https://arxiv.org/html/2610.22796#bib.bib19),[20](https://arxiv.org/html/2610.22796#bib.bib20)\], and Reformer changes attention and memory organization\[[21](https://arxiv.org/html/2610.22796#bib.bib21)\]\. Mamba, RetNet, and Hyena use recurrent or convolutional sequence mechanisms\[[22](https://arxiv.org/html/2610.22796#bib.bib22),[23](https://arxiv.org/html/2610.22796#bib.bib23),[24](https://arxiv.org/html/2610.22796#bib.bib24)\], while FlashAttention improves the execution of exact attention\[[25](https://arxiv.org/html/2610.22796#bib.bib25)\]\. The workspace complements these directions by retaining full causal self\-attention and changing the information supplied between shared\-depth applications\. Its effect is therefore evaluated through both prediction quality and the additional computation required by the write/read path\.

The experimental design follows from this architectural question\. WikiText\-103\[[26](https://arxiv.org/html/2610.22796#bib.bib26)\], enwik8\[[28](https://arxiv.org/html/2610.22796#bib.bib28)\], and TinyStories\[[27](https://arxiv.org/html/2610.22796#bib.bib27)\]provide different text distributions on which to assess the same mechanism\. Since stochastic training can affect architectural rankings\[[29](https://arxiv.org/html/2610.22796#bib.bib29)\], we evaluate the complete corpus–size grid with three seeds and score final checkpoints over all complete held\-out blocks\.

## 3Architecture

The proposed model augments one shared causal Transformer block with two complementary paths, shown in Figure[1](https://arxiv.org/html/2610.22796#S3.F1)\. The anchor preserves the original input representation as hidden states are refined; the workspace gathers and redistributes prefix information between successive applications of the block\. We first define next\-byte prediction, then the write/read operations, and finally their composition with shared\-depth refinement\.

### 3\.1Prediction task and notation

Let𝒱=\{0,…,256\}\\mathcal\{V\}=\\\{0,\\ldots,256\\\}be the 257\-symbol output vocabulary\. Symbols00–255255are byte values; symbol256256is the row separator used for WikiText only\. For a block ofL\+1L\+1symbols\(x1,…,xL\+1\)\(x\_\{1\},\\ldots,x\_\{L\+1\}\), the model consumes\(x1,…,xL\)\(x\_\{1\},\\ldots,x\_\{L\}\)and predictsyi=xi\+1y\_\{i\}=x\_\{i\+1\}at positionii, where1≤i≤L1\\leq i\\leq L\. We useL=512L=512\. Batch indices are suppressed in equations; sequences in a batch are processed independently\.

Letddbe the token\-state width,RRthe number of shared rounds,SSthe workspace slot count, andmmthe coordinates per slot\. Experiments useR=4R=4andS=4S=4\. The trainable embedding lookupE:𝒱→ℝdE:\\mathcal\{V\}\\rightarrow\\mathbb\{R\}^\{d\}and trainable position vectorpi∈ℝdp\_\{i\}\\in\\mathbb\{R\}^\{d\}produce anchoreie\_\{i\}\. Superscriptrrdenotes the state before shared roundrr, withr=0,…,R−1r=0,\\ldots,R\-1:

ei=E⁡\(xi\)\+pi,hi0=ei\.e\_\{i\}=E\(x\_\{i\}\)\+p\_\{i\},\\qquad h\_\{i\}^\{0\}=e\_\{i\}\.\(1\)Vectors are column vectors\. Brackets\[a;b\]\[a;b\]concatenate vectors\.θ\\thetadenotes all slowly learned parameters\. Affine projections include their displayed biases; weights and biases are shared across rounds unless explicitly indexed byrr\.

### 3\.2Anchoring and gated causal writes

Each round first fuses the unchanged anchor with the current token state:

uir=LNa⁡\(hir\+Fa​\(\[ei;hir\]\)\),Fa​\(w\)=A2​GELU​\(A1​w\+a1\)\+a2\.u\_\{i\}^\{r\}=\\operatorname\{LN\}\_\{a\}\\\!\\left\(h\_\{i\}^\{r\}\+F\_\{a\}\(\[e\_\{i\};h\_\{i\}^\{r\}\]\)\\right\),\\qquad F\_\{a\}\(w\)=A\_\{2\}\\,\\mathrm\{GELU\}\(A\_\{1\}w\+a\_\{1\}\)\+a\_\{2\}\.\(2\)Herew∈ℝ2​dw\\in\\mathbb\{R\}^\{2d\}is the concatenated argument,A1∈ℝd×2​dA\_\{1\}\\in\\mathbb\{R\}^\{d\\times 2d\},A2∈ℝd×dA\_\{2\}\\in\\mathbb\{R\}^\{d\\times d\}, anda1,a2∈ℝda\_\{1\},a\_\{2\}\\in\\mathbb\{R\}^\{d\}are learned parameters\. GELU is the element\-wise activation of Hendrycks and Gimpel\[[30](https://arxiv.org/html/2610.22796#bib.bib30)\]\.LNa\\operatorname\{LN\}\_\{a\}is a learned LayerNorm\[[31](https://arxiv.org/html/2610.22796#bib.bib31)\]; subscripts identify separate parameter sets, not different mathematical operations\.

With write\-normalized statezir=LNw⁡\(uir\)z\_\{i\}^\{r\}=\\operatorname\{LN\}\_\{w\}\(u\_\{i\}^\{r\}\), the write gategi​sr∈\(0,1\)g\_\{is\}^\{r\}\\in\(0,1\)and write vectorvi​sr∈\(−1,1\)mv\_\{is\}^\{r\}\\in\(\-1,1\)^\{m\}for slotssare

gi​sr=sigmoid\(\(Wgzir\+bg\)s\),vi​sr=tanh\(\(Wvzir\+bv\)s,:\)\.g\_\{is\}^\{r\}=\\operatorname\{sigmoid\}\(\(W\_\{g\}z\_\{i\}^\{r\}\+b\_\{g\}\)\_\{s\}\),\\qquad v\_\{is\}^\{r\}=\\tanh\(\(W\_\{v\}z\_\{i\}^\{r\}\+b\_\{v\}\)\_\{s,:\}\)\.\(3\)Wg∈ℝS×dW\_\{g\}\\in\\mathbb\{R\}^\{S\\times d\},bg∈ℝSb\_\{g\}\\in\\mathbb\{R\}^\{S\},Wv∈ℝS​m×dW\_\{v\}\\in\\mathbb\{R\}^\{Sm\\times d\}, andbv∈ℝS​mb\_\{v\}\\in\\mathbb\{R\}^\{Sm\}are learned\. The value projection is reshaped intoSSrows of lengthmm;\(s,:\)\(s,:\)selects a row\. Sigmoid and tanh act element\-wise\. The newly written memory for prefix1,…,i1,\\ldots,iis

Ni​sr=∑j=1igj​sr​vj​sr,Di​sr=∑j=1igj​sr,M~i​sr=Ni​srmax⁡\(Di​sr,ϵ\),ϵ=10−6\.N\_\{is\}^\{r\}=\\sum\_\{j=1\}^\{i\}g\_\{js\}^\{r\}v\_\{js\}^\{r\},\\qquad D\_\{is\}^\{r\}=\\sum\_\{j=1\}^\{i\}g\_\{js\}^\{r\},\\qquad\\widetilde\{M\}\_\{is\}^\{r\}=\\frac\{N\_\{is\}^\{r\}\}\{\\sqrt\{\\max\(D\_\{is\}^\{r\},\\epsilon\)\}\},\\quad\\epsilon=10^\{\-6\}\.\(4\)jjindexes causal sources,Ni​sr∈ℝmN\_\{is\}^\{r\}\\in\\mathbb\{R\}^\{m\}is their gated vector sum,Di​srD\_\{is\}^\{r\}their scalar gate mass, andM~i​sr∈ℝm\\widetilde\{M\}\_\{is\}^\{r\}\\in\\mathbb\{R\}^\{m\}the slot memory\. The current input position is included, but its next\-symbol target is not\. The square\-root denominator scales the accumulated value by gate mass while preserving its dependence on the amount of accumulated evidence\. Unlike a weighted average, this normalization allows aligned writes to increase the summary magnitude as the square root of prefix mass\.

The main model recomputesMi​sr=M~i​srM\_\{is\}^\{r\}=\\widetilde\{M\}\_\{is\}^\{r\}every round\. This choice lets each round revise its summary after the preceding Transformer update\. In a separate retention experiment, we replace the direct assignment with

Mi​s0=M~i​s0,Mi​sr=ρr​Mi​sr−1\+\(1−ρr\)​M~i​sr​\(r≥1\),ρr=sigmoid⁡\(ℓr\)\.M\_\{is\}^\{0\}=\\widetilde\{M\}\_\{is\}^\{0\},\\qquad M\_\{is\}^\{r\}=\\rho\_\{r\}M\_\{is\}^\{r\-1\}\+\(1\-\\rho\_\{r\}\)\\widetilde\{M\}\_\{is\}^\{r\}\\ \(r\\geq 1\),\\qquad\\rho\_\{r\}=\\operatorname\{sigmoid\}\(\\ell\_\{r\}\)\.\(5\)Hereℓr\\ell\_\{r\}is a learned scalar that controls the contribution of the preceding round; its initializationlog⁡\(0\.2/0\.8\)\\log\(0\.2/0\.8\)givesρr=0\.2\\rho\_\{r\}=0\.2\. This variant blends old and newly written summaries at the same prefix position\. Both update rules reset workspace state at input\-block boundaries\. Section[6](https://arxiv.org/html/2610.22796#S6)compares their effect under the same parameter and training\-target budgets\.

### 3\.3Slot read and shared refinement

Each slot has a learned keyks∈ℝmk\_\{s\}\\in\\mathbb\{R\}^\{m\}, initialized with independent normal coordinates of standard deviationm−1/2m^\{\-1/2\}\. The read query and attention over slots are

qir=Wq​LNq⁡\(uir\)\+bq,αi​sr=exp⁡\(\(qir\)⊤​ks/m\)∑t=1Sexp⁡\(\(qir\)⊤​kt/m\)\.q\_\{i\}^\{r\}=W\_\{q\}\\operatorname\{LN\}\_\{q\}\(u\_\{i\}^\{r\}\)\+b\_\{q\},\\qquad\\alpha\_\{is\}^\{r\}=\\frac\{\\exp\(\(q\_\{i\}^\{r\}\)^\{\\top\}k\_\{s\}/\\sqrt\{m\}\)\}\{\\sum\_\{t=1\}^\{S\}\\exp\(\(q\_\{i\}^\{r\}\)^\{\\top\}k\_\{t\}/\\sqrt\{m\}\)\}\.\(6\)Wq∈ℝm×dW\_\{q\}\\in\\mathbb\{R\}^\{m\\times d\}andbq∈ℝmb\_\{q\}\\in\\mathbb\{R\}^\{m\}are learned;ttis a dummy slot index\. Weightsαi​sr\\alpha\_\{is\}^\{r\}are nonnegative and sum to one overss\. The workspace residualcir∈ℝdc\_\{i\}^\{r\}\\in\\mathbb\{R\}^\{d\}and next hidden state are

cir=γ⁡\(Wo​∑s=1Sαi​sr​Mi​sr\+bo\),hr\+1=Tθ​\(LNb⁡\(ur\+cr\)\)\.c\_\{i\}^\{r\}=\\gamma\\\!\\left\(W\_\{o\}\\sum\_\{s=1\}^\{S\}\\alpha\_\{is\}^\{r\}M\_\{is\}^\{r\}\+b\_\{o\}\\right\),\\qquad h^\{r\+1\}=T\_\{\\theta\}\\\!\\left\(\\operatorname\{LN\}\_\{b\}\(u^\{r\}\+c^\{r\}\)\\right\)\.\(7\)Wo∈ℝd×mW\_\{o\}\\in\\mathbb\{R\}^\{d\\times m\}andbo∈ℝdb\_\{o\}\\in\\mathbb\{R\}^\{d\}are learned;γ\\gammais a learned scalar initialized to0\.50\.5\. Unsubscriptedhr\+1,ur,crh^\{r\+1\},u^\{r\},c^\{r\}denote complete length\-LLsequences\.TθT\_\{\\theta\}is one pre\-normalized causal Transformer block, reused in all four rounds; its full equations appear in Appendix[A](https://arxiv.org/html/2610.22796#A1)\.

After the last round, the separately learned, bias\-free vocabulary matrixWLM∈ℝ257×dW\_\{\\mathrm\{LM\}\}\\in\\mathbb\{R\}^\{257\\times d\}produces logitsoio\_\{i\}:

oi=WLMLNf\(hiR\),Pθ\(yi=a∣x1:i\)=exp⁡\(oi,a\)∑b∈𝒱exp⁡\(oi,b\)\.o\_\{i\}=W\_\{\\mathrm\{LM\}\}\\operatorname\{LN\}\_\{f\}\(h\_\{i\}^\{R\}\),\\qquad P\_\{\\theta\}\(y\_\{i\}=a\\mid x\_\{1:i\}\)=\\frac\{\\exp\(o\_\{i,a\}\)\}\{\\sum\_\{b\\in\\mathcal\{V\}\}\\exp\(o\_\{i,b\}\)\}\.\(8\)aais a candidate target,bba dummy vocabulary index, andx1:ix\_\{1:i\}the consumed prefix\. Embedding and head weights are not tied\. Figure[1](https://arxiv.org/html/2610.22796#S3.F1)summarizes the resulting computation\. The input anchor and current state determine the causal writes and the slot query\. Their retrieved summary enters the shared block as a residual, after which causal self\-attention refines interactions between token states\. Repeating this sequence allows each round to construct its workspace from the representations produced by the preceding round\.

Figure 1:The evaluated architecture\. Blue: input anchor and output head\. Orange: gated causal prefix accumulation\. Green: learned\-key slot read\. Purple: shared refinement and the recurrent route\. Fast workspace activations are recomputed each round in the main model\.
### 3\.4Causality and computational structure

The workspace preserves the autoregressive structure of the shared block\. Initiallyhi0h\_\{i\}^\{0\}depends only onxix\_\{i\}and its position\. Assume thathjrh\_\{j\}^\{r\}depends only onx1:jx\_\{1:j\}\. Anchor fusion and slot reads operate position\-wise, prefix accumulation atiiincludes onlyj≤ij\\leq i, and causal attention masksj\>ij\>i\. Thereforehir\+1h\_\{i\}^\{r\+1\}depends only onx1:ix\_\{1:i\}\. Retention combines states at the same prefix position and preserves this property\. By induction, each output depends only on its consumed prefix\.

The same construction permits all prefix summaries within a round to be computed by causal cumulative sums\. Slot keys and projection matrices are shared learned parameters, whereas the summaries are input\-dependent activations\. As the hidden sequence changes, each round recomputes its writes and reads using those same parameters\. For a batch ofBBsequences, parallel execution materializes a distinctS×mS\\times mworkspace at each of theLLpositions, giving an activation tensor of shapeB×L×S×mB\\times L\\times S\\times m\.

For one sequence, the write projection requires work of orderL​S​m​dLSmd, and prefix accumulation requiresL​S​mLSm\. Query construction, slot addressing, and the output projection add terms of orderL​d​mLdm,L​S​mLSm, andL​m​dLmd, respectively\. With feed\-forward widthff, the shared Transformer block requires work of orderL​d2\+L2​d\+L​d​fLd^\{2\}\+L^\{2\}d\+Ldfper round\. Workspace operations are linear in context length for fixedS,m,dS,m,d, while the retained self\-attention term is quadratic\. Reusing parameters acrossRRrounds thus changes the parameter allocation while executing the block and workspace operationsRRtimes\. Section[6](https://arxiv.org/html/2610.22796#S6)examines the observed training cost alongside prediction quality\.

## 4A byte\-to\-prediction walkthrough

The interaction between causal writes and slot reads can be illustrated with the eight\-byte prefixthe catand next\-byte targets\. The input IDs are\(116,104,101,32,99,97,116,32\)\(116,104,101,32,99,97,116,32\), including the two spaces; the target byte has ID115115\. Figure[2](https://arxiv.org/html/2610.22796#S4.F2)traces this prefix through embedding, accumulation, addressing, recurrent refinement, and the output distribution\. To make the intermediate values explicit, the example uses two slots with two coordinates each and a fully specified set of parameters \(Appendix[D](https://arxiv.org/html/2610.22796#A4)\)\. These parameters define an arithmetic illustration of the architecture\.

1. 1\.Embed and anchor\.Equation[1](https://arxiv.org/html/2610.22796#S3.E1)turns each integer ID into a vector and adds a position vector\. In every round, Equation[2](https://arxiv.org/html/2610.22796#S3.E2)exposes both the unchanged anchor and the current refined state to the fusion network\.
2. 2\.Write from the consumed prefix\.At positioni=8i=8, Equation[4](https://arxiv.org/html/2610.22796#S3.E4)uses all eight inputs\. For a toy withS=m=2S=m=2, let every source have gates\(0\.25,0\.50\)\(0\.25,0\.50\)and slot vectors\(0\.4,0\.2\)\(0\.4,0\.2\)and\(0\.2,−0\.4\)\(0\.2,\-0\.4\)\. Masses areD8,1=2D\_\{8,1\}=2andD8,2=4D\_\{8,2\}=4; numerators are\(0\.8,0\.4\)\(0\.8,0\.4\)and\(0\.8,−1\.6\)\(0\.8,\-1\.6\)\. Dividing by square\-root mass givesM8,1=\(0\.5657,0\.2828\)M\_\{8,1\}=\(0\.5657,0\.2828\)andM8,2=\(0\.4,−0\.8\)M\_\{8,2\}=\(0\.4,\-0\.8\)\. The round index is omitted because these particular toy writes are constant\.
3. 3\.Address slots\.With queryq=\(1,0\)q=\(1,0\)and keysk1=\(1,0\)k\_\{1\}=\(1,0\),k2=\(0,1\)k\_\{2\}=\(0,1\), Equation[6](https://arxiv.org/html/2610.22796#S3.E6)gives scores\(1/2,0\)\(1/\\sqrt\{2\},0\)and weights\(0\.6698,0\.3302\)\(0\.6698,0\.3302\)\. This is attention over prefix summaries, not over individual tokens\.
4. 4\.Inject the read\.The weighted vector is\(0\.5110,−0\.0748\)\(0\.5110,\-0\.0748\)\. With identity output map, zero bias, andγ=0\.5\\gamma=0\.5, the residual is\(0\.2555,−0\.0374\)\(0\.2555,\-0\.0374\)\. The actual model learns the output map and gain, normalizes the sum, and applies full causal attention plus a feed\-forward network\.
5. 5\.Refine four times\.The same transformation is reused, but the real model normally writes different summaries from updated states\. In the simple toy, zero attention/feed\-forward weights isolate workspace arithmetic; position eight has hidden state approximately\(1,−1\)\(1,\-1\)after each round\.
6. 6\.Predict and score\.A specified toy head gives logit33tos,11toa,00to space, and−5\-5to each remaining symbol\. The full 257\-way softmax givesP⁡\(s\)=e3/\(e3\+e1\+1\+254​e−5\)=0\.7872P\(\\texttt\{s\}\)=e^\{3\}/\(e^\{3\}\+e^\{1\}\+1\+254e^\{\-5\}\)=0\.7872\. The single\-target loss is−log⁡\(0\.7872\)=0\.2393\-\\log\(0\.7872\)=0\.2393nats\. Training averages these losses over targets and updates slow parameters by backpropagation\. A target is never written before being predicted\.

Figure 2:Numerical walkthrough: byte IDs, writes, prefix memories, slot addressing, residual, shared rounds, and a 257\-class toy head\. Appendix[D](https://arxiv.org/html/2610.22796#A4)fully specifies the toy parameters\.
## 5Experimental protocol

We evaluate whether the causal workspace improves next\-byte prediction under a fixed parameter and training\-target budget\. Five controls separate its contribution from that of shared depth, input anchoring, and recurrent transition design\. Within each model size, all conditions use the same corpus representation, context length, optimizer recipe, and number of scored training targets\. The corpus and training specifications are given in Tables[1](https://arxiv.org/html/2610.22796#S5.T1)and[2](https://arxiv.org/html/2610.22796#S5.T2)\.

### 5\.1Five controls and parameter matching

The six rows use one training/evaluation harness:

1. 1\.Plain Transformer:four independently parameterized causal blocks and a final normalized vocabulary head\.
2. 2\.Fixed shared loop:one block applied four times,hr\+1=Tθ​\(hr\)h^\{r\+1\}=T\_\{\\theta\}\(h^\{r\}\), without anchoring or workspace\.
3. 3\.Adaptive\-exit mixture:four shared passes and a per\-position mixture of their exits\. An affine router oneie\_\{i\}supplies four softmax weights; every exit is computed\.
4. 4\.Anchored shared loop:Equation[2](https://arxiv.org/html/2610.22796#S3.E2)followed directly by the shared block, without a workspace residual\.
5. 5\.Anchored conditioned\-dense:anchoring, learned per\-round element\-wise scale/shift before the block, and a learned readout from earlier exits\. Additions initialize at zero, preserving the anchored control’s initial transition/readout\.
6. 6\.Causal workspace:Equations[1](https://arxiv.org/html/2610.22796#S3.E1)–[8](https://arxiv.org/html/2610.22796#S3.E8), recomputing each round, without retention between rounds\.

Hidden widths and feed\-forward dimensions are chosen to match approximately 918K or 101\.71M parameters, as detailed in Appendix[B](https://arxiv.org/html/2610.22796#A2)\. Every model executes four Transformer\-block applications\. Relative to the workspace model, the maximum parameter difference is 474 at the small budget and 6,058 at the larger budget\. Independent and shared\-depth architectures allocate these parameters differently: the larger plain Transformer has width 1,536, whereas the fixed shared loop has width 3,072\. The anchored and workspace models both use width 2,560 and 20 heads; adding the workspace is compensated by a narrower feed\-forward sublayer\. The comparison therefore measures the effectiveness of each complete architecture at a common parameter budget\. The controls are implemented within the same experimental framework, with their equations and dimensions specified in the appendices\.

### 5\.2Corpora and preprocessing

The evaluation spans three complementary text distributions\. WikiText\-103 raw contains Wikipedia articles and surrounding text, enwik8 provides a continuous Wikipedia\-derived byte stream, and TinyStories contains short synthetic narratives\. Their corpus descriptions are given by Merity et al\.\[[26](https://arxiv.org/html/2610.22796#bib.bib26)\], the Hutter Prize specification\[[28](https://arxiv.org/html/2610.22796#bib.bib28)\], and Eldan and Li\[[27](https://arxiv.org/html/2610.22796#bib.bib27)\]\. Table[1](https://arxiv.org/html/2610.22796#S5.T1)specifies the training pools, held\-out sources, and scored counts\. Evaluation examples are excluded from gradient updates\.

Table 1:Data protocol\. Held\-out evaluation scores every complete 513\-symbol block from the designated evaluation stream\. All models retain a 257\-class head\.The source datasets are available from[Salesforce/WikiText](https://huggingface.co/datasets/Salesforce/wikitext),[enwik8](https://mattmahoney.net/dc/enwik8.zip), and[TinyStories](https://huggingface.co/datasets/roneneldan/TinyStories)\. WikiText text rows are encoded as UTF\-8 bytes, with symbol 256 appended after each nonempty row\. enwik8 and TinyStories retain their raw file bytes without inserted separators\. A common 257\-class vocabulary is used for all conditions, keeping the output head identical in form across corpus representations\.

Each stream is partitioned into non\-overlapping blocks of 513 symbols\. A block supplies 512 input positions and their shifted next\-symbol targets; its first symbol provides context only, and any incomplete trailing block is discarded\. Hidden states, workspace summaries, and positions reset at each block boundary\. All architectures consequently use the same within\-block context, including when a boundary falls inside an article or story\. This common segmentation makes the scored training and evaluation targets directly comparable within each corpus\.

### 5\.3Training, seeds, and evaluation

Table[2](https://arxiv.org/html/2610.22796#S5.T2)specifies the two size\-dependent recipes\. Within a size, the architecture changes but the total number of scored training targets does not\.

Table 2:Training recipe within each parameter budget\. Microbatch steps count forward/backward calls; optimizer updates account for accumulation\.All models are trained with AdamW\[[32](https://arxiv.org/html/2610.22796#bib.bib32),[33](https://arxiv.org/html/2610.22796#bib.bib33)\]using the size\-specific settings in Table[2](https://arxiv.org/html/2610.22796#S5.T2)\. Each run performs 20,000 optimizer updates and processes 81\.92M scored targets\. The small models use eight sequences per update in FP32\. The larger models use microbatches of four sequences, accumulate two microsteps per update, and use BF16 autocast\. Their learning rate increases linearly to10−410^\{\-4\}over the first 2,000 updates and then remains constant\. Each run uses one A100 GPU and eight CPU cores\.

Training blocks are shuffled once and sampled without replacement until the fixed target budget is reached\. The reported endpoint is the saved final checkpoint after this schedule\. The full matrix comprises six architectures, three corpora, two sizes, and seeds 11, 23, and 37, for a total of 108 training runs; the retention study adds nine runs\. Exploratory held\-out results informed architecture development before the complete matrix was evaluated\.

Initialization and shuffled loading use a shared PyTorch random stream in the reported experiments\. Since architectures have different parameter shapes, identical seed labels identify repeated runs but do not impose identical minibatch orders across models\. All runs nevertheless process the same number of targets under the same size\-specific schedule\. Appendix[E](https://arxiv.org/html/2610.22796#A5)specifies this sampling protocol and the independent\-sampler option provided for subsequent experiments\.

The principal metric is negative log\-likelihood over the full held\-out evaluation stream and its exponential, perplexity\. We denote this loss byℒeval\\mathcal\{L\}\_\{\\mathrm\{eval\}\}\. ForKKscored held\-out targets, negative log\-likelihood \(NLL\) is

ℒeval=−1K∑n=1KlogPθ\(yn∣xprefix⁡\(n\)\),PPL=exp\(ℒeval\)\.\\mathcal\{L\}\_\{\\mathrm\{eval\}\}=\-\\frac\{1\}\{K\}\\sum\_\{n=1\}^\{K\}\\log P\_\{\\theta\}\(y\_\{n\}\\mid x\_\{\\mathrm\{prefix\}\(n\)\}\),\\qquad\\mathrm\{PPL\}=\\exp\(\\mathcal\{L\}\_\{\\mathrm\{eval\}\}\)\.\(9\)nnindexes all scored targets in complete blocks;yny\_\{n\}is the true symbol andxprefix⁡\(n\)x\_\{\\mathrm\{prefix\}\(n\)\}its within\-block input prefix\. Summed loss is divided by total targets, including the last short batch\. On raw enwik8/TinyStories, bits per scored byte equalℒeval/log⁡2\\mathcal\{L\}\_\{\\mathrm\{eval\}\}/\\log 2\. WikiText also scores artificial separators, so its metric is per symbol, not strictly per byte\. The resulting perplexity is measured per scored byte or separator symbol, rather than per word token\. All results in the following section use the complete held\-out stream at the final checkpoint\. The mean and sample standard deviation summarize the three run\-level scores for each condition; Appendix[C](https://arxiv.org/html/2610.22796#A3)reports the individual runs\.

## 6Results

### 6\.1Held\-out performance across corpora and sizes

Tables[3](https://arxiv.org/html/2610.22796#S6.T3)and[4](https://arxiv.org/html/2610.22796#S6.T4)report final\-checkpoint perplexity for every architecture, corpus, size and seed, summarized as a mean and sample standard deviation\. The causal workspace has the best mean in five of the six corpus–size comparisons\. This pattern is especially clear at 101\.71M parameters: it leads all three corpora against controls that use the same number of shared block applications and nearly the same parameter count\.

Table 3:Small\-budget held\-out PPL, mean±\\pmsample standard deviation across seeds 11, 23, 37\. Lower is better; bold marks the lowest mean per corpus\. Each run processes 81\.92M scored targets\.Table 4:Larger\-budget \(101\.71M\) held\-out PPL at the same scored\-target budget\. Means, variation and bolding follow Table[3](https://arxiv.org/html/2610.22796#S6.T3)\. These are final\-checkpoint scores\.At 0\.918M parameters \(Table[3](https://arxiv.org/html/2610.22796#S6.T3)\), the workspace reaches 3\.7455 PPL on WikiText\-103 and 2\.1732 on TinyStories\. The strongest alternative on both corpora is the anchored conditioned\-dense loop, at 3\.9481 and 2\.2225; the workspace reduces perplexity by 5\.13% and 2\.22%, respectively\. The WikiText advantage appears for all three seed labels, and the TinyStories advantage for two\. On small enwik8, the conditioned\-dense loop instead leads with 3\.8431 versus the workspace’s 3\.8968\. The workspace’s 0\.2019 standard deviation there is also larger than in its other small\-model cells\. Thus the proposed structure is the leading small\-model design on two distinct corpora, but not every byte distribution\.

At 101\.71M parameters \(Table[4](https://arxiv.org/html/2610.22796#S6.T4)\), the workspace achieves 2\.7820, 2\.7843 and 1\.6823 PPL on WikiText\-103, enwik8 and TinyStories\. The strongest competing means are 2\.8118 and 2\.8049 from anchored looping on the first two corpora, and 1\.6996 from conditioned\-dense looping on TinyStories\. Consequently, the workspace improves upon the best non\-workspace alternative by 1\.06%, 0\.73% and 1\.01%\. It leads for all three seed labels on each corpus, rather than owing its ranking to one favorable run\. On enwik8 its mean NLL corresponds to 1\.4773 bits per byte\.

The pattern across corpora and sizes is summarized in Figure[3](https://arxiv.org/html/2610.22796#S6.F3), which displays the endpoint NLLs and their variation across runs\. Each row corresponds to a corpus and each column to a parameter budget\. The larger\-model panels show the workspace advantage on all three corpora, while the small\-enwik8 panel shows the stronger conditioned\-dense mean and the larger spread of the workspace runs\. The plotted NLLs and the perplexities in Tables[3](https://arxiv.org/html/2610.22796#S6.T3)–[4](https://arxiv.org/html/2610.22796#S6.T4)express the same endpoint scores on different scales\.

Figure 3:Complete\-grid held\-out NLL\. Large markers/bars show three\-run means/sample standard deviations; small offset markers show seeds\. Rows are corpora, columns are parameter budgets\. Each point represents the final checkpoint evaluated over the full held\-out stream\.Against the plain Transformer, the workspace’s mean perplexity is lower by 16\.17–18\.37% at the small budget and 6\.71–8\.19% at the larger budget\. This comparison establishes the benefit of the complete architecture\. The shared\-loop, anchored, and conditioned\-dense rows then locate the gain more precisely: weight sharing and anchoring both improve performance, but neither consistently reaches the workspace result\. In particular, adding the conditioned\-dense transition to the anchored loop helps small models and larger TinyStories while slightly worsening larger WikiText\-103 and enwik8\. Merely adding more expressive recurrence is therefore not sufficient to explain the workspace’s three\-corpus advantage at 101\.71M\.

### 6\.2Architectural and retention ablations

Table[5](https://arxiv.org/html/2610.22796#S6.T5)isolates two successive design decisions at 101\.71M parameters\. The anchored\-loop row*removes the causal workspace path*: it has no gated slot writes, learned\-key reads, or workspace residual, but keeps the input anchor and four applications of the shared causal block\. To retain the parameter budget, its feed\-forward sublayer is wider\. Restoring the workspace produces the second row, reducing mean PPL from 2\.8118 to 2\.7820 on WikiText\-103, 2\.8049 to 2\.7843 on enwik8, and 1\.7040 to 1\.6823 on TinyStories\. Thus the workspace improves prediction on every larger\-budget corpus even when the anchored alternative reallocates its parameters to a wider feed\-forward sublayer\. The matched\-budget comparison supports the complete causal write/read pathway as an effective architectural addition\.

Table 5:Workspace\-path and across\-round retention ablations at 101\.71M\. Anchored and recomputed\-workspace rows are from the grid; retention adds nine runs at the same training budget\. Held\-out PPL, mean±\\pmsample standard deviation\.The third row changes only*how slot state persists between rounds*\. The main workspace model recomputes each prefix summary from the current round’s token states\. The retention variant replaces it with the learned convex blend in Equation[5](https://arxiv.org/html/2610.22796#S3.E5), while preserving the same shared block, write/read pathway and nominal parameter count\. Its mean PPLs are 2\.7794, 2\.7863 and 1\.6827\. Both workspace variants outperform the anchored loop on all three corpora, supporting the value of the causal summary pathway under either update rule\. The choice of update rule has a much smaller effect: relative to recomputation, retention changes mean PPL by−0\.0931%\-0\.0931\\%on WikiText\-103,\+0\.0722%\+0\.0722\\%on enwik8, and\+0\.0216%\+0\.0216\\%on TinyStories\. Figure[5](https://arxiv.org/html/2610.22796#S6.F5)A resolves these changes into individual seed\-labelled NLL differences\. The main model uses recomputation, which achieves the lower mean on two corpora and refreshes each prefix summary using the current token states\. The architectural comparison thus supports the added causal summary pathway, while the retention comparison determines how that pathway is updated across depth\. The roles of gating, learned keys, and the read projection can be examined separately through further component ablations\.

### 6\.3Training behavior and computational cost

The endpoint results can be considered alongside the course of optimization\. Figure[4](https://arxiv.org/html/2610.22796#S6.F4)shows smoothed training NLL for four representative 101\.71M architectures over the common 81\.92M\-target schedule\. The shared\-depth models generally reach lower training loss than the plain Transformer, and the workspace follows the favorable recurrent\-model trajectory\. These curves describe the periodically sampled training minibatches, complementing the held\-out endpoint comparisons in Table[4](https://arxiv.org/html/2610.22796#S6.T4)\.

Figure 4:101\.71M training NLL for four representative rows\. Each curve is a five\-observation moving average of periodically logged microbatch loss, then averaged over seeds; shading shows sample standard deviation of smoothed seed curves\. Each run contributes 101 periodically logged training observations\.Figure 5:A: retention minus recomputed\-workspace held\-out NLL; negative favors retention\. Points are seed\-labelled differences; black lines mark means\. B: observed training\-stage minutes versus held\-out PPL\. Color identifies architecture; circle/square/triangle denote WikiText/enwik8/TinyStories\. Training\-stage timing includes checkpoint I/O and a limited held\-out pass; it excludes setup and the complete held\-out evaluation\.The additional projections and prefix accumulations also affect training cost\. Figure[5](https://arxiv.org/html/2610.22796#S6.F5)B relates the observed training\-stage time to held\-out perplexity for the recorded 101\.71M runs\. The workspace reaches the lowest mean perplexity on each corpus, while its training\-stage records reflect the work of the added write/read pathway\. Timings include checkpoint I/O and the limited held\-out pass performed by the trainer, so they describe the observed runs rather than a controlled kernel\-throughput measurement\. The resulting quality–cost comparison characterizes the additional computation associated with improved prediction at the matched budgets\.

### 6\.4Contribution of recurrent depth and the workspace

The anchored shared loop provides the most direct architectural comparison because it preserves the input anchor and executes the same four shared rounds\. Adding the workspace lowers its mean perplexity in all six corpus–size settings\. At 0\.918M parameters, the reductions on WikiText\-103, enwik8, and TinyStories are 6\.41%, 0\.40%, and 2\.70%; at 101\.71M they are 1\.06%, 0\.73%, and 1\.27%, respectively\. Thus the workspace improves upon anchoring even in the small\-enwik8 setting where the conditioned\-dense architecture has the best overall mean\. This comparison identifies a consistent benefit from allocating part of the matched budget to the causal summary pathway\.

The other controls establish how this benefit relates to recurrent depth\. Applying a shared block four times improves upon the plain Transformer in all six settings, and adding the input anchor improves upon the fixed shared loop on every corpus at both sizes\. The workspace builds on these improvements by introducing content\-addressed prefix summaries\. The conditioned\-dense control instead modulates the recurrent transition and combines earlier exits\. Its ranking relative to anchoring varies across corpora, whereas the workspace improves upon the anchored loop throughout the grid and leads the complete larger\-model comparison\.

The repeated larger\-model ranking across WikiText\-103, enwik8, and TinyStories indicates that the benefit extends across the three tested text distributions\. Absolute perplexities reflect each corpus’s byte distribution and are therefore interpreted within corpus\. Together with the retention results, these comparisons support the architectural combination of anchored recurrence and a refreshed causal workspace: the anchor preserves the original input, and the workspace supplies a distinct summary of information accumulated over the available prefix\.

## 7Conclusion and future work

We introduced a causal latent workspace Transformer that combines anchored shared\-depth refinement with gated causal writes and learned\-key reads\. The workspace makes compact, prefix\-specific summaries available before each shared block application, providing an additional information pathway alongside token\-level self\-attention\. This design improves upon the anchored loop in all six tested corpus–size settings and achieves the lowest mean held\-out perplexity among six architectures in five\.

At 101\.71M parameters, the workspace leads on WikiText\-103 raw, enwik8, and TinyStories, with all corresponding seed\-labelled comparisons favoring it over the strongest non\-workspace control\. At 0\.918M parameters it leads on WikiText\-103 and TinyStories; conditioned\-dense looping is strongest on enwik8\. The retention study further shows that both tested workspace update rules improve upon the larger anchored control on every corpus, with direct recomputation giving the lower mean on two of the three\. These findings support a causal write/read pathway as an effective structural addition to looped Transformers at the evaluated parameter and training\-target budgets\.

The next architectural question concerns the individual contributions of gating, addressing, and residual injection\. Experiments that hold hidden and feed\-forward widths fixed while varying these operations can clarify how they contribute to the complete pathway\. Longer training schedules can assess the persistence of the observed advantage, while longer contexts and subword language modeling can test the mechanism under different sequence representations\. Incremental decoding and measurements of activation memory and compute\-matched quality will further characterize the trade\-offs of workspace\-assisted refinement\.

## Code, data and reproducibility

The implementation, experiment configurations, per\-seed results, and sources for the figures and tables are available at

[https://github\.com/izaznov/causal\-latent\-workspace\-transformer](https://github.com/izaznov/causal-latent-workspace-transformer)

The repository includes the six benchmark architectures, the retention variant, and the offline training and evaluation procedures\. Dataset sources are linked in Section[5](https://arxiv.org/html/2610.22796#S5)\. Appendix[E](https://arxiv.org/html/2610.22796#A5)specifies the implementation mapping, environment requirements, and commands for reproducing the study\.

## References

- \[1\]A\. Vaswani et al\. Attention Is All You Need\. NeurIPS, 2017\.[https://arxiv\.org/abs/1706\.03762](https://arxiv.org/abs/1706.03762)\.
- \[2\]M\. Dehghani, S\. Gouws, O\. Vinyals, J\. Uszkoreit, and L\. Kaiser\. Universal Transformers\. ICLR, 2019\.[https://arxiv\.org/abs/1807\.03819](https://arxiv.org/abs/1807.03819)\.
- \[3\]A\. Graves\. Adaptive Computation Time for Recurrent Neural Networks\. 2016\.[https://arxiv\.org/abs/1603\.08983](https://arxiv.org/abs/1603.08983)\.
- \[4\]D\. Hutchins, I\. Schlag, Y\. Wu, E\. Dyer, and B\. Neyshabur\. Block\-Recurrent Transformers\. NeurIPS, 2022\.[https://arxiv\.org/abs/2203\.07852](https://arxiv.org/abs/2203.07852)\.
- \[5\]J\. Geiping et al\. Scaling up Test\-Time Compute with Latent Reasoning: A Recurrent Depth Approach\. 2025\.[https://arxiv\.org/abs/2502\.05171](https://arxiv.org/abs/2502.05171)\.
- \[6\]S\. Bae et al\. Mixture\-of\-Recursions: Learning Dynamic Recursive Depths for Adaptive Token\-Level Computation\. 2025\.[https://arxiv\.org/abs/2507\.10524](https://arxiv.org/abs/2507.10524)\.
- \[7\]Z\. Dai, Z\. Yang, Y\. Yang, J\. Carbonell, Q\. V\. Le, and R\. Salakhutdinov\. Transformer\-XL: Attentive Language Models Beyond a Fixed\-Length Context\. ACL, 2019\.[https://arxiv\.org/abs/1901\.02860](https://arxiv.org/abs/1901.02860)\.
- \[8\]J\. W\. Rae, A\. Potapenko, S\. M\. Jayakumar, and T\. P\. Lillicrap\. Compressive Transformers for Long\-Range Sequence Modelling\. ICLR, 2020\.[https://arxiv\.org/abs/1911\.05507](https://arxiv.org/abs/1911.05507)\.
- \[9\]Y\. Wu, M\. N\. Rabe, D\. Hutchins, and C\. Szegedy\. Memorizing Transformers\. ICLR, 2022\.[https://arxiv\.org/abs/2203\.08913](https://arxiv.org/abs/2203.08913)\.
- \[10\]A\. Graves, G\. Wayne, and I\. Danihelka\. Neural Turing Machines\. 2014\.[https://arxiv\.org/abs/1410\.5401](https://arxiv.org/abs/1410.5401)\.
- \[11\]A\. Graves et al\. Hybrid computing using a neural network with dynamic external memory\. Nature 538:471–476, 2016\.[https://doi\.org/10\.1038/nature20101](https://doi.org/10.1038/nature20101)\.
- \[12\]A\. Goyal et al\. Coordination Among Neural Modules Through a Shared Global Workspace\. 2021\.[https://arxiv\.org/abs/2103\.01197](https://arxiv.org/abs/2103.01197)\.
- \[13\]F\. Locatello et al\. Object\-Centric Learning with Slot Attention\. NeurIPS, 2020\.[https://arxiv\.org/abs/2006\.15055](https://arxiv.org/abs/2006.15055)\.
- \[14\]A\. Jaegle et al\. Perceiver: General Perception with Iterative Attention\. ICML, 2021\.[https://arxiv\.org/abs/2103\.03206](https://arxiv.org/abs/2103.03206)\.
- \[15\]J\. Ba, G\. Hinton, V\. Mnih, J\. Z\. Leibo, and C\. Ionescu\. Using Fast Weights to Attend to the Recent Past\. NeurIPS, 2016\.[https://arxiv\.org/abs/1610\.06258](https://arxiv.org/abs/1610.06258)\.
- \[16\]I\. Schlag, K\. Irie, and J\. Schmidhuber\. Linear Transformers Are Secretly Fast Weight Programmers\. ICML, 2021\.[https://arxiv\.org/abs/2102\.11174](https://arxiv.org/abs/2102.11174)\.
- \[17\]A\. Katharopoulos, A\. Vyas, N\. Pappas, and F\. Fleuret\. Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention\. ICML, 2020\.[https://arxiv\.org/abs/2006\.16236](https://arxiv.org/abs/2006.16236)\.
- \[18\]K\. Choromanski et al\. Rethinking Attention with Performers\. ICLR, 2021\.[https://arxiv\.org/abs/2009\.14794](https://arxiv.org/abs/2009.14794)\.
- \[19\]I\. Beltagy, M\. E\. Peters, and A\. Cohan\. Longformer: The Long\-Document Transformer\. 2020\.[https://arxiv\.org/abs/2004\.05150](https://arxiv.org/abs/2004.05150)\.
- \[20\]M\. Zaheer et al\. Big Bird: Transformers for Longer Sequences\. NeurIPS, 2020\.[https://arxiv\.org/abs/2007\.14062](https://arxiv.org/abs/2007.14062)\.
- \[21\]N\. Kitaev, L\. Kaiser, and A\. Levskaya\. Reformer: The Efficient Transformer\. ICLR, 2020\.[https://arxiv\.org/abs/2001\.04451](https://arxiv.org/abs/2001.04451)\.
- \[22\]A\. Gu and T\. Dao\. Mamba: Linear\-Time Sequence Modeling with Selective State Spaces\. 2023\.[https://arxiv\.org/abs/2312\.00752](https://arxiv.org/abs/2312.00752)\.
- \[23\]Y\. Sun et al\. Retentive Network: A Successor to Transformer for Large Language Models\. 2023\.[https://arxiv\.org/abs/2307\.08621](https://arxiv.org/abs/2307.08621)\.
- \[24\]M\. Poli et al\. Hyena Hierarchy: Towards Larger Convolutional Language Models\. ICML, 2023\.[https://arxiv\.org/abs/2302\.10866](https://arxiv.org/abs/2302.10866)\.
- \[25\]T\. Dao, D\. Fu, S\. Ermon, A\. Rudra, and C\. Ré\. FlashAttention: Fast and Memory\-Efficient Exact Attention with IO\-Awareness\. NeurIPS, 2022\.[https://arxiv\.org/abs/2205\.14135](https://arxiv.org/abs/2205.14135)\.
- \[26\]S\. Merity, C\. Xiong, J\. Bradbury, and R\. Socher\. Pointer Sentinel Mixture Models\. ICLR, 2017\.[https://arxiv\.org/abs/1609\.07843](https://arxiv.org/abs/1609.07843)\.
- \[27\]R\. Eldan and Y\. Li\. TinyStories: How Small Can Language Models Be and Still Speak Coherent English? 2023\.[https://arxiv\.org/abs/2305\.07759](https://arxiv.org/abs/2305.07759)\.
- \[28\]M\. Hutter\. The Hutter Prize: 500,000 Euro Prize for Compressing Human Knowledge\. Project and corpus specification\.[https://prize\.hutter1\.net/](https://prize.hutter1.net/)\.
- \[29\]J\. Dodge, S\. Gururangan, D\. Card, R\. Schwartz, and N\. A\. Smith\. Show Your Work: Improved Reporting of Experimental Results\. EMNLP\-IJCNLP, 2019\.[https://arxiv\.org/abs/1909\.03004](https://arxiv.org/abs/1909.03004)\.
- \[30\]D\. Hendrycks and K\. Gimpel\. Gaussian Error Linear Units \(GELUs\)\. 2016\.[https://arxiv\.org/abs/1606\.08415](https://arxiv.org/abs/1606.08415)\.
- \[31\]J\. L\. Ba, J\. R\. Kiros, and G\. E\. Hinton\. Layer Normalization\. 2016\.[https://arxiv\.org/abs/1607\.06450](https://arxiv.org/abs/1607.06450)\.
- \[32\]D\. P\. Kingma and J\. Ba\. Adam: A Method for Stochastic Optimization\. ICLR, 2015\.[https://arxiv\.org/abs/1412\.6980](https://arxiv.org/abs/1412.6980)\.
- \[33\]I\. Loshchilov and F\. Hutter\. Decoupled Weight Decay Regularization\. ICLR, 2019\.[https://arxiv\.org/abs/1711\.05101](https://arxiv.org/abs/1711.05101)\.

## Appendix ASymbol definitions and shared\-block equations

Table 6:Main notation\. Learnable weights, biases and LayerNorm parameters are shared over rounds unless round\-indexed\.Letti∈ℝdt\_\{i\}\\in\\mathbb\{R\}^\{d\}denote one block’s input, and letHHbe the head count with widthdh=d/Hd\_\{h\}=d/H\. For heada=1,…,Ha=1,\\ldots,H, normalized inputt¯i=LN1⁡\(ti\)\\bar\{t\}\_\{i\}=\\operatorname\{LN\}\_\{1\}\(t\_\{i\}\)producesQia,Kia,Via∈ℝdhQ\_\{i\}^\{a\},K\_\{i\}^\{a\},V\_\{i\}^\{a\}\\in\\mathbb\{R\}^\{d\_\{h\}\}by learned affine projections\. CapitalKiaK\_\{i\}^\{a\}is an attention key, distinct from the held\-out target countKKand workspace keyksk\_\{s\}\. Attention probabilities and head outputs are

Ai​ja=\{exp⁡\(\(Qia\)⊤​Kja/dh\)∑k=1iexp⁡\(\(Qia\)⊤​Kka/dh\),j≤i,0,j\>i,bia=∑j=1iAi​ja​Vja\.A\_\{ij\}^\{a\}=\\begin\{cases\}\\displaystyle\\frac\{\\exp\(\(Q\_\{i\}^\{a\}\)^\{\\top\}K\_\{j\}^\{a\}/\\sqrt\{d\_\{h\}\}\)\}\{\\sum\_\{k=1\}^\{i\}\\exp\(\(Q\_\{i\}^\{a\}\)^\{\\top\}K\_\{k\}^\{a\}/\\sqrt\{d\_\{h\}\}\)\},&j\\leq i,\\\\ 0,&j\>i,\\end\{cases\}\\qquad b\_\{i\}^\{a\}=\\sum\_\{j=1\}^\{i\}A\_\{ij\}^\{a\}V\_\{j\}^\{a\}\.\(10\)The dummykkis a source position\.Ai​jaA\_\{ij\}^\{a\}is a probability;biab\_\{i\}^\{a\}is a head output, not an affine bias\. With learned output matrixWA∈ℝd×dW\_\{A\}\\in\\mathbb\{R\}^\{d\\times d\}and biasbA∈ℝdb\_\{A\}\\in\\mathbb\{R\}^\{d\},

yiblk\\displaystyle y\_\{i\}^\{\\mathrm\{blk\}\}=ti\+WA​\[bi1;…;biH\]\+bA,\\displaystyle=t\_\{i\}\+W\_\{A\}\[b\_\{i\}^\{1\};\\ldots;b\_\{i\}^\{H\}\]\+b\_\{A\},\(11\)Tθ​\(t\)i\\displaystyle T\_\{\\theta\}\(t\)\_\{i\}=yiblk\+W2​GELU​\(W1​LN2⁡\(yiblk\)\+b1\)\+b2\.\\displaystyle=y\_\{i\}^\{\\mathrm\{blk\}\}\+W\_\{2\}\\,\\mathrm\{GELU\}\(W\_\{1\}\\operatorname\{LN\}\_\{2\}\(y\_\{i\}^\{\\mathrm\{blk\}\}\)\+b\_\{1\}\)\+b\_\{2\}\.\(12\)yiblky\_\{i\}^\{\\mathrm\{blk\}\}is a block intermediate, not targetyiy\_\{i\}\.W1∈ℝf×dW\_\{1\}\\in\\mathbb\{R\}^\{f\\times d\},b1∈ℝfb\_\{1\}\\in\\mathbb\{R\}^\{f\},W2∈ℝd×fW\_\{2\}\\in\\mathbb\{R\}^\{d\\times f\}andb2∈ℝdb\_\{2\}\\in\\mathbb\{R\}^\{d\}are learned\. All dropout is zero\. Each LayerNorm on vectorz∈ℝdz\\in\\mathbb\{R\}^\{d\}uses its own learned gainλ∈ℝd\\lambda\\in\\mathbb\{R\}^\{d\}and offsetδ∈ℝd\\delta\\in\\mathbb\{R\}^\{d\}:

LN⁡\(z\)=λ⊙z−μ⁡\(z\)​𝟏ν⁡\(z\)\+η\+δ,μ⁡\(z\)=1d​∑a=1dza,ν⁡\(z\)=1d​∑a=1d\(za−μ⁡\(z\)\)2,η=10−5\.\\operatorname\{LN\}\(z\)=\\lambda\\odot\\frac\{z\-\\mu\(z\)\\mathbf\{1\}\}\{\\sqrt\{\\nu\(z\)\+\\eta\}\}\+\\delta,\\quad\\mu\(z\)=\\frac\{1\}\{d\}\\sum\_\{a=1\}^\{d\}z\_\{a\},\\quad\\nu\(z\)=\\frac\{1\}\{d\}\\sum\_\{a=1\}^\{d\}\(z\_\{a\}\-\\mu\(z\)\)^\{2\},\\quad\\eta=10^\{\-5\}\.\(13\)⊙\\odotdenotes coordinate\-wise multiplication,𝟏\\mathbf\{1\}the all\-ones vector,μ\\muthe coordinate mean,ν\\nupopulation coordinate variance andη\\etaLayerNorm epsilon\. Hereaaindexes coordinates, not heads\. Gain and offset initialize to one and zero\.

## Appendix BExact architecture dimensions

Table 7:Exact dimensions and nominal trainable counts\. FFN width isff\. Dash denotes no workspace\. Untied blocks / rounds are 4 / 1 for plain and 1 / 4 for recurrent rows\.The exit mixture usesπi=softmax⁡\(Wexit​ei\+bexit\)\\pi\_\{i\}=\\operatorname\{softmax\}\(W\_\{\\mathrm\{exit\}\}e\_\{i\}\+b\_\{\\mathrm\{exit\}\}\)and final state∑r=0R−1πi,r​hir\+1\\sum\_\{r=0\}^\{R\-1\}\\pi\_\{i,r\}h\_\{i\}^\{r\+1\}\.Wexit∈ℝR×dW\_\{\\mathrm\{exit\}\}\\in\\mathbb\{R\}^\{R\\times d\},bexit∈ℝRb\_\{\\mathrm\{exit\}\}\\in\\mathbb\{R\}^\{R\}are learned;πi,r\\pi\_\{i,r\}weights exitr\+1r\+1\. The conditioned\-dense control sends\(1\+σr\)⊙uir\+τr\(1\+\\sigma\_\{r\}\)\\odot u\_\{i\}^\{r\}\+\\tau\_\{r\}to the block, with learnedσr,τr∈ℝd\\sigma\_\{r\},\\tau\_\{r\}\\in\\mathbb\{R\}^\{d\}initialized to zero\. Its readout is

h^i=hiR\+∑r=0R−2tanh⁡\(βr\)​\(hir\+1−hiR\)\.\\widehat\{h\}\_\{i\}=h\_\{i\}^\{R\}\+\\sum\_\{r=0\}^\{R\-2\}\\tanh\(\\beta\_\{r\}\)\(h\_\{i\}^\{r\+1\}\-h\_\{i\}^\{R\}\)\.\(14\)βr\\beta\_\{r\}are zero\-initialized learned scalars\.h^i\\widehat\{h\}\_\{i\}replaceshiRh\_\{i\}^\{R\}before the final normalization/head\. These modifications change recurrent transition and readout, not the number of independently parameterized blocks\.

Parameter counts include position embeddings and the vocabulary head\. Both workspace variants allocateRRscalar retention logits to keep their nominal counts identical\. The recomputed variant does not use these logits; in the retained variant,ℓ0\\ell\_\{0\}is unused because the first round has no preceding workspace\. Projection matrices use the standard PyTorch initializers, with the slot\-key, gain, retention, and conditioned\-dense initializations specified above\.

## Appendix CComplete per\-seed scores

Table 8:All small\-budget full held\-out PPLs\. The mean is the arithmetic mean of run\-level PPLs, not exponentiated mean NLL\.Table 9:All larger\-budget full held\-out PPLs\. Matching seed labels do not establish identical minibatch order across rows\.Table 10:Retention PPL by seed\. Last column: change in mean PPL from recomputed workspace; negative is better\.Full\-precision NLL, PPL, raw\-byte bits\-per\-byte values, parameters, evaluation\-block counts, hashes, training times and run IDs for all 117 evaluations are inpaper/manuscript/results\.csv\. Source records and 108 logged trajectories are inpaper/manuscript/evidence\.json\.

## Appendix DFully specified numerical toy

The following parameter choices specify the numerical example in Section[4](https://arxiv.org/html/2610.22796#S4)\. Setd=S=m=2d=S=m=2,R=4R=4andE⁡\(xi\)\+pi=\(0\.3,−0\.1\)E\(x\_\{i\}\)\+p\_\{i\}=\(0\.3,\-0\.1\)for all eight positions; use identical embeddings and zero position vectors\. Set allFaF\_\{a\}parameters to zero and LayerNorm gain/offset to one/zero withη=10−5\\eta=10^\{\-5\}\. SetWg=0W\_\{g\}=0,bg=\(log⁡\(1/3\),0\)b\_\{g\}=\(\\log\(1/3\),0\), giving gates\(0\.25,0\.5\)\(0\.25,0\.5\)\. SetWv=0W\_\{v\}=0and

bv=\(atanh⁡\(0\.4\),atanh⁡\(0\.2\),atanh⁡\(0\.2\),atanh⁡\(−0\.4\)\)\.b\_\{v\}=\(\\operatorname\{atanh\}\(0\.4\),\\operatorname\{atanh\}\(0\.2\),\\operatorname\{atanh\}\(0\.2\),\\operatorname\{atanh\}\(\-0\.4\)\)\.atanh\\operatorname\{atanh\}is inverse hyperbolic tangent\. ChooseWq=0W\_\{q\}=0,bq=\(1,0\)b\_\{q\}=\(1,0\),k1=\(1,0\)k\_\{1\}=\(1,0\),k2=\(0,1\)k\_\{2\}=\(0,1\),Wo=I2W\_\{o\}=I\_\{2\},bo=0b\_\{o\}=0, andγ=0\.5\\gamma=0\.5, whereI2I\_\{2\}is the two\-dimensional identity\. All attention/FFN projection weights and biases are zero: their residual sublayers return the input, soTθT\_\{\\theta\}is identity\. Recompute the workspace each round\.

At position eight, the exact slot\-one weight and workspace residual are

a∗=exp⁡\(1/2\)exp⁡\(1/2\)\+1,c∗=12​\(a∗​\(0\.8,0\.4\)/2\+\(1−a∗\)​\(0\.4,−0\.8\)\)\.a\_\{\*\}=\\frac\{\\exp\(1/\\sqrt\{2\}\)\}\{\\exp\(1/\\sqrt\{2\}\)\+1\},\\qquad c\_\{\*\}=\\tfrac\{1\}\{2\}\\left\(a\_\{\*\}\(0\.8,0\.4\)/\\sqrt\{2\}\+\(1\-a\_\{\*\}\)\(0\.4,\-0\.8\)\\right\)\.a∗a\_\{\*\}is a scalar andc∗c\_\{\*\}a two\-coordinate residual\. The recurrence isur=LN⁡\(hr\)u^\{r\}=\\operatorname\{LN\}\(h^\{r\}\),hr\+1=LN⁡\(ur\+c∗\)h^\{r\+1\}=\\operatorname\{LN\}\(u^\{r\}\+c\_\{\*\}\), starting ath0=\(0\.3,−0\.1\)h^\{0\}=\(0\.3,\-0\.1\)\. It givesu0≈\(0\.999875,−0\.999875\)u^\{0\}\\approx\(0\.999875,\-0\.999875\),h1≈\(0\.999996,−0\.999996\)h^\{1\}\\approx\(0\.999996,\-0\.999996\); subsequent rounds stay nearby\. The resulting intermediate values are tabulated for all four rounds in the accompanying numerical summary\.

WriteLNf⁡\(h4\)=\(cf,−cf\)\\operatorname\{LN\}\_\{f\}\(h^\{4\}\)=\(c\_\{f\},\-c\_\{f\}\), with positive scalarcfc\_\{f\}computed using the sameη\\eta\. Choose output rows\(3/\(2cf\),−3/\(2cf\)\)\(3/\(2c\_\{f\}\),\-3/\(2c\_\{f\}\)\)fors,\(1/\(2cf\),−1/\(2cf\)\)\(1/\(2c\_\{f\}\),\-1/\(2c\_\{f\}\)\)fora,\(0,0\)\(0,0\)for space, and\(−5/\(2cf\),5/\(2cf\)\)\(\-5/\(2c\_\{f\}\),5/\(2c\_\{f\}\)\)for each other class\. This construction gives the logits used in the walkthrough through the same linear head and softmax as Equation[8](https://arxiv.org/html/2610.22796#S3.E8)\. Identical embeddings and zero attention weights make the arithmetic explicit\. In the trained models, byte\-dependent embeddings and learned attention supply the contextual representations from which workspace writes and reads are constructed\.

## Appendix EReproduction and provenance

#### Implementation map\.

baselines/workspace\_loop\.pyimplements the workspace;baselines/looped\_transformer\.pyrecurrent controls;baselines/causal\_transformer\.pythe plain model\.data/wikitext\.py,data/enwik8\.pyanddata/raw\_bytes\.pyconstruct blocks\.train\_lm\.pytrains andevaluate\_lm\.pyevaluates all held\-out batches by default\. The matrix script ishpc/gwdg/standard\_transformer\_benchmark\_matrix\.sbatch, with per\-cell scripthpc/gwdg/standard\_transformer\_benchmark\.sbatch\. They record requested resources, offline paths, sequential seeds and final\-checkpoint evaluation\.

#### Configurations and sampling\.

The small and larger recipes are specified byconfigs/standard\_benchmark\_small\.yamlandconfigs/standard\_benchmark\_medium\.yaml, respectively; retention usesconfigs/long\_workspace\_medium\.yaml\. The reported experiments use\-\-sampler\-mode historical, which shares the PyTorch random stream between initialization and shuffled loading\. The release also provides\-\-sampler\-mode independent \-\-data\-seed 11for experiments that hold the data order fixed across architectures\. Reproducing the reported protocol requires the former setting\. Exact bitwise replay additionally depends on the original hardware and package versions; executed revisions were not captured separately for each run\.

Commands below illustrate one larger\-budget WikiText workspace run, from the repository root using downloaded official shards:

```
python train_lm.py --config configs/standard_benchmark_medium.yaml \
  --model anchored_workspace_static --seed 11 \
  --train-parquet DATA/wikitext-103-raw-train-0000.parquet \
    DATA/wikitext-103-raw-train-0001.parquet \
  --heldout-parquet DATA/wikitext-103-raw-heldout.parquet \
  --output-dir runs/replay-wiki-medium-workspace-11

python evaluate_lm.py --config configs/standard_benchmark_medium.yaml \
  --model anchored_workspace_static \
  --checkpoint runs/replay-wiki-medium-workspace-11/checkpoint.pt \
  --heldout-parquet DATA/wikitext-103-raw-heldout.parquet \
  --eval-batches 0 --output runs/replay-wiki-medium-workspace-11/full.json
```

DATAis a placeholder for the local data directory; replace it\. For enwik8, substitute\-\-enwik8\-path DATA/enwik8for Parquet arguments\. TinyStories uses\-\-raw\-byte\-train,\-\-raw\-byte\-heldoutand\-\-max\-train\-tokens 250000000in training, and only\-\-raw\-byte\-heldoutin evaluation\. Repeat seeds 11, 23, 37, substitute architecture IDs, and use the small configuration for the other budget\.

#### Environment and auditing\.

Code requires Python 3\.10\+, PyTorch, NumPy, PyYAML, and PyArrow for Parquet; matplotlib only rebuilds figures\. Historical scripts load Miniforge/GCC and activate a local environment\. The archive records SHA\-256 hashes of each source evaluation output andmetrics\.json, plus matched main\-grid log hashes\. These identify the archived result records; original dataset checksums and complete package\-version locks were not recorded\. Independent replication should record these alongside the sampler state and executed code revision\.

#### Regenerate evidence displays\.

python paper/manuscript/build\_artifacts\.pyverifies the 108 main and nine retention records, seed/cell coverage, target budgets, parameter counts, NLL–PPL consistency and training\-log endpoint matches; it regenerates tables, CSV and vector figures\. The displayed curves use the archived training logs and complete held\-out scores\.tests/test\_paper\_models\.pyverifies model dimensions, causal dependence, and the workspace update equations\.

相似文章

Recurrent Looped Transformer

Hugging Face Daily Papers

论文提出 Recurrent Looped Transformer (RLT),将层划分为并行因果编码器和循环解码器,使计算路径随序列长度增长而每个 token 的成本保持固定,在奇偶校验、S_5 置换追踪和模运算等算法任务上,长度外推能力显著优于标准 Transformer。

什么是 Looped Transformers?清晰解释(8分钟阅读)

TLDR AI

Looped transformers 在多次传递中重复使用相同的层,以参数数量换取计算量,用更少的权重实现更好的推理。文章追溯了这一想法到 Universal Transformer (2018),并解释了其最初因扩展定律和时机而失败的原因。

@ZhihuFrontier: 半年前,一位知乎答主预测下一个Transformer将吸收循环、递归状态、稀疏路由……

X AI KOLs Timeline

一位知乎答主半年前的预测——下一个Transformer将吸收循环、递归状态、稀疏路由和潜在推理——随着Loop Engineering的推进,正变得越来越有现实意义。本文探讨了未来的Transformer架构如何演变为混合模型:将线性复杂度的层用于背景上下文,注意力机制用于精确推理,再加上更细粒度的稀疏性和原生的System 2推理。