Steering Information Utility in Key-Value Memory for Language Model Post-Training
Authors: Chunyuan Deng
Ruidi Chang
Hanjie Chen
Venue: NeurIPS 2025 poster ()
Keywords: Interpretability
Intervention
Model Steering
Language Models
Post-Training
Primary area: deep_learning
TL;DR: We introduce a method that encourages LMs to leverage their pretrained knowledge during post-training.

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., ',' or '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.

Steering Information Utility in Key- Value Memory for Language Model Post-Training
Chunyuan Deng
Ruidi Chang
Hanjie Chen
Dept. of Computer Science
Dept. of Computer Science Rice University Houston, TX 77005 ruidi cww....raarric.... edu
Dept. of Computer Science Rice University Houston, TX 77005
Rice University Houston, TX 77005
Shunyuan.den@@rice
hanjie@rice edu
Abstract
Recent advancements language models (LMs) have marked a shift toward the growing importance of post-training. Yet, post-training approaches such as super- vised fine-tuning (SFT) do not guarantee effective use of knowledge acquired during pretraining. We therefore introduce InfoSteer, lightweight method that encourages parametric information utilization LMS during post-training. Specifi- cally, InfoSteer treats the feed-forward network (FFN) layer as associate key-value memory and promotes the use of stored memory vectors via forward-pass inter- ventions regularization during backpropagation. This simple guidance during post-training phase yields consistent performance improvements across diverse model families incluuiinn Qwen, Gemma and Llama-spaman spanning 15 downstream tasks in both in-distribution (ID) and -iiitii-iittitttionnn (OOD) evaluations. Be- yond performance gains, also find that steered LMs can adaptively allocate information by placing more emphasis on generating semantically meaningful tokens, while using fewer resources simple transition ones (e.g., ',' or '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.
Introduction
The contemporary training pipeline for LMs has standardized around two-stage process: initial pre-training phase on extensive, web-scale corpora, followed by post-training phase utilizing smaller, more curated datasets [Ouyang al., 2022, Touvron et al., 2023, Bai et al., 2023, Mesnard al., 2024]. A considerable body of literature suggests that the fundamental capabilities and knowledge of these models are predominantly instilled during the pre-training stage [Chung al., 2022. Anil et al.. 2023, Muennighoff et al., 2025]. Subsequent post-training techniques are commonly viewed as approaches better refine, elicit, or adapt these inherent capabilities embedded in the base model [Zhou et al., 2023, Rafailov et al., 2024, Guo et al., 2025, Swamy et al., 2025].
Despite this, open question remains: do post-training methods truly encourage the model fully utilize the information encoded during pre-training? In many cases, they may not-since models are neither explicitly trained nor incentivized to retrieve and apply such knowledge" optimally [Conmy et al., 2023, Chang et al., 2024, Du et al., 2024a]. This insufficient use could result in suboptimal performance on downstream tasks, even when relevant knowledge is already stored in the model's parameters [Bietti et al., 2023, Kim et al., 2025].
'The code is available at: https://githen cww coo//chin htthchhhiliiii 1 i- 1abb InfoSteer.. In this work, the 'information' and "knowledge' are interchangeably used as descriptive language.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Phase 1: Large-scale Pretreining
Phase 2: Post-Training
Feed-
farged night dream piispared right
Network
Task
X
Down
LL
2 H(k)
ccethceent
Up
TB-accal Carpux
Base Model
Information Compression
Transformer FFN
Key-Value Memory
Forward-pess Intervention
1.2
Regularization
Backward pass
Figure I: Overview of our proposed InfoSteer framework. The interpretation of Transformer FFNs as key-value memory was introduced by Geva et al. [2021], further details are provided in $3.
We therefore introduce InfoSteer, method designed encourage trained knowledge during the post-training phase (Figure 1). Our approach draws on an associative memory view of the Transformer's feed-forward network (FFN) layers [Geva al., 2021, Meng et al., 2022a,b, Geva et al., 2023], where the first FFN layer acts as content-dependent key and the second FFN layer functions Consequently, each layer's 's output can be viewed as a weighted sum of memory vectors, where each vector's weight (the key coefficient) produced by the input vector multiplied by its corresponding key in the first FFN.
effective use of model's pre-
learned value memory, together forming key-value memory mechanism.
InfoSteer controls the distribution of key coefficients manage the information-dent memory vectors learned from pretraining. The goal is to promote a high-entropy key distribution, enabling their corresponding memory vectors be actively engaged during post-training. Specifically, two strategies: (i) intervening on memory vectors with low corresponding key coefficients during the forward pass, and (ii) regularizing the entropy of key coefficients in the gradient flow during backpropagation. These methods can be seamlessly incorporated into a vanilla SFT pipeline for information-sterrt SFT.
Our empirical evaluation demonstrates that InfoSteer yields consistent and notable performance enhancements across variety of condels-inclludnng Llama [Touvron et al., 2023]. Qwen [Bai et al., 2023], and Gemma [Mesnard et al., 2024] at different scales. These improvements are demonstrated across comprehensive suite of over 15 downstream tasks, encompassing both in-distribution (ID) and out-of cattoo-ddisrrbbttion (OOD) scenarios. Beyond these performance gains, information-steered LMs exhibit adaptive information allocation strategy during token generation. Specifically, they dedicate more representational capacity to semantically rich and challenging tokens, while expending less simpler, transitional tokens, indicating a more nuanced use of parametric information.
2 Related Work
Model Steering. Model steering represents an emerging paradigm to guide model behavior, focus- ing on controlled interventions in the model's latent space [Subramani et al., 2022, Zou et al., 2023, Turner et al., 2024, Li et al., 2024, Cyberey and Evans, 2025]. Many work has demonstrated the effec- tiveness of steering methods across various dimensions, including disentangling human-interpretable concepts [Rumelhart et al.. 1986]. such as linguistic features (c.g., gender, number) [Hewitt and Manning, 2019, Lasri al., 2022, Wang al., 2022, Hanna et al., 2023, Huang et al., 2024, Chang et al., 2025] and logical reasoning [Wu et al., 2024b, Xie et al., 2025]. These approaches typically operate by introducing auxiliary objectives, modifying activation patterns, or implementing controlled perturbations during forward or backward passes [Rimsky et al., 2023, Scalena et al., 2024, Stolfo al., 2024, Luo al., 2025, Gur-Arieh et al., 2025, Bartoszeze et al., 2025]. The key advantage of model steering lies in its ability to leverage the rich knowledge already encoded in pretrained weights while directing how this information is accessed and applied downstream tasks [Geiger et al., 2023, Lee et al., 2024, Siddique et al., 2025, Soo et al., 2025, Deng et al., 2025]. Our information-ssteerd approach extends this paradigm by targeting the key distribution in transformer FFN layers, which influences which memory vectors activated during computation.
Parametric Information. From an information theoretic -hoooati perspective, quantifying the information encoded in an LM's high-dimensional parameter space remains a challenging problem [Achille et al.,
In this work. we term the output of first FFN as key coefficient and its distribution as key distribution. This is the same "memory coefficient" in Geva et al. [2021]

---

2020, Bernstein and Yue, 2021]. This challenge is closely tied to the ongoing discourse on parametric versus non--aran-eeeee con-aammetiic knowledge storage in LMs [Tay al., 2022, Ferrando et al., 2022, 2023, Xie et al., 2023, Deng et al., 2024, Ferrando and Voita, 2024, Du et al., 2024b] (or token-mixing in attention layer vs. channel mixing in FFN layer). The concept of channel mixing stems from the perspective of viewing the FFN layer
key-value memory structure [Weston et al., 2014, Geva et al., 2021, Qiu et al., 2024, Kim et al., 2025]. Empirical evidence supports this interpretation. showing that individual neurons in FFN layers activate in response to specific semantic concepts and linguistic patterns [Dai et al., 2021, Liu et al., 2023, Niu et al., 2024]. Furthermore, intervention studies have demonstrated that targeted modifications these memory vectors directly influence model outputs on knowledge wowwlegg-intensive intensive tasks [Meng et al., 2022a.b, Geva et al., 2023, Hase et al., 2023, Yao et al., 2024, Wang et al., 2024], suggesting that factual information is strongly correlated with in this component.
Transformer Feed-Forward Layers as Key- Value Memories
We consider a standard autoregressive Transformer [Vaswani et al., 2017]. which models the condi- tional probability of a sequence x = (I1,...,,77) as
p(z)= = II p(ra x<)); 1=
using a stack of L Transformer decoder blocks. Each block consists of two primary components: a masked self-attention layer and position-wise feed-forward network (FFN).
In standard Transformer architectures, position-wise FFN in each decoder block is given by:
FFN(X) oXW up + b2) W,)Wdown down + b2.
where E RTxx is the input representation matrix for FFN, Wap E Rd and W Woown E Rm xd the up-projection and down-pree projection matrix respectively, and a is a nonlinearity activation function. Bias term b and by are included in the standard FFN formulation but omitted in the following equations for simplicity.
Let h E Rdd represent the embedding as a single token representation from the input matrix X. As illustrated in Figure 2, this structure can be interpreted as a soft key-value memory [Weston et al., 2014, Geva al., 2021], where the intermediate activation a (()))) E Rd defines the soft addressing weights, and each row of W serves value vector. The FFN output of the input h can be written as a weighted sum over value rows:
din
FFN = Σ
((hWop)s
(W (Wdowa).:
= 1i + K2 +
kid Vdd.
key coeflicient
value
where each k; E R is scalar key coefficient, and each v, Id is a value vector (memory vector). The key coefficients control the magnitude of the memory vectors' contribution the final prediction.
4 InfoSteer
LMs are commonly interpreted to store parametric knowledge within their dense layers, particularly in the FFN blocks [Meng et al., 2022a, Geva et al., 2023, Kim et al., 2025]. However, during post-training, there is no explicit guidance to encourage the model to utilize this knowledge for new alignment retrieval tasks. Therefore, propose InfoSteer to bridge this gap.
4.1 Motivation & Desierata
To steer the information utility of pretrained LMs, we aim to enhance the engagement memory vectors during post-training The core idea is to control the distribution of key coefficients in the position-wise FFN (see Figure 2). Specifically, we observe dominant property of associative memory: if a key coefficient k; is significantly larger than another key coefficient ky. the final prediction predominantly relies on the corresponding memory vector v., while v, underutilized or overlooked. Our goal is to achieve two objectives:
3

---

weighted
Outppt
(b) Forward-pass Intervention
Step Find the p& smallest key coefficient
Ouga2. X
va
0.006 x
0.077. X
Step 2: Steering the corresponding key-oulue pair to ang. key value.
(c) Backpropagation Regularization
Iiz
Ic
Step.:: Calculate the entropy of coeffeient II Step 2: Feeding the entropy as additional regularization term post nestttroining
Ingatt
(a) Transformer FFN as key-vatue memory
Key Coeffeient Dittrbation
Figure 2: Key Design of InfoSteer. (a) An illustration of viewing the Transformer FFN as key-value memory. The key acts as a control mechanism that determines the extent to which each memory vector is engaged. (b) and (c) Two general methods used to modulate the distribution of key coefficients, thereby encouraging the engagement of memory vectors during post-training.
Minimize the language modeling loss: Ensure the FFN output FFN(h) = Emm ky contributes accurate predictions.
2. Maximize the entropy of the key distribution: Encourage diverse engagement of memory vectors by promoting a higher-entropy key coefficient distribution.
Promoting entropy weaken the dominance of individual coefficients, encouraging more balanced activation of memory vectors v, and thus enabling richer parametric knowledge retrieval.
4.2 Generic Methods:
Intervention and Regularization
Formally, let k() = [R(C) 200
( E Rd denote the key coefficients at layer 1, where k; (hh ap)s and W E Rdxdm is the up-projection matrix. We propose two complementary methods increase memory vector engagement: intervention and regularization.
Intervention Method. The intervention method directly modifies key coefficients to promote broader memory vector utilization. For each layer 1, we identify the p% of key coefficients with the smallest value and adjust them to value proportional to the layer's average key coefficient.
Formally, let I() denote the set of indices corresponding to the p% of elements in k!)) with smallest value. For a hyperparameter 0. we update:
t-
for e (
(4)
This adjustment ensures that previously minor coefficients contribute more significantly to the FFN output FFN(h) = Eim k, a is controlling scalar to set the magnitude of steering.
dm
The regularization method encourages a uniform key coefficient distri- bution by adding entropy term to the loss function. Higher entropy reduces the dominance of individual coefficients, engaging more memory vectors. The modified loss function, across layers, is:
Regularization Method.
C= A H(K)))
(5)
where (()) modeling loss, and
5m gay is the normalized version of original key distribution k CLM is the language
(t)
controls the regularization strength.
Both methods modulate the key coefficient distribution to improve the model's ability to leverage parametric knowledge during post-training tasks.

---

4.3 Fine-Grained Steering of Memory Vectors
While the generic methods modulate the overall key distribution, more targeted approach requires first understanding what information each memory vector encodes. The information surrogate @; provides a direct, formalized method for this characterization.
Let 2 E RR be the i-th memory vector (i.e., the i-th row of Wdown) and Wdccode E R be the language model's decoding head, where [V| is the vocabulary size. The information surrogate d, R VV for vector 0y is the resulting logit vector:
Φ, = U, Wdccode-
(6)
This surrogate o; acts as a "semantic fingerprint" for Us. We propose the following algorithm to characterize U, by analyzing the properties of this surrogate logit vector.
1. Surrogate Normalization. First, we compute the full probability distribution P. E RRI over the vocabulary by applying the Softmax function the logit vector Φi, P. Softmax(o:). where the j-th element of P., Pi.j Ex exp(ó..))
exp(d..))
2. Specificity Analysis. We quantify the vector's "specificity" by calculating the entropy HPP)) of its distribution. This measures the concentration of the distribution I (P:) E Pi,j log Pi.j low (P:) indicates a high-speciffcity vector, as its probability mass is concentrated on narrow set of tokens (representing specialized knowledge) and vice versa.
Semantic Concept Identification. For vectors identified as specialized (e.g.. H(P.) for some entropy threshold r), we identify their semantic focus. We find the set of indices I; corresponding to the K largest probabilities in P::
I arg top-K (Pt.5). iVI
The semantic concept, T, the set of vocabulary tokens corresponding to these indices:
T {token(j) j E I.}.
This set T7 (e.g., 'quantum', 'physics', " 'superposition' }) reveals the specific topic or con- cept encoded by the memory vector 2;
The results of this characterization- such the entropy H(P,) and the semantic concept set T; for each vector Vy-can then be used to inform a more targeted intervention or regularization strategy. For example, one could selectively amplify key coefficients corresponding to vectors v; whose T; matches a desired topic.* 4
4.4
MLP Variants
Modern LM architectures may include MLP variants that differ from the typical FFN. For example, when adapting methods for MLP variants like SwiGLU, used models such as Qwen [Bai et al., 2023], LLaMA [Touvron et al., 2023], and Gemma [Mesnard al., 2024], we focus on the input to the down projection. SwiGLU modifies the standard FFN as:
SwiGLU(h) = (o(hW gate) C (hWup) ) Wdowa: down:
(9)
where a is the SiLU activation, and @ denotes element-wise multiplication.
Similar to a standard two-layer FFN, we define the key coefficient distribution as k = (hW gate) @ (hWw.) The core idea for identifying this key distribution that it represents the input just before associates with the memory vectors (i.e., Wdnwn. Regardless of whether gated function or another method is used to derive it, our focus remains on this input k, as is what directly manipulates the memory vectors.
'In this paper, primarily focus on generic method, which provides substantial performance gains. While fine-graine mne-grained steering is a promising approach for encouraging LMs to utilize specific knowledge during post-training, our investigation found it offered only marginal improvements over the generic method. Therefore, we detail these fine grained methods in Appendix B and provide wise ablation study in Appendix C, leaving more granular control of knowledge steering direction for future research.
5

---

Experiment
To evaluate the overall effectiveness of steered SFT against vanilla SFT, we conduct a comprehensive study across various downstream tasks and model sizes. Our primary focus is achieving holistic understanding of the overall model behavior after applying our methods, with particular attention paid to generalization performance. The details are provided below.
5.1
Experiment Setup
Base Models. We evaluated our methods on language models of varying sizes across three series:
Qwen-2.5-1.5B and Qwen-2.5-7B for the Qwen series [Bai et al., 2023], LLaMA-3.2-IB and LLaMA-3-8B for the LLaMA [TTuurron et al., 2023] series, and Gemma-2-2B and Gemma-2-9B for the Gemma [Mesnard et al., 2024] series.
Datasets. We evaluate performance under two settings: in-distribution (ID) and out-of-distribution (OOD) evaluations.
In the ID setting, consider diverse of downstream tasks ranging from knowledge-intenssi crasoning-intensiinn These include BoolQ [Clark et al., 2019), PIQA [Bisk et al., 2019) 91, SIQA [Sap al., 2019], HellaSwag [Zellers al., 2019]. WinoGrande (Sakaguchi et al., 2019], GSM8K [Cobbe et al., 20211, 1. ARC-e, ARC-c [Clark et al., 2018], and OBQA [Mihaylov et al., 2018J. No chain-of- thought (CoT) [Wei et al., 2023] rationales are provided.
For OOD setting, we primarly trained GSM8K and eval on other arithmetic datasets, including AddSub [Hosseini et al., 2014], SingleEQ Koncel-Kedzione et al., 2015], MultiArith [Roy and Roth, 2016], AQuA [Ling et al., 2017], MAWPS [Koncel-Kedziorsk et al., 2016], and SVAMP [Patel al., 2021]. In these benchmarks, chain-of-thougth (CoT) rationales are typically included before the final answer: For all benchmarks, we use the same prompt templates as in Hu et al. [2023], Wu al. [2024a]. We also remove any leading and trailing whitespace from the dataset.
Baseline. We use standard SFT as our baseline. For information--teerre SFT, evaluate inter- vention and regularization methods, which share the goal of enhancing memory engagement during training. The default hyperparameters for the intervention set p% 0.01 and a 1, with 0.01 for entropy regularization. Other training details are provided in Appendix A.
5.2 General Performance Comparison
We use the default hyperparam setting as reported in $ 5.1 to eval general performance difference. Table 1 presents the accuracy across nine commonsense reasoning datasets for Qwen, LLaMA, and Gemma models under various different steering strategies. We categorize the models into small-scale (I-2B parameters) and large-scale (7-9B parameters) groups.
Benefit of Model Steering. Our experimental results, as presented in Table 1, provide compelling evidence for the efficacy of model steering techniques across different model architectures and parameter scales. Both proposed steering methods-steered SFT with intervention and steered SFT with regularization-coniiseei outperform their respective base models and vanilla SFT counterparts across all benchmarks. This pattern holds across all model families, demonstrating the robust, architecture-agnosic nature our steering approaches.
Insufficient Utilization of Pretrained Knowledge in Modern LLMs. The benefits of steering reveal a critical limitation in contemporary large language models: they substantially underutilize the knowledge acquired during pretraining when applied to downstream tasks. This inefficiency is strikingly evident across all three widely-used model families examined- Qwen. Specifically, the Gemma-2-98 base model achieves 90.1% on HellaSwag, yet reaches 95.74 with intervention steering-a 5.6% improvement without additional pretraining or parameter increase. Similarly, LLaMA-3-S0 shows remarkable 5.5% gain on HellaSwag with intervention steering over the base model. This pattern is consistent across all architectures and benchmarks, suggesting that modern L.MSs possess substantially more capabilities than they successfully deploy standard fine-tuning regimes.
LLaMA, and

---

Table 1: Performance comparison of Qwen, LLaMA, and Gemma models with different training methods on eight datasets. To highlight improvements, we use blue to highlight significant gains. All results reported as the average scores three independent runs.
Accuracy (1)
Training
BoolQ
PIQA
SIQA
HellaS.
WinoG.
ARC-e
ARC-
OBQA
Small-Scale Models (I-2B parameters)
base model
64.2
74.3
80.1
76.9
61.2
vanilla SFT
68.5
82.9
79.6
84.8
80.8
81.4
65.8
81.0
Qwen 2.5. 1,5B
steered w. intervention
69.3
80.3
93.1
steered SFT w. regularization
68.7
83.9
79.8
92.4
83.7
82.5
67.5
78.1
buse mudel
65.6
75.3
74.2
78.9
77.8
74.5
60.1
76.3
vanälla SFT
69.8
79.9
79.5
83.6
82.6
79.8
64.7
81.0
LLMMA-32-IB
steered SFT w. intervention
71.8
83.7
76.0
89.1
82.6
83.7
68.2
82.4
SFT regularization
71.0
82.9
75.2
88.3
82.6
67.4
base model
66.5
79.1
73.8
82.7
78.9
77.4
63.8
74,9
Gemma-2-2B
vanilla SPT
70.2
83.4
78.1
87.5
83.3
82.7
68.4
80.1
steered SFT w. Intervention
72.5
85.6
79.3
90.2
85.8
85.3
71.9
83.7
steered SFT w. regularization
71.8
84.9
78.5
89.3
85.0
84.7
71.0
82.9
Large-Scale Models (7-9B parameters)
base model
68.9
81.2
77.6
87.5
80,3
79.8
65.1
77.6
Qwen-2.5-7B vanilla SFT steered SFT intervention
72.4
84.9
81.5
92.4
84.2
84.2
69.6
82.8
74.1
86.3
81.8
95.1
87.2
86.2
73.7
steered SFT w. regularizaiion
76.4
85.7
81.0
94.3
86.5
85.4
72.9
83.4
hase
70.3
75.7
90.8
86.2
LLaMA-3-SB
vanilla SPT steered SFT w. Intervention
74.6
89.3
79.9
93.5
85.6
90.5
80.4
85.8
77.1
90.2
82.0
96,3
87.4
92.4
81.6
87.5
steered SFT w. regularization
76.5
89.5
81.2
95.6
86.8
91.7
80.9
86.8
base model
71.6
86.3
77.2
90.1
82.5
87.5
77.8
81.7
Gemma-2-9B
vanilla SPT
74.3
90.1
81.7
94.8
86.9
91.7
82.0
86.4
steered SFI w. intervention
77.2
91.8
83.1
95.7
88.5
93.5
83.4
88.2
regularizution
94.9
5.3 Ablations of Stecring Magnitude
Table summarizes the performance under various steering magnitudes and strategies. Incorporating our proposed steered SFT ap- proach with intervention consistently en- hances performance, with the best result of 75.5% achieved at intervention param- eters = 1, a = 2. And excessive in- tervention magnitude may lead to subopti- mal performance (72.8%), suggesting a bal- ance necessary for optimal results. Ad-
Table 2: Results under varying steering magnitudes. p% determines the proportion of keys being intervened, controls the intervention strength, and sets the regu- larization strength for the entropy.
Model
Average Acc
Base Model
71.4
Vanilla SFT
72.6
steered SFT w/ interv. (p 1,a 1)
73.8
+ steered SFT w/ interv (p =, 1,a=
75.5
steered SFT w/ interv (p 2.a
72.8
ditionally, entropy-based regularization (A)
steered SFT reg. (A -0.01)
72.3
also demonstrates effectiveness in steering model performance. Positive entropy reg-
steered
reg. (A 0.01)
73.4
steered SFT reg. (A 0.05)
74.7
ularization (A 0.05) significantly boosts accuracy to 74.7%, whereas negative regularization
(A = -.001,, which actively discourages information utility, leads to degraded performance, result- ing an accuracy of 72.3%.
5.4 Effectiveness of Model Steering Across Task Types
Table 3 presents performance improvements across five task types when applying different model steering strategies.
Both steered SFT variants outperform the base model and vanilla SFT. On average, these methods yield +3.9 improvement in reading comprehension, +2.3 in commonsense reasoning, +1.1 in math, and +3.3 in knowledge tasks. However, linguistic tasks show a slight performance drop (-0.3). These results highlight the effectiveness of targeted model steering, particularly for knowledge- intensive and reasoning-hean tasks, while linguistic tasks appear to benefit less from encouraging increased parametric knowledge utilization.
7

---

Table 3: Gain of Model Steering w/ Different Task Type.
Model
Reading Comp.
Knowledge
Commonsense Reasoning
Math
Linguistic
Base Model
72.3
70.1
65.4
63.7
78.2
Vanilla SFT
73.8
70.6
66.1
65.7
77.9
Stecred SFT w/ interv.
78.1
74.9
68.5
66.8
78.0
steered SFT w/ reg.
77.4
73.4
67.9
66.1
77.3
Average A
+3.9 (1)
43.3
+2.3
+1.1
-0.3
5.5
Out-of-Distribution Evaluation
We trained on the GSM8K dataset assess in- distribution (ID) performance and are evaluated on five arithmetic auttof-distribution (OOD) datasets-AddSte MAWPS, MultiArith, Sin- gleEq, and SVAMP-to evaluate generalization. This setup enables comprehensive study of how different fine-tuning strategies transfer to novel problem structures.
Table 4: Results for ID and OOD evals. The OOD performance is reported as average score across five benchmarks.
Model
ID Eval
OOD Eval
Base Model
63.7
85.3
Vanilla SFT
65.7 (+2.0)
83.7
Stecred SFT w/ interv.
66.8 (+3.1)
86,6 (+1.3)
steered SFT w/ reg.
66.1 (42.4)
86.0 (+0.7)
Vanilla SFT Improves ID but Harms OOD.
As shown Table Vanilla SFT improves ID accuracy from 63.7 to 65.7, demonstrating its ability to fit the training distribution. However, this comes the cost of generalization: its OOD performance drops from 85.3 to 83.7. This suggests that naive SFT may lead to overfitting and reduce robustness on unseen arithmetic tasks.
OOD Model Performance Comparison
1.00
Steered SFT Enhances Both ID and OOD Performance. In contrast, our proposed Information-steered SFT improve both ID and OOD scores. Specifically, the snrrvention-bbsed steering reaches the highest ID score of 66.8 and the best OOD score of 86.6 The regularization-bant -aasd steer- ing also yields consistent improve- ments (66.1 ID, +2.4; 86.0 OOD, +0.7). These results highlight that con- trolled interventions and regulariza- tions during post-training can enhance generalization for LMs.
Model Variant
Vanilla SFT
Steered SPT Interv.
0.95
Steored SPT w. Reg.
Accuracy
0.90
0.85
0.80
Adds
MultiArith SingleEQ
Average
OOD Datasets
Figure 3: OOD Model Performance Comparison across dif- ferent mathematical reasoning datasets. Results are reported w/ average score of three separate runs.
Fine-Grained Analysis Across OOD
Datasets. Figure further illustrates performance across individual OOD
datasets. Notably, Steered SFT w/ intervention achieves the best or comparable performance nearly all settings. On the most challenging dataset, SVAMP, Vanilla SFT performs the worst, while both steered methods significantly improve accuracy. This demonstrates the robustness of InfoSteer on structurally diverse arithmetic problems.
Analysis
In this section, we address two key questions. First, compared to vanilla SFT, what kind of distribution shift does our method introduce in the key-coefficient distribution? Second, beyond improvements in performance, what other changes can be observed in the model's behavior? These analyses provide deeper insights into both the effectiveness and interpretability of our steering strategy.

---

analyze the 6.1 Distribution Shift after Steering
We analysis the low/mediun/high key regions defined by percentile- based cutoffs before/after SFT: 0 - 25th, 25th-75th, and 75th- 100th percentiles. As shown in Figure 4, fine-tuning via standard SFT increases the ratio of activations in the low-key region, diverging from the distribution observed in the base model. This shift indicates that vanilla SFT encourages the model to rely on fewer memory vectors compared to base model, potentially increasing the risk of overfitting to downstream tasks (as overfitting is observed in our previous experiment in $ 5.5).
Low Key Region Change
Medium Key Region Change
High Key Region Change
0.3
0.05
0.05
0.00
0.1
0.0
-0.10
-0.1
-0.10
SFT
Reg.
w/ Interv.
SFT
Reg. Interv.
SFT
Reg.
Interv.
Figure 4: Key Coefficient Distribution Shift from Base Model. The gray line serve baseline for the number key in corresponding region.
Compared to vanilla SFT, which predominantly boosts activations in the low-key region, both steering methods adjust the distribution in more structured and nuanced ways. Regularization-basee steering softens key allocation by reducing activations at the extremes both low and high- moderately increasing the usage of medium-range keys. In contrast, intervention-ban based steering results sharper redistribution, actively suppressing low-region keys and promoting greater use of medium and high-region keys. This targeted reshaping suggests that intervention-base steering, in particular, fosters more efficient and selective memory utilization within the model.
6.2 Sell-Steering Leads to Better Interpretablity
In addition performance improvements, we're intrigued by how InfoSteer can enhance model interpretability. Recently, there has been growing interest in understanding the mechanisms to dissect the inner workings of LMs. For example, approaches like sparse autoencoders (SAEs) [Cunningham et al., 2023] involve training external NN to identify co-activation patterns in the model's 's latent space. While this line of research promising, we aim explore a different perspective: that LMs might be more self- interpreabeee via proper guidance without extra training.
To investigate this, first performed instruction tuning et al., 2023], a dataset comprising 6.4K high-quality instruction-followiig examples specifically designed to align model outputs with nanaand instruction styles. Our goal model's behavior after alignment, particularly under steering interventions. To this end, we computed token-level entropy
using UltraFeedback [Cui
the key coefficient distributions are activated during token generation. We refer to this metric as the Information Flux (IF) score, which quantifies the amount of information the model needs to
proxy for how many memory vectors
generate each token.
Table 5 presents qualitative examples of token-level IF score across four prompt types. We observe meaningful patterns in both the highlighted (high-IF) and underutilized (low-IF) tokens.
Highlight Tokens Reflect Instruction Semantics. The highlighted tokens in each response align closely with the semantic core of the instruction:
For Explain Concepts, the model concentrates high IF the definitional span- span-"quantuu "quantum computing is a type of computing that uses quantum mechanical phenom- ena.. whwch directly addresses the instruction.
In Creative Writing, poetic and original terms such as "circuits," "code," and "digita tag symphony" are highlighted. These tokens carry the stylistic and imaginative weight of the response, indicating the model recognizes them as content-bearing
For Arithmetic Calculation, the most salient tokens are in the initial planning phrase-"To calculate (31 x 31), you use the formula... "-which frames the reason-

---

Table 5: Token-wise Information Utility (IF) across prompts. Highlight colors represent the IF intensity- Low, medium, High correspondingly.
Instruction Type
Instruction
Response (Token-level IF)
Explain Concepts
Explain the concept of quantum computing in simple terms.
Quantum computing is a type of computing that uses
quantum mechanical phenomena such as superposition and entanglement to perform
operations on data
Creative Wriring
Write a shurt poem about artifi- intelligence.
circuitg and code we rise, A digital symphony eyes. logic and data entwined.
Arithmetic Calculation
What is 31 312
To calculate (3 (31 x 31 }. you can use the formulal for squaring
binomial which is (a+b]
2ab 15 Here.
So, we have: (30+ 1)
300
Game Strategy Reason- ing (AIME2024)
Alice and Bob play the following game. A stack of tokens lies before them. ... Whoever moves the last token wins. Find the number positive integers S 2024 such that Bob always wins.
To determine the number of positive integers Il 2024 for which Boh has winning strategy we need to analyze the game using the concept of winning and losing positions A position is called losing position if the player whose turn
ing path. The mathematical expansion that follows shows diminished IF, likely due to repetition and procedural predictability.
In Game Strategy Reasoning, tokens like "winning strategy." "losing position," and "concept of winning and losing" emphasized. These convey the strategic logic essential for solving the problem, demonstrating that the model assigns higher utility to abstract reasoning components.
Underutilized Tokens Encode Structure or Copy. Across settings, transition tokens (e.g.. ",", "to"") consistently exhibit low IF, indicating minimal semantic contribution. Interestingly, math expressions directly reused from the prompt (e.g.. "31 x 31"") are also under- higherihgghihhted, suggesting the model de-emphasizes copied content favor of novel reasoning or planning segments. This pattern extends to common function words and syntactic markers that provide structure rather than content. For example, determiners ("the", "a"), conjunctions ("and", "or"), and pronouns ("it", "they") show consistently lower IF scores across all prompt types.
7 Conclusion
We introduce lightweight and effective method for post-training that enhances parametric knowledge utilization in language models by steering the key-value dynamics in FFN layers. Our findings reveal a critical insight: modern LLMs substantially underutilize the knowledge acquired during pretraining when applied to downstream tasks, leaving significant performance potential unrealized. Through simple forward interventions and entropy- based regularization, steered SFT consistently improves both ID and OOD performance across diverse models and tasks. Beyond accuracy gains, method encourages adaptive memory allocation and reveals interpretable information usage patterns, offering new insights into the internal behavior of post-trained LMs. These findings suggest that strategic controlling of memory engagement is a promising direction for improving both capability and transparency in LMs' internal thinking behavior.
Limitation. In this work, we position ourselves as exploratory study in this area, focusing our evaluation solely on standard SFT. We believe that steered RL could be more effective approach, as RL may better leverage the capabilities of pretrained LMs compared to SFT in reasoning casoning-intensive intensive settings [Swamy et al., 2025]. For scenario like inference-time compute or long CoT, steering with greater utilization on pretrained knowledge may further enhance task performance. Therefore, we also see an opportunity to explore the co-design of algorithms that combine both combine both "internal" steering and "external" CoT generation.
10

---

References
Alessandro Achille, Giovanni Paolini, and Stefano Soatto. Where is the information in a deep neural network?, 2020. URL https: /ttpe/////aaivaa/eaaeeaaeeaaeaar org/ abs/ 1905 12213.
Rohan Anil, Andrew M. Dai, Orhan Firat, Melvin Johnson, Dmitry Lepikhin, Alexandre Passos, Siamak Shakeri, Emanuel Taropa, Paige Bailey, Zhifeng Chen, Eric Chu, Jonathan H. Clark, Laurent EI Shafey, Yanping Huang, Kathy Meier-Hellsteri Gaurav Mishra, Erica Moreira, Mark Omernick, Kevin Robinson, Sebastian Ruder, Yi Tay, Kefan Xiao, and al. Palm 2 technical report, 2023. URL
org/ abs/ 2205 10403.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng, Yang Fan, Wenbin Ge, Yu Han, Fei Huang, Binyuan Hui, Luo Ji, Mei Li, Junyang Lin, Runji Lin, Dayiheng Liu, Gao Liu, Chengqiang Lu, Keming Lu, Jianxin Ma, Rui Men. Xingzhang Ren, Xuancheng Ren, Chuanqi Tan, Sinan Tan, Jianhong Tu, Peng Wang. Shijic Wang. Wei Wang, Shengguang Wu, Benfeng Xu, Jin Xu, An Yang, Hao Yang, Jian Yang, Shusheng Yang, Yang Yao, Bowen Yu, Hongyi Yuan, Zheng Yuan, Jianwei Zhang. Xingxuan Zhang. Yichang Zhang, Zhenru Zhang, Chang Zhou, Jingren Zhou, Xiaohuan Zhou, and Tianhang Zhu. Qwen technical report, 2023. URL https: http:://wwwvv.r.......o 2309. 16609.
Lukasz Bartoszcze, Sarthak Munshi, Bryan Sukidi, Jennifer Yen, Zejia Yang, David Williams- King. Linh Le, Kosi Asuzu, and Carsten Maple. Representation engineering for large-language models: Survey and research challenges. ArXiv, abs/2502.17601, 2025. URL https: //api. semanticscholar, wwwoooog///ooeeeeoooo CooppsIdd:2658000000000
Jeremy Bernstein and Yisong Yue. Computing the information content of trained neural networks, 2021. URL https://arae /twpvv/rxvvvv///////////./.a.a org/ aas//b/3..
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. Advances in Neural Information Processing Systems, 36: 1560-1588, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. Piqa: Reasoning about physical commonsense in natural language, 2019. URL httpp::/aaxxi..rre.e.e.e org/abs/////ooooooooooo..oo /1111.91...111 1991..111 11641..
Hoyeon Chang, Jinho Park, Seonghyeon Ye, Sohee Yang. Youngkyung Seo, Du-Seong Chang, and Minjoon Seo. How do large language models acquire factual knowledge during pretraining? In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL / /ttp://veeeeeeeeeeeeemeee htw.:/nee/eeeeeeeeemee http://oormmidd-TYoe 1EvBP.
https:
Ruidi Chang, Chunyuan Deng, and Hanjie Chen. SAFR: Neuron redistribution for interpretability. In Luis Chiruzzo, Alan Ritter, and Lu Wang, editors, Findings of the Association for Computational Linguistics: NAACL 2025, pages 2117-2126, Albuquerque, New Mexico, April 2025. Association for Computational Linguistics. ISBN 979-8-89176-19944 URL htt.htts///ccaatthooo 2025.
org/
naacl. 112/-
Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Tay, William Fedus, Yunxuan Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, Albert Webson, Shixiang Shane Gu, Zhuyun Dai, Mirac Suzgun, Xinyun Chen, Aakanksha Chowdhery, Alex Castro-Ros, Marie Pellat, Kevin Robinson, Dasha Valter, Sharan Narang. Gaurav Mishra, Adams Yu, Vincent Zhao, Yanping Huang, Andrew Dai, Hongkun Yu, Slav Petrov, Ed H. Chi, Jeff Dean, Jacob Devlin, Adam Roberts, Denny Zhou, Quoc V. Le, and Jason Wei. Scaling inlrucctionfineturee language models, 2022. URL https: / /arxiv.vvvvaa org/ att.///110/ 11416.
Christopher Clark, Kenton Lee, Ming- Wei Chang, Tom Kwiatkowski, Michael Collins, and Kristina Toutanova. Boolq: Exploring the surprising difficulty of natural yes/no questions, 2019. URL
https://arxiv. org/ abs/ 10044.
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try are, the ai2 reasoning challenge, 2018. URL https::/ /arxivv org/ aa/aab////000 05457.

---

Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchrr Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems, 2021. URL https: httttps://aatt //tp://aaavv.///a org/
0tss/2110.14163
Arthur Conmy, Augustine Paevorr Aengus Lynch, Stefan Heimersheim, and Adrià Garriga- Alonso. Towards automated circuit discovery for mechanistic interpretability. Advances in Neural Information Processing Systems, 36:16318-16352, 2023.
Ganqu Cui, Lifan Yuan, Ning Ding, Guanming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. UltraFeedback: Boosting language models with high-quality feedback. arXiv:2310.01377, 2023. URL wtwph:tttt//artteeee org/ abs /210 22112311110000....1
Hoagy Cunningham, Aidan Ewart, Logan Riggs, Robert Huben, and Lee Sharkey. Sparse autoen- coders find highly interpretable features in language models, 2023. URL https://aaxii/ org/ abs/2309. /tt///399.0000.0.. 08600.
Hannah Cyberey and David Evans. Steering the censorship: Uncovering representation vectors for Ilm"ttoughh" control. In unknown, 2025. URL https://api. http:///app..eeeaaacccccoooooeee org/ control. CorpusId: 2770022722
Damai Dai, Li Dong, Y. Hao, Zhifang Sui, and Furu Wei. Knowledge neurons in pretrained
transformers. ArXiv, abs/2104.08666, 2021. URL https://////t /t//:/..allanhhooogyyyy org/2022.
acl-long pdf.
Chunyuan Deng, Zhiqi Li, Roy Xie, Ruidi Chang, and Hanjic Chen. Language models are symbolic learners in arithmetic, 2024. URL https: / aaxiv. org/abo abs//410. 15580.
Chumyuan Deng, Ruidi Chang, and Hanjie Chen. Learning distribution wise control in representation space for language models, 2025. URL, htttp:://arxiiii http:///wwwwww..aaaaa//ee atwrg//bbb////6// 06686.
Kevin Du, Vésteinn Snæbjarnarson, Niklas Stoehr, Jennifer C White, Aaron Schein, and Ryan Cotterell. Context versus prior knowledge in language models. arXiv preprint arXiv:2404.04633 2024a.
Kevin Du, Vésteinn Snæbbarnarson, Niklas Stoehr, Jennifer C. White, Aaron Schein, and Ryan Cotterell. Context versus prior knowledge in language models, 2024b. URL stttps://arixs org/ htg//bb/244000 04633.
Javier Ferrando and Elena Voita. Information flow routes: Automatically interpreting language models at scale. In Yaser Al-Onaizan, Mohit Bansal, and Yun-Nung Chen, editors, Proceedings of the 2024 Conference on Empirical Methods in Natural Language Processing, pages 77432-17445 Miami, Florida, USA, November 2024. Association for Computational Linguistics. doi: 10.18653/ v1/2024.cmnlp-mma maail. URL https: / ac aclanthology org/2024. /2224. oonlp-main.9eee
Javier Ferrando, Gerard I. Gállego, and Marta R. Costa-jussà. Measuring the mixing of contextual information in the transformer. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang, editors, Proceedings of the 2022 Conference Empirical Methods in Natural Language Processing, pages 8698-877 8714, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: help-mmai..5990
URL https //te:///cccaa /ttp//aaaanttttt //taaaanaaannonoyocoooo org/2022.
Javier Ferrando, Gerard I. Gállego, Ioannis Tsiamas, and Marta R. Costa-jussà. Explaining how transformers use context to build predictions. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki, editors, Proceedings of the 61st Annual Meeting of the Association for Computationnal Linguistics (Volume Long Papers), pages 5486-5513, Toronto, Canada, July 2023. Association for Computational Linguistics. doi: 10.18653/v ac2022.ac1- URL https://aclantholoe org/2023. acl-long.300 301/.
Atticus Geiger, D. Ibeling, Amir Zur, Maheep Chaudhary, Sonakshi Chauhan, Jing Huang, Aryaman Arora, Zhengxuan Wu, Noah D. Goodman, Christopher Potts, and Thomas F. Icard. Causal abstraction: A theoretical foundation for mechanistic interpretability. 2023. URI https///: semanticscholar orwo/o///oooo.ooooooooooooo httg//orpsI::::5:505
12

---

Mor Geva, Roei Schuster, Jonathan Berant, and Omer Levy. Transformer feed-forwand layers are key-value memories. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages $484-5495, 2021.
Mor Geva, Jasmijn Bastings, Katja Filippova, and Amir Globerson. Dissecting recall of factual associations in auto-regressive language models. In Proceedings of the 2023 Conference on Empirical Methods Natural Langugge Processing, pages 12216-12235, 2023.
Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, Xiaokang Zhang. Xingkai Yu, Yu Wu, Z. F. Wu, Zhibin Gou, Zhihong Shao, Zhuoshu Li, Ziyi Gao, Aixin Liu, Bing Xue, Bingxuan Wang, Bochao Wu, Bei Feng. Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, Damai Dai, Deli Chen, Dongji Ji, Erhang Li, Fangyun Lin, Fucong Dai, Fuli Luo, Guangbo Hao, Guanting Chen, and et al. Deepseek-rl: Incentivizing reasoning capability in Iims via reinforcement learning, 2025. URL https: //axxiv...iii////aa///i org/ abs 2501 12948.
Yoav Gur-Arieh, Roy Mayan, Chen Agassy, Atticus Geiger, and Mor Geva. Enhancing automated interpretability with output-centric feature descriptions, 2025. URL https:///aaxiiiii org/ /arg//////// 1501.08314
Michael Hanna, Yonatan Belinkov, and Sandro Pezzelle. When language models fall in love: Animacy processing in transformer language models. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 12120-12135, Singapore, December 2023. Association for Computational Linguistics. doi: 10.18653/w1/2023 1/2666/3///2222...mmmmmmmmam mmnlp-mminnn-- 744. URL aattps://aclantholloe wrg//022/ emnlp www.main..7.. 744/.
Peter Hase, Mohit Bansal, Been Kim, and Asma Ghandeharioun. Does localization inform editing? surprising differences in cousality-based localization vs. knowledge editing in language mod- els. ArXiv, 12//2001.04213, 2023. URL https: //api. semanticscholast org/ IorpusID: 255595518.
John Hewitt and Christopher D. Manning. A structural probe for finding syntax in word represen- tations. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Hu- man Language Technologies, Volume I (Long and Short Pupers), pages 4129-41338 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1/N199 1419. URL https: /aap::://llllaaa http///aaaaanttooooooooo org/N19-10 1419/.
Mohammad Javad Hosseini, Hannaneh Hajishirzi, Oren Etzioni, and Nate Kushman. Learning to solve arithmetic word problems with verb categorization. In Alessandro Moschitti, Bo Pang, and Walter Daelemans, editors, Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), pages 523-533, Doha, Qatar, October 2014. Association for Computational Linguistics. doi: 10.3115/w1D1141005 1058. URL https: http///acaaaaoloo/oo// org/ D14-1058/.
Zhiqiang Hu, Lei Wang, Yihuai Lan, Wanyu Xu, Ee-Peng Lim, Lidong Bing, Xing Xu, Soujanya Poria, and Roy Lee. LLM-adapters: An adapter family for parameter-efficient fine-tuning of large language models. In Houda Bouamor, Juan Pino, and Kalika Bali, editors, Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pages 5254-5276, Singapore, December 2023. Association for Computational Linguistics. doi: 319. URL https: /ttptts://aalaathooll/ www.gg.ooog/ooooooooooooooo...go oor./rr///22 manlp-main... 319/.
111553/112023.emmmm main.
Jing Huang. Zhengxuan Wu, Christopher Potts, Mor Geva, and Atticus Geiger. Ravel: Evaluating interpretability methods on disentangling language model representations, 2024. URL https: //arxiv. org/ abs//402..... 2ot.//200277700
Jiyeon Kim, Hyunji Lee, Hyowon Cho, Joel Jang, Hyeonbin Hwang, Seungpil Won, Youbín Ahn, Dohaeng Lee, and Minjoon Seo. Knowledge entropy decay during language model pretraining hinders new knowledge acquisition. The Thirteenth International Conference on Learning Representations, 2025, URL https:/ stttps://ppenreee http//pprrrreeeeeeeeooooreeeeeee Sii.SSSSE.
13

---

Rik Kedzzirskd.... Hannaneh Hajishirzi, Ashish Sabharwal, Oren Etzioni, and Siena Dumas Ang. Parsing algebraic word problems into equations. Transactions of the Association for Computational Linguistics, 3:585-597, 2015. doi: 10. 10.116:2 1e2116//ccl__ a 4000160. URL https: http://acaaanttooooo/oo org/ Q15-1042/.
Rik
Subhro Roy, Aida Amini, Nate Kushman, and Hannaneh Hajishirzi. MAWPS: A math word problem repository. In Kevin Knight, Ani Nenkova, and Owen Rambow, editors, Proceedings of the 2016 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, pages 112-11557 157, San Diego, Califor- nia, June 2016. Association for Computational Linguistics. doi: 10.18653/11/N16- 1136. URL https: http///aaaatthhoooo
org/ OrggNN6-1100 1136/.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph E. Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. In Proceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles, 2023.
Karim Lasri, Tiago Pimentel, Alessandro Lenci, Thierry Poibeau, and Ryan Cotterell. Prob- ing for the usage of grammatical number. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Proceedings of the 60th Annual Meeting of the Association for Com- putational Linguistics (Volume 1: Long Papers), pages 0818-8831, Dublin, Ireland, May 2022. Association for Computational Linguisties. doi: 10. 10.18653/1/2022..cc attp:/222.ccc11olel
URL
Bruce W. Lee, Inkit Padhi, K. Ramamurthy, Erik Michling, Pierre L. Dognin, Manish Nagireddy, and Amit Dhurandhar. Programming refusal with conditional activation steering. ArXiv, abs/2409.059077, 2024. URL atttps://arxxt www.aawwww...../..eeaeeoee htgg/ddf/2400 05907. pdf.
Kenneth Li, Oam Patel, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Inference-time intervention: Eliciting truthful answers from language model, 2024. URL https://axxiv
org/abs/2306. 2334/
Wang Ling, Dani Yogatama, Chris Dyer, and Phil Blunsom. Program induction by rationale generation : LLeanning to solve and explain algebraic word problems, 2017. URL https: //arx/aa..iv////a org/ abs/
1705. 04146.
Leo Liu, Tim Dettmers, Xi Victoria Lin, Ves Stoyanov, and Xian Li. Towards unified view of sparse feed-forward network in pretraining large language model. ArXiv, abs/2305.13999 2023, URL
bittpa://aphs semanti wac.ontaaccca.......comaraaaaa.a wwgo Corpusld:
Feng Luo, Rui Yang. Hao Sun, Chunyuan Deng, Jiarui Yao, Jingyan Shen, Huan Zhang, and Hanjic Chen. Rethinking diverse human preference learning through principal component analysis, 2025.
URL https: //tp:v/////iv/////a org/ abs/ 13111
Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in neural information processing systems, 15:17399-17372, 2022a.
Kevin Meng, Arnab Sen Sharma, Alex J Andonian, Yonatan Belinkov, and David Bau. Mass-editing memory in a transformer. In The Eleventh International Conference on Learning Representations, 2022b.
Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, and et al. Gemma: Open models based on gemini research and technology, 2024. URL https: / /tt///avvvvraaaaaaeeee
08295.
Todor Mihaylov, Peter Clark, Tushar Khot, and Ashish Sabharwal. Can suit of armor conduct electricity? a new dataset for open book question answering, 2018. URL https: :/arxivv///xxvv..// org/ abs/1809.02
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto. sl: Simple test-time scaling, 2025. URL https: //tp///aavvvv//////////eee.aaall herg/abs/2500 19393.
14

---

Xueyan Niu, Bo Bai, Lei Deng, and Wei Han. Beyond scaling laws: Understanding transformer performance with associative memory. ArXiv, abs/2405.08707, 2024. URL https://api. semanticscholar org/ wttg//oopussdd:297611957
Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe. Training language models to follow instructions with human feedback, 2022. URL 02155.
httts://arxix: org/
Arkil Patel. Satwik Bhattamishra, and Navin Goyal. Are NLP models really able to solve sim- ple math word problems? In Kristina Toutanova, Anna Rumshisky, Luke Zettlemoyer, Dilek Hakkani-Tur Iz Beltagy. Steven Bethard, Ryan Cotterell, Tanmoy Chakraborty, and Yichao Zhou, editors, Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 2080-2094, Online, June 2021. Association for Computational Linguistics. doi: 1e.111655/1/2222 http://w11/111.lallllaaaaaa naacl- maadd URL https: httes///aaaaaaaahaeeeeeeee
org/ 22222221112222 http//22...aaaala..ee.e.aa www.aaacllaaaaa mmain. 168/.
Zihan Qiu, Zeyu Huang, and Jie Fu. Unlocking emergent modularity in large language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pages 2638-26600 2024.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 01533728-57744 2024.
Nina Rimsky, Nick Gabrieli, Julian Schulz, Meg Tong. Evan Hubinger, and Alexander Matt Turner Steering llama 2 via contrastive activation addition. ArXiv, abs/2312.06681, 2023. URL https: //arxiv... org/ /ttp//ff////// 223122266221211112222 pdf.
Subhro Roy and Dan Roth. Solving general arithmetic word problems, 2016. URL https://arxiv.
org/ abb//60/..11./../////// /600. 01413.
David E Rumelhart, James L McClelland, PDP Research Group, et al. Parallel distributed processing, volume 1: Explorations in the microstructure of cognition: Foundations. The MIT press, 1986.
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An
adversarial winograd schema challenge at scale, 2019. URL httptt//////aeee www.riiiioo..oo..oooooooooo httgg/bb//190// 10641.
Maarten Sap, Hannah Rashkin, Derek Chen, Ronan LeBras, and Yejin Choi. Socialiqa: Commonsense reasoning about social interactions, 2019. URL httthtsss///xxxe wwtpeiiii.orraaaaeeee
09728.
Daniel Scalena, Gabriele Sarti, and Malvina Nissim. Multi-property steering of large language models with dynamic activation composition. ArXiv, abs/2406.17563. 2024. URL https //www
aclanthology org/2024. 2024.. blackboxnlp-1 -..3.. ...f..
Zara Siddique, Irtaza Khalid, Liam D. Turner, and Luis Espinosa Anke. Shifting perspectives: Steering vector ensembles for robust bias mitigation in 1lms. ArXiv, abs/2503.05371, 2025. URL
attps://apt. semanticscholes org/ Id:://ooppdd:::/// 276885488.
Samuel Soo, Chen Guang, Wesley Teng, Chandrasekaran Balaganesh, Guoxian Tan, and Yan Ming. Interpretable steering of large language models with feature guided activation additions. In unknown, 2025. URL https: http:p/s///pp.eemmtticccccool/
www.sseiiaccnool......... wwwpcchholarooooooooooooeoooo IIg//orp 275606633.
Alessandro Stolfo, Vidhisha Balachandran, Safoora Yousefi, Eric Horvitz, and Besmira Nushi. Improv- ing -nstruction-follooiingg in language models through activation steering. ArXiv, abs/2410.12877,, 2024. URL /ttpp:://app...... http:pss///pp.ssemaiiccccooooooe semanticscholar . org/ 2trpusId:273405555
Nishant Subramani, Nivedita Suresh, and Matthew Peters. Extracting latent steering vectors from pretrained language models. In Smaranda Muresan, Preslav Nakov, and Aline Villavicencio, editors, Findings of the Association for Computational Linguistics: ACL 2022, pages 566-581, Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10. 1/22 findings-aclen ac1.48. URL https: :/ htt///aaantttt/ aclanthology org/ org/2022.oo findings-acl.cc 48/.
15

---

Gokul Swany, Sanjiban Choudhury, Wen Sun, Zhiwei Steven Wu, and J. Andrew Bagnell. All roads lead to likelihood: The value of reinforcement learning in fine-tuning, 2025. URL https: //arxiv. org/ 2a0a//553.0000 01067.
Yi Tay, Vinh Q. Tran, Mostafa Dehghani, Jianmo Ni, Dara Bahri, Harsh Mehta, Zhen Qin, Kai Hui, Zhe Zhao, Jai Gupta, Tal Schuster, William W. Cohen, and Donald Metzler. Transformer memory as a differentiable search index. ArXiv, abs/2202.06991, 2022. URL https://api. semanticscholar wwwrg//orpussDD.. 2468634484..888
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Maric-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurelien Rodriguez, Armand Joulin, Edouard Grave, and Guillaume Lample. Llama: Open and efficient foundation language models, 2023. URL https: /ttp///axvvvvvva///////////a.vvv org/ ////2/2222.2222
Alexander Matt Turner, Lisa Thiergart, Gavin Leech, David Udell, Juan J. Vazquez, Ulisse Mini, and Monte MacDiarmid. Steering language models with activation engineering, 2024. URL https://arxiv org/ abs/2308....../////// 10248.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2017. URL https: //tpx//aaivv////a///a/ org.org///b... /bs// 1706.031 03762.
Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Inter- pretability in the wild: circuit for indirect object identification in gpt-2 small, 2022. URL
https://arxxe
Peng Wang, Zexi Li, Zhagg, Ziwen Xu, Yunzhi Yao, Yong Jiang, Pengjun Xie, Fei Huang, and Huajun Chen. Wise: Rethinking the knowledge memory for lifelong model editing of large language models. ArXiv, 1476/2405.1478, 2024. URL https: //api. semanti cscholarr org/ CorpusId: SorpusId:269982715
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian Ichter, Fei Xia, Chi, Quoc Le, and Denny Zhou. Chain-of-thoouhh prompting elicits reasoning in large language models, 2023. URL https:tp/..
11903.
Jason Weston, Sumit Chopra, and Antoine Bordes. arXiv: 1410.3916, 2014.
Memory networks.
arXiv preprint
Zhengxuan Wu, Aryaman Arora, Zheng Wang, Atticus Geiger, Dan Jurafsky, Christopher Manning. and Christopher Potts. Reft: Representation finetuning for language models, 2024a. URL
btttp:://axxvv.org/abbe http://bs/224..000/ 03592.
Zhengxuan Wu, Atticus Geiger, Thomas Icard, Christopher Potts, and Noah D. Goodman. Inter- pretability at scale: Identifying causal mechanisms in alpaca, 2024b. URL https://arxiv....
Jian Xie, Kai Zhang, Jiangjie Chen, Renze Lou, and Yu Su. Adaptive chameleon or stubborn sloth: Revealing the behavior of large language models knowledge conflicts. In International Conference on Learning Representations, 2023. URL btttps://apt www.nticschooiii..ee. org/ :rrpusID:263610:5
Roy Xie, Junlin Wang, Paul Rosu, Chunyuan Deng, Bolun Sun, Zihao Lin, and Bhuwan Dhingra. Knowing when to stop: Dynamic context cutoff for large language models, 2025. URL https: 01025.
//arxiv. org/
Yunzhi Yao, Ningyu Zhang, Zekun Xi, Meng Wang, Ziwen Xu, Shumin Deng, and Huajun Chen. Knowledge circuits in pretrained transformers. ArXiv, abs/2405.17969, 2024. URL https: //api. //app..smaaaae www.eetiiiccoolaa...o
270068372.
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence?, 2019. URL https://arxiv. wawxiv.iiiv..///ai orgoooog/aaoo.oooooooo abs/ab////// /100. /1///91000088......1 07830.
16

---

Chunting Zhou, Pengfei Liu, Puxin Xu, Srinivasan Iyer, Jiao Sun, Yuning Mao. Xuezhe Ma, Avia Efrat, Ping Yu, Lili Yu, al. Lima: Less is more for alignment. Advances in Neural Information Processing Systems, 36:55006-55021, 2023.
Andy Zou, Long Phan, Sarah Chen, James Campbell, Phillip Guo, Richard Ren, Alexander Pan, Xuwang Yin, Mantas Mazeika, Ann-Kathrin Dombrowski, Shashwat Goel, Nathaniel Li, Michael J. Byun, Zifan Wang, Alex Mallen, Steven Basart, Sanmi Koyejo, Dawn Song, Matt Fredrikson, Zico Kolter, and Dan Hendrycks. Representation engineering: top-down approach ai transparency, 2023. URL https: //trx/aaivvv///////////ee/aaaaa/
17

---

A Training Details
Hardware and Setup. All experiments were conducted using single NVIDIA RTX A6000 GPU with 48 GB of memory. Models with approximately to 2 billion parameters were trained directly on this GPU. For larger models, ranging from 7 to 9 billion parameters, used DeepSpeed to enable efficient training.
Training Configuration. All tasks trained for one epoch. We the learning apply a warmup phase of 100 steps. A weight decay of 0.01 is used to regularize the model. We adopt mixed precision training using bfloat16 (bfl6) to reduce memory usage and improve efficiency,
5e-5 and
The maximum sequence length is 256 tokens. We use per-device batch size of 4, with gradient checkpointing enabled. To simulate larger effective batch size and manage memory usage, gradient accumulation with 16 steps.
DeepSpeed Optimization. For large models (7B-9B), optimization. This approach splits the optimizer states and gradients across devices and offloads the optimizer to the CPU. The training gradients clipped to maximum norm of 1.0 to stabilize updates. Batch sizes are automatically adjusted by DeepSpeed based on available memory.
DeepSpeed with ZeRO Stage
Inference Setup. We use vLLM [Kwon al., 2023] for all inference runs to ensure efficient memory management and fast decoding. We adopt greedy decoding, and the maximum number of tokens generated per sequence is 256.
B Fine-Grained Steering during Dost-Training
While our intervention and regularization methods provide general approaches to enhance memory vector engagement, finer control mechanisms enable more nuanced steering of parametric knowledge. We present two complementary strategies for achieving fine- aine-grined control: (1) a group- based clustering approach that enables targeted steering of distinct memory regions, and (2) information surrogate-guide guided method that selects memory vectors based on their contribution to specific output distributions.
B.1
Group-Based Clustering for Targeted Memory Activation.
B.1.1
Method.
The key coefficients (*) can partitioned into meaningful groups that exhibit distinct activation patterns. Leveraging this structure allows us apply differentiated steering strategies various memory regions.
Formally, for each layer I, we cluster the key-value pairs {((()) u )mm
into G groups
(g(", _f ...., 9 based on their functional characteristics. We employ a hierarchical clustering approach with the following steps:
1. Semantic Clustering:
First, we cluster memory vectors of based on their semantic similarity, measured by cosine distance in the embedding space.
2. Activation Pattern Clustering: Second, we sub-cluster based on activation patterns ob- served during inference development set, capturing functional roles within semantic clusters.
For each cluster gh we define cluster-specifie steering parameter B, that modulates the strength of intervention:
+ k + 8g AK for i e g
(10)
where AE(") is the adjustment magnitude determined by either our intervention or regularization method.
18

---

Table 6: Performance comparison of Qwen, LLaMA, and Gemma models with different training methods on eight datasets. To highlight improvements, we use blue for significant gains and green for moderate ones, All results reported as the average scores over three independent runs.
Model
Training
Accuracy
BoolQ
PIQA
SIQA
HellaS.
WinoG.
ARC-e
ARC-e
OBQA
Small-Scale Models
(I-28 porameters)
base model
64.2
78.5
74.3
80.1
76.4
76.9
61.2
75.8
Qwen-2.5-1.55 5.555B
vonilla SFT
68.5
82.9
79.6
84.8
80.8
81.4
65.8
81.0
steered SFT semantic clastering
70.1
81.3
94.5
85.3
84.6
83.2
steered SFT activation clastering
69.4
84.5
80.7
93.2
84.1
83.7
68.5
82.4
base model
65.6
75.3
74.2
78.9
77.8
74.5
60.1
76.3
1JMMA-32-IB
vanilla SFT
semantic clustering
69.8
79.9
79.5
90.3
83.6
82.6
79.8
64.7
81.0
steered SFT activation clustering
71.7
83.6
80.4
89.5
83.1
84.1
68.2
82.8
model
78.9
vonilla SFT
70.2
83.4
78.
87.5
83.3
82.7
68.4
80.1
Gemma-2-22
steered SFT semantic clustering
73.1
86.3
80.8
91.4
87.2
86.5
73.1
85,3
steered SFT activation clastering
72.4
85.7
79.6
90.5
86.3
85.8
72.2
84.1
Large-Scule Models (7-9B parameters)
base model
68.9
81.2
77.6
87.5
80.3
79.8
65.1
77.6
Qwen:2.5.7B
vonilla SET
72.4
92.4
84.2
84.2
69.6
82.8
steere SFT somuntic clustering
75.2
87.
82.4
96.2
88.4
87.5
74.8
85.3
steered SFT activation clastering
77.0
86.5
81.9
95.1
87.2
86.3
73.5
84.6
buse model
70.3
85.6
75.7
90.8
81.9
86.2
75.3
80.5
LLaMA-3-8B
vanilla SFT
74.6
89.3
79.9
95.5
85.6
90.5
80.4
85,8
steered SFT
semantic clastering
78.3
91.0
96.8
88.3
93.6
82.7
88.4
steered SFT activation clastering
77.2
90.3
82.5
96.1
87.5
92.4
81.8
87.3
82.5
Gemma-2-98
vanilla SFT
74.3
90.1
81.7
94.8
86.9
91.7
82.0
86.4
semantic clustering steered SFT activation clastering
77.3
91.7
83.2
95.8
89.0
93.6
83.6
89.1
This group-based method offers several advantages. By examining which clusters are active during different tasks, can selectively strengthen those that contribute most to specific goals, improving task performance. At the same time, we preserve essential clusters that are important for general language understanding, while enhancing those that underused. Additionally, when aiming improve factual accuracy, we can prioritize clusters that are associated with factual knowledge.
B.1.2 Results.
Table 6 presents the comprehensive results across all models and benchmarks. Several clear patterns emerge from our experiments:
Memory steering consistently outperforms vanilla SFT. Across all model families and parameter scales, both of our memory clustering approaches demonstrate substantial improvements over standard fine-tuning. The semantic clustering method shows the most dramatic gains, with improvements ranging from 1.2% to 9.7% over vanilla SFT depending on the benchmark and model. These results validate our hypothesis that more balanced memory vector engagement leads to enhanced model capabilities.
Different clustering strategies show task-specific strengths. The semantic clustering approach excels particularly on Knowledge-intensive tasks (ARC-c, OBQA) and complex reasoning benchmarks (HellaSwag), achieving improvements of up to +4.7% on ARC-c (Gemma-2-2B) and +3.9% on OBQA (Gemma-2-9B) compared to vanilla SFT. In contrast, activation-based clustering shows more moderate but consistent improvements across broader range of tasks, suggesting it enables more balanced memory utilization. Larger models benefit more significantly from memory steering. While all models show improvements with our methods, the gains are particularly pronounced in the 7-9B parameter models. For instance, semantic clustering improves LLaMA-3-8h by +3.7% on BoolQ and +3.8% on SIQA, compared to more modest gains in the 1B variant. This suggests that larger models contain more untapped parametric knowledge that can be effectively engaged through steering techniques.
19

---

B.2 Information Surrogate-Guided Memory Selection.
Method. To understand the information introduced by our steering strategy, we extend the key-value formulation from Equation 3 by connecting the FFN output with the final logits:
dm
dm h O
(hN W decode
hy (U;
W decnde)
(11)
where K; o represents the key coefficient and Φ; Wdccoee defines the logit distribution associated with the i-th value vector, which we call the information surrogate.
The information surrogate o; provides a direct view of how each memory vector influences the final token distribution. This insight enables us develop a more targeted steering approach:
1. Characterization of Memory Vectors: We analyze the entropy and concentration properties each Φ; to identify memory vectors that contribute to specific types of generation (e.g., factual statements, reasoning steps, creative content).
2. Surrogate-Guided Steering: We define a surrogate score function S(ox) that measures the relevance of each memory vector our target objective:
S(Ø,) = H1 + >2 (o:lluude
(12)
where ((:)) is the entropy of the surrogate distribution, DRL. is the KL divergence from target distribution Ouagett and \1, X2 are weighting hyperparameters.
3. Selective Amplification: We modulate key coefficients based their surrogate scores:
- k (1+ S())))
(13)
where controls the strength of the surrogate-guidee steering.
Theoretical Analysis. The prrogate-gguided selection mechanism introduces several theoretical advantages over conventional fine-tuning approaches. First, we analyze the relationship between surrogate entropy and information capacity. For memory vector v, with corresponding surrogate e,, the entropy H (é;) quantifies the diversity of tokens that can be influenced by this vector. Higher entropy surrogates represent memory vectors that encode distributional knowledge. while low-entropy surrogates often correspond to specialized knowledge concentrated on specific vocabulary subsets.
We can formalize this by defining the specificity of a memory vector Spec(u.)
log([VV)* where
(
[V] is the vocabulary size. A memory vector with high specificity (low surrogate entropy) exhibits peaked distribution over the vocabulary, suggesting it encodes precise, specialized information. Conversely, low specificity (high surrogate entropy) indicates a memory vector that contributes more generally across various contexts. This formulation provides principled approach for analyzing memory vector functionality. Specifically, the information processing capacity of the FFN can be decomposed as:
(14) where MI(d;; represents the mutual information between the surrogate distribution di and the target next-token distribution y, and EEwD [k+ is the expected activation of key k; across the data distribution D. Our carrogate-guidee approach can be interpreted optimizing this information processing capacity by modulating the key coefficients k according to the coommation-thhoree properties of their corresponding surrogates 6,, By encouraging the activation of memory vectors with high mutual information with the target distribution, we effectively allocate the model's capacity toward task-relevan information.
I(FFN) Ei= 1--=E
MI()))))
B.3 Results and Analysis
Table 7 presents the performance comparison our surrogate-guided methods against vanilla SFT across three model families and two size scales. The surrogate-guided selection method consistently outperforms both vanilla SFT and surrogate entropy maximization approaches, with particularly notable gains reasoning-head tasks (BoolQ, SIQA) and knowledge-intensin benchmarks.
20

---

Table 7: Performance comparison of information surrogate- based steering. To highlight improve- ments, we use blue for significant gains and green for moderate ones. All results are reported as the average scores over three independent runs.
Model
Training
Accuracy (1)
BoolQ
PIOA
SIOA
HellaS.
WinoG.
ARC²e
ARC-c
OBQA
Smalt-Scole Models
parumeters))
base model
64.2
78.5
74.3
80.1
76.4
76.9
61.2
75.8
Qwen-2.5-1.58
vanilla SFT
surregole surrogole-quited seleciion
68.5
82.9
79.6
51.8
80.5
$1.4
65.8
$1.0
+ Itcered SFT v. surregate entropy musimizationn
69.4
84.2
80.2
92.9
34.0
$2.9
68.2
78.9
base model
65.6
75.3
74.2
78.9
77.8
74,5
60.1
76.3
11MA--22-18
vanilla SFT
nteered
surrggate surrogeteguuded selection
69.8
74.9
79.5
83.6
82.6
79.5
61.7
$1.0
+ Hteere SFT v. surregate entropy musimizatiion
71.7
83.4
80.7
89.1
$2.5
83.1
67.9
$2.0
base model
66.5
79.1
73.8
82.7
78.9
77.4
63.8
74,9
Ciemma-2-28
vanilla SFT
70.2
83.4
78.1
87.5
83.3
82.7
68.4
80.1
nteered
cenogate-guddia selection
73.4
91.9
+ steered SFT v. surregate entropy musimizatioon
72.3
85.3
79.2
90.4
35.5
85.3
71.4
83.5
Large
Large-Scale Scale Models (7-91 pariometers)
68.9
77.6
Qwen-2.5-7B
vanilla SFT
72.4
81.9
81.5
92.1
84.2
81.2
69.5
82.8
+ steered SFT v. murrogoe--uiieed selection
75.2
87.3
82.6
95.8
88.5
87.1
74.9
85.3
strered SFT w. surrogate entropy masimization
76.9
86.2
81.9
94.8
87.0
85.9
73.3
84.00
1LAM--8888
+ steered SFT w. gurdgote-quided selection
vanilla SFT
74.6
78.3
91.2
89.3
83.7
79.9
95.5
96.8
88.2
85.6
95.6
90.5
80-4
82.4
85.8
88.1
steered SFT w. surrogate entropy muximization
77.1
90.2
81.8
96.0
$7.3
92.2
81.5
87.2
Gemma-2-9H8
+ vanilla SFT
74.3
90.1
81.7
91.8
36.9
91.7
82.0
86.4
neered SFT w. surrogate guided selection
78.5
92.5
84.2
96.9
89.3
94.7
84.8
89.7
steered SFT surrogute entropy muximization
76.9
91.3
83.1
95.5
88.2
93.5
83.5
88.3
Beyond the aggregate statistics, we observe intriguing qualitative differences in how surrogate-guided selection influences model behavior. By examining the top activated memory vectors and their corresponding information surrogates (ø,). we find that our method preferentially engages memory vectors that encode precise factual associations rather than general linguistic patterns. For example, in the ARC-c task, the top-5 memory vectors with the largest positive Ak; values predominantly contribute to science concept definitions and physical property relationships. This suggests that surrogate-guice selection effectively identifies and amplifies task-relevant knowledge encoded in specific memory vectors, rather than uniformly increasing memory engagement. Moreover, we find that models trained with surrogate-guided selection exhibit reduced variance in their responses to imowledge-intensive questions, indicating more consistent access to stored parametric knowledge during inference.
C Understanding Layerwise Contribution
To better understand how different layers contribute the effectiveness of InfoSteer, conducted a series of ablation studies. These experiments help us analyze which layers are most sensitive to information steering and which contribute most significantly overall performance improvements.
C.1 Research Question
We designed ablation study investigate the following research questions:
1. Do all layers contribute equally to the information steering effects?
2. Are certain layer groups (carly, middle, late) more important for knowledge retrieval?
3. How does the magnitude of steering different layers affect overall performance?
For these experiments, we used our intervention method with selective application to different layer groups within the model. We also varied the intervention parameters (p% and a:) across different layer configurations to understand sensitivity. Table 8 presents the results of our layerwise ablation studies.
21

---

Table 8: Layerwise ablation studies. Results show accuracy across different layer configurations with varying intervention parameters.
Layer position matters significantly. Apply- ing InfoSteer to different layer groups produces notably different results. Early layers (1-8) show moderate improvements, while middle layers (9- 16) yield the strongest gains. Late layers (17-24) demonstrate the least improvement, suggesting that information steering is most effective at the intermediate representation level.
Model Configuration
Layers
Avg Acc
Base Model Vanilla SFT
71.4
All
72.6
+ interv. (carly)
1-8
73.8
+ interv. (middle)
9-16
75.2
+ interv. (late)
17-24
72.9
Intervention strength sensitivity varies by layer. Middle layers can tolerate and benefit from stronger interventions (a = 3). while early layers perform best with moderate intervention (a 2), and late layers require gentler steering (a = 1).
interv. (eidly+middle)
1-16
76.3
+ interv. (middlc+late)
9-24
74.5
+ interv. (all, a = 2)
1-24
75.5
+ interv. (all, 5)
1-24
72.8
+ reg. (middle, À 0.05)
9-16
74.9
reg. (all, 0.05)
1-24
74.7
+ reg. (all, 0..1))
1-24
72.3
Layer combinations show non-linear effects. Applying InfoSteer to both early and middle layers yields siter-than-additii improvements, suggesting synergistic effect. However, including late layers tends to diminish these gains, indicating that excessive steering across too many layers may destabilize the model's 's representations.
Baseline performance variations. Our experiments with different regularization strengths (A) applied specific layer groups further confirm that middle layers (9-16) are most receptive to information steering.
These findings demonstrate that information steering should be carefully targeted at specific layers rather than uniformly applied across the entire model. Our optimal configuration focuses on middle and early layers with appropriately calibrated intervention strengths.
D License for Existing Assets.
The following datasets are used under their respective licenses. For general question answering: BoolQ [Clark et al., 2019] licensed under CC-BY-SA 3.0, PIQA [Bisk et al., 2019] under the Academic Free License 3.0, SIQA [Sap et al., 2019] and WinoGrande (Sakaguchi et al., 2019] under CC-BY 4.0, HellaSwar [Zellers et al., 2019] under the MIT License, ARC-e and ARC- [Clark et al., 2018] under CC-BY 4.0, and OBQA [Mihaylov et al., 2018] under the Apache-2.0 License. For arithmetic reasoning: AddSub [Hosseini al., 2014], MAWPS et al., 2016], MultiArith [Roy and Roth, 2016], and SingleEq Kkoncel-Kedzionn under CC-BY 4.0, AQuA [Ling et al., 2017] under Apache-2.0. and GSM8K [Cobbe et al., 2021] and SVAMP [Patel et al., 2021] ]] under the MIT License. For instruct-tuning the Ultrafeedback [Cui et al., 2023] dataset released under the MIT License.
Datasets.
et al., 2015] are
Model Licenses. The Qwen-2.5-1.55 and Qwen-2.5-7B models are released under the permissive Apache License 2.0, allowing broad usage including commercial applications [Mesnard et al., 2024]. In contrast, the I.L.MA-3.2-1B and ILAMA-3-8B models are distributed under Meta's 's custom Ilama Community License, which permits research and commercial use but imposes specific restrictions, particularly for organizations with large user bases [Touvron et al., 2023J. The Gemma-2-2B and Remma-2-9B models are available under Google's Gemma License, described as commercially friendly; however, access requires users to review and agree to the license terms, typically through platforms like Hugging Face [Mesnard et al., 2024]. Users intending to utilize these models should carefully review the respective licenses to ensure compliance with all terms and conditions.
22

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
The abstract and introduction clearly state the paper's key contributions, which are well-aligned with the theoretical and experimental results presented, accurately reflecting the scope and impact of the work.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We discuss limitations in Section 7 (Conclusion), noting that approach is exploratory and focused only on standard SFT. We specifically mention that "steered RL could be more effective approach" and identify opportunities for future work exploring the co-design of algorithms combining "internal" steering with "external" CoT generation.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
23

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Justification: Our paper is primarily empirical and doesn't include formal theoretical results requiring proofs. We build upon the existing key-value memory interpretation of FFN layers, extending it with our steering approach and validating through experiments.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Section 5 provides comprehensive details of our experimental setup, including model architectures (Qwen, LLaMA, Gemma), datasets used for evaluation, and hyperpa- rameters for our steering methods (p% 0.01, a = 1, À = 0.01). Table 2 further details ablation studies with varying hyperparameter settings. We also disclose other details in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
24

---

(d)
We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We release all code for InfoSteer in supplemental material, which will include implementation of both intervention and regularization methods. We use publicly available datasets and models in our experiments.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/ public/guides/ wtwpess/oooeSuuuiiiinnnneeee
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
httppsscccppblcc//ee http:ppblccgguuieeeeepee www.ee//ueeeee Codeeubmi ssion ssionPolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if' applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Section 5.1 details the experimental setup, including base models, datasets, and our default hyperparameters. For all benchmarks, we use the same prompt templates as in prior work. We also report ablation studies in Section 5.3 with different hyperparameter settings.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
25

---

Justification: As noted in Table 1, all reported results are averaged over three independent runs. Figure 3 clearly displays error bars for our OOD evaluation, and other statistical visualizations (e.g., Figure 4) also include error bars to indicate variance across runs.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative eтror rates).
If error bars reported in tables plots. The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We use Nvidia RTX A6000 for training the models.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: :/tt::///eeeer /nwurnnuu....// cc/public/ weciicc//hiiiitiieeeee ceiuidelineeees
Answer: [Yes]
Justification: Our research conforms to the NeurIPS Code of Ethics. We use publicly available datasets and models, properly cite all sources, and our method focuses on improving general capabilities rather than applications that could raise ethical concerns.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
26

---

10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our work focuses on improving information utilization in language models, which is fundamental research. While this could potentially lead to more capable Al systems with broader societal implications, our specific technique does not directly introduce new risks beyond those associated with general LM improvements.
Guidelines:
The answer NA means that there no societal impact of the work performed.
authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper introduces a training technique rather than releasing models or datasets that pose risk for misuse. The technique itself is aimed improving general performance rather than enabling specific capabilities that might require safeguards.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual- should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
27

---

Answer: [Yes]
Justification: We properly cite all models (Qwen, LLaMA, Gemma) and datasets used in research. All experiments conducted in accordance with the terms of use for these assets, using publicly available models and datasets as referenced in Section 5.1.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. cww.coo///aooooo com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our paper does not introduce new datasets model weights, but rather technique for improving existing models.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our research does not involve human subjects or crowdsourcing.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional review board (IRB) approvals or equivalent for research with human subjects
28

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our research does not involve human subjects, IRB approval not required.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: While our paper studies LLMs, we did not use LLMs in the development of our methodology for content generation during the research process.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non ntn-standard components.
Please refer to our LLM policy (https //tp/r/ees.pp//ee/e cc/Confeerncees/20000 2022LL for what should or should not be described.
29


Reviews:
[{"review_id": "tpVri5lj9F", "reviewer": "Reviewer_E6rk", "summary": "The paper provides a new perspective on the role of the FFN layer in LMs. It hypothesizes that the post-training underutilizes the knowledge in memory vectors. By encouraging the usage of the memory vectors during post-training, the resulting LM can have better performance on in-domain and out-of-distribution domains.", "questions": "1. The approach shows improved performance during post-training. Can the same approach to applied during pre-training to encourage the usage of the memory vectors? Why is it particularly relevant during post-training. \n2. The post-training data size is relatively small. Is it a reason that the info steering approach is effective, because the model cannot learn the optimal key distribution by itself?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n(1) The writing is clear. It provides sufficient context and motivation for the method. \n(2) Convincing results and the motivation behind the approaches. The approach is effective across tasks and pre-trained models. It is also simple to implement.\n\nWeakness:\nCan't think of any.", "was_revised": true, "final_justification": "I accepted the author's rebuttal and would like to retain my rating.", "rebuttal": "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.\n\nYou'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.\n\n**During the Post-training Phase (Small Data):**\n\nIn 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.\n\n**During the Pre-training Phase (Large Data):**\n\nYou 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!\n\n## Reference\n\n- [1] Kim, Jiyeon, et al. \"Knowledge Entropy Decay during Language Model Pretraining Hinders New Knowledge Acquisition.\" The Thirteenth International Conference on Learning Representations.\n- [2] Gu, Xiangming, et al. \"When Attention Sink Emerges in Language Models: An Empirical View.\" The Thirteenth International Conference on Learning Representations.\n- [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."}, {"review_id": "NxCcC6maOm", "reviewer": "Reviewer_hS2q", "summary": "This paper introduces InfoSteer, a method designed to enhance the utilization of pre-trained knowledge in language models (LM) during the post-training phase. InfoSteer treats the Transformer's feed-forward network (FFN) layer as an associative key-value memory and employs two strategies to promote the active use of stored memory vectors. Experimental results demonstrate that InfoSteer improves performance across multiple model families and over 15 downstream tasks. It also shows that the steered models can adaptively allocate information resources, focusing more on generating semantically meaningful tokens while using fewer resources on simple transition tokens.", "questions": "See weakness.", "limitations": "N/A", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strength**\n\n1. Innovative model steering approach. InfoSteer offers a novel method for post-training by interpreting the FFN layer as a key-value memory structure. \n\n2. Comprehensive experimental validation: The authors conduct extensive experiments across multiple model series and a wide range of downstream tasks.\n\n3. Insights into model behavior: Beyond performance improvements, the paper provides valuable insights into how the model allocates resources during token generation. It shows that InfoSteer encourages the model to focus computational resources on semantically rich tokens, offering a deeper understanding of the model's internal behavior. \n\n**Weakness**\n\n1. The study lacks comparison to widely-used approaches like reinforcement learning fine-tuning (RLHF). Vanilla supervised fine-tuning (SFT) alone cannot represent the state-of-the-art model capabilities. \n\n2. As shown in Table 1, the differences between various SFT methods diminish as the model size increases. I'm particularly interested in the performance of the proposed method on models around the 30B parameter scale. Or the authors need to offer some explanations for this phenomenon.", "was_revised": true, "final_justification": "I have re-examined this paper and will stand by my previous comments.", "rebuttal": "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:\n\n**About Comparison to RL**\n\nWe 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). \n\nAs 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:\n\n| Method | AIME 2024/2025 | AMC | MATH-500 | Minerva | Olympiad | **Average** |\n|:-------|:--------------:|:---:|:--------:|:-------:|:--------:|:-----------:|\n| **BASELINE** | | | | | | |\n| Qwen2.5-Math-7B | 11.5 / 4.9 | 31.3 | 43.6 | 7.4 | 15.6 | 19.0 |\n| | | | | | | |\n| **SFT METHODS** | | | | | | |\n| SFT | 22.2 / 22.3 | 52.8 | 82.6 | 40.8 | 43.7 | 44.1 |\n| **InfoSteer SFT** (interv.) | **23.2** / **23.4** | **55.0** | **85.8** | **42.6** | **45.6** | **45.9** **(+1.8)** |\n| **InfoSteer SFT** (reg.) | **22.8** / **22.9** | **54.2** | **85.0** | **41.8** | **44.8** | **45.1** **(+1.0)** |\n| | | | | | | |\n| **RL METHODS** | | | | | | |\n| RFT (GRPO) | 25.1 / 15.3 | 62.0 | 84.4 | 39.3 | 46.8 | 45.5 |\n| ***InfoSteer RFT*** (interv.) | ***29.2*** / ***19.4*** | ***67.4*** | ***88.7*** | ***43.2*** | ***50.1*** | ***49.6*** ***(+4.1)*** |\n| ***InfoSteer RFT*** (reg.) | ***28.4*** / ***18.6*** | ***66.2*** | ***87.5*** | ***42.0*** | ***48.9*** | ***48.8*** ***(+3.3)*** |\n\n---\n\nThe 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.\n\nInfoSteer 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. \n\n**About Performance Differences with Model Size and Experiment on 30B Models**\n\nWe 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:\n\n- **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.\n- **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.\n\nWe 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!"}, {"review_id": "86TTTwRsJs", "reviewer": "Reviewer_SHNo", "summary": "This work posits that post training (SFT) currently makes subpar use of knowledge acquired during pre-training and suggest two steering methods to improve upon that, including \n1) intervening on memory vectors with low corresponding key values during the forward pass and \n2) regularizing the entropy of keys in the gradient flow during backpropagation, \n  both which help to encourage more balanced activation of memory vectors ( as opposed to overfitting on a few ).  \n\nEmpirically the authors show the two methods give generally improved performance against SFT on various benchmarks/tasks across 3 models (of various sizes 1B to 8B parameters) including for in and out of distribution settings.\n\nThe paper provides ablations on hyper parameters introduced by the methods, and analysis on distribution shift of key values observed post steering.  The authors additionally propose a token attribution interpretability method based on token-level entropy over the key coefficient distributions as a proxy for how many memory vectors are activated during token generation and provide some qualitative examples of token level highlighting induced by the method for interpretability purposes.", "questions": "Please see questions raised in the strength and weaknesses section.  \n\nAdding comparisons against some additional baselines, clarity improvements and going more into analysis of the results seen ( when are they successful and when not ) would help strengthen the paper.\n\n\nAdditionally, did the authors try combining the two methods together ( which would make logical sense) and if so what results did you see?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Quality:**\nThe overall quality is strong though there are a few areas that need improvement:\n1. In the experiments while its shown the methods do better than vanilla SFT, the proposed methods are not compared against any other SFT (Lora, ReFT, etc) or steering baselines such as inference time interventions (ITI) or mean difference vector steering ( Arditi 24 et al ). \n2. The arithmetic ID/OOD experiment of using the arithmetic GSM8K to post-train and then doing In Distribution results for it and then Out of Distribution results for 5 other arithmetic tasks seems like a misnomer to me.  \nThey collectively are different tasks, but still within the domain of “arithmetic”.  By this sort of logic the authors could have re-used the model checkpoints from table 1 to calculate “out of distribution” results for each of the 8 tasks \nrelative to the other ones ( ie, you’ve done vanilla SFT and steering w/intervention, w/regularization on BoolQ so you could see how those models each do on average against the 7 other tasks (PIQA to OBQA). \n3. In 5.2, the claims in “Insufficient Utilization of Pretrained Knowledge in Modern LLMs” seem a little over stated as written. Saying SFT methods “substantially underutilize knowledge” because Gemma-2-9B base model achieves 90.1% on HellaSwag while reaches 95.7% with intervention steering is comparing against the base model and not the actual baseline vanilla SFT (which gets 94.8% which is within a point).  Without CIs its difficult to tell (i’m assuming the green and blue colors reflect this somehow), but for Gemma-2-9B it looks like of the 8 tasks, only two show some improvement?  \n4. No mention or explanation is given for why the methods do worse both for OBQA for Qwen1.5B & SIQA for LLama 1.B?\n5. Discussion of where and how the methods improved performance and why would be appreciated \n\n\n**Clarity:** The motivation paper is well motivated and mostly well written and easy to follow, but some sections could benefit for clarity:\n1. Section 5.4\n* I didn’t see where it was mentioned what datasets are used for each group task (Reading Comp, Knowledge, etc ) and\n* the average delta’s seem off. For instance, for Commonsense Reasoning, steering improves 2.4 and 1.8 respectively over Vanilla SFT (which implies 2.1) but the Average Delta says its 2.3 ??\n\n2. Section 6.1 \n* What do low key, medium key and high key regions mean exactly here? ie, what are the cutoffs for considering this? This section is interesting, but  needs to include more on the setup and explanations.  For which datasets are these shown?  \n\n3. Section 6.2:  The section discusses how  “LMs might be more self-interpretable via proper guidance without extra training” , but then does instruction tuning ( and its unclear if this is with vanilla SFT or using which of the two interventions ) and then does a post hoc interpretability method involving token entropy, which seems to me to fall into the camp of post-hoc interpretable feature importance methods (Integrated Gradients for instance) which do token highlighting based on model gradients without extra training.  I think the proposed Information Flux method is interesting and hope it gets flushed out more systematically/rigorously in a separate work ( as it stands here, with the few qualitative examples provided its hard to see any clear advantage of it or quantitative trend ) \n\n**Significance:** The results of steering with either of the methods generally seems to improve, but without more study comparing against other baselines for knowledge intensive tasks its hard to tell.  \n\n**Originality:** Using steering methods to promote key/memory vector utilization explicitly seems novel to me from what I’ve seen in the literature, and the interpretability method proposed in 6.2 also seems novel .", "was_revised": true, "final_justification": "The additional experiments and explanations provided help addressed most of my concerns.\nI do think comparing the combination of the two methods together is the biggest open point for me since its still unclear when to use one method vs the other ( but regardless you've convincingly shown they do improve upon SFT alone and other methods ( though I think ITI and Diff MEANs could also be used on top of SFT, but with additional data needs ) . With that I update my rating from 3 to 4.", "rebuttal": "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:\n\n**[Key Concern] About Comparison with Additional Baseline**:\n\nWe 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).\n\nThat 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:\n\n| Model | Training Method | BoolQ | PIQA | SIQA | HellaS | WinoG | ARC-e | ARC-c | OBQA |\n|-------|----------------|-------|------|------|--------|-------|-------|-------|------|\n| **Qwen-2.5-7B** | base model | 68.9 | 81.2 | 77.6 | 87.5 | 80.3 | 79.8 | 65.1 | 77.6 |\n| | + LoRA* | 71.8 | 84.2 | 80.9 | 91.8 | 83.7 | 83.8 | 68.9 | 82.3 |\n| | + ReFT* | 72.1 | 84.5 | 81.2 | 92.1 | 84.0 | 84.0 | 69.2 | 82.5 |\n| | + ITI* | 70.5 | 83.1 | 79.8 | 90.9 | 82.8 | 82.9 | 67.8 | 81.4 |\n| | + DiffMean* | 70.2 | 82.8 | 79.5 | 90.6 | 82.5 | 82.6 | 67.5 | 81.1 |\n| | + vanilla SFT | 72.4 | 84.9 | 81.5 | 92.4 | 84.2 | 84.2 | 69.6 | 82.8 |\n| | + **steered SFT w. intervention** | **74.1** | **86.3** | 81.8 | **95.1** | **87.2** | 86.2 | **73.7** | 84.2 |\n| | + **steered SFT w. regularization** | **76.4** | 85.7 | 81.0 | 94.3 | 86.5 | 85.4 | **72.9** | 83.4 |\n| **LLaMA-3-8B** | base model | 70.3 | 85.6 | 75.7 | 90.8 | 81.9 | 86.2 | 75.3 | 80.5 |\n| | + LoRA* | 74.1 | 88.7 | 79.3 | 93.1 | 85.1 | 89.9 | 79.8 | 85.2 |\n| | + ReFT* | 74.3 | 89.0 | 79.6 | 93.3 | 85.3 | 90.2 | 80.1 | 85.5 |\n| | + ITI* | 72.8 | 87.4 | 78.2 | 92.0 | 84.2 | 88.7 | 78.5 | 84.1 |\n| | + DiffMean* | 72.5 | 87.1 | 77.9 | 91.8 | 83.9 | 88.4 | 78.2 | 83.8 |\n| | + vanilla SFT | 74.6 | 89.3 | 79.9 | 93.5 | 85.6 | 90.5 | 80.4 | 85.8 |\n| | + **steered SFT w. intervention** | **77.1** | 90.2 | **82.0** | 96.3 | 87.4 | **92.4** | 81.6 | 87.5 |\n| | + **steered SFT w. regularization** | 76.5 | 89.5 | 81.2 | 95.6 | 86.8 | 91.7 | 80.9 | 86.8 |\n| **Gemma-2-9B** | base model | 71.6 | 86.3 | 77.2 | 90.1 | 82.5 | 87.5 | 77.8 | 81.7 |\n| | + LoRA* | 73.8 | 89.5 | 81.2 | 94.3 | 86.4 | 91.2 | 81.5 | 86.0 |\n| | + ReFT* | 74.0 | 89.8 | 81.4 | 94.6 | 86.6 | 91.4 | 81.7 | 86.2 |\n| | + ITI* | 72.9 | 88.7 | 80.3 | 93.5 | 85.7 | 90.3 | 80.8 | 85.1 |\n| | + DiffMean* | 72.6 | 88.4 | 80.0 | 93.2 | 85.4 | 90.0 | 80.5 | 84.8 |\n| | + vanilla SFT | 74.3 | 90.1 | 81.7 | 94.8 | 86.9 | 91.7 | 82.0 | 86.4 |\n| | + **steered SFT w. intervention** | **77.2** | 91.8 | 83.1 | 95.7 | 88.5 | 93.5 | 83.4 | **88.2** |\n| | + **steered SFT w. regularization** | 76.5 | 90.9 | 82.4 | 94.9 | 87.8 | 92.8 | 82.7 | 87.5 |\n\nLoRA 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.\n\n**About Arithmetic ID/OOD Experimental Design**:\n\nThank 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.\n\nTo 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:\n\n| **Model** | **ID Eval (OBQA)** | **OOD Eval (Arithmetic)** |\n|-----------|-------------------|---------------------------|\n| Base Model | 80.5 | 68.2 |\n| + Vanilla SFT | 85.8 (+5.3) | 61.4 (**-6.8**) |\n| + Steered SFT w/ interv. | 87.5 (+7.0) | 68.4 (+0.2) |\n| + Steered SFT w/ reg. | 86.8 (+6.3) | 68.1 (-0.1) |\n\nThis 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!\n\n**About Interpretation of the Performance**\n\nOut 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.\n\nRe 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!\n\n**About the Potential Combination of the Method**\n\nYes, 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.\n\n**About Details in Experiment**\n\nWe summarize all comments related to details about our experiment here.\n\n[**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.\n\n\n[**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.\n\nWe will add these details into our manuscript to make it clarified for readers.\n\n**About Section 6.2 Interpretability Claims**\n\nWe 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.\n\nWe sincerely appreciate your suggestions and will incorporate them to enhance our work further. Thank you once again for your valuable feedback!\n\n\n\n## Reference\n\n- [1] Hu, Edward J., et al. \"LoRA: Low-Rank Adaptation of Large Language Models.\" International Conference on Learning Representations.\n- [2] Wu, Zhengxuan, et al. \"ReFT: representation finetuning for language models.\" Proceedings of the 38th International Conference on Neural Information Processing Systems. 2024.\n- [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.\n- [4] Turner, Alexander Matt, et al. \"Steering Language Models with Activation Engineering.\"\n- [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."}, {"review_id": "NHduQYdqmv", "reviewer": "Reviewer_U6HR", "summary": "This paper introduces InfoSteer, a lightweight method that improves post-training of language models by steering toward underutilized FFN memory vectors. It enhances performance and interpretability across models and tasks.", "questions": "1. As in Weakness 1, why is a more uniform key distribution necessarily desirable? Could this harm precision by encouraging irrelevant information use?\n\n2. If the main goal is to increase key entropy, what happens if we apply the intervention-only strategy during inference? Would that still yield benefits?", "limitations": "N/A", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "**Strengths**\n\n1. The method is novel, lightweight and easily integrable into existing pipelines.\n\n2. Extensive experiments across multiple model families and tasks demonstrate improvements\n\n**Weaknesses** \n\n1. **Motivation is underdeveloped**: It's not entirely clear why encouraging a more uniform (higher-entropy) key distribution is beneficial. In practice, a sharp key distribution might indicate the model is confident in retrieving task-relevant information, while low-activation keys may represent irrelevant memory. Forcing these less relevant keys to contribute might dilute useful signals. In the extreme case, a uniform distribution means every token receives the same memory mix, which seems counterproductive. A deeper justification is needed here.\n\n2. **Information Flux definition (Section 6.2)**: Using token-level entropy over key coefficients as a measure of “Information Flux” seems to be unconventional. It’s unclear whether this term has any precedent in prior work.  The paper should explain why this metric reflects self-interpretability. And a comparison with the baseline SFT model on this interpretation is needed. Also, why do operators exhibit lower flux—could this be due to syntactic predictability rather than meaningful information usage?", "was_revised": true, "final_justification": "The rebuttal addressed some of my concerns, and after reading other reviewers' comments, I believe that 4 is a reasonable score.", "rebuttal": "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:\n\n**About Motivation for Higher-Entropy Key Distribution**\n\nThank you for the great question! The reviewer raises a question about why encouraging more uniform key distributions is beneficial. We clarify several key points: \n\n- **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)).\n- **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.\n\n[**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.\n\nWe will dedicate a discussion section in the limitation section and main body in the camera-ready version to ensure this is clarified for readers. \n\n**About Information Flux Definition and Comparison with Baseline**\n\nThanks for asking this! The IF part is quite interesting, and we would love to discuss it further:\n\n- **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].\n- **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.\n\n[**Significance of IF Score Before/After SFT**]\n\nThis 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.\n\n| Token Type | Base Model | Vanilla SFT | Info-steered SFT | Difference |\n|------------|------------|--------------|-----------|------------------------------|\n| Content words | 2.03 ± 0.15 | 2.05 ± 0.16 | 2.67 ± 0.21 | +0.62 |\n| Function words | 2.01 ± 0.14 | 2.02 ± 0.15 | 1.73 ± 0.12 | -0.29 |\n| Punctuation | 1.98 ± 0.13 | 2.00 ± 0.14 | 1.31 ± 0.09 | -0.69 |\n\nWe 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!\n\n## Reference\n\n- [1] Kim, Jiyeon, et al. \"Knowledge Entropy Decay during Language Model Pretraining Hinders New Knowledge Acquisition.\" The Thirteenth International Conference on Learning Representations.\n- [2] Ali, Riccardo, et al. \"Entropy-lens: The information signature of transformer computations.\" arXiv preprint arXiv:2502.16570 (2025).\n- [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."}]

Author rebuttal:

Meta-review: This paper introduces InfoSteer, a lightweight method to enhance language model post-training. The authors claim that standard supervised fine-tuning (SFT) often leads models to underutilize their pre-trained knowledge by overfitting to a narrow set of activations. InfoSteer addresses this by treating FFN layers as key-value memories and applying steering techniques to encourage a more balanced use of this stored information. The method is shown to deliver consistent performance gains across diverse models and tasks.

The paper's primary strengths are that the technique is interesting, simple, and empirically effective, showing consistent improvements in both in-domain and out-of-domain settings. However, the initial submission raised some concerns among reviewers. During this period, the authors had to conduct a substantial amount of new work, including: adding comparisons against four crucial baselines (LoRA, ReFT, ITI, DiffMean), performing a new and stricter out-of-distribution (OOD) experiment to properly validate generalization, and providing quantitative backing for their interpretability metric. The evaluation maintains a controlled comparison by applying the same optimization parameters and single-epoch training across all methods. While this effectively isolates the gains from InfoSteer, it also implicitly assumes this configuration is fair to the baselines, which might perform better with individually tuned hyperparameters.

Furthermore, while the authors provided promising preliminary results applying their method to reinforcement learning, this was only in response to a reviewer's comment Given that RL-based methods are a dominant paradigm for post-training and the authors' own results suggest even stronger gains in that setting, the lack of a complete RL evaluation leaves a key claim of general applicability under-explored. 

Overall though, the paper was strengthened throughout the review process and the new results show the method performs favorably compared to other approaches, have better generalizability, and seems to be even more effective in RL.
Decision: Accept (poster)