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

Thank you for your constructive review of our manuscript with a positive assessment of our work. We noted your attention to a finer-grained discussion about the key-distribution in FFNs and clarification re IF score. We will your concerns as follows:

About Motivation for Higher-Entropy Key Distribution

Thank you for the great question! The reviewer raises a question about why encouraging more uniform key distributions is beneficial. We clarify several key points:

• Not Uniform Distribution, But Less Extreme Low-Entropy Distribution: We first want to clarify that our motivation for promoting a higher-entropy key distribution is not to achieve a perfectly uniform distribution, as that would indeed be another extreme. Instead, our objective is to introduce a "damp factor" to the extreme low entropy distribution introduced in vanilla SFT. We think the extreme low entropy is harmful and leads to detrimental underutilization of parametric knowledge (only a very small portion of memory vectors are used). In practice, we find after training with infosteer the entropy is close to that of the pre-trained models, while vanilla SFT would have a very low key distribution (H(pretrained LM) ~= H(Info-steered SFT) >> H(vanilla SFT)).
• Learning with Intervention, Not Intervention Alone: Following up the clarification, our main goal is to not leverage the steering-only method to work (re Q2), instead, the steering specifically designed to help or assist post-training methods like SFT (SFT -> Info-steered SFT). We know that the issue of SFT is that it will lead models to mimic the pattern of a small curated post-training datasets so easily overfit. With the help of the steering method, it feels like adding a weighted damp factor to prevent the model blindly fit this small curated dataset pattern but learn in a more “global” loss landscape.

[Empirical Evidence] A straightforward experiment is that vanilla SFT performs worse on OOD tasks, but info-steered SFT generalizes better to OOD situations. The results across three model families on 15 downstream tasks prove the effectiveness of our methods.

We will dedicate a discussion section in the limitation section and main body in the camera-ready version to ensure this is clarified for readers.

About Information Flux Definition and Comparison with Baseline

Thanks for asking this! The IF part is quite interesting, and we would love to discuss it further:

• Metric Definition and Precedent. You are correct that Information Flux is our proposed metric. We define it as $\sum_{l}{H(\hat{k}^{(l)})}$ where $\hat{k}^{(l)}$ is the normalized key coefficient distribution at layer $l$. This measures the diversity of memory vectors engaged during token generation—higher IF score indicates more memory vectors are engaged during the process. This name may be new, but using cumulative entropy across layers is indeed having some precedent literature[2][3].
• Why This Reflects Self-Interpretability. The IF score reveals which tokens require diverse vs. narrowed knowledge access. The case you mention in W2 (syntactic predictability) is precisely the point—for those syntactic tokens, models do not need many knowledge sources (low IF score), so they easily make the prediction. For some hard tokens, the model needs to "think harder" to use more memory vectors to make decisions. We think the IF score reveals how LMs deal with different types of token generation. Though it seems not quite comparable to mech interp methods like SAE/circuits/neurons, we think it provides a different angle to understand what happens in the model itself.

[Significance of IF Score Before/After SFT]

This is a very good question to ask! We indeed have these scores and report the results when noticing a difference. Here are 1K samples from Ultrafeedback and we highlight the significance below.

We believe this represents a clear difference before/after Info-steered SFT, so we report this interesting case study in our main paper. We will add this table to the camera-ready version to quantitatively highlight the difference. Thank you once again for your valuable feedback!

Reference

• [1] Kim, Jiyeon, et al. "Knowledge Entropy Decay during Language Model Pretraining Hinders New Knowledge Acquisition." The Thirteenth International Conference on Learning Representations.
• [2] Ali, Riccardo, et al. "Entropy-lens: The information signature of transformer computations." arXiv preprint arXiv:2502.16570 (2025).
• [3] Chang, Hoyeon, et al. "How do large language models acquire factual knowledge during pretraining?." Proceedings of the 38th International Conference on Neural Information Processing Systems. 2024.

We thank you for your detailed and constructive review of our manuscript! We are also grateful for your specific suggestions, which we believe will significantly strengthen the paper. Specifically, we want to address key feedbacks below:

[Key Concern] About Comparison with Additional Baseline:

We first kindly clarify that InfoSteer is not a replacement for post-training methods like LoRA or ReFT, nor is it a test-time algorithm like ITI or mean difference vector steering. Instead, InfoSteer is a steering technique that can be applied to enhance post-training methods (SFT -> info-steered SFT).

That said, we acknowledge the value of empirical comparisons and have included additional baselines* (LoRA[1], ReFT[2], ITI[3] and Diffmean[4][5]) for 7B scale model in our results:

LoRA and ReFT perform similarly to vanilla SFT, with slight variations depending on the task. ITI and DiffMean (inference-time methods) generally perform worse than training-based approaches. Info-steered SFT methods consistently outperform all baselines, demonstrating the value of information steering during post-training. We will incorporate these baseline into the paper to strengthen its position.

About Arithmetic ID/OOD Experimental Design:

Thank you for mentioning this! We agree this terminology could be clearer. We think if we look at those data even if they are mostly arithmetic, the question narrative and domain may be quite different. But we also agree with reviewers that, esp in the LLM era, the term “distribution” should be in a stricter inspection.

To address the reviewer's concern more directly, we implemented an even stricter evaluation strategy than suggested. We took Llama-3-8B checkpoints specifically trained on OBQA and evaluated them on all five arithmetic datasets to test cross-domain transfer:

This stricter evaluation reveals that vanilla SFT dramatically hurts cross-domain generalization (-6.8%), while adding our steering methods maintain the overall generalization performance. We will add these results into the main body, thank you for your suggestions!

About Interpretation of the Performance

Out of 96 experiments (8 benchmarks × 3 model families × 2 model sizes × 2 intervention methods), we observe consistent gains in 94 results, suggesting that language models may not fully realize their potential from pre-training alone. However, we agree that we should adopt a more humble and scientific tone in our paper. We will revise it to use modest language that objectively reports the gains our method achieves.

Re the failure cases you mentioned, we examined those two benchmarks and found that the entropy of key distributions before steering in FFNs was abnormally higher (around 2.97) for these models compared to others (around 2.03). We believe this is an interesting phenomenon that may be worth investigating further—specifically, whether considering prior key entropy could enable more fine-grained steering. We will add this discussion to the main body of the paper. Thank you for raising this excellent question!

About the Potential Combination of the Method

Yes, we have tried to add two pytorch hooks and work together, in some preliminary experiments we find it showed marginal improvements over the individual method so we do not explore that further.

About Details in Experiment

We summarize all comments related to details about our experiment here.

[Key Distribution Shift] We report the average results across all 15 datasets previously mentioned in this analysis section. The low/medium/high key regions are defined by percentile-based cutoffs: 0-25th, 25th-75th, and 75th-100th percentiles, respectively. Due to page limits, we do not disclose all information in the main body. We will add a detailed description and analysis for this section in the appendix.

[Dataset Picked in Different Domain] Based on the dataset characteristics, we categorize them as follows: Reading Comprehension includes BoolQ (yes/no questions based on passages); Knowledge encompasses ARC-e and ARC-c (science knowledge questions) and OBQA (factual knowledge questions with reference material); Commonsense Reasoning covers PIQA (physical interaction reasoning), SIQA (social interaction reasoning), and HellaSwag (sentence completion requiring commonsense); Math includes all arithmetic datasets: AddSub, MAWPS, MultiArith, SingleEQ, SVAMP, and GSM8K (various forms of arithmetic and math word problems); and Linguistic uses the linguistic portion of WinoGrande (pronoun resolution requiring syntactic understanding). When averaging the score there are some cumulated rounding errors (as you noticed) which leads to confusion of this, we will report the larger precision score to avoid the confusion.

We will add these details into our manuscript to make it clarified for readers.

About Section 6.2 Interpretability Claims

We are glad you find this section interesting (as do we!). You are correct that we should at least dedicate a section to give them more attention, but due to limited page space we can only squeeze this one in. A quick response to your question is: both methods see the phenomenon, but the one using entropy regularization shows the most significant results, and we report these qualitative results in the paper. We also have rigorous quantitative experiments to highlight their statistical significance, would love to include that into the camera-ready for better readability.

We sincerely appreciate your suggestions and will incorporate them to enhance our work further. Thank you once again for your valuable feedback!

Reference

• [1] Hu, Edward J., et al. "LoRA: Low-Rank Adaptation of Large Language Models." International Conference on Learning Representations.
• [2] Wu, Zhengxuan, et al. "ReFT: representation finetuning for language models." Proceedings of the 38th International Conference on Neural Information Processing Systems. 2024.
• [3] Li, Kenneth, et al. "Inference-time intervention: eliciting truthful answers from a language model." Proceedings of the 37th International Conference on Neural Information Processing Systems. 2023.
• [4] Turner, Alexander Matt, et al. "Steering Language Models with Activation Engineering."
• [5] Arditi, Andy, et al. "Refusal in language models is mediated by a single direction." Proceedings of the 38th International Conference on Neural Information Processing Systems. 2024.

Yeah that's a very interesting take if we could combine the two methods to boost the performance further. Also, it would be extremely interesting to see if all those test-time algorithms like ITI could be used to boost post-training (tho we do need to create contrastive pairs as you mentioned). That would be quite beneficial to the whole community.

Thank you again for your constructive feedback and sincere discussion with us!

We appreciate your thoughtful and constructive review of our manuscript. We’ve noted your main concerns regarding this paper mainly focusing on SFT instead of RL and correlation between performance gain and model size. We try to address your concerns as follows:

About Comparison to RL

We want to kindly clarify that infosteer is not a post-training method compared to RL, but a plug-and-play method to steer both RL and SFT those post-training methods. SFT can be seamlessly incorporated as information-oriented SFT (SFT -> info-steered SFT), which is the same as RL (RL -> info-steered RL).

As stated in the limitation section, in this work as an explorer of this steering method in post-training we primarily focus on the SFT. But we also understand that recently RL become a hot topic which people find it is effective on verifiable domain like math, We also add a preliminary experiment to use infosteer on RL with:

---

The RFT training was conducted using GRPO for 500 steps with 8 rollouts per prompt. The training dataset consisted of approximately 46k diverse math reasoning problems from a length-filtered subset of the OpenR1-Math-220K dataset, with each problem equipped with demonstrations generated by DeepSeek-R1.

InfoSteer consistently enhances both SFT (+1.0 to +1.8 points) and RL methods (+3.3 to +4.1 points) across all mathematical reasoning benchmarks, with InfoSteer RFT achieving the best overall performance of 49.6 average score. While these experiments are preliminary and cannot make a conclusive justification before conducting whole pack of rigorous ablation, we believe info-steered RL is a very interesting and promising direction for our future work.

About Performance Differences with Model Size and Experiment on 30B Models

We appreciate your observation on the different performance with different model sizes. We may not have the resources to reproduce all pack of results on 30B models in this short period, but we are glad to offer some of the discussion on this Q here:

• Higher Baseline Performance in Larger Models Make Less Room for Improvement: Larger base models generally exhibit significantly higher baseline performance. The "gap" for improvement through post-training techniques naturally narrows. Thus we think a smaller percentage gain on a very high baseline can represent a substantial absolute improvement in capabilities. For instance, a 5.5% gain on a LLaMA-3-8B model with an already high baseline (e.g., 90.8% on HellaSwag for vanilla SFT vs. 96.3% with steered SFT w. intervention) is still a highly significant and impactful improvement, showcasing InfoSteer's ability to unlock latent potential even from highly optimized models.
• Substantial Gains Remain Consistent: Even with higher baselines, our method continues to deliver meaningful improvements. 48 out of 48 configs (8 benchmarks x 3 8B models x 2 methods) deliver consistent gain. 37 out of 48 configs achieve +1.5% gain. And OOD exp is also done on 8B models, validating our methods’ effectiveness on extended distribution. All results are reported on avg score of three separate runs.

We will elaborate on these enhancements in the revised manuscript to provide readers with a comprehensive understanding of our findings. Thank you once again for your valuable feedback!

Dear Reviewer hS2q,

We appreciate your thoughtful and constructive review of our manuscript. Re your key concerns about infosteer on RL, we have added additional experiments for info-steered RL during this period. Our preliminary results see that info-steered RL achieves even stronger performance compared to info-steered SFT. We have also included additional baselines (see our rebuttal to Reviewer SHNo). We believe your suggestions further strengthens our work's position, and steering RL may be an interesting and valuable direction for future work.

As the rebuttal period draws to a close, we appreciate knowing whether our response adequately addresses your concerns. Your support during this period would mean a lot to us!

Thank you for your recognition of our work! We are glad that you find our approach easy to implement and effective across tasks and pre-trained models.

You've raised two questions regarding the applicability of our approach during pre-training and its particular relevance to post-training, especially related to the size of the post-training data. We think these two questions are inherently connected and we will address them together.

During the Post-training Phase (Small Data):

In the post-training phase we typically involve much smaller, more curated datasets (yes your take is correct in Q2!) aimed at aligning the model with specific downstream tasks or human preferences. Here, the key concern is overfitting caused by SFT on this small dataset. Standard SFT on limited data can cause the model to aggressively mimic the patterns of this curated dataset, potentially leading to an undesirable shift in its key distribution towards even lower entropy as we discussed in Sec 6.1. This can result in the underutilization of the rich, generalized knowledge acquired during pre-training, harming the model's ability to generalize, especially to OOD scenarios.

During the Pre-training Phase (Large Data):

You also raise an interesting question about whether the steering method could be applied during the pre-training phase. Previous research has observed that the entropy of key distributions also tends to decrease over pre-training steps [1]. Additionally, researchers have also noticed the phenomenon of attention matrix collapse or attention sink (low "key" entropy in attn matrix)[2][3]. Many studies aim to resist this phenomenon to enable better use of attention parameters. We think this research shares similar intuition with ours but focuses more on the attention module. We would be glad to view this direction as an interesting future work. Thank you for the great questions!

Reference

• [1] Kim, Jiyeon, et al. "Knowledge Entropy Decay during Language Model Pretraining Hinders New Knowledge Acquisition." The Thirteenth International Conference on Learning Representations.
• [2] Gu, Xiangming, et al. "When Attention Sink Emerges in Language Models: An Empirical View." The Thirteenth International Conference on Learning Representations.
• [3] Voita, Elena, et al. "Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned." 57th Annual Meeting of the Association for Computational Linguistics. ACL Anthology, 2019.