Physics of Language Models: Part 4.1, Architecture Design and the Magic of Canon Layers

AI summary

Overview

• Research area: Neural architecture design for language models; specifically, controlled empirical study of sequence architectures (Transformers, linear attention, state-space models) using synthetic pretraining tasks.
• Technical level: Intermediate. The core ideas (local convolution, residual connections, attention vs. recurrence) are accessible, while some results reference bits-per-parameter capacity measurement and state-space specifics.
• Scope: The paper introduces a class of lightweight components called Canon layers, which add horizontal information flow across neighboring tokens in any sequence architecture, and evaluates them on five synthetic pretraining tasks plus a 1.3B-parameter, 100B-token real pretraining run.

What This Paper Is About

Comparing language model architectures is hard because real-world pretraining at academic scale is dominated by noise, and benchmark differences of a couple of percent can come from random initialization rather than architecture. The paper builds controlled synthetic pretraining tasks that isolate individual capabilities — reasoning depth, reasoning breadth, knowledge capacity, knowledge manipulation, and hierarchical structure — and uses them to identify a missing architectural ingredient. That ingredient is the Canon layer, a small module that mixes information between adjacent tokens and, when added, substantially changes which architectures look strong.

Key Contributions

1. Five synthetic pretraining tasks ("the playground"). Depo (reasoning depth), Brevo (reasoning breadth), Capo (knowledge capacity), Mano (knowledge manipulation), and Lano (hierarchical language structure), each designed to isolate an atomic skill and each guided by four stated design criteria (not shallow, emphasis on mental thinking without Chain-of-Thought, avoid emphasis on length generalization, relevance to real-world skills).
2. Canon layers as a general architectural primitive. Lightweight horizontal residual links that compute weighted sums of nearby token representations and can be inserted at four points in a block: Canon-A (before attention), Canon-B (inside attention, after Q/K/V projections), Canon-C (before MLP), and Canon-D (inside MLP).
3. Systematic evaluation across architecture families. Canon layers are tested on Transformers (Llama with RoPE and with NoPE), linear attention (GLA), and state-space/recurrent models (Mamba2 and GDN), with ablations over which Canon positions are enabled.
4. A controlled, "apple-to-apples" architecture comparison. Equipping all models with full Canon layers enables fair comparisons between Transformers and linear models, plus a reported academic-scale real-world check at 1.3B parameters and 100B tokens.

Main Findings

• Initial landscape without Canon layers. Linear attention (GLA) performed weakest overall; Mamba2 excelled at knowledge tasks (Capo, Mano) but lagged in reasoning; GDN improved on Mamba2's reasoning and occasionally surpassed Llama(RoPE) on Brevo, but not on other tasks.
• Canon layers boost Transformer reasoning. In the synthetic playground, Canon layers improved reasoning depth (200–400%), reasoning breadth (30%), and knowledge manipulation length (30%), attributed to enhanced hierarchical learning dynamics, with minimal computational overhead.
• Reviving NoPE. Adding Canon layers transformed NoPE models into strong performers, often matching or surpassing RoPE(+Canon); Canon layers outperformed positional fixes such as ALiBi and H-ALiBi, and reducing or removing RoPE improved length generalization.
• Sublayer ablations. Canon layers contributed cumulatively across positions Canon-A/B/C/D, independent of the attention or MLP components; residual Canon improved training efficiency; minimal parameter tuning was required without compromising stability.
• Gated MLP and MoE. Canon layers recovered some knowledge capacity lost in gated MLP or mixture-of-experts architectures, via improved training efficiency and stability.
• Universal boost for linear models. Across GLA, Mamba2, and GDN, Canon layers consistently enhanced in-context reasoning (Depo/Brevo), knowledge (Mano), and structural reasoning (Lano). For GLA, Canon lifted reasoning depth from 1 to 4-hop, doubled reasoning breadth and knowledge length, and surpassed Mamba2.
• Mamba2's conv1d is partial Canon-B. Removing Mamba2's built-in conv1d dropped performance to GLA level, while replacing it with full Canon yielded further improvements. GDN benefited least, as its gating and delta updates already capture part of Canon-like behavior.
• Canon is not tied to attention or SSMs. Canon-ACD alone often matched conv1d/Canon-B, showing horizontal context flow is universal rather than specific to linear-attention or SSM sub-layers.
• Modern refinements may largely replicate Canon-like mixing. Most linear-model performance (for Mamba2/GDN) was achievable with a simple GLA+Canon design.
• Transformers vs. linear models (controlled). Linear models showed roughly 40% higher knowledge capacity, but Transformers reached 2–4× greater reasoning depth and stronger structural reasoning.
• Root cause of shallow linear reasoning. Linear models fell short not from insufficient memory — each layer's recurrent state is described as vastly over-provisioned — but from cumulative compression and retrieval errors, pinpointing memory dynamics as the bottleneck.
• Academic-scale real-world pretraining (Result 12). At 1.3B parameters, 100B tokens, context length 4096, noise was high and resolution limited, making many comparisons statistically unreliable. Still: Canon layers markedly improved NoPE and GLA (raising them to match RoPE and Mamba2/GDN respectively); removing conv1d reduced Mamba2 to GLA level; linear models lagged full Transformers on retrieval-heavy tasks even with Canon; and all models failed 2-hop reasoning even in short 100-token contexts. Reducing or removing RoPE improved long-context generalization when Canon layers were present.
• Token-copying sanity check. On a 500-token copying task, 1-layer RoPE required hidden size d ≥ 128, while 2-layer RoPE or 1-layer RoPE + Canon achieved 100% with d = 16.
• Cost of Canon layers. Fewer than 0.45% additional parameters for GPT2-small; for a 1.3B-parameter Llama with Canon-ABCD, parameters increased by 0.0063%. Runtime overhead on an H100 GPU with a naive implementation (PyTorch bf16, flash attention, causal conv1d kernels) was 12.4%, 14.1%, and 20.8% for forward, backward, and generation respectively; for Canon-AC these reduced to 5.8%, 5.8%, and 7.0%.

Methodology in Plain English

The author starts from the observation that on real pretraining data, architecture differences are buried under noise: at academic scale (1.3B parameters, 100B tokens) small accuracy swings (2–4%) can come from random initialization, and models fail even the simplest 2-hop reasoning. The response is to decompose intelligence into atomic skills and build synthetic data generators for each one, so every experiment has infinite clean data, controlled difficulty, and a single skill under test.

Five task families are built. Depo asks a model to follow k-hop chains through a random directed permutation entirely internally, with parameters controlling permutation size N and depth K (Depo1: N = 225/300/375, K = 8, one-to-two-token names from a 50-token vocabulary, 2048-token contexts; Depo2: N = 75/100/125, K = 16, five-to-seven-token names from a 4-token vocabulary; reduced K = 4 setups used for weaker models). Brevo asks the model to output all vertices reachable from a query in a directed acyclic graph, in topological order from the leaves, with graph size N (Brevo1: N = 70/90/110 in 1024 tokens; Brevo2: N = 30/40/50 in 1536 tokens, degree at most 4). Capo measures knowledge capacity in bits per parameter using fake biographies in many paraphrases, each biography exposed only 100 times during pretraining, spanning 50K to 2M biographies and 2×10⁶ to 10⁸ total bits. Mano asks for mental multi-step modular arithmetic without Chain-of-Thought, e.g. evaluating ((a×b)+(c−d)) mod 23 with operands in [0,22], with maximum expression lengths L = 10, 13, and 16. Lano tests hierarchical structure using context-free grammars with token-level ambiguity requiring global dynamic-programming resolution, extending the earlier cfg3f (lengths 100–500, 512-token contexts) with cfg3j and cfg3k (up to 200–1000 tokens, 1536-token contexts).

Architectures are standardized by size for fairness — Llama variants with 12 or 8 layers and hidden dimensions 768 or 512 (12 or 8 heads), denoted 12L768D, 12L512D, 8L768D, 8L512D, where 12L768D matches GPT-2-small — and translated to comparable-parameter GLA, Mamba2(mlp), and GDN configurations (GLA using a 4d² + 8d² budget, Mamba2(mlp) and GDN using 6d² + 6d²). For Capo, sizes are denoted ℓ-h (ℓ layers, hidden size 64h, h heads). All architectures share identical training settings, fixed random seeds, and identical data sequences, and results are reported as "3×4" mini scaling laws over three data scales and four model sizes.

The Canon layer itself is a 1-d causal convolution of kernel size 4 with an explicit residual: h′ₜ = hₜ + conv1d([hₜ, hₜ₋₁, hₜ₋₂, hₜ₋₃]), implemented with the open-source causal_conv1d library. The author notes that even fixed random weights work well, that simple scalar weighting (e.g., hₜ + 0.4hₜ₋₁ + 0.2hₜ₋₂ + 0.1hₜ₋₃) can degrade performance, and that more expensive variants such as input-dependent dynamic convolutions are not studied. Hybrid architectures (Griffin, Samba, GDN-H1/H2, sliding-window attention) are deliberately excluded so that individual modules can be attributed.

Why This Matters

Impact on research. The paper argues that architectural comparison at academic scale is currently unreliable — models fail 2-hop reasoning and benchmark differences can be generated by random seeds — and offers a cheap, reproducible alternative. It also reframes several existing design elements (Mamba2's conv1d, parts of GDN's gating and delta updates) as instances of Canon-like mixing, suggesting that some recent architectural progress may be re-deriving a single mechanism rather than adding new computation. The synthetic playground is proposed as a way to predict how architectures will behave as training pipelines improve.

Real-world applications:

• Selecting backbone architectures for language models under fixed compute budgets, using skill-isolated benchmarks rather than noisy perplexity.
• Improving long-context behavior: the paper reports that reducing or removing RoPE improves length generalization when Canon layers are present.
• Making small or on-device models more capable at multi-step manipulation of stored facts (the Mano-style skill), which matters for tool-free question answering.
• Recovering knowledge capacity in gated MLP and mixture-of-experts designs, which are widely used to scale models efficiently.

Industry relevance. The reported overhead is small (0.0063% extra parameters for a 1.3B model with Canon-ABCD; 12.4%/14.1%/20.8% runtime for forward/backward/generation with a naive implementation), and integration is described as a few lines of code across Transformers, linear attention, and state-space models. The author reports that preliminary larger pretrains (1–8B models on 1–2T tokens) align with the findings here, and that recent large hybrids — Falcon-H1 (May 2025, 32B, Mamba2 plus full attention) and Qwen3-Next (Sep 2025, 80B, GDN plus full attention) — validate Mamba2 and GDN as representative base linear models.

Future Directions

• Scaling beyond academic scale. The paper plans to extend Canon layer study beyond the academic scale, noting preliminary results from 1–8B models on 1–2T tokens that align with the reported findings and make several synthetic trends clearly observable.
• Runtime optimization. The author states that further runtime efficiencies are possible — for example, consolidating multiple Canon operations across layers — but remain outside this paper's scope.
• Canon variants. Dynamic convolutions with input-dependent weighting are described as possible but unstudied, since it remains unclear whether the additional cost is justified.
• Hybrid designs. The paper proposes Canon-equipped Transformer–linear hybrids as a path to combine deep reasoning with linear efficiency, building on the finding that linear models suffer from cumulative compression and retrieval errors rather than insufficient memory.
• Improved training pipelines. The paper suggests synthetic tasks like Depo may better predict model potential as data curation and RL-based post-training improve, and that synthetic benchmarks could forecast which architectures scale to more advanced tasks.

Target Audience

Architecture researchers and engineers working on sequence models — especially those comparing Transformers, linear attention, and state-space models — will get the most from this paper. It is also relevant to practitioners who need cheap, reliable architecture selection (including MoE and gated-MLP designers), and to academic groups with limited compute, since the synthetic playground is presented as an economical way to get meaningful signal where large-scale pretraining comparison is infeasible. Readers interested in evaluation methodology, rather than architecture specifically, will find the critique of perplexity-based comparison and academic-scale benchmark noise directly useful.

Authors’ abstract

Understanding architectural differences in language models is challenging, especially at academic-scale pretraining (e.g., 1.3B parameters, 100B tokens), where results are often dominated by noise and randomness. To overcome this, we introduce controlled synthetic pretraining tasks that isolate and evaluate core model capabilities. Within this framework, we discover CANON LAYERS: lightweight architectural components -- named after the musical term "canon" -- that promote…