Balancing Memory Pathways: Analyzing and Improving Memory Utilization in Hybrid LMs

• Restricting recurrent access reduced accuracy from 30.2% to 20.2%, while restricting attention reduced it from 30.2% to 5.1%.Standard SFT improved overall performance but mostly by improving attention use, leaving recurrent use limited.Source: §1 · p. 2, Figure 1 · p. 2
• 3.9% and 5.2% average absolute accuracy gains across six QA datasets came from recurrent auxiliary loss on Qwen3.5-4B and Nemotron-H-4B, respectively.The loss increased recurrent-only accuracy while largely preserving attention-only accuracy.Source: §1 · p. 2, Figure 1 · p. 2
• Across two models, recurrent auxiliary supervision improved success rates by 3.5% and 9.3% on TextWorld Quest and Treasure, and 22.0% and 13.5% on BabyAI PutNextLocal and GoToObjMaze.BabyAI trajectories averaged 134.0 steps, compared with 42.2 in TextWorld.Source: §1 · p. 2
• With retries, auxiliary-loss models improved cumulative recovery rates over standard SFT by 16.6% on QA and 19.2% on agentic tasks.On QA, gains were 30.3% for multi-evidence questions versus 2.8% for single-evidence questions.Source: §1 · p. 3
• In agentic tasks, auxiliary-loss agents produced trajectories shorter than any successful demonstrations 6.0% of the time on average, with larger gains for longer demonstrations.The imbalance and benefits also generalize to attention-based LMs combining different memory types.Source: §1 · p. 3

2.1 RECURRENT-ATTENTION HYBRID LMS

• Attention layers access current and preceding tokens directly, while linear recurrent layers update a fixed-size state with computation linear in sequence length and constant memory.Hybrid models interleave the layers, often using relatively few attention layers alongside a larger recurrent portion.Source: §2.1 · p. 3

2.2 ROLES AND STRENGTHS OF ATTENTION AND RECURRENT MEMORY PATHWAYS

• Prior work finds attention especially important for precise retrieval and recurrent components important for language modeling, but does not test whether hybrids effectively use both strengths.The section frames effective pathway use and its improvement as the open question.Source: §2.2 · p. 3, §2 · p. 3
• Ablations show attention contributes more to retrieval QA, while recurrent layers contribute more to language modeling.For Qwen3.5, attention ablation increased NLL more on retrieval QA (5.3 vs. 1.5), while recurrent ablation had the larger language-modeling effect (1.3 vs. 3.7).Source: §2.2 · pp. 3–4, Figure 2 · p. 4
• Attention-only retention was higher for verbatim than non-verbatim questions (63% vs. 46%) and for single-evidence than multi-evidence questions (60% vs. 50%).The authors relate this pattern to attention’s role in recovering precise details and its lesser role in abstraction or integration across context.Source: Figure 3 · p. 4, §2.2 · p. 4

• Experiments use Qwen3.5-4B with 24 Gated DeltaNet and 8 full-attention layers, and Nemotron-H-4B-Instruct with 24 Mamba2 and 4 full-attention layers.Source: §3 · p. 4
• All models are trained with LoRA on every layer, rank 16, scaling factor 32, dropout 0.05, and AdamW at learning rate 1e−4.Auxiliary-loss variants use half as many optimization steps to match standard SFT’s total training compute.Source: §3 · p. 4
• QA training pools four datasets into 12k training and 2k validation examples; evaluation covers six datasets and 8.2k test examples averaging 84k context tokens.Agentic evaluation uses TextWorld Quest and Treasure and BabyAI PutNextLocal and GoToObjMaze, with training on Quest and MiniBossLevel, respectively.Source: §3 · pp. 4–5

4 STANDARD SFT FAVORS ATTENTION OVER RECURRENCE

• Overall, SFT improves performance without balancing pathway use: attention reliance rises while recurrent-pathway use changes little.Source: §4 · p. 5, §4.2 · p. 6

4.1 DISENTANGLING TWO MEMORY PATHWAYS

• The pathway test splits past from current input, blocking cross-segment attention for recurrent-only and resetting recurrent state at boundaries for attention-only.Source: §4.1 · p. 5

4.2 MEMORY PATHWAY UTILIZATION AFTER SFT

• After SFT, Qwen3.5 QA full accuracy rose from 30.2% to 34.6%, while recurrent-only accuracy changed from 5.1% to 5.0%; Figure 4 also covers agentic tasks.Source: Figure 4 · p. 6, §4.2 · p. 6
• Across question types, attention-only retention increases; verbatim and single-evidence questions gain 25.5% on average versus 19.2% for non-verbatim and multi-evidence questions.Source: §4.2 · p. 6, Figure 5 · p. 6

5.2 RESULTS

• 38.5% versus 34.6% average QA accuracy for Qwen3.5, and 36.2% versus 31.0% for Nemotron-H, under SFT+Lrec versus standard SFT.Gains were particularly large on Qasper, NarrativeQA, and MINTEval; recurrent-only accuracy rose from 5.0% to 22.8% while attention-only accuracy was largely preserved, from 25.0% to 23.7%.Source: Table 1 · p. 7, §5.2 · p. 7

5.1 RECURRENT AUXILIARY SUPERVISION

• The auxiliary pass masks attention to preceding context while recurrent layers process the full sequence, so that preceding information can reach the answer through recurrent states.The objective jointly optimizes unrestricted SFT and the recurrent auxiliary loss, with λ controlling the auxiliary supervision strength.Source: §5.1 · p. 7

5.2 RESULTS

• SFT+Lattn gains were smaller than SFT+Lrec gains, and recurrent-only accuracy remained much lower when attention was encouraged instead.The average absolute improvements were 2.1% for SFT+Lattn and 4.6% for SFT+Lrec; SFT+Lattn raised attention-only accuracy from 25.0% to 28.0%, while recurrent-only accuracy was 11.3%. Restricted losses alone also underperformed standard SFT.Source: §5.2 · p. 7, Figure 6 · p. 8
• Across TextWorld and BabyAI, recurrent auxiliary supervision improved success rates more than attention-focused auxiliary supervision.The recurrent objective's average absolute gains over SFT were 3.5% and 9.3% on TextWorld Quest and Treasure, and 22.0% and 13.5% on BabyAI PutNextLocal and GoToObjMaze, respectively. BabyAI trajectories were longer on average than TextWorld trajectories.Source: Table 2 · p. 8, §5.2 · p. 8

Learning from Failure Trajectories

• 16.6% average absolute recovery gain on QA and 19.2% on agentic tasks; the QA gain was larger for multi-evidence than single-evidence questions.Across attempt budgets, multi-evidence recovery rose from 27.3% to 57.6%, compared with 14.8% to 17.6% for single-evidence questions.Source: §6 · p. 9, Figure 7 · p. 9

Learning from Successful Trajectories

• Auxiliary supervision improved the rate of generating successful trajectories shorter than the shortest demonstration by an average of 6.0%.Gains were larger when demonstrations were longer relative to the optimal path.Source: §6 · p. 9, Figure 7 · p. 9

Generalizes to Attention-Based LMs

• 24.1% accuracy with both memories under standard SFT was below text-only accuracy of 25.9%; auxiliary supervision raised accuracy to 28.6%.The per-example Oracle reached 32.0%, and the improvement with auxiliary supervision accompanied increased latent-memory use while largely preserving textual-memory use.Source: Figure 8 · p. 9, §6 · pp. 9–10

7 CONCLUSION

• Hybrid LMs tend to rely predominantly on attention despite the complementary strengths of recurrent layers.Standard SFT further increases attention use while yielding little improvement in recurrent memory use.Source: §7 · p. 10
• An auxiliary loss encouraging recurrent-pathway use improves utilization of both memory pathways and overall performance.Source: §7 · p. 10
• The reported imbalance and auxiliary-loss benefits generalize across recurrent-attention LMs, QA and agentic tasks, and attention-based LMs combining different memory forms.Source: §7 · p. 10

A.1 VERBATIM VS. NON-VERBATIM QUESTIONS

• Questions are classified as verbatim when the normalized gold answer appears in the source document, and non-verbatim otherwise.About 34.6% of test examples were verbatim; an equal number of non-verbatim examples were randomly sampled.Source: §A.1 · p. 14

A.2 SINGLE- VS. MULTI-EVIDENCE QUESTIONS

• Questions are classified as single-evidence or multi-evidence using dataset evidence annotations, with equal numbers sampled for the two categories.The test set had approximately 30.2% single-evidence and 39.7% multi-evidence questions with annotations; around 2.5k examples per category were sampled.Source: §A.2 · p. 14

B.1 QUESTION ANSWERING

• Question answering uses token-level F1 for NarrativeQA and Qasper, and exact match for DROP, HotpotQA, and MINTEval; large test sets are subsampled to 1.5k.Models train for five epochs with early stopping; MuSR uses exact match after the model generates one of the provided answer options.Source: §B.1 · p. 14

B.2 AGENTIC SETUP

• Agentic evaluation measures the fraction of games completed within the step limit, using 100 TextWorld steps and official BabyAI limits.TextWorld uses 1,000 training, 150 validation, and 200 test games; retry experiments allow up to seven attempts, and agents train on gold trajectories for five epochs with early stopping.Source: §B.2 · p. 14

C.1 EFFECT OF THE AUXILIARY LOSS WEIGHT λ

• 78.5% is standard SFT (λ = 0); every tested auxiliary-loss weight improves Qwen3.5 TextWorld Quest success, with λ = 0.25 achieving the best performance.The authors use λ = 0.25 in all experiments unless otherwise specified.Source: §C.1 · p. 14

C.2 AGENTIC TASKS

• Across λ values, auxiliary loss beats standard SFT on TextWorld Quest and improves cumulative recovery by 19.2% on average as retries accumulate.With auxiliary loss, recovery continues to improve as attempts accumulate, whereas standard SFT plateaus.Source: §C.2 · p. 15, Figure 10 · p. 15, Figure 9 · p. 15

C.3.1 TEXTUAL & LATENT MEMORY.

• In Qwen3-4B-Instruct, textual summaries and latent KV-prefix memories are complementary: textual memory fares better on verbatim or single-evidence questions, while latent memory is more competitive on non-verbatim or multi-evidence questions.Latent memory is constructed from layer-wise key-value representations at appended learnable memory tokens.Source: §C.3.1 · pp. 15–16, §C.3 · p. 15, Figure 11 · p. 16

C.3.2 STANDARD SFT FAVORS TEXTUAL MEMORY OVER LATENT MEMORY

• Standard SFT with both memories scores 24.1%, below text-only at 25.9% and the Oracle by 7.9%, despite complementary information across the memories.The two memories are jointly available during training without additional supervision to coordinate their use.Source: §C.3.2 · p. 16

C.3.3 ADDING AN AUXILIARY LOSS TO ENCOURAGE LATENT MEMORY USE IMPROVES PERFORMANCE

• The auxiliary term makes the non-designated memory uninformative for the target, encouraging reliance on the designated pathway alongside the standard hybrid-memory objective.For text-role examples, the correct textual memory is retained while an irrelevant latent memory is substituted; for latent-role examples, the correct latent memory is retained and an irrelevant textual memory is substituted.Source: §C.3.3 · p. 17
• Role assignment can follow single-memory correctness (OUTCOME), random assignment, the verbatim distinction (VERBATIM), or its reverse (REVERSE).The VERBATIM criterion assigns verbatim examples to text and non-verbatim examples to latent memory.Source: §C.3.3 · p. 17
• Without auxiliary supervision, replacing latent memory with an unrelated one slightly improves accuracy (∆Latent = −2.4), indicating that standard SFT largely ignores it.The table defines ∆Latent and ∆Text as accuracy drops after replacing the corresponding memory with one from an unrelated context; higher values indicate greater use.Source: §C.3.3 · p. 17, Table 3 · p. 17
• 24.1% to 28.6%: VERBATIM role assignment yields the best accuracy and largest latent-memory-use increase, from −2.4 to +4.7, while textual-memory use is largely preserved (+11.4 to +10.4).VERBATIM assigns verbatim examples to the text role and non-verbatim examples to the latent role.Source: §C.3.3 · pp. 17–18, Table 3 · p. 17
• LATENT-AUX raises accuracy to 26.1% while reducing textual-memory use; TEXT-AUX adds little accuracy, and UNROUTED improves neither memory use nor accuracy.The authors conclude that gains are largest when each example is assigned to the memory suited to it.Source: §C.3.3 · p. 18