Learning Distribution-wise Control in Representation Space for Language Models

Thank you for your insightful feedback on our submission! We’re encouraged by your recognition of our experimental results and value your suggestions for improving clarity and rigor. Below, we address your concerns and outline our revision plan.

About Main Concern: Method Clarity & Training Process

• Clarification: The MI formulation in Section 3.3 (Eq. 2) was intended as a high-level conceptual framing to motivate why interventions are useful in general (i.e., they aim to preserve/enhance task-relevant information) and in early layers (DPI in information theory in appendix A). It is not the objective function directly optimized during training. We agree with you that this will lead to confusion and will revise it in camera-ready.
• Actual Training Objective: Our method is trained end-to-end by minimizing the standard next-token prediction cross-entropy loss of the entire system (frozen base LM + trainable stochastic intervention layers between transformer blocks), identical to previous methods like RED/LoFiT/ReFT. Gradients flow from the loss, back through the reparameterized stochastic intervention networks (see psudo alg below).

pseudocode algorithm for learnable intervention

# for X_input, Y_target in training_batches:
# --- Forward Pass --- 

# 1. forward pass before layer l
Z = LM_pre(X_input)              # Activations at layer l

# 2. Predict intervention distribution params
mu = Net_mu(Z)
sigma = softplus(Net_sigma(Z))   # Ensure sigma > 0

# 3. Stochastic intervention via reparameterization
epsilon = sample_gaussian_noise(shape=mu.shape) # Sample N(0,I)
Z_intervened = mu + sigma * epsilon # Differentiable sampling

# 4. forward pass after layer l
Logits = LM_post(Z_intervened)

# --- Loss & Backward Pass ---
# 5. Calculate standard LM loss and backpropagate only on intervention network

We sincerely apologize for the lack of clarity in the method description and training process. This is a critical point, and we will revise the relevant sections (primarily Sections 3 and 4) to provide a much clearer explanation. We will make the following changes accordingly:

• Explicitly State Training Objective. We will dedicate space in Section 4 (Methodology) to explicitly state that the training objective. We will clearly define the overall system architecture (frozen LM + trainable intervention networks) and specify which parameters are updated during training.
• Add Reference with MI Discussion. We will relocate the MI-based motivation to a distinct subsection, referencing its grounding in variational information bottleneck (VIB) theory [1][2]. This will be clearly separated from the training objective to enhance readability and avoid confusion.

About Theoretical Claims & Justification for Stochasticity

• Key Motivation. Previous research found that the effect is an intervention can be controlled by multiplying a constant to control its magnitude (more positive or less positive). Then why don’t we directly learn the distribution to represent that region? Consequently, we use stochastic intervention to reflect this idea and find it very helpful.
• Key Benefits. The reviewer’s intuition is correct – a primary benefit is enhanced exploration of the intervention space during training. By sampling interventions from a learned distribution $N(μ, σ²)$ instead of applying a single deterministic transformation, the model is exposed to a wider range of related intervention effects.
• Potential Benefits of Test-Time Stochasticity. Test-time stochasticity might offer robustness to slight variations in input representations (See Figure 6 as D-ReFT is much more robust) or act as a cheap form of ensembling if multiple forward passes were considered.

About Terminology and Presentation

Thank you for your valuable feedback! We generally agree with the reviewer and will make the following changes:

• We will replace “point-wise/distribution-wise intervention” with “deterministic/stochastic intervention”.
• We will add non-intervened base model performance for Table 2 and Table 3.
• We will add concrete examples to illustrate “low/high-level control” and include a sentence illustrating “concept space.”
• We will replace “optimal” and other descriptive languages with more measured ones to maintain scientific neutrality.

Overall, we are sincerely grateful for the reviewer’s insightful feedback to help us improve our clarity. We believe that the updated presentation of methodology + the original experiment results will strengthen this work and contribute meaningfully to the intervention community.

Reference

• [1] Tishby, Naftali, and Noga Zaslavsky. "Deep learning and the information bottleneck principle." 2015 ieee information theory workshop (itw). Ieee, 2015.
• [2] Alemi, Alexander A., et al. "Deep variational information bottleneck." International Conference on Learning Representations (2017)

We appreciate your thoughtful and constructive review of our manuscript. We’ve noted your main concerns regarding the statistical significance of our results and the need for clarification on our ablation choices. Thus, we try to address your feedback as below:

About eval setting & statistical significance

To begin, we’d like to clarify that our evaluation setup generally follows eval settings on what's already been established in prior work, such as RED, LoFiT, or ReFT. Consistent with the previous literature, we report the average score from three independent runs. We also agree with the reviewer that including additional metrics like standard deviation would be valuable—please refer to the table below for general results.

With the inclusion of variance metrics, our D-ReFT method demonstrates consistent and statistically meaningful improvements over the ReFT baseline across all tested models (pls also refer to Reviewer dkHP's rebuttal for additional results on Qwen/Gemma models). We’ll expand on the variance details in the revised paper to validate this point.

To further bolster reviewer’s confidence in our results, we’d like to highlight a few additional points:

• All reported numbers for prior work are fully optimized. We meticulously hypertuned all baseline methods to ensure a fair comparison (see Appendix B for details).
• For methods where replication proved challenging, we opted to directly cite the numbers reported in their original papers.
• Unlike some prior approaches that tuned parameters on the test set, we adhered to a separate dev-set hyperparam tuning process to maintain rigor.

We’re happy to provide further details if the reviewers feel additional transparency would strengthen the manuscript.

About ablation with epsilon

We acknowledge that the phrasing in our manuscript might have caused some confusion for the reviewer. To clarify, for each scaling factor $\lambda$, they effectively sample from $N(0, \lambda^2I)$, with the variance scaled accordingly. When the scaling factor is 0, this collapses to the original ReFT (deterministic).

The reviewer suggests that different scaling factors (or init) would ultimately result in the same learned variance. We respectfully disagree. Prior work on intervention methods like ReFT demonstrates that different initializations (LoReFT vs DiReFT) yield markedly distinct outcomes (different init cannot lead $\mu$ to rescale itself). Our experiments also show that introducing varying degrees of randomness at the initial learning stage leads to very different optimization trajectories.

On the presentation front, we appreciate the reviewer’s feedback regarding Figure 4. Initially, we did present the results as described (in a single figure with standard deviations). However, upon closer analysis, we find an interesting correlation between learned variance and accuracy. To emphasize this discovery, we split the data into two figures. We welcome your thoughts on this presentation choice and are open to reverting to the original single-figure format if preferred.

About results on instruction tuning

To address your inquiry about instruction tuning, we provide the following updated results.

$^*$ Number taken from ReFT. Three separate runs for each method.

Our proposed D-ReFT method achieves great performance with a win-rate of 87.19%, surpassing ReFT (85.27%) and other baseline. This superior performance demonstrating that controlled stochasticity during optimization leads to better generalization on instruction-following tasks.

Thank you for catching the typo. We will correct it and improve the paper writing thoroughly.

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

About comparison with LoFiT

We agree with the reviewer that LoFiT is a significant contribution to the representation fine-tuning literature, as we highlighted in our introduction. LoFiT, ReFT, and our method share some same benchmarks—SIQA, ARC-c, and BoolQ (for commonsense reasoning), along with SVAMP (for math). Drawing from existing benchmark results and our own replication efforts, LoFiT and ReFT are tied across these four benchmarks (60.5 for ReFT, 60.7 for LoFiT, and 62.3 for D-ReFT). As for the remaining 11 benchmarks, we’re actively addressing them and plan to incorporate the results into the main body of our work to further solidify its standing in the intervention literature.

About testing on simple tasks

Thank you for your suggestion. Given time constraints, it may be challenging for us to run a fresh batch of experiments on these tasks. However, recent literature, such as RAVEL[1] and Axbench[2] reports scores for learnable intervention methods on simple tasks like entity recognition and concept detection. These results show that learnable interventions, such as ReFT, significantly outperform traditional methods like DiffMean[3] on these simple tasks, achieving over 90% accuracy in some settings. As a result, we may expect the intervention research community to shift toward tackling more complex and challenging tasks.

About the standard deviation

In this paper, we follow the evaluation settings of ReFT, LoFiT, and RED to report the average score across three different random seeds. Generally, the std for ReFT and D-ReFT is similar (avg. std. 2.71e-3 and 2.95e-3 for ReFT and D-ReFT, respectively) across settings.

About the scaling factor of epsilon

The ablation study in section 6.2 varies the scaling factor (from 0 to 3.0 with a step of 0.2) applied to $ϵ ~ N(0,I)$, not the specific values of ϵ. For each scaling factor $\lambda$, they effectively sample from $N(0, \lambda^2I)$ - still centered at 0 but with scaled variance. At scaling factor 0, this reduces to the original deterministic ReFT (no stochasticity). As they increase the scaling factor, they're increasing the variance of the noise distribution, which affects how widely the model explores the neighborhood around the learned mean. By doing this we could maintain the stochastic nature while controlling the amount of randomness.

About the presentation & code

Thank you for providing valuable suggestions regarding paper presentation! We will add a specific section in the appendix to introduce different intervention methods like ReFT to give readers outside this field better background information. We are also working on cleaning the codebase and will release all code and model checkpoints (approximately 70 for all ablations for one model) to make our research accessible to all users.

Reference

• [1] Huang, Jing, et al. "Ravel: Evaluating interpretability methods on disentangling language model representations." Association for Computational Linguistics (2024).
• [2] Wu, Zhengxuan, et al. "AXBENCH: Steering LLMs? Even Simple Baselines Outperform Sparse Autoencoders." arXiv preprint arXiv:2501.17148 (2025).
• [3] Marks, Samuel, and Max Tegmark. "The geometry of truth: Emergent linear structure in large language model representations of true/false datasets." Conference of Language Modeling (2024).

Thank you for your constructive feedback! We sincerely appreciate the time and effort you dedicated to reviewing our paper and providing thoughtful comments. We have carefully reviewed your concerns and questions and recognize that your primary focus is the generalization to other series models and motivation of experiment design. Below, we address each point in detail.

About model diversity

We appreciate the feedback regarding our claims' generalizability beyond Llama models. Previously we followed the setting in ReFT for better comparison, but we realize the limitation of this design. In response, we have conducted additional experiments on Qwen and Gemma models on arithmetic reasoning tasks.

For layer-wise setting

Qwen-2.5-7B (28 layers)

Gemma-3-12B (42 layers)

For all-layer setting:

$^*$ For Qwen-2.5-7B, it’s the top 8 layers. For Gemma-3-12B, it’s the top 11 layers.

These new findings back up our original observations (early layers perform better + a mixed strategy works well) and show they apply across different model types. We’ll update the main paper by adding these results to Table 2 and Table 3, strengthening our overall analysis.

About the motivation for clamping

The key motivation for implementing clamping is to prevent numerical instability issues that may arise from introducing large stochasticity during training (similar to variational methods[1][2]). Regarding potential limitations, theoretically, the latent distribution we learn will deviate slightly from the standard Gaussian due to clamping. We also find without clamping performance is slightly better, but not too much as the extreme values may occur very rarely during sampling.

About the setting in robustness eval

During our preliminary studies, we have tried both synonym replacement (using WordNet) and paraphrase generation (using back translation). However, our empirical analysis revealed that these semantically-preserving transformations produced insufficient perturbation magnitude to effectively discriminate between intervention methodologies, so we implemented a more challenging delete-N setting to provide stronger attack.

About the standard deviation

Generally, the std for ReFT and D-ReFT is similar (avg. std. 2.71e-3 and 2.95e-3 for ReFT and D-ReFT, respectively) across settings.

About the writing and presentation

We acknowledge the reviewer’s valuable feedback concerning several aspects of the paper’s writing and presentation that require improvement. We plan to incorporate those suggestions and revise claims to accurately reflect the scope of experimental validation (model families tested) and ensure terms like "significant" are moderated.

Reference

• [1] Alemi, Alexander A., et al. "Deep variational information bottleneck." International Conference on Learning Representations, 2017.

• [2] Zhu, Zhiyu, et al. "Narrowing Information Bottleneck Theory for Multimodal Image-Text Representations Interpretability." International Conference on Learning Representations, 2025