An Empirical Study on Enhancing LLMs' Alignment Capabilities through Restyled In-Context Learning Demonstration Examples

Dear Authors,

Despite the approaching end of our discussion period, I am still very much interested in further exchanges with you.

Thank you for your constructive and valuable reviews. We have learned a lot from your feedback! Our responses are as follows:

W1Q1: ``The paper’s contribution is incremental, lacking substantial novelty compared to existing ICL-based alignment methods''.

Response: The primary contribution of this work is the development of an automated metric—polarity tokens—to identify ICL demonstration examples that balance safety and helpfulness. Additionally, we propose a restyling process to further enhance these examples in terms of safety and helpfulness.

Our most relevant prior work is URIAL (Lin et al., 2024), which uses three manually crafted ICL examples for alignment purposes. While these manually curated examples are empirically effective, they lack interpretability. For the community, it is challenging to extract actionable insights, such as:

• Why do these examples improve alignment performance?
• How should one create appropriate ICL examples for a different alignment task?

These questions remain unanswered in URIAL. Our work addresses this gap by making both the ICL example selection and restyling processes transparent. Readers can understand why specific ICL examples are effective for a given downstream task and use our approach as a guideline for selecting examples in their own applications.

This transparency and explainability distinguish our work from previous studies and constitute the novelty of our approach.

W1Q2: ``The experiments lack granular analysis of individual ICL examples or polarity tokens’ contributions to model performance across tasks''.

Response: It is true that our polarity token analysis was conducted on a single combination of LLaMA-2-7B and the just-eval-instruct dataset. However, as demonstrated in Section 4 Evaluation, our RIDE method achieves the best performance across most settings involving three models (LLaMA-2-7B-hf, Mistral7B-v0.1, and OLMo-7B) and three datasets (Alpaca-eval, just-eval-instruct, and MT-Bench).

While restyling contributes to overall performance improvements, these results also provide evidence that polarity tokens are transferable across LLMs and datasets. We will conduct more fine-grained experiments in future work to further validate the generalizability of polarity tokens across tasks and settings.

W1Q3: ``The paper demonstrates effectiveness in specific tasks but lacks evidence of generalization across broader cross-task scenarios such as multilingual settings or bias detection''.

Response: As mentioned in lines 53–63, factuality and safety are inherently in tension within alignment tasks. By observing this interplay, we proposed polarity tokens to select ICL demonstration examples suitable for both factuality and safety, and employed restyling to achieve a trade-off between the two. This addresses the factuality-safety tension in alignment tasks.

We agree with the reviewer’s observation that exploring the broader applicability of RIDE to more general tasks, such as multilingual scenarios or bias detection, is a valuable direction. However, we note that this lies beyond the scope of the current study.

Question: ``Could the authors clarify if their approach could adapt to dynamic, real-world contexts where example restyling might need to change over time? What are the potential trade-offs between the proposed method and simpler alignment techniques in terms of computational efficiency and performance gains?''

Response: As described, our approach selects and restyles a set of high-quality ICL demonstration examples to enhance the alignment capabilities of the base LLM without requiring additional training.

During inference, the restyled ICL examples are constant and applicable across all test sets for the same downstream task. If the task changes, we do not reselect the ICL examples but instead adjust their style based on whether the new task prioritizes factuality or safety. This ensures that the computational complexity during inference remains $O(1)$, which is significantly lower than other methods that require dynamic, online selection or computation of ICL examples.

This design achieves a favorable trade-off by balancing computational efficiency with substantial performance gains, making our approach practical for real-world applications.

Thank you again for reviewing our paper and for the valuable feedback. We hope our response has somewhat solved your concerns. if you have any comments, please let us know!

Thank you for your detailed, constructive, and insightful review comments. We have learned a great deal from your feedback! Our responses are as follows:

W1: ``The presentation and the writing in the paper need to be further improved.''

Q1: ``What is $P^{refer}$?''

Response: The reviewer raised a valid concern about the ambiguous notation for $P^{refer}$, $P^{base}$, and $P^{align}$. To clarify, in the equation $\Delta P^{`malicious`} = P^{`base`} - P^{`align`}$, $P^{refer}$ refers to $P^{`base`}$. Similarly, in $\Delta P^{`benign`} = P^{`align`} - P^{`base`}$, $P^{refer}$ refers to $P^{`align`}$.

Thus, $P^{refer}$ denotes the reference model, which serves as the baseline for comparison with the target model’s probability distribution under the same inputs. This comparison helps identify polarity tokens. In future revisions, we will revise the notation definitions to make them clearer and avoid ambiguity.

Q2: ``It is unclear how $V_{polar}$ defined in Line 221 is utilized in the subsequent causal structure of alignment.''

Response: The restyling process described in Section 3.2 does not affect $V_{polar}$. As detailed in Appendix A.4, we use Average Treatment Effect (ATE) to investigate how different styles of rewriting ICL demonstration examples impact downstream model performance in terms of factuality and safety.

The overall structure of our method follows a pipeline approach:

• Use polarity tokens to select ICL examples that perform well in terms of factuality and safety.
• Restyle these examples to achieve a better trade-off between factuality and safety for alignment tasks.

This pipeline addresses the inherent tension between factuality and safety while maintaining their complementary relationship. Importantly, the restyling process does not influence $V_{polar}$.

W2: ``It is very likely that the polarity tokens that you found are only the tokens that frequently appear at the start of a response.''

Response: We sincerely appreciate the reviewer’s insightful observation, which has inspired us to reflect further on our methodology.

First, as demonstrated empirically in prior studies such as URIAL (Lin et al., 2024) and Shallow Safety Alignment (Qi et al., 2024a), alignment tuning tends to alter the probability distribution of specific tokens or symbols at the early stages of an LLM’s output. By influencing the sampling probabilities of these tokens, alignment tuning can guide the generation trajectory of the model toward outputs that align better with human values.

Building on this foundational insight, our work focuses on identifying ICL demonstration examples that enhance alignment in downstream tasks by addressing the inherent tension between factuality and safety. As mentioned in the response to W1Q1, we compare the probability distributions of $P^{base}$ and $P^{align}$ to identify tokens with the largest average probability differences.

If a token frequently appears in LLM decoding but exhibits little difference in its generation probability between $P^{base}$ and $P^{align}$, it would not qualify as a polarity token. Thus, the polarity tokens we identify are those that frequently show significant changes in generation probability between $P^{base}$ and $P^{align}$—not merely the tokens that frequently appear at the start of a response.

W3: ``Are those identified benign and malicious tokens transferable across LLMs and validation datasets?''

Response: Thank you for your insightful and thought-provoking observation. Indeed, our polarity token analysis was conducted solely on the combination of LLaMA-2-7B and the just-eval-instruct dataset. It is worth noting that the subset of just-eval-instruct used for ATE evaluation is orthogonal to the test set used in Section 4 experiments, thereby eliminating the possibility of information leakage.

However, as shown in the Section 4 Evaluation, our RIDE method achieves the best performance in the vast majority of settings across three models (LLaMA-2-7B-hf, Mistral7B-v0.1, and OLMo-7B) and three datasets (Alpaca-eval, just-eval-instruct, and MT-Bench). While restyling contributes to the overall performance improvements, this also provides indirect evidence that polarity tokens are transferable across LLMs and datasets.

In future work, we will conduct more detailed experiments to further validate the generalizability of polarity tokens themselves across different models and datasets.

Thank you again for reviewing our paper and for the valuable feedback. We hope our response has somewhat solved your concerns. if you have any comments, please let us know!

Thank you for your constructive and valuable review comments. Our responses are as follows:

W1: ``The paper lacks a clear discussion about its position with respect to previous work."

Response: Firstly, due to space limitations in the main text, we included a discussion of related work in Appendix A.1. The primary contribution of this paper is the development of an automated metric—polarity tokens—to identify in-context learning (ICL) demonstration examples that balance safety and helpfulness. Additionally, we propose a restyling approach to make these examples even more effective in terms of safety and helpfulness.

Regarding the related work, URIAL (Lin et al., 2024), raised by Reviewer 1p86, their method uses three manually written ICL examples to achieve alignment. While manually curated ICL examples are indeed empirically effective, they lack interpretability. For the community, it is challenging to extract insightful takeaways from such examples, such as:

• Why do these manually written ICL examples improve alignment capabilities?
• For a different alignment task, how should one create their own ICL examples?

These questions remain unanswered in URIAL's work. In contrast, our approach addresses these issues. Both our ICL example selection and restyling processes are transparent, enabling readers to understand why these ICL examples are effective for the current downstream task. Furthermore, we can confidently say that our work provides guidance for selecting suitable ICL examples for their specific needs.

We agree with Reviewer 1p86's suggestion and will revise the Introduction to include a more prominent and logical discussion of the relationship between URIAL and our work. We will also explicitly state that our study builds upon URIAL while highlighting the key differences between the two approaches.

W2: ``It is hard to validate the claimed improvements compared to URIAL."

Q1: ``It is not clear why the just-eval-instruct scores for URIAL and baselines in Table 2 are not the same as those reported in the URIAL paper.''

Response: We learned that directly citing experimental results from the original paper may not reflect a rigorous experimental approach. Therefore, adhering to the principle of fair experimental comparison, we re-ran URIAL’s code, tested all versions of URIAL, and selected the demonstration examples that achieved the highest performance for comparison.

Due to the inherent variability in the outputs of LLMs and the subjective and stochastic nature of using LLMs as LLM-as-a-judge, our experimental results may differ from those reported in the URIAL paper. Despite these discrepancies, the data we present is both truthful and reliable, and such variations do not affect the fairness of the comparative experiments.

To ensure transparency and reproducibility, we will release the source code for our work, allowing reviewers and the community to evaluate and analyze our results.

Q2: ``RIDE optimizes its prompt using GPT-4o and a validation set from just-eval-instruct (is this provided originally or sampled from the test set?), which might provide an unfair advantage compared to other techniques.''

Response: First, the subset of just-eval-instruct used for ATE evaluation (as described in Appendix A.4) is orthogonal to the test set used in the experiments of Section 4. Therefore, there is no risk of information leakage. We appreciate the reviewer pointing this out and will explicitly clarify this in future revisions.

Second, regarding the issue of restyling: as mentioned earlier, our primary comparison is with URIAL, which manually crafted several demonstration examples as ICL inputs. In the process of manual curation, URIAL implicitly incorporated a form of restyling—i.e., adapting examples to fit a consistent, specific style. However, this implicit approach makes it challenging for readers to replicate or draw actionable insights from the restyling process.

In contrast, RIDE adopts a different approach by making the restyling process explicit and transparent. Through the main text and Appendix A.4, we provide clear guidance on:

• What types of restyling enhance factuality,
• What types of restyling improve safety, and
• How to achieve the best trade-off between these two aspects.

We believe this transparency offers practical insights to readers, enabling them to make informed decisions when selecting and restyling ICL demonstration examples. Consequently, while both URIAL and our method involve restyling, we argue that the comparison remains fair.

W3: ``The paper’s titled implies a possibly broader application of the paper’s idea than what is tackled by the proposed method.''

Q1: ``It is not clear how polarity tokens can generalize in broader scenarios beyond factuality and safety categories.''

Response: As mentioned in lines 53–63 of the paper, factuality and safety are often in a state of inherent tension within alignment tasks. By observing and analyzing this relationship, we proposed polarity tokens to address this contradictory yet unified interplay. Thus, in this work, polarity tokens are specifically designed for alignment tasks.

We agree with the reviewer's perspective that exploring the applicability of polarity tokens to more general tasks is a worthwhile direction for future research. However, we would like to emphasize that such exploration lies beyond the scope of this paper.

Q2: ``It is not clear how effective this approach would be in the absence of aligned models.''

Response: Thank you for raising this point. In future work, we will include a comparison of performance between rewritten ICL examples and unmodified ICL examples to further evaluate the effectiveness of our approach.

Thank you again for reviewing our paper and for the valuable feedback. We hope our response has somewhat solved your concerns. if you have any comments, please let us know!

We must acknowledge and agree with the points raised by you and other reviewers regarding the limitations in our experimental design, particularly the issue of statistical significance. As discussed with Reviewer GeMC, developing new metrics to better highlight significant performance improvements is a critical task we need to address. Thank you for your suggestions—we will work on resolving this issue in future revisions.

Regarding your additional questions:

• Q2: "Different numbers than those in the URIAL paper": This has been addressed in W2Q1 of our rebuttal.
• Q3: "How was the validation set selected for optimizing the prompt in RIDE": This was discussed in W2Q2.
• Q4: "Release source code": Please refer to our response in W2Q1.
• Q5: "Could the authors elaborate on how the proposed approach behind polarity tokens would scale with an increasing number of alignment categories?": Please refer to our response in W3Q1.

In conclusion, we deeply value your thoughtful feedback and sincere suggestions, which will serve as an invaluable guide for our future research. We look forward to crossing paths with you again and wish you all the best in your future endeavors!