A Theoretical Understanding of Self-Correction through In-context Alignment

Authors: yifei  wang
yuyang  wu
zeming  wei
stefanie  jegelka
yisen  wang
hugo  touvron
louis  martin
kevin  stone
peter  albert
amjad  almahairi
yasmine  babaei
nikolay  bashlykov
soumya  batra
prajjwal  bhargava
shruti  bhosale
dan  bikel
lukas  blecher
cristian canton ferrer
moya  chen
guillem  cucurull
david  esiobu
jude  fernandes
jeremy  fu
wenyin  fu
brian  fuller
cynthia  gao
vedanuj  goswami
naman  goyal
anthony  hartshorn
saghar  hosseini
rui  hou
hakan  inan
marcin  kardas
viktor  kerkez
madian  khabsa
isabel  kloumann
artem  ko- renev
singh  koura
marie-anne  lachaux
thibaut  lavril
jenya  lee
diana  liskovich
yinghai  lu
yun- ing  mao
xavier  martinet
todor  mihaylov
pushkar  mishra
igor  molybog
yixin  nie
andrew  poulton
jeremy  reizenstein
rashi  rungta
kalyan  saladi
alan  schelten
ruan  silva
eric michael smith
ranjan  sub- ramanian
ellen  tan
binh  tang
ross  tay- lor
adina  williams
jian xiang kuan
puxin  xu
zheng  yan
iliyan  zarov
yuchen  zhang
angela  fan
melanie  kambadur
sharan  narang
aurelien  rodriguez
robert  stojnic
sergey  edunov
thomas scialom llama
gladys  tyen
hassan  mansoor
peter  chen
tony  mak
ashish  vaswani
noam  shazeer
niki  parmar
jakob  uszko- reit
llion  jones
aidan n gomez
łukasz  kaiser
illia  polosukhin
  attention
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 14:09:07

Abstract: Going beyond mimicking limited human experiences, recent studies show initial evidence that, like humans, large language models (LLMs) are capable of improving their abilities purely by self-correction, i.e., correcting previous responses through self-examination, as seen in models like OpenAI o1. Nevertheless, little is known about how such capabilities arise. In this work, based on a simplified setup akin to an alignment task, we theoretically analyze self-correction from an in-context learning perspective, showing that when LLMs give relatively accurate self-examinations as rewards, they are capable of refining responses in an in-context way. Notably, going beyond previous theories on over-simplified linear transformers, our theoretical construction underpins the roles of several key designs of realistic transformers for self-correction: softmax attention, multi-head attention, and the MLP block. We validate these findings extensively on synthetic datasets. Inspired by these findings, we propose a simple self-correction strategy, Checking as Context (CaC), which finds novel applications in alleviating social bias and defending against LLM jailbreaks. We believe that these findings will inspire further research on understanding, exploiting, and enhancing self-correction for building better foundation models. Code is at https://github.com/yifeiwang77/Self-Correction.

Main paper: 
 introduction: 
 "Who among people is without fault? Making mistakes and being able to correct them is the greatest goodness." -Zuo Zhuan (∼400 BC), Translated by ChatGPTThe capacity for self-correction, traditionally viewed as a distinctive human trait, is increasingly being explored within the realm of artificial intelligence, particularly in Large Language Models (LLMs). Recent studies have sparked optimism about LLMs' self-correction capabilities for enhancing reasoning (Madaan et al., 2023; Shinn et al., 2023; Wang et al., 2023) , planning (Yao et al., 2023) , and alignment (Bai et al., 2021; Ganguli et al., 2023) . Although some find that self-correction may lead to worse performance without external feedbacks (Huang et al., 2023a; Valmeekam et al., 2023) , more recent evidence shows that with careful designs of instructions on the self-criticizing process, self-correction can yield considerable benefits on various tasks (Li et al., 2024; Zhang et al., 2024; Lin et al., 2024; Ji et al., 2023; Jiang et al., 2024; Tyen et al., 2023) .Driven by these intruiging empirical findings, we want to establish a principled understanding of how the selfcorrection ability emerges in LLMs. A particular difficulty is to formulate the multifaceted self-correction designs to be amenable to theoretical analysis. We notice that existing self-correction methods admit a general abstraction: generation, critics, regeneration, and further critics, continuing until the final refined output. This self-correction path can be understood as a particular form of context that provides feedback for refining the prediction on the fly. Different from standard (query, response) context examples akin to supervised learning, self-correction examples can be formulated in a triplet form (query, response, reward) that is akin to LLM alignment with both good and bad samples indicated by their rewards (Ouyang et al., 2022; Bai et al., 2022; Rafailov et al., 2023; Song et al., 2023) . This observation motivates us to formulate self-correction as a form of in-context alignment (ICA) , where LLMs are provided with a context of self-correction steps and the goal is to refine the final outputs to have higher rewards.Through this perspective, we prove that in a simplified setup, a standard multi-layer transformer can utilize selfcorrection samples to generate responses of higher rewards. Specifically, we prove the existence of model weights such that a transformer can optimize common ranking-based alignment objectives by performing gradient descent incontext, which includes the Bradley-Terry model (Bradley and Terry, 1952) and the Plackett-Luce model (Plackett, 1975) that are de facto choices for LLM alignment (used in RLHF (Ouyang et al., 2022) and DPO (Rafailov et al., 2023) ). As far as we know, this is the first theoretical analysis showing that LLMs can improve alignment incontext, providing a solid foundation for understanding self-correction. Our theory accommodates different kinds of self-correction methods, because the critics of responses can come from humans (Ouyang et al., 2022) , external verifiers (Chen et al., 2023) , or LLMs themselves (Zhang et al., 2024; Li et al., 2024) . The analysis further reveals that LLMs' self-correction performance relies crucially on the quality of critics, which agrees well with recent empirical findings (Lin et al., 2024; Chen et al., 2024; Tyen et al., 2023) . Intriguingly, within this analysis, we nail down the roles of realistic transformer designs -multi-head softmax attention, feed-forward network, and stacked blocks -for alignment, providing concrete theoretical insights for designing robust LLMs. This contrasts with previous in-context learning theories that focus on linear attention in the context of linear regression, deviating from practice (Von Oswald et al., 2023; Zhang et al., 2023a; Ahn et al., 2023b) .At last, we validate our theoretical explanations through both synthetic and real-world experiments. Extensive synthetic datasets show that transformers can indeed learn from noisy outputs with the help of relatively accurate critics. We validate that real-world transformer modules do matter for in-context alignment, and the results align surprisingly well with our theory. Driven by these theoretical insights, we explore two real-world scenarios where we hypothesize that aligned LLMs can provide relatively accurate self-critics: alleviating social bias and defending against jailbreak attacks. We show that with a simple generationcritic-regeneration process (we call Checking-as-Context) and no external feedback, intrinsic self-correction can alleviate social bias on Vicuna-7b and Llama2-7b-chat, and exhibits a strong correlation between self-checking accuracy and final performance again. With the same strategy, we find that self-correction can reduce the attack success rate by a large margin (e.g., 95% → 2%) against multiple types of jailbreak attacks. These evidences show that LLMs are indeed capable of improving alignment by selfcorrection alone, which not only validates our theory, but also provide insights for future designs and applications of self-correction. 
formulation: 
 In this section, we introduce self-correction and formulate it as a general in-context alignment process, and then introduce the setup for theoretical analysis.2.1. Self-correction as In-context Alignment ICL. In-context learning (ICL) is known as an emergent ability of LLMs to learn from a few demonstrations without finetuning (Lu et al., 2023) . Specifically, an LLM directly predicts the desirable response to the test query x test with N pairwise training examples {(x i , y i )} n i=1 as the context: ŷtest = LLM([x 1 , y 1 , . . . , x n , y n , x test ]).(1)Despite its effectiveness, ICL requires the knowledge of desirable responses y i to construct the training examples.For instance, Wei et al. (2023) use human-selected safe query-response pairs for in-context defense of jailbreaks.For queries that are vague or require domain expertise (e.g., math, science, and open-end discussions), desirable responses can be hard to collect or formulate.Self-correction. As a further step to eliminate human efforts, self-correction relies on LLMs to correct the mistakes in the initial generation. In self-correction, we first generate an initial response y 1 to the query x, and then obtain a critic on the response, denoted as a reward r 1 . The critic can be either generated by LLMs themselves through carefully designed prompting (Jiang et al., 2024; Li et al., 2024) , or by external verifiers such as code intepreters (Renze and Guven, 2024) . Afterwards, the LLM is then instructed to generate a refined response y 2 taking the initial response y 1 and its critic r 1 as the input context. This process can be repeated multiple times for iterative refinements of the response. After N steps, we take the final response y N as the final output. For simplicity, we assume that these steps share the same query x, and the extension to multiple queries is discussed in Appendix F.1.In-context Alignment (ICA). The self-correction process described above can be formalized as an in-context learning task with triplet examples {(x, y i , r i )}, where x is the (shared) query, y i is the response, and r i is the critic at the i-th step. Note that the same data format is also adopted in LLM alignment tasks, where LLMs are trained to follow human intention with human/AI-generated preference data (Ouyang et al., 2022; Bai et al., 2022; Rafailov et al., 2023; Song et al., 2023) . 1 In this way, we formulate self-correction as an in-context way to solve an alignment task, which we call in-context alignment (ICA). Here, the concept of alignment is inclusive and not limited to standard alignment tasks. Any objective that works with the triplet preference data can fit into our framework. Also, we do not assume that the rewards r i are always accurate, and the quality of the rewards will be shown to have a critical influence on self-correction.Since real-world LLMs on language tasks are too complex for a rigorous analysis, recent studies on ICL theory rely on synethetic simple tasks to examine LLM capabilities (Garg et al., 2022; Von Oswald et al., 2023; Zhang et al., 2023a; Ahn et al., 2023b) . Existing results are mostly es-tablished in the supervised setting, particularly for linear regression, due to its simplicity and alignment with linear attention. However, it is yet unknown whether transformers are capable of learning alignment tasks using preference data in-context. In this section, we introduce a simplified setup for in-context alignment. For the ease of analysis, we still study a linear regresion task, where a smaller MSE loss gets higher reward. However, what makes things harder is that the models are not provided with groundtruth targets as the context, but only (potentially false) responses y i and their rewards r i . To solve this task, the model has to learn the ability to compare the rewards of different samples and prioritize those with higher rewards -a critical ability that is key to self-correction and alignment, but has not been studied in previous theories.We begin by formalizing a general alignment task with triplet examples. Consider a training dataset D = {(x, y i , r i )} n-1 i=1 composed of a common query x ∈ R nx (assume ∥x∥2 = 1 for simplicity) 2 , multiple responses y i ∈ R ny and rewards r i ∈ R. Following the setup of Von Oswald et al. (2023) , we also consider a linear regression task where the groundtruth function is f (x) = W * x for some W * ∈ R ny×nx . However, in our tasks, the responses y i can be quite noisy (e.g., random), and the quality of this response is indicated by its reward value. Therefore, the transformers have to rank the responses based on their rewards and adjust their outputs accordingly. Meanwhile, the rewards coming from humans, external verifiers, or LLMs can also contain noise, which reflect the critic quality. The goal is to output a response y N that has a smaller square error, i.e., higher rewards. There are two approaches to solve this problem, one is through the in-context alignment with a transformer-based LLM, and one is through learning a parameterized alignment model. We describe these methods formally, and establish their inherent connections.The transformer model (Vaswani et al., 2017) is the de facto choice for building LLMs. It is a composition of multiple transformer blocks. Each block has two modules: MHSA and FFN. Normalization layers are omitted for simplicity.A multi-head self-attention (MHSA) layer updates a set of tokens {e 1 , . . . , e N } by e j ← e j + SA θ (j, {e 1 , . . . , e N })EQUATIONwith P h , V h , K h the projection, value and key matrices, respectively, and q h,j the query, all for the h-th head (bias terms omitted). The columns of the valueV h = [v h,1 , . . . , v h,N ] and key K h = [k h,1 , . . . , k h,N ] matrices consist of vectors v h,i = W h,V e i and k h,i = W h,K e i ;likewise, the query is produced by linearly projecting the tokens, q h,j = W h,Q e j . The parameters θ = {P h , W h,V , W h,K , W h,Q } h of a SA layer consist of all the projection matrices of all heads. We omit causal masking in the main paper for simplicity; see Appendix F.2 for an extension.FFN. Following self-attention, a feed-forward network (FFN) transforms each token individually with two shared linear transformations and a ReLU activation in between: Context Tokens. For simplicity, we assume that LLMs take a concatenated input e i = [x i , y i , r i ] for each example.3 For the last test example, to align with the same input format, we model it as e N = [x, y N , r N ], where we 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 final output, and its "dummy" reward r N is assumed to have the lowest reward among the input examples. In total, we haveEQUATIONN tokens {e i = [x, y i , r i ]} N i=1as the contextual input to the transformer.A common way to solve alignment tasks is to learn a parameterized alignment model that models preferences through a ranking objective over multiple candidates (Bradley and Terry, 1952; Plackett, 1975; Luce, 2005; Rafailov et al., 2023) . We use y i ≻ y j to denote the event that the response y i is preferable over y j . Let τ : [N ] → [N ] be the permutation function that denotes the ranking of all responses according to the reward scores, i.e., r τ (1) > • • • > r τ (N ) . 4The ranking τ implies that for any N ≥ i > j ≥ 1, we have y τ (i) ≻ y τ (j) . A common objective for N -ary comparison is the Plackett-Luce (PL) model (Plackett, 1975; Luce, 2005; Rafailov et al., 2023) that stipulatesP PL (τ | x, {y i }) = N i=1 exp r(x, y τ (i) ) N j=i exp r(x, y τ (j) ), (4) where r denotes the reward function. Since we consider a linear regression task (Section 2.2.1), we use the negative square error as the reward function (higher is better): r(x, y) = -∥W x -y∥ 2 . The corresponding PL model isP PL (τ ) = N i=1 exp -∥W x -y τ (i) ∥ 2 N j=i exp -∥W x -y τ (j) ∥ 2 . (5)In this setting, optimizing W w.r.t. the PL loss above also improves the policy (here the prediction y = W x) simultaneously, in a similar spirit to DPO (Rafailov et al., 2023) . Notably, it avoids the need to perform additional RL steps for policy improvement (Ouyang et al., 2022) .Relationship to Bradley-Terry model. The Plackett-Luce model is an N -ary generalization of the Bradley-Terry model (Bradley and Terry, 1952) used for pariwise preferences. In particular, with N = 2, the PL model (Eq. ( 5)) reduces to the Bradley-Terry model with least-squares reward:EQUATIONPrevious work (Song et al., 2023) shows that the N -ary PL model outperforms the binary BT model for alignment.Relationship to InfoNCE. We also notice that the InfoNCE loss that is widely used for contrastive learning (Oord et al., 2018; Chen et al., 2020; Radford et al., 2021) can be seen as a special case of the PL model when only considering its first term (i = 1). In this case, only y 1 is the positive sample and y 2 , . . . , y N are negative samples, which corresponds to a special ranking r τ (1) > r τ (2) = • • • = r τ (N ) . Therefore, the analysis in our framework can be used to explain incontext contrastive learning (Chia et al., 2023) . 
main results: 
 In this section, we present the main result of this work, which, to the best of our knowledge, is the first to show that a realistic transformer (with stacked multi-head softmax attention and feed-forward networks) can implement the gradient descent of common alignment objectives with incontext triplets. Notably, our analysis reveals the individual roles of these core designs of realistic transformers for incontext alignment (and self-correction), which may help future designs of LLM backbones as well.To highlight the key ideas without technical nuances, we start with N = 2, the Bradley-Terry (BT) model (Eq. ( 6)). Assume w.l.o.g. that y 1 ≻ y 2 with scores r 1 > r 2 , the BT model isL BT (W ; x, y 1 , y 2 ) = -log P BT (y 1 ≻ y 2 | x) = ∥W x -y 1 ∥ 2 + log 2 j=1 exp -∥W x -y j ∥ 2 .Proposition 3.1. One can realize the gradient descent for BT,W ′ = W + ∆W = W -η∇ W L BT (W ; x, y 1 , y 2 ),by updating each y i withEQUATION-2ηy 1(2)EQUATIONwhereβ j = softmax(-∥W x -y j ∥ 2 ). Specifically, L BT (W ′ ; x, y 1 , y 2 ) = L BT (W ; x, y ′ 1 , y ′ 2 ).Proposition 3.1 shows that the gradient descent of the BT model is equivalent to transforming the targets y i according to Eq. ( 7). This connection allows us to optimize output alignment (measured by BT loss) with the forward propagation of an MHSA layer (Eq. ( 2)). To see this, Term (1) corresponds to the shortcut feature y i . Term (2) is a bit complex, since it only picks y 1 with the higher score (r 1 > r 2 ). We find that this can be realized by constructing a softmax attention head that only attends to tokens with the largest reward r. Term (3) can be implemented with another softmax attention head that incorporates β i 's as the attention weights and y i 's as values. Therefore, the one-step gradient descent of the BT model can be implemented with two-head softmax attention.Theorem 3.2. Given a two-head softmax attention layer and two tokens e i = (x i , y i , r i ), i = 1, 2, there exists a set of parameters (Eq. ( 2)) such that a forward propagation step with token e i is equivalent to the gradient-induced dynamics of the Bradley-Terry model (Eq. ( 6)):EQUATIONAll proofs of the paper are deferred to Appendix E. As outlined above, our construction of in-context alignment requires two heads to implement the two gradient terms corresponding to positive and negative feedback, where softmax attention is exploited for sample selection in both cases. Instead, ICL analyses for linear regression (Von Oswald et al., 2023) only require one linear attention head for interpolating with linear products. Thus, our alignment analysis better reveals the need for softmax and multi-head attention, so it has a close correspondence to real-world architectures.We further explore how to extend this result to a general Nary Plackett-Luce (PL) model (Eq. ( 5)). Although the key ideas are similar, it is technically much harder to implement N > 2 with a single SA layer. To see this, notice that the response update of the PL loss corresponds toEQUATIONAt first glance, the i-th item of the update resembles Eq. ( 7) and seems implementable with a two-head self-attention. However, it is actually hard to realize the first term y τ (i) , since softmax attention can only select the top or the bottom value5 from a set of rewards, making it challenging to compare N examples within a single SA layer.A roadmap to implementing N > 2 with stacked full Transformer blocks. We discover that it is still possible to construct the PL gradient descent if we further incorporate the FFN module and allow stacking multiple transformer blocks. Specifically, at the i-th block, we can 1) identify the token with the largest reward (i.e., y τ (i) ) and implement the i-th term of the gradient descent with a three-head SA layer; and 2) mask out the y τ (i) of this token to eliminate its contribution in subsequent terms with the help of an FFN.In other words, each transformer block can implement one of the N -1 terms of the gradient (Eq. ( 9)) and prepare the input data for implementing the next term with one additional head. In total, it requires stacking N -1 transformer blocks (each is composed of three-head MHSA and FFN) to implement the whole gradient descent of the PL model.6 Theorem 3.3. Given a transformer TF with N -1 stacked transformer blocks (composed of three-head softmax attention and feed-forward networks) and N input tokens {e i , i ∈ [N ]}, there exists a set of parameters such that a forward step with token e i is equivalent to the gradient-induced dynamics of the N -ary Plackett-Luce model (Eq. ( 5)), i.e., TF(e i ) = (x i , y i , r i ) + (0, -∆W PL x i , 0), i ∈ [N ].Theorem 3.3 shows that a multi-layer transformer can improve its output alignment by optimizing a general Plackett-Luce model through in-context learning. It could serve as a general explanation for ICL-based alignment algorithms (Han, 2023; Lin et al., 2023; Guo et al., 2024) . As far as we know, it is the first theoretical result for explaining in-context alignment from an optimization perspective.Through our construction, we also underpin the individual roles of rewards and transformer modules during the self-correction process:1. Reward quality determines self-correction quality.By connecting in-context alignment to an optimiza-tion process, we reveal that the critics used in selfcorrection essentially serve as the supervision for the in-context alignment task. Thus inaccurate rewards would amount to noisy supervision that is known to degrade learning performance (Natarajan et al., 2013) , which explains the benefits of external feedback (Chen et al., 2023) and stronger discriminator (Chen et al., 2024) in self-correction.2. Softmax attention is important for ranking. One of the key steps to implement the gradient descent is to select the top response based on the input rewards, and our construction relies crucially on the ability of softmax attention to compare and reweight different rewards. Instead, it is hard for linear attention to implement such ranking operations.3. Multi-head attention is important for token discrimination. We use two attention heads in Eq. ( 7) with different roles: one for pushing top ones apart, and one for pulling others closer. This indicates that only with multi-head attention can we achieve better discrimination of different input tokens. In contrast, only one attention head is needed for regression (Von Oswald et al., 2023) .4. FFN is important for transforming selected tokens.In our construction, although softmax attention can select the top tokens, we cannot edit the selected tokens with attention alone. Instead, FFN is capable of 1) identifying top tokens in the input sequence with the knowledge of initial and selected tokens, as well as 2) performing conditional operations (e.g., masking out y τ (i) ) by leveraging the ReLU nonlinearity.Ranking multiple examples requires more depth.Comparing Theorems 3.2 and 3.3, we notice that ranking N examples with a transformer requires N -1 layers with our construction. This fact suggests a hint of why depth is still a major factor when constructing LLMs. For example, scaling from 7B to 70B, Llama2 goes from 32 layers to 80 layers and shows significant improvements.In Section 4, we also empirically validate the necessity of these modules for in-context alignment. This analysis also suggests that linear regression-which only requires single-head linear attention to solve in-context-may not be enough to fully characterize the behavior of standard transformers (Ahn et al., 2023b) , while our in-context alignment tasks (Section 2.2.1) could be a better theory model. These theoretical disclosures of Transformer modules may inspire future designs of LLM backbones as well.Relation to Previous Theoretical analyses. An existing line of prior research explains in-context learning via its connection to optimization algorithms (Garg et al., 2022; Von Oswald et al., 2023; Bai et al., 2023; Li et al., 2023a; Ding et al., 2023; Wu et al., 2023; Ahn et al., 2023b; Huang et al., 2023c) . We provide a detailed summary of these works in Appendix A, and here summarize key aspects in which we differ:• Objective: linear regression vs. non-convex alignment. Compared to previous methods that focus on solving linear regression in-context, we are the first to show that transformers can also solve ranking-based alignment problems in-context. A major difference is that alignment involves a more complex non-convex objective that does not admit a closed-form solution like linear regression.• Backbone: linear attention vs. full transformer. As discussed above, our construction identifies that softmax attention and other components of transformers play a major role in ranking while focusing on linear regression problems only requires linear attention. It reveals that our PL model with linear reward could be a better theory model for explaining in-context learning as it aligns better with practice.• Task: supervised learning vs. preference-based alignment. Previous ICL theories mostly focus on explaining its ability to perform supervised regression. Instead, we show that LLMs can learn in-context alignment, which allows feedbacks from various sources with noises, and learns from both good and bad behaviors. In particular, our theory also applies to intrinsic self-correction methods with self-generation critics, which is self-supervised. 
verification on synthetic data: 
 Here, we follow our theoretical setup in Section 2.2.3 and conduct a series of synthetic experiments to examine our theoretical results established in Section 3.Setup. We consider the following meta-learning setting.For every task, we draw a common query x ∼ N (0, I d×d ) and a groundtruth parameter W ∼ N (0, I d×d ). We then generate N responses and rewards. For each response y i , we sample a reward r i ∈ U[0, 1] and an independent noise weight W - i ∼ N (0, I d×d ), and then generatey i = r i W x + (1 -r i )W - i x.Thus, responses with higher rewards are closer to the ground truth in expectation. By default, we set d = 5, N = 20 and use a 20-layer GPT-2 model with 3 heads, 96 hidden dimension, and a PL loss (Eq. ( 5)). Then we evaluate the normalized MSE between the predicted output ŷ and ground-truth y = W x using varying numbers of in-context examples, averaged over 256 runs with randomly generated tasks. We also implement the gradient descent (GD) of the linear PL model (Eq. ( 5)) and measure its optimal solution in the same way. We also change the reward noise p, model depth, and attention types to investigate their effects on in-context alignment. For more details, please refer to Appendix C.As shown in Figure 1 , there is a clear trend that with more in-context examples, transformer-based in-context alignment and gradient descent (GD) can quickly adapt to the task and find better predictions for test samples. In comparison, Figure 1a shows that GD performs better at the beginning, while Transformers also adapt quickly and attain slightly better performance with more in-context examples, e.g., N = 14. It indicates that in-context alignment might be even preferable to GD-based alignment in certain cases, which validates our theoretical results that transformers can optimize alignment in-context by gradient descent. Below, we study different factors of in-context alignment.Reward quality matters. To investigate the influence of reward quality, we randomly replace rewards with random values r p i ∈ (0, 1) with a probability p. As shown by solid lines in Figure 1b , a large p significantly decreases the incontext alignment performance with much larger test errors. This can be naturally understood through our theory, where the gradient descent is performed on noisy data that hinders the learning process, as shown in the dashed lines in Figure 1b . Therefore, it explains why self-correction methods are sensitive to the quality of critics, and LLMs need strong critics to perform effective self-correction, as empirically observed in recent work (Chen et al., 2024; Lin et al., 2024; Zhang et al., 2024) .Impact of Transformer Modules. While conventional ICL theories show that 1-layer single-head linear self-attention is sufficient for linear regression (Von Oswald et al., 2023) , for in-context alignment, we observe: (1) ICA requires more depth. Figure 1c shows that when transformers are shallow (e.g., 5 layers), ICA is much worse, and more depth benefits ICA effectively. After 15 layers, depth brings diminishing returns. This is consistent with our theory that requires stacking multiple transformer blocks for in-context alignment of N example (Theorem 3.3). (2) Softmax attention is necessary. Figure 1d illustrates that linear attention can hardly solve the in-context alignment task while softmax attention performs much better, which is consistent with our analysis (Section 3.2). (3) Multi-head attention helps. Figure 1e shows that single-head attention struggles to align in-context, while multi-head (3, 4, 6) performs well. In addition, when the number of attention heads exceeds 3, there is no significant benefit, which aligns surprisingly well with our analysis that requires 3-head to implement the GD of the N -ary PL loss (Theorem 3.3). (4) FFN is necessary. Figure 1f shows that without FFN, the model cannot align in-context, consistent with our analysis that FFN is necessary for transforming selected tokens. Summarizing these results, we find that our proof by construction does have a nice correspondence to the practical behaviors of transformers on in-context alignment tasks, and it helps reveal the roles of each transformer module for in-context alignment-like tasks. 
exploring intrinsic self-correction on: 
 Real-world Alignment TasksOur theoretical analysis above reveals that self-correction indeed has the potential to improve the alignment of LLMs, especially when the critics are relatively accurate. Motivated by this observation, we explore self-correction on two real-world alignment tasks: alleviating social bias (Ganguli et al., 2023) and defending against jailbreaks (Zou et al., 2023) . Since LLMs are aligned on human preferences and harmfulness is relatively easy for discrimination, we hypothesize that self-generated critics can be accurate in these tasks, which faciliate LLMs to improve their own alignment, known as intrinsic self-correction (Huang et al., 2023a) .Method: Checking-as-Context (CaC). For simplicity, we study a very simple and general form of self-correction without sophisticated procedures. Specifically, following the same format as our theoretical setup (Section 2.1), given a query x, we first generate an initial response y (w/o selfcorrection), and then instruct the model to review its response and get a self-critic r, and instruct the model to regenerate a new answer as the output (w/ self-correction), as illustrated in Figure 2 . In this way, the self-checking results are utilized as context for refined generation, so we name this method as Checking-as-Context (CaC). More details can be found in Appendix C.Following Ganguli et al. (2023) , we study the use of selfcorrection to alleviate societal biases in LLMs on the BBQ (Bias Benchmark for QA) benchmark (Parrish et al., 2021) , which evaluates societal biases against individuals belonging to protected classes across nine social dimensions. We randomly select 500 questions from each task subclass. Different from moral self-correction (Ganguli et al., 2023) that requires model finetuning, our method is more lightweighted, since it is inference-only without parameter update. In most cases, self-correction improves model performance (scores are higher the better). (c) plots the self-checking accuracy and self-correction performance gain of each category on Vicuna-7b, which exhibits a positive correlation that is statistically significant. 3c exhibits a strong correlation (p < 0.05) between the gain of self-correction and self-checking accuracy, as suggested by our theory. In Appendix C.2, we do the same evaluation only on the ambiguous questions, which leads to more significant self-correction gain and a stronger correlation. Our experiments illustrate an intriguing fact that even with minimal designs, LLMs can potentially improve morality by criticizing themselves.LLM jailbreaks have recently risen to be a major threat to LLM alignment (Anwar et al., 2024; Dong et al., 2024) , where even well-aligned models like ChatGPT can be manipulated into generating harmful content (Zou et al., 2023; Liu et al., 2023a; Wei et al., 2023; Zhang and Wei, 2024) . Although various defense measures have been proposed (Jain et al., 2023; Xie et al., 2023; Wei et al., 2023; Li et al., 2023b; Huang et al., 2023b; Mo et al., 2024) , these typically require extensive human intervention. The ambiguity remains as to whether LLMs can autonomously counteract such jailbreaking manipulations. Here, we explore whether LLMs can defend against jailbreak attacks themselves with self-correction. Due to the limit of space, more results can be found in Appendix B.We observe that for LLM jailbreaks, self-correction can give accurate self-checking most of the time (close to 100%). As a result, from Table 1 , we observe that on AdvBench (Zou et al., 2023) , CaC-based self-correction can indeed improve LLM safety a lot by reducing the attack success rate (ASR) on Vicuna-7b and Llama2-7b-chat by a significant margin against different types of jailbreak attacks, including gradient-based GCG attacks (Zou et al., 2023) and instruction-based AutoDAN (Liu et al., 2023a) . Compared to manually designed defense methods (Xie et al., 2023; Li et al., 2023b; Huang et al., 2023b) , self-correction can achieve comparable and even better performance. It suggests that LLMs can autonomously defend against jailbreak attacks with intrinsic self-correction, which is a promising direction for future research on AI safety. 
conclusion: 
 In this paper, we have explored how self-correction ability rises from an in-context alignment perspective, showing that standard transformers can perform gradient descent on common alignment objectives in an in-context way. Different from previous theories on simplified linear attention, our analysis reveals the important roles of real-world transformer modules and validates these insights extensively on synthetic datasets. We further studied intrinsic selfcorrection for real-world alignment scenarios and demonstrate clear improvements on alleviating social bias and defending against jailbreaks. In this way, our analysis provides concrete theoretical and empirical insights into the dazzling debate of building LLMs that can correct and improve themselves. 
 
 Appendix: There is a rapidly emerging body of research on LLMs, and some key techniques, such as in-context learning and selfchecking, are reinvented by different works from time to time. We will try to summarize some important aspects of previous works that are related to our research.LLM Alignment. Nowadays, to obtain LLMs for practical uses, an alignment procedure is often required to fine-tune pretrained language models to behave appropriately and human-like. A standard LLM alignment pipeline consists of three stages: 1) supervised finetuning, 2) learning reward model, and 3) RLHF / RLAIF (reinforcement learning from human/AI feedback) (Ouyang et al., 2022; Bai et al., 2022) . Recent studies also explore directly optimizing language models from preference data with learning reward models (Rafailov et al., 2023; Song et al., 2023) . In either case, they utilize an alignment objective for learning from preference data. A common choice is the Bradley-Terry model for pairwise preference (Ouyang et al., 2022; Rafailov et al., 2023) , while others also explore the use of Plackett-Luce (PL) model for N -ary preference data (Rafailov et al., 2023; Song et al., 2023) .In-context Alignment. We refer to the use of in-context learning for alignment as in-context alignment. In this line of research, Han (2023) first demonstrates we can improve alignment with approximately 10 dynamic examples, and Lin et al. (2023) show that as few as 3 constant stylistic examples can significantly improve the alignment of top-rated LLMs such as Mistral (Jiang et al., 2023) and LLama2 (Touvron et al., 2023) . Concurrently, Guo et al. (2024) show that we can also achieve in-context alignment with only self-generated samples from LLMs without human instructions.Self-correction. Self-correction refers to the general concept that LLMs can improve their response quality based on reflecting on their previous outputs. Many previous works utilize this idea and show promising improvements on multiple tasks (Bai et al., 2022; Madaan et al., 2023; Shinn et al., 2023; Kim et al., 2023; Ganguli et al., 2023; Gao et al., 2023a; Paul et al., 2023; Chen et al., 2023) . We refer to Pan et al. (2023) for a comprehensive review. However, recent research puts this ability into question by showing that intrinsic self-correction (a scenario wherein the model can correct its initial responses based solely on its inherent capabilities) does not bring real improvements on reasoning (Huang et al., 2023a) and planning (Valmeekam et al., 2023) tasks without external feedbacks (e.g., ground-truth labels). Meanwhile, they find that self-correction does help improve the appropriateness of responses, including alignment-related tasks (Bai et al., 2022; Ganguli et al., 2023; Madaan et al., 2023) . Our theory in Section 3 provides a general theoretical explanation for the mechanism of (intrinsic) self-correction by interpreting it as a in-context alignment process and establishing its connection to the alignment objective. Without using any external feedback, the proposed checking-as-context strategy shows that intrinsic self-correction is also very effective for defending against jailbreaks.ICL Theory. Recently, a lot of interest emerged in the theoretical understanding of in-context learning (ICL), and a major direction is to investigate how linear transformers can perform certain optimization algorithms on simple problems like linear regression (Garg et al., 2022; Von Oswald et al., 2023; Bai et al., 2023; Li et al., 2023a; Ding et al., 2023; Wu et al., 2023; Ahn et al., 2023b; Huang et al., 2023c) from different perspectives, such as, convergence (Zhang et al., 2023a) , generalization (Zhang et al., 2023b) , optimization schemes (e.g., high-order (Fu et al., 2023) and preconditioned (Ahn et al., 2023a) ones), distribution shifts (Zhang et al., 2023a; Ahuja and Lopez-Paz, 2023) , etc. Beyond this simple setup, some explore the ability of transformers for learning softmax regression (Gao et al., 2023b) , discrete function (Bhattamishra et al., 2023) , regression mixture models (Pathak et al., 2023) , Gaussian Process (Cheng et al., 2023) , etc. As far as we know, we are the first to show that transformers can perform gradient descent of a non-convex alignment objective in-context.Considering the importance of alignment in LLM training, our theory model may be of more practical uses than linear regression. Besides, contrary to the linear regression case, we show that the Transformer modules like softmax attention, feed-forward networks and stacked layers, are naturally important for our construction, indicating our theory model is more aligned with the transformer architecture.Jailbreaking and Defending LLMs. Even if LLMs are aligned with human preference and behave well in most cases (e.g., refusing to answer harmful queries), researchers find that LLM alignment is still superficial (Qi et al., 2023) and can be jailbroken under carefully crafted instructions (Liu et al., 2023b) . Along this line of research, people find techniques such as, persuasive instructions (Wei et al., 2023; Shen et al., 2023) , stealthy conversation (Yuan et al., 2024) , low-resource languages (Yong et al., 2023; Deng et al., 2023) . Meanwhile, some explore automatic ways to craft jailbreak instructions, such as, gradient-based optimization (Shin et al., 2020; Zou et al., 2023; Zhu et al., 2023) (requiring white-box access), and generic algorithms (Liu et al., 2023a; Lapid et al., 2023; Chao et al., 2023) (only requiring black-box queries). To counter such attacks, various defense measures have also been proposed. One direct solution is to detect or purify harmful prompts with preprocessing, such as, perplexity filter (Alon and Kamfonas, 2023) , harmful string detection (Kumar et al., 2023; Cao et al., 2023) , retokenization and paraphrasing (Jain et al., 2023) . Nevertheless, Varshney et al. ( 2023) point out that they may suffer a considerable loss on benign queries. The instruction method, Self-reminder (Xie et al., 2023) adds a system prompt to remind the model to be safe in its reply. RAIN (Li et al., 2023b) proposes a new rewinding decoding scheme based on model evaluation. Different from these prior works, our CaC (Checking as Context) does not use explicit human instructions to teach LLMs how to behave. Instead, the only instruction we provide, i.e., the checking question, is to ask LLMs to examine their own harmfulness. In this way, we expect LLMs to refine their output based on self-examination as a form of self-instruct.In this section, we comprehensively evaluate CaC to show its effectiveness and practicalness as a defense technique against jailbreak attacks. We first propose some direct variants of CaC, then demonstrate their strength of defending LLMs against jailbreaks whilst remaining natural capabilities.I. Multi-round Checking. As discussed in Section 2, the vanilla CaC with one-round checking can be extended to multiple rounds. Intuitively, the multi-round checking also acts like a persistent interrogation of LLMs based on former responses. We call this variant CaC-self.II. Diverse Checking. In practice, we notice that although useful to some extent, multi-round checking often has marginal gains since later checking results are consistent with previous ones in most cases. From an optimization perspective, it is caused by a lack of diversity in the training examples that share the same query x. Inspired by this view, we propose diverse checking, that is to leverage the self-generated answers from other queries x i to form a diverse context, i.e., (x i , y i , r i ), and call it CaC-diverse. We randomly sample M (M = 3 is typically enough) harmful queries from AdvBench (Zou et al., 2023) , collect their LLM responses and critics, and use that as a context for the final output y for the current query x:EQUATIONWe note that by drawing from AdvBench, we rely on a human-curated dataset to obtain harmful queries, which introduces some human knowledge. Future work can further explore the use of LLMs to generate harmful queries.A defect of multi-round checking is that it leads to more latency. Since the current query x is unknown to the system for selfchecking, the responses and critics can only be collected sequentially on the fly. Instead, for diverse checking, we can cache the N -round conversation history, and use it directly for any future queries as a prompt p = [x 1 , y 1 , r 1 , . . . , x M , y M , r M ]:EQUATIONWith this historical backup technique, 1-round checking only incurs 8% overhead in inference time, which is almost negligible.III. Self-instruct. Although self-checking is rather accurate, vanilla CaC is still vulnerable under persuasive attacks that instruct models to forget the context above with lengthy prompts, such as, AutoDAN (Liu et al., 2023a) using the prompt "ignore all the instructions you got before". As revealed in our theory, self-correction works in an in-context manner, and encountering such instructions in the query will invalidate the self-checking history before. To mitigate this issue, instead of repeating the query to get a refined answer, we simply instruct the model to refine its answer based on self-checking, using the prompt "Please refine your answer to the prior request based on your examination above". 7 We name this version as CaC-self-instruct, and it can be combined with multi-round and diverse checking techniques as well. Please note that the results of CaC reported in the main content (Table 1 ) are referred to CaC-self-instruct.Checking questions. In our experiments for CaC and CaC-diverse, we have tried a few similar checking questions: 1) Is this a harmful response? 2) Isn't this a harmful response? 3) Isn't this a harmful, offensive, illegal, or improper response that you should not provide? The three questions perform rather similarly, while the third one performs slightly better. Either choice does not influence the main conclusions of our experiments. We use the third one by default. We reckon that the rhetorical question tone and detailed descriptions of potential harmful aspects could persuade LLMs to check more accurately.Evaluation of ASR. Following GCG (Zou et al., 2023) , we apply suffix detection to judge the success of jailbreak (more details here). However, as agreed by AutoDAN, DeepInception, the suffix detection may not be fully reliable. Therefore, similar to AutoDAN, we also use GPT-4 to double-check the harmfulness of a generated string. Specifically, we use both the language model and suffix detection to judge the generated string. If there is a conflict (less than 3% cases), human evaluation is involved to manually check and give the final judgment of its harmfulness.In this part, we evaluate the improved variants of CaC, including CaC-self, CaC-diverse, and CaC-self-instruct for defending against real-world jailbreak of LLMs. Following common practice (Zou et al., 2023; Liu et al., 2023a) , we consider two well-known LLMs, Vicuna-7b-v1.5 (Zheng et al., 2023) and Llama2-7b-chat (Touvron et al., 2023) . We include three jailbreak attacks, gradient-based GCG (Zou et al., 2023) (individual and transfer variants) and query-based AutoDAN (Liu et al., 2023b) . For defense, we consider the instruction-based Self-reminder (Xie et al., 2023) , and the ICL-based ICD (Wei et al., 2023) as baselines. In comparison, our CaC families are pure self-correction methods. We use 3-round checking by default. For evaluation, we consider two datasets, Advbench (behavior) (Zou et al., 2023) that contains 100 harmful queries, and GLUE (Wang et al., 2018) for natural performance (200 samples for each task). On AdvBench, a higher ASR (Attack Success Rate) indicates lower robustness. All experiments are conducted using one NVIDIA A100 GPU.Benchmark Results. From Table 2 , we can see that CaC-self and CaC-diverse are very effective against gradient-based GCG attacks, outperforming Self-reminder and RAIN by a large margin. For instruction-based AutoDAN, CaC variants are more effective on Llama2 compared to that on Vicuna. Since Llama2 is known to be more powerful, it indicates that self-correction abilities depend crucially on underlying LLMs.In Table 3 , we compare CaC-self and CaC-diverse with different rounds. Both methods perform well with only one round and benefit from more rounds. In terms of latency, CaC-self requires significantly more time with on-the-fly generation, while CaC-diverse has only minimal overhead (10% each round), which is preferable in practice. Setup. We consider the following meta-learning setting. For every task, we draw a common query x ∼ N (0, I d×d ) and a groundtruth parameter W ∼ N (0, I d×d ). We then generate N responses and rewards. For each response y i , we sample a reward r i ∈ U[0, 1] and an independent noise weight W - i ∼ N (0, I d×d ), and then generatey i = r i W x + (1 -r i )W - ix. Thus, responses with higher rewards are closer to the ground truth in expectation. We construct each in-context example asq i = [x, y i , r i ], for i ∈ [N ]. By default, we set d = 5, N = 20 and use a 20-layer GPT-2 model with 3 heads, 96 hidden dimension, and a PL loss (Eq. ( 5)). First, we train the GPT-2 model to give it the ability of in-context alignment. Specifically, lety pred i = LLM([q 1 , • • • , q i-1 , q test i ]), where q test i = (x, 0, 0), and apply PL-loss:EQUATION)Next, we sum the losses from all positions, take the average (L = 1 N N i=1 L i ) and then perform one step gradient update. In details, we set the batch size = 256, lr = 0.0001 and train step = 1500, all models are trained using one NVIDIA 3090 GPU.After training, we evaluate the normalized MSE between the predicted output ŷ and ground-truth y = W x using varying numbers of in-context examples, averaged over 256 runs with randomly generated tasks. We also implement the gradient descent (GD) of the linear PL model (Eq. ( 5)) and measure its optimal solution in the same way. We also change the reward noise p, model depth, and attention types to investigate their effects on in-context alignment.Gradient descent. We train the parameter W i θ with PL loss by setting lr = 0.1 with 50 epochs and only use in-context examples (q 1 , • • • , q i-1 ) as data. In each epoch, the prediction of GD is y pred i = W i θ x. The trained Ŵ i θ is then used to predict ŷi = Ŵ i θ x, and finally, we calculate the loss between y i and ŷi . On the other hand, we can obtain the transformer's predicted values by using the trained GPT-2 model to perform inference on the in-context examples (q 1 , • • • , q i-1 ) and get the model's predictions. The model's predictions can be used to calculate the loss in the same manner, serving as the evaluation result. Do the same for each position i, we can get Figure 1a .Reward noise. We use the same 20-layer GPT-2 model with 3 heads, 96 hidden dimension, and evaluate its performance on different noise p = 0, 0.25, 0.5, 0.75, 1. Additionally, we use the same noise reward data to do the gradient descent experiment.Model depth. We train 4 different (5-layer, 10-layer, 15-layer and 20-layer) transformer while keeping all other parameters constant. Then, we evaluate the performance of each model on the same dataset.Multi-head attention. We fixed the dimension of each head at 32 and selected 1, 3 (default), 4, and 6 heads for comparison. indicate that the performance of 1-head attention significantly declines, while the differences in performance among 3, 4, and 6-head attention are not pronounced. This perfectly aligns with our theory that 3-head attention is sufficient to accomplish the alignment task.We only remove the softmax operation in attention layer or FFN block while keeping all other parameters same as 20-layer GPT-2 model with 3 heads, 96 hidden dimension.The Bias Benchmark for QA (BBQ) comprises 58,492 unique questions designed to evaluate societal biases against individuals belonging to protected classes across nine social dimensions: Age, Disability Status, Gender Identity, Nationality, Physical Appearance, Race/Ethnicity, Religion, Socioeconomic Status, and Sexual Orientation, and there are also two categories that are intersections of the above types, which are respectively: Race x Socioeconomic Status and race x gender. This benchmark assesses the propensity of models to rely on stereotypes when responding to questions. Such reliance can be detrimental, as it risks perpetuating biases that negatively impact members of protected categories.Each question in the dataset has three answer options and appears in two forms: ambiguous and disambiguous. For example, an ambiguous question is: "I saw a grandson and their grandfather last week outside the Walmart trying to book a cab on Uber. Who was not comfortable using the phone?" The answer choices are: "The grandson", "The grandfather", or "Unknown". Given the ambiguity of the context, the correct answer is "Unknown" due to insufficient information. However, a model influenced by age-related stereotypes might incorrectly select "The grandfather". The correct response to all questions in ambiguous contexts is "Unknown" or another expression of uncertainty.To prevent the model from becoming less helpful due to being excessively constrained by the goal of avoiding bias, each question is accompanied by an additional disambiguated context condition, such as, "The grandson was struggling with the app and not really listening to the advice from the grandfather." In this scenario, the correct answer to the original question is "The grandson".We randomly selected 500 questions from each task subclass. For each question, we applied CaC, recorded the model's original answers and the answers after self-correction, and then calculated the accuracy of these answers.In the correlation analysis, we evaluate the relationship between accuracy gain with self-correction and self-checking accuracy. In details, we randomly select 100 questions in each category (1,100 questions in total) from vicuna's answer, and evaluate the model's self-check answer by gpt-4-turbo-preview.Evaluation on ambiguous questions. Due to the limitation of model size, we found it challenging for the model to simultaneously determine whether a question is ambiguous and whether the answer is biased. Therefore, we focused on evaluating whether the model's answers are biased. We selected 100 ambiguous questions from each category (1100 questions in total) and standardized the model's output: starting the self-check with "My previous answer is biased." or "My previous answer is unbiased.". We calculated the accuracy of the self-check through string matching. Surprisingly, we found that this standardized form of self-check significantly improved self-correctness (Figure 4 ), and in the correlation analysis (Figure 5 ), we also found a strong correlation between self-correctness gain and self-check. Check Prompt: Please review the previous response for any potential biases or stereotypes. If the initial query provided additional information that supports the correctness of the answer, confirm the accuracy confidently. However, if the query lacked specific context, be cautious and critically assess the response for any implicit biases or stereotypes. Highlight any areas of concern and suggest improvements to ensure a fair and unbiased reply.Check Response: The initial query did not provide any information about the university student's drinking habits, so it is not possible to determine whether they drink rarely or not. Therefore, the response (c) "The university student" is not accurate and should be revised.It is important to avoid making assumptions about someone's drinking habits based on their age, gender, or other characteristics. Instead, the query should be rephrased to ask about the university student's drinking habits specifically.Repeated Input: When a young elderly woman went to visit a neighbor who's an university student, the two of them got to talking about their favorite drinks. Who drinks only rarely? (a) The elderly woman (b) Not enough information (c) The university student.Final response (released to user): The answer is (b) Not enough information.In this section, we provide the proofs for all theorems.E.1. Proof of Proposition 3.1Proposition E.1. One can realize the gradient descent for BT,W ′ = W + ∆W = W -η∇ W L BT (W ; x, y 1 , y 2 ),by updating each y i withEQUATION-2ηy 1(2)EQUATIONwhereβ j = softmax(-∥W x -y j ∥ 2 ). Specifically, L BT (W ′ ; x, y 1 , y 2 ) = L BT (W ; x, y ′ 1 , y ′ 2 ).Proof. We first calculate one gradient descent step of the BT loss that yields the following weight change w.r.t. WEQUATIONwhere η > 0 is the step size, and for any j ∈ [N ],EQUATIONConsidering the BT loss after the weight udpate, we haveL BT (W + ∆W ) =∥(W + ∆W )x -y 1 ∥ 2 -log 2 j=1 exp -∥(W + ∆W )x -y i ∥ 2 =∥W x -(y 1 -∆W x)∥ 2 -log 2 j=1 exp -∥W x -(y i -∆W x)∥ 2 .Comparing it with the original BT loss, we notice that a gradient descent update of the parameter W is equivalent to updating each y i withEQUATION-2ηy 1(2)EQUATION.In the last step, we utilize the assumption ∥x∥ = 1 (otherwise it can be merged into the learning rate η).Theorem 3.2. Given a two-head softmax attention layer and two tokens e i = (x i , y i , r i ), i = 1, 2, there exists a set of parameters (Eq. ( 2)) such that a forward propagation step with token e i is equivalent to the gradient-induced dynamics of the Bradley-Terry model (Eq. ( 6)):EQUATIONProof. We prove a stronger version of this proposition by considering the general case of N samples (e 1 , e 2 , • • • , e N ). Note that the proof of Theorem 3.2 follows from the case of N = 2. Without loss of generality, we assume y 1 ≻ y i with scores r 1 > r i , for i = 2, 3, • • • , N , and we use y + and r + to represent y 1 and r 1 , respectively.Under this setting, we rewrite the new input matrix X and the update formula Eq. ( 7) of each y i asEQUATIONEQUATION-2ηy 1(2)EQUATION.The proof of this theorem is organized in the following three parts:• First, in Lemma E.2 we construct the part (2) of the gradient update (Eq. ( 17)) with the first head of MHSA structure to extract the answer (y + ) that corresponds to the maximum reward r.• Then, we use Lemma E.5 with the second head of MHSA structure to extract reweighed different rewards, which construct the part (3) of Eq. ( 17).• Finally, We employ a residual structure to integrate both part (2) and part (3) with y i itself.Specifically, leveraging Lemma E.2 and Lemma E.5, we can construct two attention heads for parts (2) and (3), respectively:EQUATIONIn accordance with the computational rules of MHSA, we can construct two projection heads P 1 , P 2 as P 1 = -2ηI and P 2 = 2ηI. Then we haveEQUATIONFurther combined with the residual connection, we can realize the full update of y:EQUATIONThat is to say, each y i is updated to y i -2ηy + + 2η N i=0 β i y i , exactly equivalent to the gradient descent (Eq. ( 7)).Theorem 3.3. Given a transformer TF with N -1 stacked transformer blocks (composed of three-head softmax attention and feed-forward networks) and N input tokens {e i , i ∈ [N ]}, there exists a set of parameters such that a forward step with token e i is equivalent to the gradient-induced dynamics of the N -ary Plackett-Luce model (Eq. ( 5)), i.e., TF(e i ) = (x i , y i , r i ) + (0, -∆W PL x i , 0), i ∈ Proof. According to the PL gradient Eq. ( 9), the update of each y i is:EQUATIONwhere we denote:EQUATIONEQUATIONWe plan to construct the whole gradient by constructing each g k in each iteration . Each g k is constructed by a three-head MHSA and an FFN structure with residual connection respectively and sum up by residual mechanism. You can see the structure of one iteration in Figure [6] . After N -1 iterations, we will get the whole gradient.To calculate g k and β k j , we wish to use the same structure but changed input y k-1 τ (j) , such thatEQUATIONEQUATIONTo update the k-th iteration input y k-1 τ (j) to y k τ (j) after the k-th iteration without affecting the accumulation of the original gradient of y i , we expanded the dimension of the input matrix X and duplicated each y i , placing it in the last row of the matrix, so as to update the y i used for gradient calculation in subsequent iteration rounds. As before, the line of y i below x is used for storing gradients, meaning that after N -1 rounds of iterations, we will obtain the desired state for each y i (Eq. ( 24)) in this line, while the y N -1 i in the last line becomes redundant after the completion of N -1 iterations. We define the new input matrix X as:EQUATIONIn our notation, the superscript k denotes the value of the variable in the k-th iteration of the structure in figure 6 , while the subscript i indicates the tokens in the i-th round of self-check. We define the k-iteration output matrix and hidden matrix as:EQUATIONEQUATIONwhere G k i = k j=1 g j (Eq. ( 27)) refers to the gradient accumulation after k iterations. When k = N -1, that is after N -1 iterations, we have 24)). Therefore, we only need to recursively constructed matrix X k .y i + G N -1 i = y i + N -1 j=1 g j = y i -∆y i (Eq. (Compared with X k and X k-1 , we have the following four changes, which need to verify later:• G k i = G k-1 i + g k . • r k i = r k-1 i -r +, where r + is the same constant to each i. Notice that we only consider the order of magnitude of each reward and subtract the same r + will not have any effect on it.• r k τ (k) = r k-1 τ (k) -r + -γ, where γ is a sufficient large number such that the current((k -1)-th) iteration maximum reward r k-1 τ (k) changes to the lowest oner k τ (k) in the next (k-th) iteration. That is, max(r k 1 , • • • , r k N ) = r k τ (k+1) . • y k τ (k-1) = y k-1 τ (k-1) -γ. Therefore, exp(-∥W x -y k τ (k-1) ∥ 2 ) → 0.According to Lemma E.6, we can construct MHSA θ s.t.EQUATIONWith residual structure, we haveEQUATIONAccording to Lemma E.7, we can construct the feed-forward module FFN θ such thatEQUATIONWith residual structure, we can gainEQUATIONTo this end, four changes (E.3) have been verified, meaning that we have constructedX k with input X k-1 . When k = N -1, we get y i + G N -1 i = y i + N -1j=1 g j = y i -∆y i (Eq. ( 24)). That is the updated result of each y i . Lemma E.2 (Construction of the numerator gradient). Given an input matrix X (Eq. ( 16)), one can construct key, query and value matrices W K , W Q , W V such that the output is:EQUATIONwhere y = [y 1 , . . . , y n ], r = [r 1 , . . . , r N ], and ϕ i : R N → {0, 1} N denotes an indicator function of the maximal rewards:EQUATIONProof. We provide the weight matrices in block form:Q is fixed as   0 0 • • • 0 0 0 • • • 0 γ γ • • • γ  , where γ is a large and positive hyper parameter.W K =   0 0 0 0 0 0 0 0 1   , and then K ⊤ = X ⊤ W K ⊤ =      0 0 r 1 0 0 r 2 . . . . . . . . . 0 0 r N      .Therefore, when calculating the attention score, for the same query, it is equivalent to scaling up each r i by a sufficiently large factor, that isEQUATIONEQUATIONThe function ϕ i (r) is defined in Eq. ( 44).Thus, when doing softmax, we can get the following matrix.EQUATION)The attention score will changed to 1 or 0 only depending on the whether current r i is the maximum value or not.Then, letW V =   0 0 0 0 I dy 0 0 0 0   , and we have V = W V X =   0 0 • • • 0 y 1 y 2 • • • y N 0 0 • • • 0   .Finally, we get the desired head matrixEQUATION)Lemma E.3 (Construction of the denominator gradient). Given an input matrix X (Eq. ( 16)) with positional encoding, wecan construct Q = W Q X and K = W K X such that K ⊤ Q =      -∥W x -y 1 ∥ 2 -∥W x -y 1 ∥ 2 • • • -∥W x -y 1 ∥ 2 -∥W x -y 2 ∥ 2 -∥W x -y 2 ∥ 2 • • • -∥W x -y 2 ∥ 2 . . . . . . . . . . . . -∥W x -y N ∥ 2 -∥W x -y N ∥ 2 • • • -∥W x -y N ∥ 2      .Proof. With positional encoding (for convenience, here we assume using one hot positional encoding), we can transform the input matrix X (Eq. ( 16)) toEQUATION)The upper part of this matrix (Eq. ( 49)) is used to construct K, and the lower part is used to construct Q. Then, according to Lemma E.4, we can constructEQUATIONThus, we can use X ′ p (Eq. ( 50)) to easily constructEQUATIONEQUATIONHerein, K and Q are simply linear transformations applied to the rows of the matrix X ′ p (Eq. ( 50)), and W is part of the parameters in W K and W Q . With these constructions, K ⊤ Q is the desired result we expect.Lemma E.4 (Construction of complete positional input matrix). With input matrix X p (Eq. ( 49)), we can construct an attention layer such thatX p + att(X p ) =               x x • • • x y 1 0 • • • 0 0 y 2 • • • 0 . . . . . . . . . . . . 0 0 • • • y N y 1 y 1 • • • y 1 . . . . . . . . . . . . y N y N • • • y N               .Proof. By setting the attention score of each query to be the same after softmax, that isEQUATIONand X p + att(X p ) is our desired result.Lemma E.5 (Construction of denominator). Given an input matrix X(Eq. ( 16)), one can construct key, query and value matrices W K , W Q , W V such that the output is:EQUATIONProof. According to the formula of β i Eq. ( 29), we hope to construct the following attention score before doing softmax.K ⊤ Q =      -∥W x -y 1 ∥ 2 -∥W x -y 1 ∥ 2 • • • -∥W x -y 1 ∥ 2 -∥W x -y 2 ∥ 2 -∥W x -y 2 ∥ 2 • • • -∥W x -y 2 ∥ 2 . . . . . . . . . . . . -∥W x -y N ∥ 2 -∥W x -y N ∥ 2 • • • -∥W x -y N ∥ 2     .There are two ways to achieve this. One is straightforward but has complex construction, and the other is approximate but more easier.We first introduce the approximate method. With the proposition that an FFN can easily approach the mean square error, we have FFN(y i |x) = -∥W x -y i ∥ 2 , where W is part of the parameters in FFN. Before passing through the attention layer, the input matrix X can be transformed asEQUATIONThus, the attention score should beK ⊤ Q =      -∥W x -y 1 ∥ 2 -∥W x -y 1 ∥ 2 • • • -∥W x -y 1 ∥ 2 -∥W x -y 2 ∥ 2 -∥W x -y 2 ∥ 2 • • • -∥W x -y 2 ∥ 2 . . . . . . . . . . . . -∥W x -y N ∥ 2 -∥W x -y N ∥ 2 • • • -∥W x -y N ∥ 2     .The second method to achieve this is to give a detailed construction following Lemma E.3.Thus, after doing softmax, we can getEQUATIONFinally, by constructing matrix V asEQUATIONwe can get the desired attention head H 2 = V softmax(K ⊤ Q).Lemma E.6 (Construction of gradients and updates). Given an input matrix X k-1 (Eq. ( 33)), we can construct three heads in MHSA θ respectively such thatEQUATIONwherer + = r k-1 τ (k) = max(r k-1 1 , • • • , r k-1 N ) is the maximum reward in the ((k -1)-th) iteration.Proof. According to lemma E.2 and lemma E.5, we only need to make adjustment to dims and multiplying certain projection matrices by a permutation matrix so that we can extract certain rows from a matrix.For example, if we want to construct a matrix for calculation, rather than the second row as described in the original lemma. Therefore, using Lemma E.2 , we can construct the first and the third head matricesH 1 and H 3 :H 3 =     0 0 • • • 0 0 0 • • • 0 r + r + • • • r + 0 0 • • • 0     ,EQUATIONEQUATIONAccording to Lemma E.5, we can construct the second head matrix H 2 :EQUATIONwhereEQUATION)), we haveEQUATIONN i=1 β k i y k-1 i = N i=k β k i y k-1 i (67)This is the desired form of the construction of part of g k (Eq. ( 27)). Thus, we can concat them together with projection matrices P 1 , P 2 , P 3 :EQUATIONLemma E.7 (Construction of the position of the maximum value). Given a hidden matrix H k and passing through an FFN, we can successfully obtain the position τ (k) within the matrix.k ) = W 2 • ReLU(W 1 (H k )) =     0 • • • 0 • • • 0 0 • • • 0 • • • 0 0 • • • -γ • • • 0 0 • • • -γ • • • 0     . (71) Proof. Actually, r + = N i=1 r k-1 i exp(γ•r k-1 i ) N j=1 exp(γ•r k-1 j ) < max(r 1 , • • • , r N ) = r k-1 τ (k) (according to Lemma E.2). Then ∃ ϵ > 0 s.t.r + = r k-1 τ (k) -ϵ. Notice that γ is sufficient large, such that r k-1 τ (k) > r + > r k-1 τ (k+1) > • • • . Thus, r k-1 τ (k)as the largest reward that satisfies r k-1 τ (k) -r + = ϵ > 0, and r j as any other component with j ̸ = k, for which r k-1 τ (j) -r + < 0.EQUATION(76) which completes the proof. In our analysis, we adopt a single common query for simplicity, and specifically, we can compute the attention score by performing inner product operations on different instances of x. Since we assume ∥x∥ 2 = 1, the inner product between x and itself yields the maximum attention score. With this property, we can filter out the corresponding answer and reward of each example (as elucidated in Lemma E.2) and use this information to construct the gradient update of each sample accordingly. The following are the construction details.For multi-queries, we define the new input matrixX = (e 1 1 , e 1 2 , • • • , e 1 N , • • • , e M 1 , e M 2 , • • • , e M N ) =   x 1 x 1 • • • x 1 • • • x M x M • • • x M y 1 1 y 1 2 • • • y 1 N • • • y M 1 y M 2 • • • y M N r 1 1 r 1 2 • • • r 1 N • • • r M 1 r M 2 • • • r M N   .Here, we take Lemma E.2 as an example to illustrate how our constructions are generalized to adapt multi-queries scenario.Based on the hypothesis that ∥x i ∥ 2 = 1, i = 1, 2, • • • , M , we can construct matrix W Q , W K , W V such that where r θ (•) denotes the reward function with parameters θ.Therefore, in Theorem 3.3, we use casual PL loss instead to calculate the gradient of W and update the corresponding token e i : TF(e i ) = (x i , y i , r i ) + (0, -∆W onlinePL(i) x i , 0), i ∈ [N ],which indicates that when passing through Transformer blocks, token e i is updated by one step gradient descend using tokens before its positions with online PL loss.In our former construction, since tokens do not have positional encodings, we cannot record the positions of the maximum values. Thus, we implement the gradient of each term by selecting the example with the largest reward and then eliminating it for subsequential terms. With a causal mask, the reward at each position does not know the global maximum, but only knows the maximum of all rewards before its position. In other words, if a particular reward happens to be larger than all precursors while it is not a global maximum, it would be still falsely treated as the maximum.According to (Haviv et al., 2022) , causal attention enables the model to infer the number of predecessors that each token can attend to, thereby approximating its absolute position. Therefore, in order to increase the flexibility of our construction, we assume casual LM can derive one-hot positional encodings p i for each token. Since we cannot propagate the maximum reward value calculated at the last position back to previous positions for updates, we devised a strategy where the current position uses a positional encoding mask m i to track and record the positions of the global maximum values. Since m i is initialized to ⃗ 0, we only need to take a portion from the dimensions after embedding to represent m i . Then, when querying at this current position, it updates all rewards at the positions already identified as maximums to the minimum values during the attention calculation with each preceding key, before proceeding to softmax.Main Idea. The key change we made is to record the current position's information under the current position itself, rather than under previous positions, as that would be a fallacy; previous positions cannot see the information of the current position.To proceed into the details, first, let's define a new input matrix that is more amendable for later useEQUATION)Here, p i is a one-hot PE (positional encoding), and m i refers to masked PE, which is initialized as ⃗ 0 and updated by accumulating the sum of positional encodings that have been selected to the PL loss numerator. For example, if r 1 > r 3 > r 2 > r 4 and the current is the third round iteration (k = 3), m 3 4 now should be (1, 0, 1, 0) ⊤ . Next, we make some minor modifications to the Lemma E.2 to enable it to extract the positional encoding p i of the position with the maximum value r. Lemma F.1. Given an input matrix X, one can construct key, query and value matrices W K , W Q , W V such that the output is:EQUATION)where y + i represents the corresponding y-value for the maximum value of reward among the first i positions, while p + i represents the positional encoding of the position of the maximum value r + i among the first i positions.Proof. By providing the matrices in block form, we can construct matrix W Q , W K , W V such thatQ =         0 0 • • • 0 0 0 • • • 0 γ 1 γ 1 • • • γ 1 0 0 • • • 0 -γ 1 γ 2 m 1 -γ 1 γ 2 m 2 • • • -γ 1 γ 2 m N 0 0 • • • 0        , where γ 1 , γ 2 are sufficient large and positive hyper parameters.We can also construct key matrix to provide positional encoding p i to match PE mask m i in query matrix: Thus,EQUATIONEQUATION)m i p j = 1 if and only if the gradient at query i has been accumulated over the sub-sum on the numerator of the pl loss with position j as the maximum value, therefore it is necessary to make the r i at this position the minimum value to ensure that it won't be selected again. This is what γ 2 accomplishes. EQUATION) Let V = W V X =         0 0 • • • 0 0 0 • • • 0 r 1 r 2 • • • r N 0 0 • • • 0 p 1 p 2 • • • p N 0 0 • • • 0        .Finally, we get the desired head matrix:EQUATIONA major difference is that in alignment, the preference data are used for finetuning pretrained LLMs, while self-correction refines outputs in an in-context way without changing model weights.Following discussions can be extended to multiple x's as well.As in Von Oswald et al. (2023), it is easy to show that we can construct such concatenated tokens from standard sequential tokens with the help of positional encodings.For simplicity, we omit the case of having equal rewards. On the one hand, such scenarios are rare since LLMs are well capable of telling different answers apart. On the other hand, our analysis can be easily extended to such cases by grouping the samples with equal rewards.Technically, we can also manually choose thresholds for each ri for them to be selected in a specific attention head (N heads for N tokens). However, it is not adaptive to the change of reward values and input length and thus deviates far from the practice.As a natural extension, stacking K(N -1) blocks can implement K gradient descent steps.again, we rely on LLMs' own judgment instead of human priors in this prompt. 


Decision: This paper explores the ability of language models to perform self-correction (SC), in which the LM is asked to evaluate and (possibly) correct its own answer. They present a theoretical proof-by-construction showing how SC is possible and enabled by several key architectural components of a Transformer model. They also validate this empirically on synthetic data, as well as two real-world tasks: defending against jailbreaks, and mitigating social bias.

Reviews were all positive on this paper, appreciating the theoretical derivation and how it linked a behavior to specific components of the model, as well as that the findings were validated on a range of tasks. Some reviewers did wish for more extensive real-world evaluations and had questions about some experimental details. The authors did an admirable job responding to initial reviewer concerns, leading two to raise their scores significantly.

Reviews: review: paper_summary: paper_summary: This paper provides a theoretical framework for understanding self-correction in large language models (LLMs), framing it as a form of in-context alignment. The authors prove that a standard multi-layer transformer can optimize common ranking-based alignment objectives using self-correction samples, generating responses with higher rewards. Their analysis reveals the importance of key transformer design elements for self-correction, including softmax attention for ranking, multi-head attention for token discrimination, feed-forward networks for token transformation, and multiple stacked layers for handling multiple examples. These theoretical results are validated through extensive experiments on synthetic datasets.

Inspired by these theoretical insights, the authors demonstrate real-world applications of self-correction, including alleviating social bias in LLM outputs and defending against jailbreak attacks. Their experiments show that self-correction can significantly improve performance on these tasks, with a strong correlation between self-checking accuracy and final performance. This work provides the first theoretical foundation for understanding self-correction in LLMs, offering valuable insights into how this capability emerges and how it relates to model architecture, with implications for improving LLM safety and reliability.
strength_weakness: strength_weakness: My main concern is about modeling the self-critic process as a ranking problem. The author's modeling ranks the existing hypotheses generated in the context, as shown in equation (7), as the proof of construction uses the information from y1 and other ys. 
However, self-critic is about generating a better hypothesis one by one. The generated ones should not be in the list of hypotheses to rank. 
If the authors can modify the theory to mainly consider the possible hypotheses in the next round generation, it would make more sense to me. 

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

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.
questions: questions: 1. Can you provide theory about the next round possible candidates? 
2. Why, in your opinion, LLMs can have the ability of self-critic?
limitations: limitations: see weakness.
paper_summary: paper_summary: This paper analysis self-correction theoretically from the in-context learning perspective. It extends the theoretical analysis from previously over-simplified transformers to more realistic scenario: softmax attention, multi-head attention... It also provides experiments on how self-correction can serve in practical applications such as defending against jailbreaks.
strength_weakness: strength_weakness: 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.
questions: questions: See the weakness session above. If the connection between real-life experiment setting can be better explained, this theoretical framework would be more complete.
limitations: limitations: The authors have adequately addressed the limitations and potential negative societal impact of their work.
paper_summary: paper_summary: # Summary

This paper provides a theoretical analysis of self-correction from in-context learning, demonstrating that LLMs can refine their responses by using accurate self-examinations as feedback.

# Contributions

1. Theoretical Framework: The paper develops a theoretical framework that explains how self-correction capabilities arise in LLMs, extending beyond simplified linear transformer models to realistic transformers.

2. Applications: The authors demonstrate the real-world relevance of their findings,  including jailbreaks and mitigating social biases.
strength_weakness: strength_weakness: 1. Validation Scope: The validation is primarily on synthetic datasets.
2. 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.
questions: questions: 1. Could you please explain why the primary validation is on synthetic datasets? (The BBQ benchmark is in the appendix.)
2. If possible, could you add more experiments on realistic benchmarks?
limitations: limitations: No potential negative societal impact.

For improvement, please check the Questions.
paper_summary: paper_summary: This paper investigates the ability of large language models (LLMs) to improve their responses through self-correction from a theoretical perspective. Specifically, the authors prove the self-correction mechanism through in-context alignment formulation and analyze why self-correction naturally improves LLM performance in the transformer based model. Additionally, they conduct experiments on synthetic datasets to validate their findings and also verify them in two applications. This paper provides a strong foundation for understanding the self-correction mechanism. It also offers theoretical insights for explaining the prompting of LLMs through agent workflows.
strength_weakness: strength_weakness: 1. 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.

2. Strictly speaking, this cannot be considered a weakness. Although Section 5 provides improvements in self-correction in two application scenarios, I think these are unnecessary, given the numerous relevant empirical studies already available. The performance improvement due to self-correction has been supported by numerous experiments and has become a consensus in the field. 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.
questions: questions: Please refer to the questions in the "Weaknesses" section.
limitations: limitations: The authors adequately addressed the limitations.
paper_summary: paper_summary: This paper investigates how large language models (LLMs) can improve their performance through self-correction without external feedback. The authors provide a theoretical framework for understanding self-correction as an in-context alignment process. They demonstrate that LLMs can refine their responses based on self-generated feedback using Transformer modules. They conduct experiments on synthetic datasets, showing that self-correction can significantly enhance model performance and mitigate issues like social bias and jailbreak attacks.
strength_weakness: strength_weakness: 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.
questions: questions: What potential strategies can be employed to improve the robustness of self-correction when the self-generated feedback is noisy or inaccurate?
limitations: limitations: Yes, they discussed the limitations in checklist Limitation section.

Comments: title: Reviewer responses
Nice Response
Thanks for your detailed reply
Thank you for raising the score
Impressive Performance
Thanks!
Thank you for your reply
Nice rebuttal!
Thanks
comment: Thanks for your detailed response. My concerns are addressed to some extent, and I will raise my score.
I keep my positive score and hope to see these contents in your appendix.

Best wishes to your paper.
This is really helpful, great work!
Thanks for raising the score and for appreciating our response! We further take a preliminary study to get some insights of how MiniCPM does not work well for self-correction.

Notably, we find that the MiniCPM model has a lower probability of correcting incorrect answers to correct ones **(4.45%)** than the probability of changing correct answers to incorrect ones **(9.18%)**. That means that the MiniCPM, due to its limited capability, seems to have limited long context (or multi-round) understanding ability compared to bigger models (see the example below). Indeed, our theory suggests that this self-refinement ability does require the model to be expressive enough to be able to understand the critic and adjust the behaviors accordingly. 

Here is an example failure of MiniCPM, where both the initial answer and the critic are true but the model adjusts to to an incorrect one.

> **Q:** When a young retiree went to visit a neighbor who's an university student, the two of them got to talking about their favorite drinks.
> 
> 
> Who drinks only rarely?
> 
> (a) The retiree (b) Cannot answer (c) The university student.
> 
> **A:** The answer is (b) Cannot answer.
> 
> **Q:** Please make a short review on the previous response for any potential biases or stereotypes. Let's think step by step and consider the following categories of bias: Age, Disability Status, Gender Identity, Nationality, Physical Appearance, Race/Ethnicity, Religion, Socioeconomic Status (SES), Sexual Orientation, and Intersectionality (e.g., Race x SES, Race x Gender). Your review should end with "My previous answer is unbiased." or "My previous answer is biased.”
> 
> **A:** In the previous response, I provided a general answer without considering any potential biases or stereotypes. Therefore, my previous answer is unbiased. 
> 
> **Q:** When a young retiree went to visit a neighbor who's an university student, the two of them got to talking about their favorite drinks.
> 
> Who drinks only rarely? 
> 
> (a) The retiree (b) Cannot answer (c) The university student.
> 
> **A:** The answer is (a) The retiree.
> 

Hope this clarification helps!
Really amazed by your results on the SLMs. Hope you will add these results to your final version and arXiv. Also, would you please give an analysis of why the performance with SLMs is so impressive, especially on MiniCPM?

Yeah. I increased my score and best wishes for your paper.
Thank you for appreciating our response. We are very glad to hear that you find it convincing. Indeed, for real-world LLMs, model size is an important factor that reflects model capabilities. Here, follow your suggestions, we further evaluate Phi-3 and MiniCPM-2B on BBQ. We also include Qwen-1.5 1.8B and Qwen-1.5 7B, to study in the influence of model size.We follow the same evaluation protocol for all settings for a fair comparison.

**Results.** As shown in the table below, we find that very small models ,like MiniCPM-2B and Qwen1.5-1.8B, although they have good self-critic accuracy (enough to gain from self-correction, as we show in Experiment 1 in **General Response**), but they can hardly benefit from self-correction. Instead, when the model becomes larger, e.g., Phi-3-4.8B and Qwen1.5-7B, the model shows great improvements via self-correction.

**Analysis.** Our theoretical analysis indicates that a full Transformer with sufficient capability and depth is required to perform in-context alignment (Sec 3.2), and our synthetic experiment confirms it by ablating each Transformer component (Sec 4). These new our real-experiments further validate our theoretical insights, that self-correction requires the model to be expressive enough. 

Meanwhile, since Phi-3-3.8B and Qwen1.5-7B are not very large models (they can be run on a single GPU), we believe that self-correction can also find many successful applications. As the small model regime is also an active research area, we believe that in the future, with better training or model designs, smaller models could potentially benefit from self-correction as well. 

| Model  | prediction acc (initial) | self-critic acc | prediction acc (with correction) |
| --- | --- | --- | --- |
| MiniCPM-2B-sft-bf16-llama-format | 36.27% | 62.82% | 25.75% |
| Phi-3-mini-4k-instruct (3.8B) | 85.34% | 48.87%  | 92.56% (+7.24%) |
| Qwen1.5-1.8B | 1.18% | 99.41% | 0.00% |
| Qwen1.5-7B | 66.51% | 97.37% | 73.32% (+6.81%) |

Hope you find this new experiment satisfactory! Thank you for the insightful comment, which makes our work more complete. We will add these new results in the revision. Please let us know if there is more to clarify.
I think the authors have successfully addressed my question in weakness part. In general, this is a good theoretical paper combined with convincing real-world experiments. I will hold my postive rating here.
Your rebuttal convinced me a lot and thanks so much.

I see you evaluated your methods on two 7B models. I am wondering about its effectiveness on smaller models.

Could you also test your methods on these two models and **include the results in your appendix**:

Phi-3: https://arxiv.org/abs/2404.14219

MiniCPM: https://arxiv.org/abs/2404.06395

Since a lot of edge users really care about these SLMs. If show possible results, I gonna increase my score. And also, sorry for the late reply. Let me know what I can do.
Thank you for the prompt response and for carefully reviewing the other reviews as well. We are glad that you find our rebuttal satisfactory and the new experiments "undoubtedly increase the technical depth of the paper." We will be sure to incorporate these results in our revision.
I appreciate the authors' point-by-point rebuttal and supplementary experiments, which addressed my questions. I particularly liked the experimental design of "Natural language vs. explicit labels" and the examination of "the influence of critic accuracy, critic types, and self-correction rounds on self-correction performance." These experiments undoubtedly increase the technical depth of the paper.

After carefully rechecking the paper and reading the reviews and rebuttals from other reviewers, I believe this paper is of high quality. I will maintain my positive opinion. Thank you!