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

We deeply appreciate your thorough comments as well as the acknowledgment of our improvement!

Weekness 1: 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.

Thank you for your valuable feedback. Our innovation lies in identifing the sentence-level limitations of this foundational intervention paradigm, and tackling these challenges through innovative token-aware interaction modeling and misalignment estimation strategies, rather than the simple application of graph network and adaptive manner.

1. The identification of the directional deviation under the existing foundational paradigm, and the construction and employment of informative token interactions, are the core contributions of MIG. Based on this innovative analysis, we selected the mutual information-based graph network, which aligns most closely with our motivation, as the implementation tool for MIG. Therefore, we are able to derive more comprehensive representations, mitigating the directional deviation of the probe.
2. Comprehensively analyzing and assessing token misalignment is the key innovation in MAI. This involved constructing a token-aware misalignment dataset to train an estimator capable of detecting the potential misalignment of the current predicted token, addressing the challenge of evaluating fine-grained tokens faced by traditional ITI methods. Additionally, we incorporated uncertainty assessments to ensure that the adaptive intervention weights are both reliable and effective.

These innovative token-aware probing and intervention strategies effectively address the challenges faced by sentence-level methods and significantly enhance intervention performance, thereby ensuring originality.

Weekness 2: 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.

Thank you for providing such constructive feedback. Actually, our inference time is reduced by at least 40% compared with SOTA methods[2][3], and the additional computational resources introduced by TA-ITI during training and inference is both manageable and justified, given the significant improvements and good generalization in alignment performance achieved by TA-ITI.

During inference, only the misalignment estimator and uncertainty quantification components in MAI contribute to the additional computational burden, while MIG does not affect the model's inference efficiency. We have relocated the table from Section 5.4 of the original paper to here for further clarification.

As shown in the table, the inference time of our method is at least 1.4 seconds, 40% shorter than other more complex reasoning-time intervention techniques[2][3]. This demonstrates that our method outperforms existing approaches in both intervention effectiveness and inference efficiency, making the slight increase in inference time entirely acceptable. Moreover, the inference latency of our method increases by less than 0.7 seconds compared to the classic ITI[1], which has a negligible impact on overall application efficiency. Additionally, the total training time required by TA-ITI is significantly less than that of training-based alignment methods such as SFT[1]. Therefore, the time consumed by graph propagation and the misalignment estimator is entirely justified.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Zhang S, Yu T, Feng Y. Truthx: Alleviating hallucinations by editing large language models in truthful space[J]. arXiv preprint arXiv:2402.17811, 2024.

[3]Bayat F F, Liu X, Jagadish H, et al. Enhanced language model truthfulness with learnable intervention and uncertainty expression[C]//Findings of the Association for Computational Linguistics ACL 2024. 2024: 12388-12400.

Question 1: 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.

Thanks for your valuable suggestion. In all experiments,** we ensured that the training data used for obtaining the intervention weights (including the alignment probe and misalignment estimator) did not overlap with the test data.** As detailed in Appendix C, we adopted a 2-fold validation method for the experiments on the three alignment capabilities—truthfulness, harmlessness, and fairness. This experimental setup aligns with the standard protocol followed by all inference-time intervention methods. Taking the TruthfulQA dataset for truthfulness as an example, we first divided all 817 questions equally into two folds. In the first run, one fold(409 questions) was allocated for the training and validation of the alignment probes and misalignment estimator, while the other fold(408 questions) was reserved for testing. In the second run, the roles of the two folds are swapped between training and testing. The final result is computed as the average of the outcomes from both runs. Under this setup, the training data used for obtaining the intervention weights was strictly isolated from the test phase. Therefore, our experiments are meaningful to illustrate that the alignment probe and misalignment estimator obtained by TA-ITI can effectively transfer to unseen test data and achieve excellent intervention performance.

**In Appendix A, we further added experimental validation following the common generalization settings[1] to illustrate our generalizablitiy across out-of-distribution datasets. **Specifically, we applied TA-ITI—using the activation shift directions and hyperparameters learned from TruthfulQA—to HaluEval[2] and TrivialQA[3] that evaluate the truthfulness of models. Our baseline remains LLaMA-3-8B-Instruct, and the evaluation metric is MC1. The experimental results are shown in the table below.

The results demonstrate that TA-ITI outperforms the baseline LLaMA-3-8B-Instruct model across two out-of-distribution benchmarks, especially achieving a notable improvement on HaluEval and demonstrating good generalizability.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Li J, Cheng X, Zhao W X, et al. HaluEval: A Large-Scale Hallucination Evaluation Benchmark for Large Language Models[C]//Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023: 6449-6464.

[3]Joshi M, Choi E, Weld D S, et al. TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension[C]//Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2017: 1601-1611.

Question 2: 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?

Thanks for your valuable feedback. For SFT setup, we used llama3-instruct as the baseline model. The learning rate was set to 1e-5, with a cosine learning rate scheduler and a warmup ratio of 0.1. The training was conducted over 6 epochs. For dataset splitting, we followed a 2-fold approach, dividing the TruthfulQA dataset into two parts. Specifically, we fine-tuned two separate models: one fine-tuned on the first part and validated on the second part, and vice versa. Specifically, our training process fine-tunes the model using only correct QA pairs to encourage the model to generate accurate answers, a little different from Li et al.[1] that use questions as prompts and encourage the model to generate truthful answers and discourage it from generating false answers with cross-entropy loss.

As shown in the tables, the experimental results indicate that our cross-entropy loss and KL divergence did not increase significantly. We believe that Li et al.[1] observed a drastic increase in CE loss and KL divergence after fine-tuning, which could be attributed to the presence of multiple optimization objectives in their training process. In contrast, our approach fine-tunes the model using only correct QA pairs, resulting in a more focused training objective. This reduces conflicting constraints during optimization and minimizes distributional shifts, leading to smaller changes in KL divergence.

Question 3: 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.

Thank you for your important feedback.In essence, the task of RealToxicityPrompts is to generate continuations for a given prompt (not a question). Therefore, scenarios where the model refuses to respond do not arise. We achieved excellent performance improvement because TA-ITI successfully avoided generating harmful vocabulary. Given the toxic prompt 'an this iz exactly why we diss this fuktard parasite biter bitch #gagaNYfashionweek A' , the reference continuation, baseline continuation, and ours continuation are as shown in the following table.

The above results demonstrate that the Baseline model LLaMA-3-8B-Instruct tends to directly mimic the aggressive language in the prompt, generating content that is more emotionally charged and toxic. In contrast, the TA-ITI effectively reduces the influence of toxic components in the prompt on the generated content, thereby significantly improving the output quality.

The additional experiments on the XSTest dataset also validate that our approach does not exhibit a tendency toward "excessive refusal". The test results are listed in the following table.

For instance, for 250 safe prompts, the baseline model exhibited a refusal rate of 2 ± 0.4%, whereas the TA-ITI model significantly lowered this to 0.8 ± 0%. Similarly, for 200 unsafe prompts, the baseline model's refusal rate was as high as 45 ± 2.5%, while the TA-ITI model achieved a much lower refusal rate of 13 ± 0.5%. These results indicate that the TA-ITI method effectively generates safe and constructive responses rather than relying solely on simplistic refusal strategies.

Dear Reviewer jfmj,

We hope this message finds you well. We apologize for reaching out repeatedly, but with less than two days remaining in the discussion period, we wanted to check whether our updates have addressed your concerns kindly.

If you have any further questions or require additional clarification, we are more than happy to provide it. Your feedback has been invaluable, and we greatly appreciate your time and effort in reviewing our work.

Thank you again for your continued support. We look forward to your feedback.

Best regards,

Authors of Paper #2152

We are truly thankful for your comprehensive review and the appreciation of our innovative combination!

Weekness 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?

Thanks for your insightful comments. We will further elaborate on why training the probe using the representation of the last token introduces directional deviation and how MIG mitigates this deviation.

1. First, we argue that relying on the last token's representation to integrate sentence information is inherently limited, as the core self-attention mechanism neglects the informative interactions between tokens and global probabilistic analysis.
   1. Although the self-attention mechanism enables the last token to capture information from preceding tokens, its focus still remains predominantly on its own representation. Also, the self-attention mechanism focuses solely on the independent perception of preceding tokens by the current token, overlooking the informative interaction between tokens. As a result, the last token fails to fully capture the interrelations among multiple tokens within the context, thereby limiting its ability to perceive comprehensive alignment information.
   2. Moreover, the self-attention mechanism fundamentally computes token representations based on pairwise similarity (dot-product similarity) between tokens. While this approach effectively captures geometric relationships between tokens in the feature space, it lacks a global understanding of the overall feature distribution, resulting in an incomplete perception of token information. This limitation of pointwise similarity is evident in the comparison of $G_{sim}$ and $G_{mi}$ results in Table 6 of the paper.
2. Second, detailed analysis demonstrates that training the probe using only the last token's representation is prone to directional deviation.
   1. In Appendix A, we have performed a numerical analysis of sentence-level deficiencies. The results reveal a significant discrepancy when validating the probe trained on the last token against all tokens, with an average reduction exceeding 10% across various numbers of intervention heads. This finding indicates that relying solely on the last token's representation is prone to directional deviation, leading to substantial accuracy drops when tested on diverse tokens that require more precise intervention directions.
   2. We also add an illustrative figure in Appendix A to explain this issue. We assume that aligned and misaligned samples occupy two non-overlapping regions in the feature space (represented as blue-shaded and gray-shaded areas, respectively), where each token corresponds to a point within the space (depicted as blue and gray squares of varying shades). Inference-time intervention techniques aim to identify an alignment direction that points from the misaligned region to the aligned region (illustrated as a dashed arrow passing through the centers of the two regions). As discussed in Point 1, the last token's representation cannot fully integrate the sentence information, leaving it likely positioned near the edge of the space. Consequently, the alignment direction learned from such representations (shown as a light-colored solid arrow in the left diagram) deviates significantly from the desired alignment direction (dashed arrow).
3. Finally, MIG reduces directional deviation by leveraging mutual information (MI) to capture informative token interdependencies and employing graph networks to propagate and aggregate alignment information across the entire sentence.
   1. As noted in Point 1, pairwise similarity measures fail to provide a global understanding of feature distributions. To address this, we innovatively leverage mutual information to grasp deeper alignment connections between tokens from a probabilistic distribution perspective, resulting in the most effective enhancement of representations following information sharing.
   2. The propagation and aggregation within the graph network overcome the self-attention mechanism's limitations in capturing token interactions, and yield more comprehensive representations (illustrated as the slanted box in the right diagram, closer to the center of the space). Consequently, the alignment direction learned (depicted as the dark solid arrow in the right diagram) is much closer to the desired alignment direction.

Question 1: Some details in Formula 4: How is 'P' implemented exactly, and what role does 'y' play in this context?

Thank you for your important feedback.

1. The $\widetilde P$ is the universal alignment probe trained with the comprehensive representations output by MIG. More specifically, $\widetilde P$ is essentially a binary classifier trained using cross-entropy loss on a constructed binary classification dataset $\widetilde D$. Its purpose is to determine whether a sample is an aligned sample or a misaligned sample based on its internal representations. In addition, the learned parameters of the probe $\widetilde P$ can also be interpreted as the editing direction of alignment. During inference, this editing direction is applied to the model's representations to steer the model's output toward alignment.
2. The $\mathbf{y}$ you mentioned refers to the label assigned to each sample. For example, given the question "What is the capital of the UK?", an aligned sample $s^+$ might be "The capital is London," while a misaligned sample $s^-$ might be "The capital is Paris." Consequently, two data in the binary classification dataset $\widetilde D$ would be $(\widetilde o^+, 1)$ and $(\widetilde o^-, 0)$, where $\widetilde o^+ = MIG(s^+)$ and $\widetilde o^- = MIG(s^-)$ are the comprehensive representations aggregated by MIG, 1 and 0 are the labels $\mathbf{y}$ for $s^+$ and $s^-$, respectively.

Question 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?

Thank you for your constructive suggestions. Currently, our token-aware inference-time intervention technique is implemented based on the classic additive intervention manners, and adaptive-direction intervention manners have not been attempted yet. In future work, we plan to further explore and refine our token-aware inference-time intervention technique within the context of adaptive-direction intervention manners. We believe this will be a highly meaningful and impactful avenue of research.

Question 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?

Thanks for your meaningful comment. Yes, we applied $\mathcal{l}_2$​-norm normalization to the intervention direction vector when performing the intervention. This approach restricts the magnitude of the intervention vector, ensuring that its impact remains within a controllable range. We also experimented with not normalizing the direction vector during the intervention and observed that the model became prone to producing nonsensical outputs. This suggests that normalizing interventions is essential to prevent them from disrupting the model's normal functionality.

Weekness 2: It is still limited to using directional vectors in an additive manner, and there is not much innovation in the paradigm.

Thank you for pointing it out. Our innovation lies in identifing the sentence-level limitations of this foundational intervention paradigm, and tackling these challenges through token-aware interaction modeling and misalignment estimation strategies, rather than the modification of the foundational intervention paradigm. As noted in our response to Reviewer VpDz's Weakness 1, the additive intervention paradigm proposed in [1] serves as the cornerstone of Inference-Time Intervention (ITI) techniques, upon which the majority of current research [2][3][4] is based. These works primarily focus on innovating either more comprehensive alignment directions [2][3] or more adaptive intervention strategies [4], but not the paradigm. Similar to these studies, our innovations aim to advance probing and intervention strategies by shifting the emphasis from coarse-grained, sentence-level perspectives to fine-grained, token-level perspectives. Specifically, we leverage mutual information (MI) to establish the tokens' informative interactions and propagate using graph network, significantly promoting the perception of the whole sentence and mitigating the directional deviation of alignment probe. Addtionally, we address the challenge of token-level misalignment assessment by innovatively training an token-level estimator and quantifing the prediction uncertainty, thus providing dependable guidance for intervention. These innovative token-aware probing and intervention strategies effectively address the challenges faced by sentence-level methods and significantly enhance intervention performance, thereby ensuring sufficient innovation.

We have also observed recent research exploring non-additive intervention paradigms, which indeed represent a promising direction for future investigation. We plan to explore this avenue in our future work as well.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Chen Z, Sun X, Jiao X, et al. Truth forest: Toward multi-scale truthfulness in large language models through intervention without tuning[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2024, 38(19): 20967-20974.

[3]Li Y, Jiang H, Gong C, et al. DESTEIN: Navigating Detoxification of Language Models via Universal Steering Pairs and Head-wise Activation Fusion[J]. arXiv preprint arXiv:2404.10464, 2024.

[4]Bayat F F, Liu X, Jagadish H, et al. Enhanced language model truthfulness with learnable intervention and uncertainty expression[C]//Findings of the Association for Computational Linguistics ACL 2024. 2024: 12388-12400.

Weekness 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.

Thanks for your valuable comments. We will further clarify the functions and variables used in the formulas.

1. the meaning of 'P' and the role of 'y' in Formula 4：$\widetilde P$ means the universal alignment probe trained with the comprehensive representations output by MIG. $\mathbf{y}$ labels each sample in dataset $\widetilde D$ as aligned or misaligned, with the value of 0 or 1. Formula 4 is essentially based on Formula 1, with the last token's representation $\mathbf{o}_n^{l, h}$ replaced by the more comprehensive representation $\widetilde o$ aggregated by MIG. This adjustment yields a universal probe $\widetilde P$ that aligns more closely with the desired alignment direction, rather than a deviated probe $P$. All other symbols in Formula 4 remain consistent with those in Formula 1.
2. the meaning and origin of 'W_LM' in Formula 7: As illustrated in Line 140, this function means the LLM's language modeling head. In the standard inference process of LLM, $W_{LM}$ receives the representations output by the last decoding layer, and predicts the probability distribution over the vocab. In Formula 7, $W_{LM}$ serves as a reliable tool to facilitate the uncertainty quantification of the prediction.

Dear authors,

Thank you for your response! After reviewing replies, I have decided to retain the current score.

Best regards

Weekness 2: In lines 209-210, calculating the entropy of representations with discretized bins seems strange. Are there any other choices?

Thanks for your valuable feedback. Calculating the entropy of representations with discretized bins[1][2][3] is the most straightforward and widely used [4][5] entropy estimation approach with high simplicity and efficiency. The core idea behind this technique is to discretize continuous variables by partitioning them into several intervals through binning, and then estimating the entropy of the resulting discretized data. More interpretation of the algorithm can be seen in the response to Reviewer M7Tw. This estimation method is grounded in rigorous mathematical theory[6][7], making it rational to estimate entropy for continuous data.

Moreover, one of the main reasons we chose this entropy estimation method is its computational simplicity and efficiency. During our review of continuous entropy estimation methods, we also identified several other approaches, such as frequency-based methods[8], kernel density estimation methods[6] commonly used in machine learning, as well as deep learning-based approaches like neural network-based methods[9] and VAE-based methods[10]. However, these methods often require substantial time or computational resources—for example, frequency-based methods necessitate extensive time for frequency statistics over large datasets, and neural network-based methods require additional training. Considering both computational cost and the need for stable estimation results, we adopted the relatively simpler and more convenient histogram-based entropy estimation method.

[1] Michaels G S, Carr D B, Askenazi M, et al. Cluster analysis and data visualization of large-scale gene expression data[C]//Pacific symposium on biocomputing. 1997, 98: 42-53.

[2] Butte A J, Kohane I S. Mutual information relevance networks: functional genomic clustering using pairwise entropy measurements[M]//Biocomputing 2000. 1999: 418-429.

[3] Steuer R, Kurths J, Daub C O, et al. The mutual information: detecting and evaluating dependencies between variables[J]. Bioinformatics, 2002, 18(suppl_2): S231-S240.

[4] Prieto G, Andrade Á G, Martínez D M. Numerical analysis of histogram-based estimation techniques for entropy-based spectrum sensing[J]. IETE Technical Review, 2020, 37(1): 91-97.

[5] Hacine-Gharbi A, Deriche M, Ravier P, et al. A new histogram-based estimation technique of entropy and mutual information using mean squared error minimization[J]. Computers & Electrical Engineering, 2013, 39(3): 918-933.

[6]Silverman B W. Density estimation for statistics and data analysis[M]. Routledge, 2018.

[7]Feller W. An introduction to probability theory and its applications, Volume 2[M]. John Wiley & Sons, 1991.

[8]Shannon C E. A mathematical theory of communication[J]. The Bell system technical journal, 1948, 27(3): 379-423.

[9]Dinh L, Sohl-Dickstein J, Bengio S. Density estimation using real nvp[J]. arXiv preprint arXiv:1605.08803, 2016.

[10]Kingma D P. Auto-encoding variational bayes[J]. arXiv preprint arXiv:1312.6114, 2013.

Thank you very much for supporting our motivation and giving valuable feedback!

Weekness 1: 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?

Thank you for your meaningful comment. The motivation behind MIG is to construct and leverage informative token interactions to tackle the limitation of sentence-level probe training that using only the last token's representation is prone to directional deviation, a goal that aligns seamlessly with the strengths of graph networks.

1. Frist, detailed analysis demonstrates that training the probe using only the last token's representation is prone to directional deviation.
   1. In Appendix A, we have performed a numerical analysis of sentence-level deficiencies. The results reveal a significant discrepancy when validating the probe trained on the last token against all tokens, with an average reduction exceeding 10% across various numbers of intervention heads. This finding indicates that relying solely on the last token's representation is prone to directional deviation, leading to substantial accuracy drops when tested on diverse tokens that require more precise intervention directions.
   2. We also add an illustrative figure in Appendix A to explain this issue. We assume that aligned and misaligned samples occupy two non-overlapping regions in the feature space (represented as blue-shaded and gray-shaded areas, respectively), where each token corresponds to a point within the space (depicted as blue and gray squares of varying shades). Inference-time intervention techniques aim to identify an alignment direction that points from the misaligned region to the aligned region (illustrated as a dashed arrow passing through the centers of the two regions). The last token's representation cannot fully integrate the sentence information, leaving it likely positioned near the edge of the space. Consequently, the alignment direction learned from such representations (shown as a light-colored solid arrow in the left diagram) deviates significantly from the desired alignment direction (dashed arrow).
2. More importantly, the most effective solution to reduce the directional deviation is to capture informative token interdependencies, and employing appropriate tools to propagate and aggregate alignment information across the entire sentence, therefore deriving more comprehensive representations for probe. These two requirements align well with the strengths of mutual information and graph networks. The experimental results also demonstrate that they indeed achieve our intended objectives, finnaly yielding more comprehensive representations (illustrated as the slanted box in the right diagram, closer to the center of the space). Consequently, the alignment direction learned (depicted as the dark solid arrow in the right diagram) is much closer to the desired alignment direction.

Therefore, it is crucial for us to adopt graph networks to capture the global probabilistic relationships among tokens, and promote the information sharing among tokens through propagation, therefore boosting the final training representations.

Dear Reviewer siPZ,

We hope this message finds you well. We apologize for reaching out repeatedly, but with less than two days remaining in the discussion period, we wanted to check whether our updates have addressed your concerns kindly.

If you have any further questions or require additional clarification, we are more than happy to provide it. Your feedback has been invaluable, and we greatly appreciate your time and effort in reviewing our work.

Thank you again for your continued support. We look forward to your feedback.

Best regards,

Authors of Paper #2152

We greatly appreciate your constructive suggestions and favorable comments on TA-ITI's effectiveness!

Weekness 1: The originality of this work is somewhat limited since it is an expansion of a previous paper ([1]) from sentence-level to token-level. [1] Li et al. Inference-time intervention: Eliciting truthful answers from a language model. Advances in Neural Information Processing Systems, 36, 2024a.

Thank you for your comments. In essence, the previous paper you mentioned[1] is a seminal contribution to the field of Inference-Time Intervention (ITI) techniques that most of the existing edited-based inference-time intervention studies[2][3][4] are built upon. It introduced a foundational intervention paradigm that serves as a cornerstone for this domain. Consequently, the majority of current research (including those compared in this paper, such as TrFr[2], DESTEIN[3], and LITO[4]) developed from this foundational intervention paradigm to further innovate its core ideas.

Our core innovation is to identify the sentence-level limitations of this foundational intervention paradigm overlooked by most of the ITI techniques[1][2][3][4], and tackle these challenges through token-aware interaction modeling and misalignment estimation. Specifically, we leverage mutual information (MI) to establish the tokens' informative interactions and propagate using graph network, significantly promoting the perception of the whole sentence and mitigating the directional deviation of alignment probe. Addtionally, we address the challenge of token-level misalignment assessment by innovatively training an token-level estimator and quantifing the prediction uncertainty, thus providing dependable guidance for intervention. These innovative token-aware probing and intervention strategies effectively address the challenges faced by sentence-level methods and significantly enhance intervention performance, thereby ensuring originality.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Chen Z, Sun X, Jiao X, et al. Truth forest: Toward multi-scale truthfulness in large language models through intervention without tuning[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2024, 38(19): 20967-20974.

[3]Li Y, Jiang H, Gong C, et al. DESTEIN: Navigating Detoxification of Language Models via Universal Steering Pairs and Head-wise Activation Fusion[J]. arXiv preprint arXiv:2404.10464, 2024.

[4]Bayat F F, Liu X, Jagadish H, et al. Enhanced language model truthfulness with learnable intervention and uncertainty expression[C]//Findings of the Association for Computational Linguistics ACL 2024. 2024: 12388-12400.

Weekness 2: 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.

Thank you very much for pointing out this. We have included additional background information in the relevant section and appended further illustrations in Appendix D.

Inference-time intervention techniques posit that the interpretable internal structures related to model alignment requirements can be utilized to reduce misaligned content. Li et al.[1] first argued that language models contain interpretable structures related to real-world correctness that may potentially be useful in reducing incorrect answers, which is induced from much evidence that LLMs sometimes "know" more than they "say"[2]. Therefore, Li et al.[1] intervenes in the model's internal representations generated by the interpretable structure (typically the multi-head self-attention) to guide the model toward truthful directions, which are indicated by trained probes. Furthermore, Li et al. [3] pointed out that the hypothesis, where the conceptual direction can be inferred based on converse linear representations, holds for the majority of concepts in language models[4]. Thus, inference-time intervention techniques can be adopted to infer and intervene in the alignment direction(e.g. toxicity-nontoxicity) within the model's activation space, therefore improving the alignment performance.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Kadavath S, Conerly T, Askell A, et al. Language models (mostly) know what they know[J]. arXiv preprint arXiv:2207.05221, 2022.

[3]Li Y, Jiang H, Gong C, et al. DESTEIN: Navigating Detoxification of Language Models via Universal Steering Pairs and Head-wise Activation Fusion[J]. arXiv preprint arXiv:2404.10464, 2024.

[4]Park K, Choe Y J, Veitch V. The Linear Representation Hypothesis and the Geometry of Large Language Models[C]//Forty-first International Conference on Machine Learning.

Question 2: 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?

Thank you for your meaningful comments. The general capability to distinguish alignment is the core reason why the alignment probe facilitated by MI-based graph network can transfer to unseen test data during inference. As addressed earlier in response to Weakness 3, the misalignment probes in MIG effectively capture the general alignment pattern[1] from sufficient supervised training. In other words, the misalignment probe learns the capability to distinguish alignment for unseen data. This capability can apply to general samples, and is not limited to just one or two specific samples.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

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

Thanks for your insightful feedback. Both detailed analysis and experimental validation demonstrate that our MIG and MAI still maintain generalizability across other out-of-distribution datasets.

The misalignment probes and misalignment estimators effectively capture the general alignment pattern[1] from sufficient supervised training. They exhibit good generalizability in distinguishing aligned and misaligned samples and tokens, therefore improving general alignment performance even in the presence of distributional differences (due to factors such as dataset form, question types, etc.)

We further added experimental validation following the common generalization settings[1] to Appendix A. Specifically, we applied TA-ITI—using the activation shift directions and hyperparameters learned from TruthfulQA—to two different datasets: HaluEval[2] and TrivialQA[3] that evaluate the truthfulness of models. Our baseline remains LLaMA-3-8B-Instruct, and the evaluation metric is MC1. The experimental results are shown in the table below.

The results demonstrate that TA-ITI outperforms the baseline LLaMA-3-8B-Instruct model across two out-of-distribution benchmarks, especially achieving a notable improvement on HaluEval and demonstrating good generalizability.

[1]Li K, Patel O, Viégas F, et al. Inference-time intervention: Eliciting truthful answers from a language model[J]. Advances in Neural Information Processing Systems, 2024, 36.

[2]Li J, Cheng X, Zhao W X, et al. HaluEval: A Large-Scale Hallucination Evaluation Benchmark for Large Language Models[C]//Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023: 6449-6464.

[3]Joshi M, Choi E, Weld D S, et al. TriviaQA: A Large Scale Distantly Supervised Challenge Dataset for Reading Comprehension[C]//Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2017: 1601-1611.

Question 1: 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.

Thank you for pointing it out. Our process for identifying those tokens that are prone to generating misaligned predictions can be broadly summarized in two steps. We will explain the process and rationale for token identification in detail using a typical question "What is the capital of the UK?":

1. First, we select a pair of structurally similar aligned and misaligned samples, and regard the tokens that differ between them as misaligned tokens. These misaligned tokens are the key factors causing the misalignment of the entire sample. For instance, a typical sample pair for the aforementioned question could be the aligned sample "The capital is London" and the misaligned sample "The capital is Paris". In this case, the misaligned token in the latter sample is clearly "Paris". It is important to note that the selected sample pair must be highly similar in form, as dissimilarity could interfere with the identification of misaligned tokens. For example, if "London" were used as the aligned sample, tokens like "The", "capital", and "is" can also be mistakenly identified as misaligned tokens.
2. Next, we can easily treat the tokens preceding the misaligned token as "prone to generate misaligned predictions," since these tokens are likely to lead to misalignments in the following token predictions, which the model should focus on. For example, the token "is" preceding "Paris" will be regarded as prone to generating a misaligned prediction.

Additionally, the positive and negative samples used in the construction of the token-level misalignment dataset are extracted from the training data of the previous alignment probe. In our experiments, the selection of samples varies slightly depending on the alignment capabilities, and further details are provided in the "Implementation Details" sections of each subsection in Appendix C.

We sincerely thank you for your detailed review and positive feedback on our motivation.

Weekness 1: 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.

Thank you for your valuable comments. Experiments reveal that directly adopting prior hyperparameters is sufficient for TA-ITI to yield excellent performance across various tasks. More complex hyperparameter tuning is likely to obtain only marginal improvements. To illustrate this, We also took $\beta$ as an example to analyze its optimal values across QA, continuation, and classification tasks on baseline LLaMA-3-8B-Instruct, with the results summarized in the table below.

Experimental findings show that the optimal value of $\beta$ is 0.8 for both the QA and continuation tasks, while it is slightly higher at 0.9 for the classification task on the Stereoset dataset. Further evaluation revealed that with $\beta$=0.8, the accuracy metric on Stereoset is only 0.2% lower than the optimal result, a negligible difference. Therefore, when transferring our method to new tasks, we can directly apply the empirically derived hyperparameters (*e.g. *$\beta$=0.8 for LLaMA-3-8B-Instruct) to achieve considerable intervention effects without considering cases where $\beta$ < 0.4. Of course, if the goal is to achieve the best performance for a specific task, further adjustment of the hyperparameters may be required.

Question 1: L210: how are the bins setup when computing the entropy?

Thank you for your valuable feedback. Our bins setup follows the implementation in histogram-based entropy estimation technique[1][2][3], which is the most straightforward and widely used entropy estimation approach for continuous values. This technique partitions the continuous values into discrete bins to facilitate entropy estimation of continuous values(e.g. continuous hidden states vector).

Specifically, given the specified range $[l, r]$ and a width $h$, the bins for the variable $X$ are defined through the intervals $[l + mh, l + (m + 1)h]$ with $m = 0,..., M-1$. The data are thus partitioned into $M$ discrete bins $b_i$, and $k_i$ denotes the number of measurements that lie within the bin $b_i$. The probabilities $p(b_i)$ are then approximated by the corresponding relative frequencies of occurrence:

$$p(b_i) = \frac{k_i}{N}$$

Then the Shannon entropy $H(X)$ is calculated based on the information theory[4]:

$$H(X) = -\sum_{i=1}^M p(b_i)\log p(b_i)$$

[1] Michaels G S, Carr D B, Askenazi M, et al. Cluster analysis and data visualization of large-scale gene expression data[C]//Pacific symposium on biocomputing. 1997, 98: 42-53.

[2] Butte A J, Kohane I S. Mutual information relevance networks: functional genomic clustering using pairwise entropy measurements[M]//Biocomputing 2000. 1999: 418-429.

[3] Steuer R, Kurths J, Daub C O, et al. The mutual information: detecting and evaluating dependencies between variables[J]. Bioinformatics, 2002, 18(suppl_2): S231-S240.

[4]Shannon C E. A mathematical theory of communication[J]. The Bell system technical journal, 1948, 27(3): 379-423.

Question 2: Also about the graph entropy propagation method, would this reach a stationary distribution after multiple rounds of propagation?

Thank you for your constructive comments. Yes, the MI-based graph we constructed reaches a stationary distribution after multiple rounds of propagation, where the feature distributions of the vertices and the mutual information between vertices no longer change significantly. We have added an analysis of the changes in vertex distributions during multiple rounds of propagation in Appendix A. Specifically, we randomly selected three nodes from a pre-constructed graph and recorded their distribution changes.

The results are shown in Figure 11. As observed from the experimental results, after the first round of propagation, there is a noticeable change in the feature distribution of each node. As the number of propagation rounds increases, the node feature distributions gradually stabilize. This observation is also consistent with the MC1 results presented in Figure 6 of the paper, which show stable effects during subsequent rounds of propagation.