Language models behavior calibrated through hidden state corrections
HeadEdit: Calibrating Language Model Behavior Through the Frozen Unembedding Matrix
Computation and Language
Summary
Language models sometimes act in ways that aren’t desired, like refusing simple requests or making mistakes, even after efforts to align them. The authors noticed that the information needed to fix these behaviors can still be found inside the model’s internal representations, even when the output looks wrong. They created a method called HeadEdit that adjusts how the model’s output is generated by changing the way internal signals are decoded, without changing the model’s parameters or needing manual fixes. This method works across different tasks and models, improves behavior, and can be reused after further training.
What this means in practice
- •For ai model developers: Calibrate and improve language model responses without retraining or manual output tweaks.
- •For chatbot designers: Adjust chatbot output behaviors dynamically by modifying internal decoding representations rather than model weights.
Authors
Zirui He, Haiyan Zhao, Jingyu Hu, Yinghao Wu, Chenxi Yuan, Yingcong Li, Yandong Bai, Mengnan Du
Abstract
Alignment does not eliminate behavioral errors in language models. Models may still refuse benign requests, call unnecessary tools, or yield to false user claims. Current methods mitigate such errors as a computation problem, and rarely explore if the desired behavior is already encoded in the model's representation. Motivated by the observation that behavior-relevant information remains linearly decodable from the final hidden state even when the resulting logits produce the undesired behavior, we introduce HeadEdit, a gradient-free method that calibrates model behavior through the unembedding matrix. HeadEdit extracts a low-rank behavioral subspace from paired completions and uses each prompt's coordinates within it to generate a vocabulary-wide correction, thereby implementing implicitly adaptive steering without manually specified target tokens or parameter updates. HeadEdit improves all nine experimental settings across three tasks and three model families, with negligible inference overhead and no systematic loss of general capabilities. It also reveals a connection to gradient-based alignment. HeadEdit's low-dimensional representation partly predicts how preference tuning changes output logits on unseen prompts. The subspace learned from the model can also be reused after tuning, improving performance without re-extracting or retuning. These results show that HeadEdit provides a practical, lightweight, and interpretable way to calibrate model behavior through the unembedding matrix.