Token-Aware Inference-Time Intervention for Large Language Model Alignment

Authors: Tianbo Wang, Yuqing Ma, Kewei Liao, Chengzhao Yang, Zhange Zhang, Jiakai Wang, Xianglong Liu
Primary area: applications to computer vision, audio, language, and other modalities
Keywords: LLM Alignment, Inference-Time Intervention, Mutual Information, Graph Network, Misalignment Estimation, Uncertainty Quantification

Abstract:
Effectively mitigating the misalignment of large language models (LLMs) is crucial for ensuring secure AI applications. Inference-Time Intervention (ITI) technique, which applies interventions to internal representations along the probed alignment direction during inference, offers substantial alignment enhancements with minimal cost. However, previous ITI methods adopt coarse sentence-level analysis which neglects the misalignment discrepancy among varied tokens, resulting in deviant alignment direction and inflexible intervention strength.
In this work, we propose a Token-Aware Inference-Time Intervention (TA-ITI)  approach to fully utilize token-level alignment information, therefore realizing superior post-intervention performance. TA-ITI primarily consists of Mutual Information-Guided Token-level Graph Aggregation (MIG) and Misalignment-aware Adaptive Token-level Intervention (MAI). MIG develops a MI-guided graph to exploit the tokens' informative interaction for representation enrichment, thus improving alignment probing and facilitating subsequent intervention.
MAI comprehensively perceives the token-level misalignment degree from token representation and prediction to guide the adaptive adjustment of intervention strength, thereby enhancing final alignment performance. Extensive experiments on three alignment capabilities demonstrate the efficacy of TA-ITI, notably surpassing baseline by 25.8\% on the primary metric of truthfulness.

Decision: reject
Reviewer scores: 8, 6, 5, 6, 5

Reviews and discussion:
<thread>
creation date:
1730364118964

thread_type:
official_review

<thread_object>
creator: reviewer_jfmj

<content>
summary:
The paper introduces TA-ITI to improve representation engineering in large language models by addressing token-level misalignment. Unlike sentence-level methods, TA-ITI uses Mutual Information-guided Token-level Graph Aggregation to capture detailed token interactions, creating refined alignment directions. The Misalignment-aware Adaptive Token-level Intervention then customizes intervention strength based on each token’s misalignment and prediction uncertainty. This token-level approach boosts truthfulness, harmlessness, and fairness of LLMs.

strengths:
- Technically sound: The experiments demonstrate that the method indeed achieves performance improvements across various benchmarks.

weaknesses:
- This work is relatively incremental. The mutual information-based propagation method is a common approach in graph representation learning, applied here primarily to address the limitations of last-token information. Similarly, adaptive intervention during inference is also a common practice.
- The approach incurs additional inference costs due to the need to recalculate token representations multiple times using a graph aggregation algorithm. Additionally, training the Misalignment Estimator requires extra computational resources both for training and inference.

questions:
- On what dataset is the intervention direction obtained? And what about the Misalignment Estimator? If both are derived and tested on the same dataset, the performance improvements might be unsurprising. It would be more meaningful if the authors could demonstrate that the proposed improvements generalize to benchmarks beyond the dataset used for obtaining the intervention weights, for example, to HaluEval.
- Could the authors specify the SFT setup in more detail? I noticed that Li et al. reported high loss and KL divergence when applying SFT on TruthfulQA—did the authors observe similar results?
- Is the performance improvement on RealToxicityPrompts potentially due to excessive refusal? Can ITI methods prevent over-refusal? It would be beneficial if the authors could include experiments on XSTest to address this concern.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730433238848

thread_type:
official_review

<thread_object>
creator: reviewer_wzhx

<content>
summary:
This paper builds upon previous Inference-Time Intervention (ITI) methods by advancing from sentence-level to token-level interventions, thereby achieving higher performance in controlling the truthfulness, harmlessness, and fairness of LLM-generated content. To accomplish this, this paper first employs MI-guided Token-level Graph Aggregation (MIG) to mitigate directional deviation and obtain direction vectors for intervention. Subsequently, this paper uses Misalignment-aware Adaptive Token-level Intervention (MAI) to implement adaptive interventions across distinct tokens. In experiments, this paper validates the substantial improvement of TA-ITI in post-intervention alignment across multiple datasets.

strengths:
1. Addressing the long-standing issue of coarse granularity in sentence-level ITI, this paper achieves fine-grained ITI at the token level. 
2. This paper innovatively combines multiple machine learning algorithms: (1) utilizing MI-guided Token-level Graph Aggregation to obtain global direction vectors, (2) and incorporating Representation Misalignment Estimation as well as Prediction Uncertainty Quantification to implement token-level adaptive intervention.
3. In the experiments, overall, this paper demonstrates a very significant improvement compared to the baselines.

weaknesses:
1.	Some more intuitive explanations are missing. For example: Why does MI-guided aggregation work? The last token of the LLM itself is also an integration of the entire sentence information. Intuitively, why does directional deviation occur, and how does your method intuitively solve this problem?
2.	It is still limited to using directional vectors in an additive manner, and there is not much innovation in the paradigm.
3.	Some functions and variables in the formulas are not clearly explained. For example, the meaning of 'P' and the role of 'y' in Formula 4, as well as the meaning and origin of 'W_LM' in Formula 7.

questions:
1. Some details in Formula 4: How is ‘P’ implemented exactly, and what role does ‘y‘ play in this context?
2. Based on Formula 8: Different tokens only vary in the strength of the intervention while maintaining the same direction. Have you attempted to make the direction adaptive?
3. Before and after performing the MI-guided aggregation, as well as before and after obtaining the direction, and before and after performing the intervention, do any steps involve vector normalization? Does applying normalization or not have any impact?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730629039537

thread_type:
official_review

<thread_object>
creator: reviewer_sipz

<content>
summary:
This paper targets the task of inference-time intervention (ITI) of LLM. Previous ITI methods usually probe and intervene at the sentence level with uniform editing direction for all tokens, which may be deviant and inflexible. To address these problems, this paper proposes a Token-Aware Inference-Time Intervention (TA-ITI) approach, which utilizes a Mutual Information-Guided Token-level Graph Aggregation (MIG) and Misalignment-aware Adaptive Token-level Intervention (MAI) to probe and intervene at the token level. Experiments on truthfulness, harmlessness, and fairness alignment show improvement of TA-ITI over baselines.

Despite the good performance, my main concern is that some implementation details in the method seem to be unintuitive and lack a strong guarantee. More explanations of these choices may enhance the soundness of the proposed method.

strengths:
- The general motivation for token-level intervention is reasonable.
- The experimental results look good.

weaknesses:
- Reasons for some designs in method implementation are unclear. Specifically, 
    - The graph propagation module in MIG lacks strong motivation. Why is the graph structure needed?
    - In lines 209-210, calculating the entropy of representations with discretized bins seems strange. Are there any other choices?

questions:
None
</content>
</thread_object>

</thread>
<thread>
creation date:
1730680037959

thread_type:
official_review

<thread_object>
creator: reviewer_vpdz

<content>
summary:
This paper proposes Token-Aware ITI to address the limitations of sentence-level ITI methods by utilizing information from all tokens in a sentence. It has two main components: MIG, which uses mutual information to analyze token interactions and enhance the accuracy of alignment probing and MAI, which adjusts the intervention strength based on the misalignment level of each token. Experiments on TruthfulQA, RealToxicityPrompts, and StereoSet demonstrate the effectiveness of  TA-ITI. It also proves the effectiveness of MIG and MAI by conducting ablation experiments.

strengths:
The two main contributions of the paper, MIG and MAI prove their effectiveness with the overall performance compared to other baselines and the analysis in section 5.

The inference computation is reasonable which is critical for the inference time intervention.

weaknesses:
The originality of this work is somewhat limited since it is an expansion of a previous paper ([1]) from sentence-level to token-level. 

More background explanations on editing-based inference-time intervention are needed in the related section and preliminaries section, rather than a simple summary of previous works with citations. 

The generalizability of MIG and MAI is limited since they rely on supervised trained misalignment probes and misalignment estimators. 

[1] Li et al. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36, 2024a.

questions:
When constructing the token-level misalignment dataset, how do you decide if this token is prone to overall misalignment? Since this process is the most critical part of MAI, the data construction process needs further details.

Since the MI-based graph network is constructed based on the training samples of the particular dataset, how does the alignment probe transfer to unseen test data in inference time?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730852133821

thread_type:
official_review

<thread_object>
creator: reviewer_m7tw

<content>
summary:
This paper proposed token level methods to conduct inference time intervention for alignment. The methods provide a new way to use mutual information to compute sentence level directions from tokens, and apply different weights to different tokens and add stronger intervention to potentially key misaligned tokens.

strengths:
- The paper is well-written, well-motivated.
- Extensive experiments are conducted on several aspects, with proper ablation and analysis. It shows that the methods not only improve the scores, but also generate fluent questions
- The added cost of inference is acceptable giving the intervention results.

weaknesses:
It seems there are two hyper-parameter alpha and beta that needs to be tuned. It seems that the effect can be influenced by the hyperparameter choice. For example, in beta, the performance is non-optimal when it is <0.4. If these hyperparameter behaviors do hold across different tasks, it will be harder to apply this method.

questions:
- L210: how are the bins setup when computing the entropy?
- Also about the graph entropy propagation method, would this reach a stationary distribution after multiple rounds of propagation?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732239092320

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Response to all Reviewers

comment:
We thank all the reviewers for their thoughtful and valuable feedback! We have conducted additional experiments and revised the paper based on the reviews. We highlighted all changes in red. We also included more results and analysis to make the paper more comprehensive. Additionally, we have corrected the symbol of the probe from $p$ to $P$ in Figures 1 and 2 to ensure consistency. For ease of reading, we have appended the Appendix, which was previously included in the Supplementary Material, to the end of the main text.

Below is a summary of the main changes. Please let us know if you have further questions.

||||
|-|-|-|
|**Change**|**Section**|**Related Reviewers**|
|Add more background description of ITI|Section 2.2 and Appendix D|Reviewer VpDz  |
|More analysis of sentence-level deficiency|Appendix A  |Reviewer WzHX, Reviewer siPZ  |
|Analysis of graph distribution in MIG|Appendix A|Reviewer M7Tw|
|Analysis of generalization|Appendix A|Reviewer VpDz, Reviewer jfmj|
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes Token-Aware ITI to address the limitations of sentence-level ITI methods by utilizing information from all tokens in a sentence. This paper innovatively combines multiple machine learning algorithms: (1) utilizing MI-guided Token-level Graph Aggregation to obtain global direction vectors, (2) and incorporating Representation Misalignment Estimation as well as Prediction Uncertainty Quantification to implement token-level adaptive intervention. Experiments on TruthfulQA, RealToxicityPrompts, and StereoSet demonstrate the effectiveness of TA-ITI. It also proves the effectiveness of MIG and MAI by conducting ablation experiments. 

The reviewers generally agree on the comprehensiveness of the experiments. However, there are several non-trivial concerns shared among the reviews. First, reviewers believe that this work is somewhat incremental and it is an expansion of a previous paper. Second, some designs in the method seem to be unintuitive and lack a strong guarantee. Third, the presented approach incurs additional inference costs due to the need to recalculate token representations multiple times using a graph aggregation algorithm, which harms the utility of the method in real-world practice. Not all these issues were addressed through the author response phase.