Competing Learning in Hybrid Architectures

Summary (Overview)

• Identifies "Competing Learning" phenomenon: In hybrid language models combining linear recurrence (e.g., KDA, SSMs) and softmax attention, different mixer types compete to learn the same functions, with earlier-learned capabilities suppressing learning signals for other layers
• Key finding: Recurrent layers (KDA) learn retrieval functions faster than softmax attention under typical data distributions (where short-distance patterns dominate), causing attention layers to fail at learning long-range retrieval—leading to suboptimal division of labor
• Proposes Better Layer Division (BLD): A staged training procedure that freezes recurrent layer weights during early pretraining (~1/10 of initial tokens) to give attention layers a "first-mover advantage" in learning retrieval
• Empirical validation: BLD improves performance across models from 180M to 1B parameters on language modeling loss, common-sense reasoning (CSR), recall-intensive tasks, RULER, and LongPPL benchmarks
• Transferability confirmed: BLD works with multiple recurrent architectures (KDA, Gated DeltaNet, Gated Linear Attention) and produces more retrieval-specialized attention heads

---

Introduction and Theoretical Foundation

Background

Modern LLMs increasingly adopt hybrid architectures that combine:

• Linear recurrence (e.g., state space models, linear attention): Efficient, expressive mixing via fixed-size state compression
• Softmax attention: Strong contextual information retrieval capabilities

Representative models include Kimi K3, GLM-5.3-Flash, Qwen3.8-Flash-Next, and Nemotron-3-Ultra.

Prior work on labor division

The ideal division of labor: softmax attention specializes in long-distance retrieval; recurrent layers specialize in short-distance mixing.

Problem identification

Despite this architectural intention, standard pretraining leads to Competing Learning: recurrent layers' inductive bias toward short-distance patterns (from low-initialized forget gates) lets them learn retrieval faster, which suppresses attention's learning of the same capability.

---

Methodology

1. Synthetic Task: "Spaced Retrieval"

Input: $ X₁ X₂ ··· Xₜ ⊔···⊔ (g spaces) $
Output: X₁ X₂ ··· Xₜ

• Each $X_i$ is a digit from {0,...,9}; $g$ is the gap (64 or 8192); $t = 5$
• Deliberately simple to isolate retrieval behavior from natural language complexity

2. Layer Retrieval Ability Score Estimation

Model notation: Each layer computes:

$$U_l(x) = \text{Mixer}_l(H_{l-1}(x)), \quad R_l(x) = H_{l-1}(x) + U_l(x), \quad H_l(x) = R_l(x) + \text{MLP}_l(R_l(x))$$

Retrieval score: For input pair $(x_A, x_B)$ differing only in source digit ($a$ vs. $b$):

$$C^{A \leftarrow B}_l = m(z(x_A, j); a, b) - m(\tilde{z}^{A \leftarrow B}_{l,j}; a, b)$$

where $m(z; a, b) = [z]_a - [z]_b$ is the logit margin, and $\tilde{z}$ is the logit after replacing layer $l$'s mixer output at position $j$ with that from the paired input. Higher scores indicate greater layer contribution to retrieval.

3. BLD Training Procedure

• Uses DCLM corpus, Qwen3 tokenizer, 16,384-token context
• AdamW for early phase; Muon with auxiliary AdamW afterward
• Token budget follows Chinchilla scaling (~20× non-embedding params)

---

Empirical Validation / Results

1. Competing Learning Evidence (Synthetic Experiments)

Freezing experiments: Freezing KDA layers at various training steps causes attention layers' retrieval ability to increase sharply (Figure 4), supporting suppression hypothesis.

Unfreezing dynamics:

• Freezing before step 125 → attention retains retrieval advantage after unfreezing
• Freezing after step 125 → attention is outcompeted again once KDA resumes

Gradient evidence: KDA freezing amplifies attention gradient norms ~10× relative to baseline; unfreezing mitigates this amplification.

Data distribution effects: As short-distance fraction ($p$) increases, KDA dominates retrieval; attention dominates only when long-distance examples overwhelmingly outnumber short-distance ones.

2. LLM Pretraining Results (Table 1, selected sizes)

Key observations:

• BLD consistently achieves lower LM loss across all scales
• Substantial gains on recall-intensive tasks and RULER (up to ~3× improvement at 180M scale)
• Modest but consistent CSR improvements at most scales
• LongPPL improvements are dramatic (e.g., 57.22 → 17.96 at 250M)

3. Transferability Results (Table 2)

BLD transfers effectively to GDN and GLA architectures—recall improvements are particularly striking (GDN: 17.33 → 32.05).

4. Labor Division Evidence

Retrieval head analysis (α values, following Xiao et al., 2025):

• Vanilla model: Only 1 high-α retrieval head
• BLD model: 3+ high-α retrieval heads, showing stronger functional specialization
• Masked-head accuracy experiments confirm BLD's retrieval heads are more critical for task performance

5. Scaling Laws

Power-law fitting: $L_m(N) = E + A_m(N/10^9)^{-\alpha_m}$ reveals BLD maintains its loss advantage beyond evaluated scales (up to 3B non-embedding parameters), with fitted parameters $E = 2.38867$, $A_V = 0.45960$, $\alpha_V = 0.34649$, $A_B = 0.45266$, $\alpha_B = 0.34502$.

---

Theoretical and Practical Implications

Theoretical significance

1. New understanding of training dynamics: The Competing Learning phenomenon reveals that learning signals are not independent across architectural components—early advantages create self-reinforcing cycles that determine specialization

2. Architectural inductive biases matter: KDA's low-initialized forget gates create short-distance bias:

   $$S_t = (I - \beta_t k_t k_t^T)\text{Diag}(\alpha_t)S_{t-1} + \beta_t k_t v_t^T$$

   The cumulative product of small gate values rapidly decays historical information, explaining KDA's faster learning of short-distance patterns.

3. Gradient dynamics insight: Freezing dominant components amplifies gradients in subordinate components progressively (not abruptly), revealing an "evolving optimization dynamics" mechanism.

Practical implications

1. Training recipe: BLD is a simple, effective intervention—freeze recurrent blocks for first ~9% of tokens, then unfreeze—that requires no architectural modification

2. Scalability: Improvements persist from 180M to 1B parameters and extrapolate to larger scales via scaling laws

3. Architecture-agnostic: Works across KDA, GDN, and GLA recurrent architectures

4. Benchmark improvements: Particularly dramatic gains in long-context and retrieval tasks (RULER, LongPPL), with consistent improvements in CSR and language modeling

---

Conclusion

This work identifies and characterizes the Competing Learning phenomenon in hybrid architectures, where recurrent and attention layers compete for retrieval functionality during pretraining, with recurrent layers typically winning due to their short-distance inductive bias. The proposed BLD training procedure addresses this by temporarily freezing recurrent layers to give attention layers a first-mover advantage, resulting in:

• Lower training loss
• Better common-sense reasoning
• Substantially improved retrieval capabilities
• Better long-context modeling
• Greater attention head specialization for retrieval

Future directions: The authors note that BLD is "only a preliminary approach," suggesting deeper understanding of Competing Learning could inspire more advanced architectures and training strategies. Key open questions include:

• How does Competing Learning manifest in other capability domains beyond retrieval?
• Can more sophisticated scheduling of layer freezing/unfreezing yield further gains?
• How do these dynamics change at much larger scales (e.g., 10B+ parameters)?