Steering Information Utility in Key-Value Memory for Language Model Post-Training

TL;DR

This paper introduces InfoSteer, a lightweight post-training method that guides language models to better utilize their stored knowledge, improving performance and interpretability across various models and tasks.

Contribution

The paper proposes InfoSteer, a novel approach that treats FFN layers as key-value memory and encourages their use during post-training, enhancing model performance and interpretability.

Findings

Consistent performance improvements across multiple models and tasks.

Steered models allocate information more efficiently, focusing on meaningful tokens.

Enhanced ability to adapt to in- and out-of-distribution data.

Abstract

Recent advancements in language models (LMs) have marked a shift toward the growing importance of post-training. Yet, post-training approaches such as supervised fine-tuning (SFT) do not guarantee the effective use of knowledge acquired during pretraining. We therefore introduce InfoSteer, a lightweight method that encourages parametric information utilization in LMs during post-training. Specifically, InfoSteer treats the feed-forward network (FFN) layer as associate key-value memory and promotes the use of stored memory vectors via forward-pass interventions or regularization during backpropagation. This simple guidance during post-training phase yields consistent performance improvements across diverse model families -- including Qwen, Gemma and Llama -- spanning 15 downstream tasks in both in-distribution (ID) and out-of-distribution (OOD) evaluations. Beyond performance gains, we also find that steered LMs can adaptively allocate information by placing more emphasis on generating semantically meaningful tokens, while using fewer resources on simple transition ones (e.g., `\texttt{,}' or `\texttt{and}'). Our work underscores that vanilla post-training does not fully exploit the potential gained during pre-training, and that steering LMs in latent representation space offers a promising approach to enhance both performance and interpretability. The code is available at: https://github.com/chili-lab/InfoSteer.