Source-linked AI summary
Generalization or Memorization: Dynamic Decoding for Mode Steering
Xuanming Zhang
TL;DR
LLMs can shift unpredictably between robust generalization and brittle memorization, undermining reliable deployment. The paper formalizes these modes with the Information Bottleneck and proposes Dynamic Mode Steering, an inference-time probe-and-steer approach; its examples show improved reasoning steps and more accurate medical facts after steering.
Problem
LLMs can unpredictably switch between robust generalization and brittle memorization, creating a challenge for safe and reliable deployment.
Source: §1 · p. 1Method
Results
The conclusion’s examples report that DMS correctly executes the reasoning steps in a sales calculation and steers away from memorized but inaccurate folk wisdom toward more accurate medical facts.
Source: §7 · pp. 7–8Takeaways & Limitations
In the illustrated cases, steering is associated with procedural integrity in reasoning and more accurate factual responses.
Source: §7 · pp. 7–8Standard probes reveal correlations but may not identify features that causally contribute to the output, motivating causal activation patching to choose the probing and intervention layer.
Source: §4.1 · p. 4
Abstract
from arXiv · showhide
Large Language Models (LLMs) exhibit a troubling duality, capable of both remarkable generalization and brittle, verbatim memorization of their training data. This unpredictability undermines their reliability in high-stakes applications. In this work, we propose a unified framework to understand, identify, and control these distinct reasoning modes. First, we introduce a theoretical model based on the Information Bottleneck (IB) principle, formalizing generalization as the learning of a compressed, task-relevant representation and memorization as a failure to compress. Building on this theory, we develop Dynamic Mode Steering (DMS), a novel inference-time algorithm which comprises two components: (1) a lightweight, causally-grounded linear probe that identifies the model's instantaneous reliance on memorization, and (2) a dynamic activation steering mechanism that nudges the model's computation towards pre-identified generalization circuits. We frame DMS as a form of adaptive, self-contrastive decoding. Experiments on reasoning and faithfulness tasks demonstrate that DMS significantly improves logical consistency and factual accuracy, thereby offering a principled approach to enhancing LLM reliability.
1. Introduction
LLMs can generalize from learned principles or brittlely reproduce memorized training fragments, making their behavior unpredictable and deployment less reliable. The paper proposes an Information Bottleneck account of these modes and introduces Dynamic Mode Steering to shift inference toward generalization.
1. Introduction
- LLMs can unpredictably switch between robust generalization and brittle memorization, undermining safe and reliable deployment.The paper frames generalization as applying abstract principles to novel problems and memorization as retrieving specific training examples without learning their underlying process.Source: §1 · p. 1
- Grokking illustrates that memorization and generalization can be distinct solutions, with generalization emerging after a model has memorized its training data.The paper uses this delayed transition to motivate a formal theory and a practical mechanism for controlling reasoning modes.Source: §1 · p. 1
- The paper proposes an Information Bottleneck model to formally distinguish generalization from memorization and analyze LLM learning dynamics.The framework provides a mathematical language for the distinction.Source: §1 · p. 2
- Dynamic Mode Steering is a training-free, inference-time algorithm that steers computation away from memorization-associated pathways and toward generalization circuits.It is introduced as a way to dynamically manage the model’s reasoning mode.Source: §1 · p. 2
2. Related Work
The related work situates the paper at the intersection of grokking, Information Bottleneck analysis, and inference-time control. Its stated contribution is to use the IB objective to define reasoning modes and make activation steering adaptive to a real-time mode signal.
Generalization, Memorization, and Grokking
- The paper proposes the Information Bottleneck as a unifying framework for explaining grokking as an optimization toward more information-theoretically efficient representations.Prior work links grokking to changes such as reduced feature rank and identifies model size, data size, and weight decay as relevant factors.Source: §2 · p. 2
Information-Theoretic Analysis of Neural Networks
- Unlike prior IB work described here as an analytical tool, this framework uses the IB objective to define and distinguish the computational goals of generalization and memorization.The cited lineage includes work on fitting and compression phases and bounds connecting input compression to generalization.Source: §2 · p. 2
Inference-Time Decoding and Activation Steering
- The paper frames activation steering as self-contrastive control, made dynamic and adaptive by a real-time signal from a targeted mode probe.The related work contrasts this with decoding methods that manipulate output distributions and other adaptive methods that respond to confidence or intermediate signals.Source: §2 · p. 2
3. Theoretical View
The paper uses the Information Bottleneck to frame reasoning modes as a trade-off between task utility and representational complexity. Generalization compresses away irrelevant input information, whereas memorization retains it; the paper then motivates inference-time control of the model’s trajectory through its available computations.
3.1. The Information Bottleneck Principle
- The Information Bottleneck seeks a representation that compresses the input while retaining information relevant to a target, with β controlling the utility–compression trade-off.In a neural network, the representation Z corresponds to the activations of an intermediate layer.Source: §3.1 · pp. 2–3
3.2. Formalizing Reasoning Modes
- The information plane frames generalization and memorization as distinct solutions to the complexity–utility trade-off.The paper presents these modes as different strategies for navigating the IB landscape.Source: §3.2 · p. 3, Figure 1 · p. 3
- Generalization is modeled as a compressed representation that discards irrelevant input details while retaining information needed to predict the target.The representation approximates a minimal sufficient statistic and captures the underlying data structure.Source: §3.2 · p. 3
- Memorization is modeled as an IB-suboptimal, high-complexity representation that retains training-example details and can fail on novel inputs.It may fit training data well while preserving noise and irrelevant information rather than the underlying data structure.Source: §3.2 · p. 3
- The paper decomposes representation complexity into task-relevant information and information about the input not explained by the target.The latter can include noise, stylistic details, and instance-specific artifacts; retaining it is presented as a reason memorization is suboptimal and performs poorly on unseen data.Source: §3.2 · p. 3
- The model is described as containing circuits for both modes, while inference follows an activation-space trajectory that determines which circuits are engaged.A model may default to a memorization-based path even when it has the capacity for a generalization-based solution.Source: §3.2 · p. 3
- The paper proposes nudging the model’s state during inference to select among latent capabilities in the trained model.This is presented as a control mechanism bridging the gap between the model’s available modes and its chosen inference trajectory.Source: §3.2 · p. 4
4. Methodology
DMS uses a closed-loop process to identify whether the model is relying on memorization and steer its computation toward generalization. Its probe is trained on internal activations, while causal activation patching selects where to probe and intervene.
4. Methodology
- DMS first identifies the model’s current reasoning mode, then intervenes to steer computation toward generalization.The paper presents this as a two-stage, closed-loop inference-time process.Source: §4 · p. 4
4.1. Stage 1: Mode Identification
- A lightweight linear probe predicts memorization from activations, producing an inference-time score for the model’s current reliance on memorization.The probe is a logistic regression classifier; its simplicity is intended to reduce the risk of learning capabilities independent of the host model.Source: §4.1 · p. 4
- Probe labels contrast prompts eliciting verbatim recall with prompts requiring novel reasoning, classifying repetitive outputs as memorization and diverse outputs as generalization.Prompts without clear low or high output diversity are discarded from the probe’s training signal.Source: §4.1 · p. 4
4.2. Activation Patching
- Activation patching replaces corrupted-run activations with those from a clean, correct run, identifying the layer where patching flips the answer to the generalized response.That causally critical layer is used for both probe training and steering intervention.Source: §4.2 · pp. 4–5, §4.1 · p. 4
4.3. Stage 2: Mode Control
- At the selected layer, DMS computes a steering vector from the difference between mean generalization and memorization activations, then adds it to the residual stream.The intervention’s magnitude scales with the real-time memorization score, so it is negligible when that score is near zero and strongest when it is near one.Source: §4.3 · p. 5
- DMS is framed as adaptive, self-contrastive decoding: it contrasts the model’s default computation with a steered pro-generalization path rather than a separate model.The steering vector itself contrasts the two modes before being applied to the model’s internal state.Source: §4.3 · p. 5
5. Experiment
The experiments evaluate DMS on reasoning, commonsense, and faithfulness tasks using Llama-3 8B and 70B models. DMS improves reasoning and factual accuracy, outperforming the reported decoding baselines across both model scales.
5.1. Experimental Setup
- DMS was evaluated with Llama-3 8B and 70B on GSM8K (Pass@1 accuracy), HellaSwag commonsense reasoning, and TruthfulQA (% True & Informative), against contrastive decoding.Source: §5.1 · pp. 5–6
5. Experiment
- Dynamic Mode Steering provides substantial improvements in complex reasoning and factual accuracy, outperforming all baseline methods across both model scales.Source: §5.1 · p. 6
5.2. Reasoning Task
- 68.3 Pass@1 (+6.2%) for Llama-3 8B and 86.7 Pass@1 (+5.2%) for Llama-3 70B on GSM8K; the 70B result exceeds Contrastive Decoding by 3.6%.HellaSwag improvements are described as more modest but consistent.Source: Table 1 · p. 6, §5.2 · p. 6
5.3. Faithfulness Task
- On TruthfulQA, DMS achieves the highest scores, improving over baseline by 6.4% for the 8B model and 5.4% for the 70B model.TruthfulQA rewards answers that are both factually correct and comprehensive.Source: Table 2 · p. 6, §5.3 · p. 6, §5.1 · p. 6
6. Discussion
The ablations identify a layer-specific performance peak and a non-monotonic effect of steering strength. The qualitative analysis uses examples to illuminate how DMS improves model behavior.
6.1. Ablation Study
- DMS reached its best GSM8K performance at layer 22, with a clear peak around the mid-to-post layers and degraded performance at early or post layers.The authors interpret this as support for targeting the causally identified junction point.Source: §6.1 · pp. 6–7, Figure 2 · p. 7
- As steering strength α increases from zero, performance first benefits and then declines beyond an optimum, consistent with overly aggressive steering disrupting learned computations.The authors suggest tuning α because aggressive intervention can push activations into an out-of-distribution region; the intervention magnitude is also scaled by probe confidence.Source: §6.1 · p. 7, Figure 3 · p. 7
6.2. Qualitative Analysis
- The qualitative analysis presents examples to illuminate the mechanisms by which DMS improves model behavior.Source: §6.2 · p. 7
7. Conclusion
DMS combines an Information Bottleneck framework with closed-loop, inference-time probing and activation steering toward generalizable circuits. On Llama-3, it outperforms strong baselines on GSM8K reasoning and TruthfulQA factual faithfulness, with examples illustrating improved reasoning and factual responses.
7. Conclusion
- On the Llama-3 family, DMS significantly outperforms strong baselines on GSM8K complex reasoning and TruthfulQA factual faithfulness benchmarks.The conclusion describes DMS as a training-free, inference-time closed loop that probes memorization reliance and steers toward generalization circuits.Source: §7 · p. 7
- For the knuckle-cracking question, DMS says cracking knuckles does not cause arthritis, unlike the baseline's claim that it can cause arthritis.The conclusion attributes the baseline response to a memorized but factually incorrect folk belief.Source: §7 · p. 7
- In the bakery example, the baseline uses Monday's 45 muffins for Wednesday and answers $495, while DMS uses Tuesday's 30 and answers 405.The authors describe the baseline mistake as an attentional slip and DMS's response as preserving the step-by-step calculation.Source: §7 · p. 8
Impact Statement
The impact statement links control of LLM reasoning modes to safety and reliability, and presents DMS as a step toward a more engineering-driven approach to alignment.
Impact Statement
- Preferentially activating generalization circuits offers a path toward AI systems less likely to regurgitate memorized falsehoods or fail on novel inputs.Source: Impact Statement · p. 8
- The authors frame analyzing and controlling internal computations as a step toward a more rigorous, engineering-driven approach to model alignment.Source: Impact Statement · p. 8
- The ability to understand, identify, and control LLM reasoning modes is presented as having implications for AI safety and reliability.Source: Impact Statement · p. 8
A.1. Implementation Details
Implementation details describe how probe labels were formed, how intervention layers were identified, and how DMS hyperparameters were tuned.
A.1. Implementation Details
- Probe training used benchmark training splits, with GSM8K activations labeled generalization and TruthfulQA activations labeled memorization after output-diversity filtering.Only samples showing clear high or low diversity were retained for the respective classes.Source: §A.1 · p. 11
- Causal activation patching on held-out GSM8K and TruthfulQA prompts identified layer 22 for Llama-3 8B and layer 55 for 70B, which were then used for probing and steering.Source: §A.1 · p. 11
- The steering strength α and interpolation factor λ were grid-searched on a small reasoning-and-faithfulness validation set, with α = 1.4 reported as a robust balance.Source: §A.1 · p. 11
A.2. Theoretical Foundations
The foundations recast generalization as task-relevant compression and memorization as retained task-irrelevant information, then motivate steering activations toward generalization states. Future work identifies unsupervised behavior-vector discovery, multi-attribute control, and a deeper account of the connection between the Information Bottleneck and activation geometry.
A.2. Theoretical Foundations
- The IB Lagrangian trades task information against input compression and decomposes to (1 − β)I(Z;Y) − βI(X;Z|Y), penalizing task-irrelevant detail.Source: §A.2 · p. 11
- A trained model contains potential generalization and memorization pathways, and each input follows a deterministic activation trajectory that may default to superficial-feature-driven memorization.Source: §A.2 · p. 12
- Neural collapse supplies a geometric target: reducing within-class variance discards instance-specific information, while vg = µg − µm pushes states from the memorization centroid toward the generalization centroid.Source: §A.2 · p. 12
A.3. Future Work
- Future work calls for unsupervised behavior-vector discovery, simultaneous multi-attribute control, and deeper study of how the IB trade-off is realized in transformer activation geometry.Source: §A.3 · p. 12