Causal Latent Workspaces for Looped Transformers

Ilia Zaznov

cs.CLOct 16, 2026

Causal Latent Workspaces for Looped Transformers

Authors: Ilia Zaznov

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.

Figures & tables

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.

Figure 2: Numerical walkthrough: byte IDs, writes, prefix memories, slot addressing, residual, shared rounds, and a 257-class toy head. Appendix D fully specifies the toy parameters.

Corpus	Training source and representation	Held-out evaluation stream
WikiText-103 raw	Official wikitext-103-raw-v1 training Parquet shards; UTF-8 bytes; blank rows skipped; symbol 256 appended to each nonempty row.	Official held-out split: 2,235 complete blocks, 1,144,320 scored targets.
enwik8	Official 100,000,000-byte file; first 90,000,000 bytes form the training pool. Raw bytes; no inserted separator.	Next 5,000,000 bytes: 9,746 blocks, 4,989,952 scored targets. Final 5,000,000 bytes reserved and not reported.
TinyStories	TinyStoriesV2-GPT4-train.txt ; first 250,000,000 bytes form the training pool. Raw file bytes, without inserted separators or story-wise parsing.	Entire TinyStoriesV2-GPT4-valid.txt : 43,864 blocks, 22,458,368 scored targets.

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.

Setting	Small (0.918M)	Larger (101.71M)
Context / output vocabulary	512 / 257	512 / 257
Scored targets per run	81,920,000	81,920,000
Microbatch / accumulation	8 / 1	4 / 2
Microbatch steps / updates	20,000 / 20,000	40,000 / 20,000
Effective sequences per update	8	8
Learning rate	$3\times 10^{-4}$	$10^{-4}$

Table 2: Training recipe within each parameter budget. Microbatch steps count forward/backward calls; optimizer updates account for accumulation.

Architecture	WikiText-103	enwik8	TinyStories
Plain Transformer	$4.5884\pm 0.0751$	$4.7063\pm 0.0558$	$2.5923\pm 0.0718$
Fixed shared loop	$4.1668\pm 0.1420$	$4.4215\pm 0.2663$	$2.3389\pm 0.0030$
Adaptive-exit mixture	$4.1305\pm 0.0216$	$4.4937\pm 0.2993$	$2.4100\pm 0.0631$
Anchored shared loop	$4.0022\pm 0.0201$	$3.9125\pm 0.0482$	$2.2336\pm 0.0738$
Anchored conditioned-dense	$3.9481\pm 0.0069$	$\mathbf{3.8431\pm 0.0510}$	$2.2225\pm 0.0533$
Causal workspace (ours)	$\mathbf{3.7455\pm 0.0042}$	$3.8968\pm 0.2019$	$\mathbf{2.1732\pm 0.0151}$

Table 3: Small-budget held-out PPL, mean

\pm

sample standard deviation across seeds 11, 23, 37. Lower is better; bold marks the lowest mean per corpus. Each run processes 81.92M scored targets.

Architecture	WikiText-103	enwik8	TinyStories
Plain Transformer	$3.0300\pm 0.0519$	$3.0178\pm 0.0217$	$1.8033\pm 0.0171$
Fixed shared loop	$2.8889\pm 0.0077$	$2.8810\pm 0.0025$	$1.7416\pm 0.0017$
Adaptive-exit mixture	$2.8950\pm 0.0033$	$2.8803\pm 0.0014$	$1.7395\pm 0.0030$
Anchored shared loop	$2.8118\pm 0.0063$	$2.8049\pm 0.0081$	$1.7040\pm 0.0025$
Anchored conditioned-dense	$2.8136\pm 0.0075$	$2.8057\pm 0.0022$	$1.6996\pm 0.0032$
Causal workspace (ours)	$\mathbf{2.7820\pm 0.0032}$	$\mathbf{2.7843\pm 0.0010}$	$\mathbf{1.6823\pm 0.0015}$

Table 4: Larger-budget (101.71M) held-out PPL at the same scored-target budget. Means, variation and bolding follow Table 3 . These are final-checkpoint scores.

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.

Architecture	WikiText-103	enwik8	TinyStories
Anchored shared loop	$2.8118\pm 0.0063$	$2.8049\pm 0.0081$	$1.7040\pm 0.0025$
Causal workspace (ours)	$2.7820\pm 0.0032$	$2.7843\pm 0.0010$	$1.6823\pm 0.0015$
Workspace + loop retention	$2.7794\pm 0.0092$	$2.7863\pm 0.0082$	$1.6827\pm 0.0020$

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

\pm

sample standard deviation.

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.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$\mathcal{V},x_{i},y_{i}$	Vocabulary, input symbol, next-symbol target
$L,R,S,m,d,f,H$	Context, rounds, slots, slot width, state width, FFN width, heads
$B,K$	Microbatch sequences and scored held-out-target count
$e_{i},h_{i}^{r},u_{i}^{r},z_{i}^{r}$	Anchor, recurrent state, fused state, write-normalized state
$g_{is}^{r},v_{is}^{r}$	Scalar write gate and value vector
$N_{is}^{r},D_{is}^{r}$	Gated prefix sum and scalar gate mass

Appendix

Table 6: Main notation. Learnable weights, biases and LayerNorm parameters are shared over rounds unless round-indexed.

Budget	Architecture	$d$	$H$	$f$	$S\times m$	Parameters
Small	Plain	128	4	506	–	918,504
	Fixed loop	256	4	761	–	918,009
	Exit mixture	256	4	759	–	918,011
	Anchored	215	5	860	–	918,050
	Conditioned dense	215	5	856	–	918,049
	Workspace	215	5	758	$4\times 32$	918,030

Appendix

Table 7: Exact dimensions and nominal trainable counts. FFN width is

f

. Dash denotes no workspace. Untied blocks / rounds are 4 / 1 for plain and 1 / 4 for recurrent rows.

Dataset	Architecture	Seed 11	Seed 23	Seed 37	Mean
WikiText-103	Plain	4.62059	4.50256	4.64196	4.58837
	Fixed loop	4.10625	4.06510	4.32895	4.16677
	Exit mixture	4.14598	4.13962	4.10585	4.13048
	Anchored	4.02385	3.99855	3.98409	4.00216
	Conditioned dense	3.95518	3.94761	3.94147	3.94809
	Workspace	3.75042	3.74290	3.74327	3.74553

Appendix

Table 8: All small-budget full held-out PPLs. The mean is the arithmetic mean of run-level PPLs, not exponentiated mean NLL.

Dataset	Architecture	Seed 11	Seed 23	Seed 37	Mean
WikiText-103	Plain	3.08994	2.99986	3.00033	3.03004
	Fixed loop	2.89454	2.89208	2.88015	2.88893
	Exit mixture	2.89524	2.89818	2.89159	2.89500
	Anchored	2.81900	2.80726	2.80927	2.81184
	Conditioned dense	2.81796	2.80494	2.81803	2.81364
	Workspace	2.77861	2.78226	2.78506	2.78198

Appendix

Table 9: All larger-budget full held-out PPLs. Matching seed labels do not establish identical minibatch order across rows.

Dataset	Seed 11	Seed 23	Seed 37	Mean PPL change (%)
WikiText-103	2.76913	2.78694	2.78209	-0.0931
enwik8	2.78836	2.77730	2.79322	+0.0722
TinyStories	1.68210	1.68488	1.68106	+0.0216

Appendix

Table 10: Retention PPL by seed. Last column: change in mean PPL from recomputed workspace; negative is better.

Autoregressive Language Modeling

Autoregressive language modeling predicts sequences one step at a time, using prior tokens to generate text, speech, code, and other data.

See all 131 papers

Sep 1, 2026

Looped Transformers under the Jacobian Lens: Does the Global Workspace Survive Recurrence?

Recent work identifies a mid-depth band of verbalisable, causally potent representations in a standard feedforward transformer --- a functional analogue of a global workspace. Whether the same workspace functionality emerges when depth is implemented through recurrence rather than a stack of distinct layers remains unknown. Looped and depth-recurrent transformers provide a direct test of this question because they reuse the same weights across depth. We extend the Jacobian lens to iterated architectures using a virtual-unrolling adapter. We apply the full workspace suite --- lens fitting, readout, and eleven causal experiment families --- to Ouro-2.6B (48 layers looped 4 times, deeply supervised) and Huginn-0125 (a 4-layer core recurred 16 times, trained for latent reasoning), using Qwen3.6-27B (64 untied layers) as the standard baseline. We find that a workspace forms in the iterated part of each architecture, but that recurrence changes how it can be accessed. Ouro reconstructs workspace content in every loop, and linear transport cannot carry that content across loop boundaries; writes and ablations must therefore span every remaining loop. Huginn carries content forward across all sixteen recurrences, while reads, writes, and ablations act only within a sliding window of roughly two recurrences. Whether newly injected content can be verbalised tracks explicit per-iteration supervision; whether existing content can be steered does not.

Wenlong Wang, Fergal Reid

Fin AI Research

Transformer Interpretability Cognitive Modeling Recurrent Transformers

Ask PDF

Jun 9, 2026

Parallel Causal Associative Fields: Gated Sparse Memory for Long-Context Language Modeling

Transformers achieve strong language modeling performance by providing direct token-to-token communication paths, but causal self-attention scales quadratically with context length. Recurrent and state-space models reduce this cost, yet compress history into sequentially updated fixed-size states. This paper studies a third primitive: a parallel content-addressed memory over causal successor records. The proposed Parallel Causal Associative Field (PCAF) writes local records from a context window into hash buckets, retrieves a bounded candidate set for the current query, forms a sparse cache distribution over successor tokens, and mixes that cache with a parametric local language model through a learned gate. The resulting model maintains sparse long-context access while avoiding a single fixed recurrent state bottleneck. We evaluate PCAF under full autoregressive pretraining on WikiText-103 and PG-19 using a distributed Google Cloud TPU v4-32 pod. At 303M parameters and context length T = 2048, PCAF-semantic reaches 36.31 perplexity on WikiText-103 and 52.45 perplexity on PG-19, compared with 47.49 and 53.84 for a matched dense Transformer. PCAF-semantic simultaneously processes 0.61-0.62M tokens/s across the TPU pod, versus 0.43M tokens/s for dense and local attention baselines. Supporting 41M-parameter multi-seed sweeps and single-GPU component ablations show that the associative cache, retrieval capacity, and learned gate materially affect the speed-quality trade-off.

Muhammad Ahmed

Independent Researcher

Memory-Augmented Language Models Long-Context Language Model Inference Long-Context Language Modeling Autoregressive Language Modeling Associative Memory

Ask PDF

Oct 1, 2026

The Surprising Effectiveness of Shared Memory in Looped Transformers

Looped Transformers apply the same layers several times per token, adding compute to improve quality without more parameters. Each recursion, however, writes its own key-value cache, so memory still grows with compute. Inference-time techniques can shrink this cache at a cost in quality. We pretrain looped language models to share memory: only the first recursion writes a cache, and later recursions read it while keeping a short window of their own. Surprisingly, we find that sharing memory does not cost quality and instead improves it. At 150M-1B parameters, our Looped Prediction Transformer (LPT) and its hybrid variant set a new quality-memory frontier for looped models: with five recursions, the hybrid lowers validation perplexity on FineWeb-Edu by 1.12-1.82 relative to a same-size standard Transformer while using 76-79% less context memory. Through an extensive analysis, we investigate why memory sharing helps. Shared and local memory develop different representations, and later recursions attend mostly to the shared memory, which also acts as a gradient highway to the first recursion.

Giovanni Monea, Keshav Ramji, Yousef El-Kurdi +4

Cornell University · IBM Research

Memory-Efficient Inference Recurrent Transformers

Ask PDF

Abstract

Autoregressive Language Modeling

Explore similar work

Looped Transformers under the Jacobian Lens: Does the Global Workspace Survive Recurrence?

Parallel Causal Associative Fields: Gated Sparse Memory for Long-Context Language Modeling

The Surprising Effectiveness of Shared Memory in Looped Transformers