Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Generalization or Memorization: Dynamic Decoding for Mode Steering

Xuanming Zhang

TL;DR

The paper addresses the instability of LLMs arising from the coexistence of generalization and memorization. It introduces an Information Bottleneck (IB)–based theoretical framework that treats generalization as an IB-optimal, compressed representation and memorization as an IB-suboptimal, high-entropy representation, connected by a two-stage Dynamic Mode Steering (DMS) procedure. DMS uses a lightweight linear probe to detect memorization reliance and an offline-computed steering vector $v_g$ to adjust activations at a causally identified layer, scaled by a real-time memorization score $m$ and steering strength $ abla$. Experiments on GSM8K, HellaSwag, and TruthfulQA with Llama-3 models show that DMS improves logical consistency and factual accuracy compared to baselines, with ablations confirming the importance of the causal layer and tuned steering strength. These results offer a principled, training-free path to enhance LLM reliability and safety by actively promoting generalization during inference.

Abstract

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.

Generalization or Memorization: Dynamic Decoding for Mode Steering

TL;DR

The paper addresses the instability of LLMs arising from the coexistence of generalization and memorization. It introduces an Information Bottleneck (IB)–based theoretical framework that treats generalization as an IB-optimal, compressed representation and memorization as an IB-suboptimal, high-entropy representation, connected by a two-stage Dynamic Mode Steering (DMS) procedure. DMS uses a lightweight linear probe to detect memorization reliance and an offline-computed steering vector to adjust activations at a causally identified layer, scaled by a real-time memorization score and steering strength . Experiments on GSM8K, HellaSwag, and TruthfulQA with Llama-3 models show that DMS improves logical consistency and factual accuracy compared to baselines, with ablations confirming the importance of the causal layer and tuned steering strength. These results offer a principled, training-free path to enhance LLM reliability and safety by actively promoting generalization during inference.

Abstract

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.
Paper Structure (41 sections, 11 equations, 3 figures, 2 tables)

This paper contains 41 sections, 11 equations, 3 figures, 2 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Generalization, Memorization, and Grokking.
  4. Information-Theoretic Analysis of Neural Networks.
  5. Inference-Time Decoding and Activation Steering.
  6. Theoretical View
  7. The Information Bottleneck Principle
  8. Formalizing Reasoning Modes
  9. Generalization as an IB-Optimal Solution:
  10. Memorization as an IB-Suboptimal Solution:
  11. From Compression to Generalization:
  12. Methodology
  13. Methodology
  14. Stage 1: Mode Identification
  15. Memorization-Eliciting Prompts ($P_M$):
  16. ...and 26 more sections

Figures (3)

  • Figure 1: The information plane illustrating generalization and memorization modes, framing the two reasoning modes as distinct solutions within the IB optimization landscape.
  • Figure 2: Impact of Intervention Layer on DMS Performance. The plot illustrates the Pass@1 accuracy of the DMS algorithm on GSM8K as a function of the transformer layer at which the probing and steering interventions are applied for Llama-3 8B.
  • Figure 3: Impact of Steering Strength ($\alpha$) on DMS Performance.