A Single-Layer Model Can Do Language Modeling

Technical Summary: A Single-Layer Model Can Do Language Modeling

Problem Statement
Modern language models typically scale capacity by increasing depth, stacking multiple layers where each layer maintains its own state (e.g., per-layer KV caches in Transformers or per-layer matrix states in Mamba, GDN, RWKV, and xLSTM). This distributes computation and memory across the stack. In contrast, biological neural networks rely heavily on recurrence—revisiting the same circuitry over time—rather than stacking. The paper investigates the viability of this biological shape for language modeling: specifically, how far a single, shallow recurrent layer can go when it revisits a single state vector at every step, rather than distributing state across a deep stack.

Methodology: Grounded Prediction Networks (GPN)
The authors propose Grounded Prediction Networks (GPN), an architecture defined by a single state vector revisited at every time step. The core loop consists of three operations:

1. Grounding: The previous predicted state ($s^p_{t-1}$) is fused with the current observed token ($x_t$) via a gated residual to produce a grounded state ($s^g_t$).
2. Prediction: The grounded state is processed to predict the next state ($s^p_t$). This prediction is decoded into a token distribution for the next step.
3. Memory (GPN+M): The authors extend the base GPN to GPN+M by adding a centralized matrix memory ($M$). This memory is updated via a gated delta rule and read in parallel with a Feed-Forward Network (FFN). Both the FFN (offline parameters) and the memory read (online parameters) act as additive predictors to the state update.

Unlike standard recurrent models that maintain a state per layer, GPN uses one wide, shallow layer with a single shared matrix memory. Depth is achieved through time (revisiting the loop) rather than space (stacking layers). The state vector serves a dual role: it acts as a forecast for the next step and as an encoding of the history observed so far, with "grounding" completing the encoding for the next cycle.

Key Contributions

1. Architecture: The introduction of GPN and GPN+M, which utilize a single state vector, one FFN, and one centralized matrix memory revisited at every step, contrasting with the stacked per-layer states of Transformers and linear RNNs.
2. Conceptual View: A framing where the state's prediction of the next step is viewed as part of the encoding of that step, with a single vector playing both roles.
3. Empirical Evidence: Performance benchmarks at 130M parameters showing that a 1-layer GPN+M achieves competitive perplexity compared to deep baselines, with a 2-layer variant narrowing the gap further.
4. Observability: Because the entire working context is a single vector, the authors can directly inspect its geometry without needing to dig through residual streams. They identify emergent structures not explicitly designed into the architecture.

Results
Evaluated on the FineWeb-Edu dataset (130M parameters, 2048 context window):

• 1-Layer GPN+M: Achieves a perplexity (PPL) of 18.06. This is within 13% of a 12-layer Transformer++ (PPL 16.05) and 18% of a 10-layer Gated DeltaNet (GDN) (PPL 15.34).
• 2-Layer GPN+M: Achieves a PPL of 16.95, closing the gap to 6% of the Transformer++ and 11% of the GDN.
• LAMBADA Task: The shallow models show lower accuracy on context-dependent tasks compared to deep baselines, highlighting a cost in per-token compute depth.
• No Memory (GPN): A version without matrix memory performs significantly worse (PPL 23.51), confirming the necessity of the associative memory component.

Emergent Structures (Analysis)
The single-vector state allows for direct geometric analysis, revealing three structures that emerged spontaneously without explicit regularization:

1. High-Rank Geometry: The state vector is wide ($d=2496$) and utilizes its full width. It exhibits near-decorrelated dimensions and no "dead" dimensions, with a stable rank of ~40.
2. Persistent Default-Token Direction: The mean of the state vector ($\bar{s}$) is a large, fixed offset that decodes to a distribution dominated by frequent function tokens (e.g., "the", "and", ","). This direction persists even when an explicit bias is added to the decoder, suggesting it serves structural roles beyond simple decoding (e.g., anchoring RMSNorm).
3. Fast/Slow Memory Split: The 15 architecturally identical memory heads spontaneously split into two pools during training: a "fast" pool (6 heads) with half-lives of 1–3 tokens, and a "slow" pool (9 heads) with half-lives ranging from 5 to 322 tokens. This creates a two-timescale memory system without explicit architectural design.

Significance and Limitations
The paper claims that a single shallow recurrent layer can perform language modeling effectively, reaching within ~13% of deep baselines at matched parameter counts. The significance lies in demonstrating that "depth" can be replaced by time-revisited recurrence and that a single vector can support complex, emergent memory structures (fast/slow splits) and geometric properties (persistent anchors) that are typically hidden in deep residual streams.

However, the authors are modest about the current state of the work:

• Not Production-Ready: The model runs sequentially in Python with no parallelization along the sequence dimension, making training costs significantly higher than Transformers.
• Scale: Results are based on a single seed at 130M parameters with no scaling sweeps.
• Descriptive Nature: The analysis of emergent structures (like the default-token direction) is observational; the paper does not claim causal proof of their utility via ablation.
• Context Costs: The shallow model pays a measurable cost on tasks requiring deep context (like LAMBADA), indicating that while the approach is viable, deep stacks still hold advantages for specific context-heavy predictions.

The work serves as a proof of concept for the "grounded prediction" loop and the viability of single-vector state models, offering a new lens for observing how language models organize information.