Research map / Large language models / Instruction and fine-tuning
Steering Information Utility in Key-Value Memory for Language Model Post-Training
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.
Proposes InfoSteer, which treats transformer FFN layers as key-value memories and steers key-coefficient distributions during SFT via forward-pass intervention on low-coefficient keys or entropy regularization, improving ID and OOD performance across Qwen, Llama, and Gemma models.
Problem
Post-training with SFT does not guarantee that knowledge acquired in pretraining is effectively used; models may underutilize parametric memory on downstream tasks.
Approach
Interpret FFN layers as soft key-value memories and promote a higher-entropy key-coefficient distribution during SFT, either by boosting the smallest p% key coefficients in the forward pass or by adding an entropy regularizer to the loss, so more memory vectors are engaged.
Key results
- Steered SFT with intervention on Gemma-2-9B raises HellaSwag from 94.8 (vanilla SFT) to 95.7, and the base model achieves 90.1 on HellaSwag, per the paper's text.
- On the GSM8K-trained ID/OOD setup, intervention steering reaches 66.8 ID and 86.6 OOD versus 65.7 ID and 83.7 OOD for vanilla SFT.
- Intervention with p=1, alpha=2 gives the best average accuracy of 75.5 in the magnitude ablation, compared with 72.6 for vanilla SFT.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use intervention-based InfoSteer instead of vanilla SFT | Gemma-2-9B · HellaSwag | accuracy: 94.8 -> 95.7 | Helped |
| Use intervention-based InfoSteer instead of vanilla SFT | Qwen-2.5-7B · ARC-c | accuracy: 69.6 -> 73.7 | Helped |
| Use intervention-based InfoSteer instead of vanilla SFT | arithmetic reasoning post-training · Arithmetic OOD (average across five benchmarks) | average score: 83.7 -> 86.6 | Helped |
| Use vanilla SFT instead of the base model | arithmetic reasoning post-training · Arithmetic OOD (average across five benchmarks) | average score: 85.3 -> 83.7 | Hurt |
| Increase intervention strength from α=1 to α=2 at p=1 (ablation) | InfoSteer intervention · Average Acc (steering magnitude experiment) | accuracy: 73.8 -> 75.5 | Helped |
| Increase intervention proportion to p=2 and strength to α=5 (ablation) | InfoSteer intervention · Average Acc (steering magnitude experiment) | accuracy: 73.8 -> 72.8 | Hurt |
| Set entropy regularization strength to λ=0.05 instead of λ=0.01 (ablation) | InfoSteer regularization · Average Acc (steering magnitude experiment) | accuracy: 73.4 -> 74.7 | Helped |
| Reverse entropy regularization to λ=-0.01 (ablation) | InfoSteer regularization · Average Acc (steering magnitude experiment) | accuracy: 73.4 -> 72.3 | Hurt |
| Apply intervention only to middle layers 9-16 rather than all layers (ablation) | InfoSteer intervention · Average Acc (layerwise) | accuracy: 72.6 -> 75.2 | Helped |
| Apply intervention only to late layers 17-24 rather than all layers (ablation) | InfoSteer intervention · Average Acc (layerwise) | accuracy: 72.6 -> 72.9 | No clear effect |
| Apply intervention to early and middle layers 1-16 rather than all layers (ablation) | InfoSteer intervention · Average Acc (layerwise) | accuracy: 72.6 -> 76.3 | Helped |
| Use InfoSteer intervention instead of vanilla SFT | task-type evaluation · Linguistic | score: 77.9 -> 78.0 | Hurt |
Surprising findings
- Vanilla SFT lowers OOD arithmetic accuracy from 85.3 to 83.7 while improving ID accuracy, suggesting overfitting.
- Negative entropy regularization (lambda = -0.01) degrades performance below vanilla SFT, consistent with the claim that engaging more memory vectors helps.
Limitations
- Evaluation limited to standard SFT; steered RL not explored.
- Fine-grained steering methods gave only marginal gains over the generic method.
- Layer and hyperparameter choices (p, alpha, lambda) are tuned on evaluation settings with limited analysis of their sensitivity.
Experimental setup
Datasets: BoolQ PIQA SIQA HellaSwag WinoGrande GSM8K ARC-e ARC-c OBQA AddSub SingleEQ MultiArith
Benchmarks: BoolQ PIQA SIQA HellaSwag WinoGrande ARC-e ARC-c OBQA AddSub MAWPS MultiArith SingleEQ
Metrics: accuracy average accuracy
Baselines: vanilla SFT base model
Builds on: transformer feed-forward layers as key-value memories supervised fine-tuning
Closest papers by meaning
- Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory (NeurIPS 2026 · similarity 0.89)
- Language Model Memory and Memory Models for Language (NeurIPS 2026 · similarity 0.88)
- Once Read is Enough: Domain-specific Pretraining-free Language Models with Cluster-guided Sparse Experts for Long-tail Domain Knowledge (NeurIPS 2024 · similarity 0.88)
- Mitigating Memorization in Language Models (ICLR 2025 · similarity 0.87)
- Improving Instruction-Following in Language Models through Activation Steering (ICLR 2025 · similarity 0.87)
- Generative Adapter: Contextualizing Language Models in Parameters with A Single Forward Pass (ICLR 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.