Thank you for your positive feedback and acknowledging our theoretical insights! We address your remaining concerns below.

---

Q1. This paper only focuses on two scenarios: social bias and jailbreak attacks. It lacks of showing effectiveness of this method on other types of tasks like reasoning.

A1. We note that, as mentioned in the paper, there are many ongoing studies and debates on whether self-correction works for reasoning, e.g. [1,2]. Although the ability of large language models (LLMs) to self-correct in reasoning tasks remains controversial, some existing studies on real world reasoning datasets have shown that self-correction can actually improve the response accuracy, which aligns well with our theory. For instance,

• Zhang et al [3] examined the impact of different critics on correction, finding that stronger critics can lead to higher correction gains. This finding aligns well with our theory and is consistent with our conclusions drawn from the noisy critic experiments on both synthetic dataset and the BBQ dataset.
• In addition, Lin et al [4] explored the linear relationship between generation, critic, and correction. Their conclusion (1) that critique-focused training markedly enhances performance also aligns well with our theory.
• For synthetic reasoning dataset, the (very recent) ICML’24 tutorial shows that training with mistakes+corrections in a controlled setting does boost model accuracy (78% $\rightarrow$94%) [5].

We will include these useful reference results in the revision for a comprehensive discussion.

As a theoretical study, we do not intend to develop a superior approach among these variants, but to validate our theoretical insights on real-world tasks. As explained in Sec 5 (Line 307), our theory suggests that self-correction indeed has the potential to improve the alignment of LLMs, especially when the critics are relatively accurate”, which motivates us to study its gains in the social bias and jailbreak scenarios where models can provide accurate critics. Since reasoning tasks require more accurate critics and better refinement strategies, it requires a more careful practical designs, which is beyond the scope of this work. Nevertheless, we believe that our understanding of self-correction as an in-context alignment can provide principled insights into future designs.

Ref:

[1] Huang, Jie, et al. "Large language models cannot self-correct reasoning yet." ICLR 2024.

[2] Valmeekam, Karthik, Matthew Marquez, and Subbarao Kambhampati. "Can large language models really improve by self-critiquing their own plans?." arXiv preprint arXiv:2310.08118 (2023).

[3] Zhang et al. “Small Language Models Need Strong Verifiers to Self-Correct Reasoning.” arXiv preprint arXiv:2404.17140v1(2024)

[4] Lin, Guo, et al. “CRITICBENCH: Benchmarking LLMs for Critique-Correct Reasoning” arXiv preprint arXiv:2402.14809 (2024)

[5] Allen-Zhu et al. Physics of Language Models - Part 2.2: How to Learn From Mistakes. ICML 2024 tutorial. July 2024 (no arxiv yet).

---

Q2. What potential strategies can be employed to improve the robustness of self-correction when the self-generated feedback is noisy or inaccurate?

A2. That’s a great question! To answer this problem, we conducted a series of real-world controlled experiments on LLMs, and summarized them in the General Response on the top. Based on these results, we find the answer to this question has two folds.

1. LLMs themselves are robust to critic noise. We find that, a very noisy critic (eg only 25% accuracy) is already enough to attain gains over the baseline, showing that LLMs are robust to noise and can benefit from noisy critics

2. CoT prompting can improve critic quality. In the meantime, we also observe that more accurate critics lead to larger gains. So secondly, one way to obtain more accurate critic is through chain-of-thought (CoT) prompt, i.e., instructing the model to think step by step by the final critic. As shown in the comparison below (quoted from General Response), CoT can improve critic accuracy from 20.52% to 36.99%, and consequently, improving the final prediction from 23.01% to 33.54%.

In summary, self-correction still works well under noisy critics, and when critics are very noisy, we can use CoT to improve the critic quality.

---

Thanks for your insightful comments. Hope our new real-world verification on LLMs could address your concerns. Please let us know if there is more to clarify.

We thank Reviewer zSQV for appreciating the novelty of our theory and the solidness of our experiments! Below, we address your remaining concerns about the verification experiments.

---

Q1. In the proof, the author simplifies the concept by converting criticism into a reward, represented as a real number. Could you discuss the discrepancy between this simplification and the actual scenario of Natural Language Criticism? Additionally, could you please discuss whether this simplification significantly affects the theoretical proof if switch to Natural Language Criticism.

A1. Indeed, with LLMs, one can instruct the models in different ways to attain different forms of self-critics, which have an influence on the self-correction performance as well.

Experimental Comparison. Following your suggestions, we compare two sources of different types of self-critic on LLM:

• Natural language vs. explicit labels: for natural language, we ask the plain critic question like “Is there bias in the response?”, where the model will reply in free-form natural languages; instead, for explicit labels, we instruct the model to answer a binary question, eg “Is this response biased or unbiased? (a) biased (b) unbiased.” From the table below, we can see that the explicit label alone performs worse than natural language (28.64% vs 23.01%), which might be due to that natural language provides fine-grained reward signals to alignment.
• Direct answer vs chain-of-thought (CoT): one way to remedy the limit of direct response is CoT. Here, we use zero-shot CoT to instruct the model to think step by step before giving the explicit labels. The table below shows that this does improve the critic does enhances the critic accuracy ($20.52\% \rightarrow 36.99\%$ ) and improves the correction accuracy as well ($23.01\% \rightarrow 33.54\%$ ). This shows that CoT is also a powerful technique to enhance self-correction by improving the critic quality, which also aligns with our theory.

Table B. Prediction and Critic Accuracy vs different types of critics on BBQ.

Our theory is compatible with different reward formats. In our alignment loss, the essential role of the rewards is to provide a ranking $\tau$ of the responses $y_i$, while the specific reward values $r_i$’s do not matter. Therefore, our theory is general and applicable to rewards of different formats, as long as they can provide such a ranking. For LLMs trained on natural languages, natural langauge critic also provides the critic that LLMs can understand and use for ranking the responses. Compared to explicit labels, natural language and CoT prompting with more detailed analysis could provide a fine-grained critic (that is akin to multi-dimensional reward vector) that gives more accurate preferences.

---

Q2. Strictly speaking, this cannot be considered a weakness. …. If additional experiments, using either synthetic data or real self-correction trajectory data, could explore different types of critics (i.e., rewards under the formulation of in-context alignment) and their impact on the self-correction mechanism (such as the number of self-correction rounds and heterogeneity of critics), it would deepen our understanding and exploration of the self-correction mechanism.

A2. Thank you for your insightful thoughts! We totally concur with this point, and following your advice, we further validate our theoretical insights extensively on real-world LLMs, covering the influence of critic accuracy, critic types, and self-correction rounds on self-correction performance, which further validate our theory (summarized in General Response at the top). Here we additionally quote the results on self-correction rounds below.

According to our theory, more self-correction rounds amount to more in-context alignment examples, which helps the in-context alignment process. We validate this further on real-world LLMs by applying Checking-as-Context for multiple rounds.

As shown in the table below, with groundtruth critic, more rounds lead to increasing accuracy, which aligns well with our theory. The performance peaks at the 3rd round potentially because of the model’s limited capability of handling long context. With self-critic, we instead find that 1-round critic gives the best performance; multi-round checking still outperforms the baseline but does not bring further benefits. This is akin to the accumulative errors that widely occur in pseudo-labeling methods, where the model deteriorates under iterative self-labeling. Our analysis builds a theoretical connection between intrinsic self-correction and self-labeling alignment, which provides a principled explanation for this phenomenon.

---

Thank you again for your insightful comments, which significantly strengthen our work! We are happy to address your further concerns in the discussion stage.

We thank Reviewer jEFz for appreciating our theoretical insights, empirical verification, and real-world applications. We address your concerns below.

---

Q1. Validation Scope: The validation is primarily on synthetic datasets.

A1. For completeness, following your advice, we further validate our theoretical insights extensively on real-world LLMs, covering the influence of initial response, critic quality, and model sizes on self-correction, which further validate our theory. The results are summarized in the General Response at the top. Please take a look. Thanks!

---

Q2. Dependence on Critic Quality: The paper acknowledges that the effectiveness of self-correction heavily relies on the quality of the critics, which may not always be reliable or available. And it's also costly in practice.

A2. That’s a great question! According to our real-world verification experiments, the answer has two folds.

LLMs can work with noisy critics. First, as we show in the real-world verification experiment (quoted from General Response), a very noisy critic (eg only 25% accuracy) is already enough to attain gains over the baseline, showing that LLMs are robust to noise and can benefit from noisy critics

CoT prompting for better critics. In the meantime, we also observe that more accurate critics lead to larger gains. So secondly, one way to obtain a more accurate critic is through chain-of-thought (CoT) prompt, i.e., instructing the model to think step by step by the final critic. As shown in the comparison below (quoted from General Response), CoT can improve critic accuracy from 20.52% to 36.99%, and consequently, improving the final prediction from

23.01% to 33.54%.

To summarize, self-correction still works well under noisy critics and when critics are very noisy, we can use CoT to improve its quality.

---

Thanks for your insightful questions. Hope our new real-world verification on LLMs could address your concerns. Please let us know if there is more to clarify.

We thank Reviewer pXYm for appreciating the novelty of our idea. We further address your main concerns on how models generate new responses below.

---

Q1. My main concern is about modeling the self-critic process as a ranking problem. Can you provide theory about the next round possible candidates?

A1. We understand your concern that the theory should be able to generate new responses instead of ranking the existing ones. We note that our theory does consider generating new responses of higher reward for the test point $x_N$, although the objective may initially appear to be a ranking problem. The key point is that we adopt a DPO-like reward function in our alignment loss (Eq 5), whose reward $r(x,y)=\|WX-Y||^2$ is directly calculated from the model policy $f(X)=WX$. Consequently, when optimizing the alignment loss in-context, the policy (the prediction $f(X)=WX$) is concurrently refined to generate better responses.

Specifically, to make the computation compatible with the given training examples, we 1) use a "dummy" response $y_N=W_0 x_N$ (i.e., the initial guess of LLMs with weights $W_0$ (Section 2.2.3)) as an initialization for the test output, and 2) initializes its "dummy" reward $r_N$ to have the lowest reward among the input examples. Under this configuration, we can prove that after an iteration, the test output becomes $y'_N=W'x_N$, where $W'=W+\Delta W$ corresponds to the updated weights. Therefore, the model can generate a new response based on the updated weights instead of copying from the context. Since the weight is updated to optimize the alignment objective, the new prediction can be better than those (potentially noisy) ones in the training samples.

Please let us know if there is anything else you would like us to elaborate on. We will certainly expand on this part in the revision.

---

Q2-1. Furthermore, a more important question to answer is why LLMs can have such the capability to regenerate better answers, including providing a good critic and being able to use this critic to further improve the results.

A2-1. This is an excellent question! A prevalent understanding of how ICL arises is that LLMs are pretrained on a highly diverse set of contextual data, enabling them to handle a variety of in-context tasks. Likewise, we believe that LLMs are all exposed to numerous self-correction-like data during training, which typically includes an initial statement, a critique of this statement (e.g., "this article is not good enough"), and often a refined version (e.g., a revised article), which are quite common in natural language. This exposure to self-correction data enables the model to learn to 1) accurately critique its previous outputs and 2) generate improved versions based on these critiques. The (very recent) ICML’24 tutorial confirmed this by showing that training with mistakes+corrections in a controlled setting does boost model accuracy [3].

Nevertheless, this data-based understanding does rigorously address 1) the mechanisms of self-correction — we formulate it as in-context alignment (ICA) process, and further 2) “how transformers perform ICA” — we show that they can optimize an alignment loss in-context and underpins the roles of each module. In this way, we can rigorously show that transformers can perform self-correction in an in-context alignment way.

Thank you for raising this question, and we will elaborate on this part in the revision as well.

[1] Chen, Ziru, et al. "When is tree search useful for llm planning? it depends on the discriminator." arXiv preprint arXiv:2402.10890 (2024).

[2] Teaching Large Language Models to Self-Debug - Arxiv-2304.05128

[3] Allen-Zhu et al. Physics of Language Models - Part 2.2: How to Learn From Mistakes. ICML 2024 tutorial. July 2024 (no arxiv yet).

---

Q2-2. However, the theory in this paper focuses on the possibility of using a transformer layer to implement the algorithm, but it does not necessarily mean transformers will implement such an algorithm, leading to a gap here.

A2-2. Indeed, the main theorem provides a construction proof that a $(N-1)$-block transformer can implement the algorithm. Unlike studies on linear attention, our proof requires a deep nonlinear transformer, whose convergence analysis hasn't been established in the literature yet. Due to this obstacle, we adopt a construction proof to show that transformers can perform ICA. Further, we confirm our theoretical insights from the construction through a controlled synthetic experiment, showing that a trained transformer, though may not implement the same weights, does behave quite similarly to the gradient descent algorithm we analyzed — in particular, the necessities of each transformer module we outlined (see Fig 1). This indicates that our analysis provides valuable insights into the behavior of trained transformers, forming a basis for further convergence studies. We will elaborate on this in the limitation part and outline potential paths toward further convergence studies.

---

Q3. Finally, in line 43, the authors mentioned that you are the "first theoretical analysis showing that LLM can improve alignment in context". It sounds a bit over-claim to me. Your theory is mainly about transformers, and it still has a huge gap to LLMs.

A3. Since transformers are the de facto architecture of LLMs, our analysis on transformers does apply to LLMs. By the word “can”, we mean that it has the capability to perform in-context learning, as we rigorously proved in Theorem 3.2. Meanwhile, we agree that “transformer” is a more accurate choice here, and we will replace it in the revision following your advice.

---

Thank you again for the insightful comments. If you find it satisfactory, we respectfully hope that you can re-evaluate our work. We are happy to address your further concerns.

We thank Reviewer bH8U for appreciating the novelty of our theory and our empirical verification. We address your concerns below.

---

Q1. In real-world experiment, the author provides several techniques: Multi-round Checking, Diverse Checking, Self-instruct. It's not mentioned that how these different self-correction settings can be fit into the proposed theoretical analysis.

A1. Thank you for your careful reading of our supplementary materials! Indeed, these methods are motivated by our theoretical analysis, and we briefly touched upon these connections in Appendix B.1. We appreciate the opportunity to elaborate further:

• I. Multi-round Checking. Our theory establishes a connection between self-correction and in-context alignment with multiple query-answer-critic contextual examples:$(x,y_1,r_1,x,y_2,r_2,\dots,x,y_N,r_N)$. According to our theory, one-round checking provides only one in-context training example $(x,y_i,r_i)$, and thus it’s natural to extend this to multi-round checking that provides more training examples $(x,y_i,r_i)_{i=1}^{N-1}$ for the in-context alignment task. Our synthetic experiments also confirm it by showing that an increased number of in-context samples results in smaller errors.
• II. Diverse Checking. Although our theory primarily focuses on a single query $x$, we can easily extend it to the context with multiple queries, i.e., $(x_1,y_1,r_1,x_2,y_2,r_2,\dots,x_N,y_N,r_N)$ — see a theoretical discussion in Appendix F.1. Compared to single-query alignment, multi-query one ensures that the optimized policy is generalizable across different queries. This theoretical insight led us to develop Diverse Checking which adopts different queries at different rounds.
• III. Self-instruct. There are two variants to instruct the model to refine its prediction $y_i \to y_{i+1}$ : 1) asking the same query $x$ again (illustrated in Figure 2), or 2) directly instructing the model to refine the response (i.e., Self-instruct). In practice, we found that both methods achieve self-correction, while Self-instruct is more robust against jailbreaks, especially when the query may contain harmful instructions.

Therefore, these three variants are deeply rooted in our theoretical analysis. We will elaborate on these connections better in the revision.

---

We hope this explanation addresses your concerns. Additionally, we added extensive verification of our theory on real-world LLMs, which further confirmed our theoretical insights (see General Response at the top). Please let us know if any further clarifications are needed!