Research map / Architectures and efficiency / Attention and transformers

On the Importance of Gating: Memorization vs. In-Context Learning in State Space Models

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: State space and recurrent sequence models · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Shows theoretically and empirically that gating in Mamba-style SSMs biases training toward in-weights memorization over in-context retrieval (strong gating = longer plateau), while strong gating aids length generalization, and tunes the gating init in tool-calling fine-tuning.

Problem

Linear-time SSMs such as Mamba lag Transformers on in-context learning and retrieval; the common explanation is an expressivity gap, but the authors ask whether training dynamics and gating initialization cause memorization shortcuts instead.

Approach

Analyzes a simplified gated linear-attention (single-head S6-like) model with a global gating parameter lambda; shows that strong gating (small decay factor alpha) makes memorization dominate initially and delays escape from the plateau, while decay induces locality and enables length generalization. Validates the trade-off on synthetic retrieval, Horn-clause logic retrieval, and BFCL tool calling

Key results

• On multi-token retrieval with N=512, strong gating (large offset) never leaves the memorization baseline (~0.43 sequence accuracy) within training, while weaker gating converges to perfect accuracy progressively faster.
• Strong gating yields better length generalization on best-last multi-token retrieval and on Horn clause retrieval, where weak-gated models degrade beyond training pool sizes up to K=128.
• On BFCL with a 1.5B Mamba fine-tuned on tool calling, weaker gating (lower scale) scores moderately better and hallucinates less than stronger gating, while all Mamba models length-generalize better than the 2B Transformer baseline.

Surprising findings

• Weakening the gating of a pretrained Mamba model before SFT can improve tool-calling accuracy over the pretrained gating.
• Gating does not appear to influence length generalization in the BFCL tool-pool range probed, with all Mamba variants beating the Transformer.

Limitations

• Theory uses a simplified single-head model with identity key/query/value matrices and orthonormal or idealized inputs.
• Tool-calling experiments use only a 1.5B Mamba and 2B Transformer and a single BFCL split; gating scale only probed in a limited range.

Experimental setup

Datasets: Nemotron-CC-HQ Nemotron Agentic tool-calling split BFCL python simple

Benchmarks: BFCL python simple synthetic retrieval multi-token key-value retrieval Horn clause logical rules retrieval

Metrics: accuracy hallucination rate escape time error probability

Baselines: Transformer

Builds on: Mamba Mamba 2 Gated Linear Attention BM25

Closest papers by meaning

• Understanding and Improving Length Generalization in Recurrent Models (ICML 2025 · similarity 0.90)
• Quantifying Memory Utilization with Effective State-Size (ICML 2025 · similarity 0.89)
• Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free (NeurIPS 2025 · similarity 0.88)
• Memorization Sinks: Isolating Memorization during LLM Training (ICML 2025 · similarity 0.88)
• Mitigating Memorization Where It Happens (NeurIPS 2026 · similarity 0.87)
• Theoretical Foundations of Deep Selective State-Space Models (NeurIPS 2024 · similarity 0.87)