Competing Learning in Hybrid Architectures

AI summary

Overview

Research area: Natural Language Processing / large language model pre-training, specifically hybrid architectures that mix linear-time recurrent layers (state space models, linear attention) with softmax attention.

Technical level: Advanced. The paper assumes familiarity with token mixers, attention heads, pretraining schedules, and perplexity-based evaluation, though its central idea is explained intuitively.

Scope: The paper identifies and analyzes a training-dynamics phenomenon it calls Competing Learning, then proposes a pretraining intervention, Better Layer Division (BLD), that improves how hybrid models divide retrieval work between recurrent and attention layers, validated on models from 180M to 1B parameters.

What This Paper Is About

Hybrid language models stack two kinds of layers: efficient recurrent layers (which compress context into a fixed-size state) and softmax attention layers (which are better at exact recall from context). The intended design is a division of labor where attention handles long-range retrieval and recurrence handles everything else. The authors show that standard pretraining does not produce this division: recurrent layers learn retrieval first, and by doing so they remove the training signal that attention layers would need to learn retrieval themselves. The goal is to diagnose this competition and design a training schedule that gives attention layers the early advantage they need.

Key Contributions

1. Identification of Competing Learning: The authors characterize a phenomenon in which recurrent and softmax-attention layers compete to learn the same retrieval function, and layers that learn it earlier reduce the retrieval-specific prediction error, weakening the learning signal available to the other layers.
2. A layer-attribution methodology: They introduce a synthetic "Spaced Retrieval" task plus a layer retrieval ability score that measures how much an individual layer contributes to retrieval, by swapping that layer's output at the prediction position with the output from a counterfactual input and measuring the change in answer margin.
3. Causal evidence via freezing: By freezing recurrent (KDA) layers at many different training steps, they show attention layers begin learning retrieval almost immediately once recurrence stops updating, and provide matching gradient-magnitude evidence.
4. Better Layer Division (BLD): A pretraining pipeline that freezes recurrent layers during an early stage of pretraining to steer the division of labor toward attention, evaluated on KDA-based hybrids from 180M to 1B parameters and transferred to GDN and GLA hybrids at 400M.

Main Findings

• Recurrence outcompetes attention by default. In a 3.6M-parameter hybrid with two KDA layers and two attention layers trained on Spaced Retrieval with a 1:1 mixture of gap sizes (g = 64 and g = 8192), KDA layers become dominant in retrieval ability when no intervention is applied. Attention layers largely fail to learn retrieval.

• Freezing recurrence flips the winner. When all KDA layers are frozen at various training steps and training continues, attention layers' retrieval ability increases sharply immediately after each freezing intervention. In the unfreezing experiment, runs where KDA was frozen earlier than update 125 ended with higher attention retrieval ability than runs where KDA stayed frozen for the whole run, while runs frozen later than update 125 ended lower. The changes appear almost immediately after KDA is unfrozen.

• Gradient evidence matches the behavioral evidence. With KDA frozen, the attention layers' gradient norm reaches roughly 10× that of a baseline checkpoint that was never frozen, and unfreezing KDA substantially attenuates this amplification. The increase is progressive rather than abrupt, which the authors read as a self-reinforcing optimization mechanism rather than an immediate side effect of freezing.

• Data distribution decides the winner. Training the same synthetic hybrid with varying fractions of short-distance (g = 64) versus long-distance (g = 8192) examples shows attention layers become more important as the long-distance proportion increases, whereas KDA dominates unless long-distance patterns appear overwhelmingly more often. The authors attribute this to KDA's forget gates being initialized with small values, giving the architecture a short-distance inductive bias. Since short-distance patterns are more frequent in natural language, standard pretraining pushes hybrid models toward the suboptimal outcome.

• BLD improves pretraining across scales. Comparing BLD to vanilla pretraining on KDA hybrids (Attention-KDA-KDA-KDA pattern, 3:1 KDA-to-attention) trained on DCLM with a 16,384-token context:

  • LM loss improves at every size: 1B goes from 2.8744 (vanilla) to 2.8680 (BLD); 390M from 3.1611 to 3.1441; 180M from 3.5903 to 3.5702.
  • Recall (six BASED tasks) gains are large: 1B from 36.50 to 38.14; 390M from 26.73 to 31.27; 180M from 10.39 to 14.12.
  • RULER (57 task–length combinations, 1K–16K) improves: 1B from 32.75 to 36.04; 180M from 2.84 to 8.82.
  • LongPPL drops sharply: 180M from 258.47 to 108.64; 1B from 7.69 to 6.50.
  • CSR improves at four of six sizes (e.g., 1B 48.69 to 49.45; 390M 41.91 to 44.25) but is marginally lower for BLD at 180M (37.49 vs 37.76) and 510M (44.49 vs 44.80).

• BLD creates more retrieval heads. In an analysis of 510M hybrid models following the DuoAttention-style α-value method (trained for 1,000 steps), BLD attention heads show higher α values than vanilla heads. When heads are masked in ascending α order, BLD accuracy drops sharply only after removing the final three heads, whereas the vanilla model degrades mainly when its single highest-α head is removed — suggesting three retrieval heads versus one.

• Gains transfer to other recurrent mixers. At 400M with a 3:1 recurrence-to-attention ratio, GDN Recall rises from 17.33 to 32.05 and RULER from 12.33 to 25.42; GLA Recall rises from 24.58 to 30.12 and RULER from 23.98 to 28.54.

• Better context utilization and scaling. Position-wise loss on Books3 and BookCorpus2 up to 16K tokens is lower for BLD, especially at later positions, and power-plus-constant scaling-law fits suggest the loss advantage persists beyond the evaluated scales.

Methodology in Plain English

The authors work in two stages.

Diagnosis on a synthetic task. They build "Spaced Retrieval": the model sees a five-digit key, a gap of spaces, and must reproduce the digits. The gap is either short (g = 64) or long (g = 8192), which lets them control whether a pattern is short-range or long-range. They train small 3.6M-parameter hybrids (KDA–attention–KDA–attention) on 50,000 distinct keys and evaluate on 2,000 held-out pairs per gap size.

To see which layer is actually doing the retrieval, they use a counterfactual swap. Take two inputs differing only in the source digit that must be recalled. Run the first input normally, then re-run it while replacing one layer's output at the prediction position with that layer's output from the second input. If the model's preference for the correct answer collapses, that layer was carrying the retrieval load. Repeating this for every layer and pair gives a per-layer retrieval score over training time.

With this measuring stick, they run three interventions: freezing the KDA blocks at many different update counts and watching what attention does; unfreezing KDA later to see whether the new division sticks; and varying the mix of short- and long-gap training examples.

Intervention at scale. Based on the diagnosis, they propose BLD: freeze the recurrent blocks for roughly the first 1/10 of the pretraining that will later be unfrozen, so attention layers get a head start in learning retrieval. Then unfreeze and train normally. They pretrain KDA hybrids from 180M to 1B total parameters on DCLM with the Qwen3 tokenizer at 16,384-token context, using token budgets of about 20× non-embedding parameters, AdamW early and Muon with auxiliary AdamW after unfreezing — identical schedules for BLD and vanilla. They then evaluate zero-shot on common-sense reasoning, six recall-intensive BASED tasks, RULER, WikiText-2 perplexity, and LongPPL, plus a head-level α-value analysis and a repeat of the recipe with GDN and GLA layers.

Why This Matters

The paper reframes a design assumption. Hybrid architectures are built on the premise that attention will handle retrieval while recurrence handles throughput, but the authors show this division does not arise automatically — it must be engineered during training. This is relevant to any team training or fine-tuning hybrid models, and it connects to prior work (Qiao et al., 2026) showing efficient-attention modules can weaken optimization pressure on full-attention layers.

Real-world applications:

• Long-document question answering and summarization, where exact recall of a passage thousands of tokens back determines answer quality.
• Retrieval-augmented generation, where retrieved passages must be quoted or reasoned over precisely rather than paraphrased.
• Long-context agents and code assistants that must track identifiers, function names, or earlier tool outputs across a large working context.
• Efficient serving at scale, since hybrids lower inference cost; BLD aims to recover retrieval quality without giving up that efficiency or changing the architecture.

Industry relevance: The paper targets KDA, the recurrence used in Kimi K3 and GLM-5.3-Flash, and evaluates KDA, GDN, and GLA hybrids — the same family of designs adopted in open-weight frontier models such as Kimi K3, GLM-5.3-Flash, Qwen3.8-Flash-Next, and Nemotron-3-Ultra. Because BLD is a training-schedule change with no architectural modification or inference overhead, it is directly deployable by anyone pretraining such a model. The authors release code at https://github.com/thunlp/competing-learning.

Future Directions

• Smarter schedules than a fixed freeze. BLD freezes recurrence for about 1/10 of the unfrozen pretraining portion. The synthetic experiments show the outcome is highly sensitive to when freezing starts relative to a critical point (around update 125 in their setup), which suggests adaptive or ability-triggered schedules could work better than a fixed fraction.
• Whether the competition extends beyond retrieval. The paper studies retrieval specifically. It is an open question whether the same suppression dynamic governs other capabilities that multiple mixer types could learn, and whether it appears between other component pairs (for example, among heads or inside mixture-of-experts layers).
• A mechanistic theory. The authors describe a self-reinforcing loop in which the current parameter configuration determines how learning signals are allocated, but they characterize it empirically. A formal account would let practitioners predict the winner from data statistics and initialization rather than measuring it.
• Better division-of-labor objectives. BLD is described by its authors as "only a preliminary approach." Combining the training-time intervention with architectural or objective-level changes — and testing whether the gains compound — is left open.

Target Audience

Most useful to researchers and engineers working on LLM pretraining and architecture design, particularly those building or training hybrid recurrent-attention models. It is also valuable to anyone studying training dynamics, feature specialization, or the formation of retrieval heads, and to practitioners deciding how to schedule freezing, staged training, or layer-wise learning rates. Readers without background in attention variants and pretraining practice will find the high-level argument accessible, but the experimental sections assume familiarity with perplexity metrics, RULER, and BASED-style evaluation.

Authors’ abstract

Many modern language models have adopted hybrid architectures that incorporate both linear-time recurrence and softmax attention. Modern recurrence layers provide a highly expressive mixing mechanism while ensuring high throughput by compressing contextual information into a fixed-size state, but underperform softmax attention in contextual information retrieval. In this paper, we reveal a Competing Learning phenomenon where different kinds of mixers compete to learn the same function, and when any layer learns a function, it suppresses the corresponding learning signal for other layers. Layer attribution analysis reveals that under typical data distributions, modern recurrence learns retrieval functions faster than softmax attention, causing attention to fail to learn retrieval in most cases. Stopping weight updates in recurrence in the middle of training helps attention learn retrieval. Based on these results, we propose a staged training procedure for pre-training hybrid architectures on natural language. Empirical results on models up to 1B parameters show that our method results in substantial improvement in language modeling loss, common-sense reasoning, as well as a wide range of retrieval-related tasks across different context lengths.