Causal latent workspaces improve looped transformer language models
Causal Latent Workspaces for Looped Transformers
Computation and Language
Summary
Transformers are powerful AI models that process sequences step-by-step, but reusing the same part repeatedly can limit how information is passed along. The authors introduce a method where each step writes and reads a special summary space that helps keep track of important past information while processing new data. This approach improves the model’s performance on language tasks by keeping useful context without needing more parameters. It works better than traditional methods across several different text datasets.
What this means in practice
- •For natural language processing engineers: Enhance language models by integrating causal latent workspaces to improve prediction accuracy with fewer parameters.
- •For embedded ai developers: Deploy efficient language models that maintain context through looped Transformer blocks with minimal parameter increase.
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.