**Overhead** Smaller models are a preferred choice
for resource-constrained deployments, where the
inference cost matters the most. BiAlign does
not introduce any additional cost during inference.
The additional computational overhead only occurs
once during model training. To quantify the
increase in computational overhead caused by the
ranking loss, we use DeepSpeed Flops Profiler
(Rasley et al., 2020) to calculate the training FLOPs
of Output-Align and BiAlign, which are 3.3×10 [17]
and 7.6×10 [17] respectively (about 2.3 times).
Therefore, we further design two experiments to
compare BiAlign and Output-Align under the same
training FLOPs: ( _i_ ) we combine the original ICL
training examples with the sampled subset data
and conduct Output-Align on the combined data
(roughly the same FLOPs as BiAlign), which
performs (29.5) similarly to the original OutputAlign method (29.4), verifying the superiority of
BiAlign; ( _ii_ ) we reduce the training epochs of
BiAlign from 4 to 2 (roughly the same FLOPs
as Output-Align) and assess the final checkpoint.
There is no significant performance degradation
(from 31.0 to 30.8), which also demonstrates that
BiAlign can outperform baselines under the same
training FLOPs.
**Different Ranking Loss Formulations** In the
right part of Equation 3, we employ the rank()
function to represent the relative ranking of the
model’s preference scores instead of relying on
the scores themselves. This approach is grounded
in the idea that the primary objective of input
preference alignment is to match the rankings
of the subset scores, rather than their specific
values. By focusing on rankings, we can reduce the
impact of potential variations in score magnitudes,
allowing the model to prioritize the relative ranking
of preferences. We further conduct experiments
with an alternative ranking loss formulation that
does not incorporate rank(), while maintaining all
other implementation details. The average results
reported in Table 5 underscore the importance of
using rank() for alignment.
**Connection with In-Context Pretraining** Shi
et al. (2024) propose In-Context Pretraining (ICP)
which pretrains language models on a sequence of
related documents. BiAlign mainly differs from it
in the following two aspects: ( _i_ ) ICP focuses on
the pretraining stage while BiAlign is specifically
designed for more lightweight supervised finetuning. ( _ii_ ) The objective of ICP is to design more
effective pretraining data. In contrast, BiAlign
leverages distillation to improve the capabilities
of the student model. Therefore, BiAlign can be
seamlessly integrated with ICP to further improve
the ICL ability.
**Effect of Demonstration Numbers** As mentioned in Section 4.1, each constructed ICL training
sample contains 4 ≤ _k_ ≤ 10 demonstration
examples, which could enhance the model’s
ability to generalize to different numbers of
demonstrations. To investigate the effect of
demonstration numbers in source tasks, we further
conduct training on examples containing a fixed
number _k_ ∈ {5 _,_ 8 _,_ 10} of demonstrations. The
average results of the 5 target tasks are reported
in Table 6. We can see that training with a fixed
number of demonstrations results in performance
degradation to a certain degree, justifying our
training set construction strategy.
**Preference Score Consistency** As illustrated
in Section 3.2, _L_ [rank] enables the student model
to match the relative ranking of the preference
scores for different ICL demonstrations with that
of the teacher model. To verify this, we report the
_preference score consistency_ comparison between
BiAlign and Output-Align in Figure 4. Specifically,
we randomly select 500 examples from MMLU
(see results on other datasets in Appendix A.7).
For each example, we randomly sample 5 subsets


Default (4 ≤ _k_ ≤ 10) 5 8 10
BiAlign **31.0** 30.8 30.4 30.5
Table 6: Average results (%) of BiAlign with different
_k_ (demonstration number) for constructed ICL training
samples.
from the set of all demonstrations and obtain
their preference scores using different models.
The preference score consistency of different
methods is then calculated as the proportion of the
highest/lowest scoring subsets that are consistent
between the corresponding student model and
the teacher model. From the results, we can
see that BiAlign can indeed achieve much higher
preference score consistency than Output-Align,
indicating the effectiveness of _L_ [rank] .
In addition, for interested readers, we show
further justification of input preference alignment,
more cross-task generalization experiments, the
results with different subset sampling methods,
different numbers of subsets and different source
task selections, the analysis of KL divergence
calculation, training steps and additional training
data, the influence of ranking loss weight, the
effect of contrastive pair selection, and a case
study of model output in Appendix A.8 ∼ A.18,
respectively.

Improving In-context Learning via Bidirectional Alignment

Chengwei Qin
Wenhan Xia
Fangkai Jiao
Shafiq Joty


In this work, we have introduced Bidirectional
Alignment (BiAlign) that can improve the ICL
capabilities of student models by aligning the
input preferences between student and teacher
models in addition to aligning the token-level
output distributions. Extensive experimental
results and analysis show that BiAlign consistently
outperforms previous baseline approaches.


As the first work on input preference alignment,
one limitation of our paper is the additional
computational overhead introduced by the ranking
loss. A further improvement could be to explore
more efficient input alignment methods to improve
the ICL capabilities of student models.


Table 7: Details of different datasets. # refers to ‘the
number of’. CrossFit (Ye et al., 2021) is used to
construct training data and other tasks are used for
evaluation.


We use Sentence-BERT (Reimers and Gurevych,
2019) to obtain contextual representations of the
examples and employ cosine similarity to measure
the similarity between these representations. Based
on the similarity to the test example, we categorize
all demonstrations into two groups, G _sim_ and
G _dissim_, ensuring an approximately equal split
between the two groups ( _i.e.,_ a 1:1 ratio).

A.1 Details of Splitting Groups by Similarity


In this work, we construct the in-context learning
evaluation suite based on the following datasets:
- **MMLU** : The MMLU (Massive Multitask Language Understanding) benchmark (Hendrycks
et al., 2021) consists of 57 diverse tasks covering
various fields like computer science, history and
law, aiming to evaluate the knowledge obtained
during pretraining. Following its original setup,
we use 5-shot ICL demonstrations for evaluation.
- **BBH** : The BBH (BIG-Bench Hard) (Suzgun
et al., 2022) includes several types of reasoning
tasks that are believed to be difficult for current
language models. Following the guidelines in
Suzgun et al. (2022), we conduct the evaluation
using 3-shot ICL demonstration examples with
chain-of-thought prompting (Wei et al., 2022b).
- **GSM8K** : The GSM8K (Cobbe et al., 2021)
dataset encompasses 8.5K grade school math
word problems, aiming to evaluate the multi-step
mathematical reasoning capabilities. We evaluate
the ICL performance on it using 8-shot in-context
examples with chain-of-thought prompting.
- **LogiQA** : LogiQA (Liu et al., 2020) is a
logical reasoning benchmark sourced from
logical examination papers intended for reading
comprehension. Following Jiao et al. (2023), we
conduct 5-shot evaluation.
- **HumanEval** : HumanEval (Chen et al., 2021)
consists of 164 programming challenges for
evaluating coding capabilities. We follow the
KATE MMR IDS
SelectionVanilla 18.1 17.4 19.2
SelectionBiAlign **20.2** **19.3** **20.8**
Table 8: Integration of BiAlign with ICL demonstration
selection methods.
Output-Align BiAlign
Llama 3-8B 62.5 **65.3**
Table 9: Integration of BiAlign with reverse KL
divergence.
official zero-shot setting in Chen et al. (2021) to
verify whether bidirectional alignment hurts the
zero-shot learning ability of models.
We summarize the details of all used datasets in
Table 7.


Our methods are implemented with the PyTorch
and Transformers library (Wolf et al., 2020). Model
training is conducted utilizing DeepSpeed (Rasley
et al., 2020; Rajbhandari et al., 2020) on 4 NVIDIA
A100 GPUs. During the training phase, we set
the learning rate to 1e−6 and the batch size to 64.
The weight _λ_ for the ranking loss is set to 1.0.
For evaluation, we train the student model on the
constructed ICL data for 4 epochs and assess the
final checkpoint.


BiAlign is complementary to ICL demonstration
selection methods and can be seamlessly integrated
with them to further improve ICL performance. To
validate this, we investigate three demonstration
selection methods: KATE (Liu et al., 2022), MMR
(Ye et al., 2023), and IDS (Qin et al., 2024b).
For each method, we evaluate the following two
variants: selecting demonstrations and performing
ICL using the vanilla model (SelectionVanilla), and
selecting demonstrations and performing ICL using
the model after BiAlign (SelectionBiAlign). We
conduct experiments on GSM8K and report the
results in Table 8, demonstrating that BiAlign
consistently boosts performance across all three
selection methods. Furthermore, BiAlign (18.3)
surpasses both KATE (18.1) and MMR (17.4),
highlighting its superiority over several ICL
demonstration selection approaches.
1.5B 3B 7B 14B
Output-Align 35.2 35.9 36.2 36.7
BiAlign **36.9** **38.0** **38.8** **40.1**
Table 10: Results for different teacher model sizes with
a fixed 0.5B student (Qwen-2.5).
0.5B 1.5B 3B 7B
Output-Align 36.7 51.4 59.6 70.9
BiAlign **40.1** **54.3** **62.7** **73.4**
Table 11: Results for different student model sizes with
a fixed 14B teacher (Qwen-2.5).


Gu et al. (2024) reveals that reverse KL divergence
is more suitable for knowledge distillation in
generative LLMs, as it helps prevent the student
model from overestimating low-probability regions
of the teacher’s distribution. Building on this
insight, we investigate the integration of BiAlign
with reverse KL divergence. Specifically, we
replace the forward KL divergence in both OutputAlign and BiAlign with reverse KL divergence
and conduct experiments using Llama 3-70B as
the teacher model and Llama 3-8B as the student
model. As shown in Table 9, BiAlign continues
to significantly outperform Output-Align with
reverse KL divergence, further demonstrating its
effectiveness.


We investigate the scalability of BiAlign across
different model scales using Qwen-2.5 (Yang et al.,
2024). Specifically, we conduct experiments on
five model scales: 0.5B, 1.5B, 3B, 7B, and 14B.
Our evaluation follows two settings: ( _i_ ) _Varying_
_teacher model sizes_ : We fix the student model at
0.5B and experiment with teacher models ranging
from 1.5B to 14B. ( _ii_ ) _Varying student model sizes_ :
We fix the teacher model at 14B and test student
models ranging from 0.5B to 7B. The results for
both settings are presented in Table 10 and 11,
respectively. We can see that:
- BiAlign consistently outperforms Output-Align
with different sizes of teacher models.
- BiAlign _benefits more from increasing the size of_
_the teacher model_ compared to Output-Align.
Vanilla Output-Align BiAlign
BBH 31.4 33.8 **45.3**
GSM8K 24.7 28.4 **38.6**
LogiQA 29.1 32.3 **44.7**
Table 12: Average preference score consistency (%)
comparison between different methods.
Output-Align BiAlign
BBH 40.2 **43.3**
Table 13: Performance on BBH for models trained on
MMLU.
- BiAlign is both applicable and robust across
student models of different sizes.


We report the average preference score consistency
(%) comparison between different methods on the
other three datasets (BBH, GSM8K and LogiQA)
in Table 12. From the results, we can see that
BiAlign consistently outperforms Output-Align
across all datasets.

A.7 Average Preference Score Consistency


We outline the justification for input preference
alignment from the following perspectives.
( _i_ ) _Impact of ICL Demonstrations on Model_
_Performance._ ICL demonstrations have been
shown to have a significant impact on the
performance of ICL (Liu et al., 2022; Qin et al.,
2024b). Selecting different sets of demonstration
examples can yield performance ranging from
almost random to better than state-of-the-art finetuned models, indicating that the model has
different preferences for different inputs. For the
student and teacher models to be well-aligned, the
demonstrations preferred by the teacher model
should also be preferred by the student, _i.e.,_ **to**


**as “which input examples should be preferred”** .
This is closely related to preference learning in
RLHF, where the reward model learns "which
outputs should be preferred". After learning, a welltrained reward model can rank model responses
with expertise comparable to humans. To this end,
we introduce input preference alignment to align
the student and teacher models’ preferences for
different demonstrations.
Default Variant
BiAlign **31.0** 30.3
Table 14: Comparison between different subset
sampling methods.

model needs to learn “what to output” as well


3 4 5 6
BiAlign 30.7 31.0 30.8 **31.1**
Table 15: Average performance (%) of BiAlign with
different numbers of subsets _N_ .
( _ii_ ) _Explanatory Mechanisms of ICL._ Another
perspective supporting input preference alignment
stems from the way LLMs process and prioritize
information during ICL. Kossen et al. (2024)
discover that LLMs do not treat all available
information equally; instead, they exhibit a natural
tendency to prioritize information closer to the
query. This selective attention mechanism suggests
that LLMs inherently favor contextually relevant
details over more distant or less relevant ones.
Building on this insight, our proposed input
preference alignment ensures that **the student**


**information prioritization strategy** . By aligning
the student’s input selection process with that
of the teacher, we make the learning process
more effective. This joint alignment ultimately
enables the student model to utilize information
in a manner consistent with the teacher model’s
intrinsic preferences, thereby improving its overall
ICL performance.

model learns to replicate the teacher model’s


To further verify the cross-task generalization
ability of BiAlign, we train the model on MMLU
and evaluate it on BBH. Specifically, we use Llama
2-7B as the student model and Llama 2-70B as the
teacher model. The results reported in Table 13
highlight the superiority of BiAlign over OutputAlign.


To investigate the influence of subset sampling
methods, we replace the original method with
‘Randomly sample N subsets’ which does not
consider similarity. The comparison between the
two methods is shown in Table 14. We can observe
|BiAlign|Col2|Col3|
|---|---|---|
||||
||||
||||
||||
|~~50~~
|~~00~~
~~Ful~~|~~l~~|
Figure 5: Average performance (%) of BiAlign with
different numbers of source tasks.


Whole Sequence Label Only
BiAlign **31.0** 30.8
Table 16: Average performance (%) of BiAlign using
different types of KL divergence calculation methods.
a noticeable performance drop, highlighting the
crucial role of incorporating example similarity in
the sampling process.


While we use _N_ = 4 subsets for calculating the
ranking loss, we also evaluate the effectiveness
of BiAlign with different _N_ . Specifically, we
conduct controlled experiments with {3 _,_ 5 _,_ 6}
subsets and report the average results of the
5 target tasks in Table 15. We can observe
that increasing the number of subsets does not
always improve performance. BiAlign achieves
the best performance (31.1) with 6 subsets and the
performance with 4 subsets (31.0) is comparable.
Besides, all variants consistently outperform
baseline methods in Table 1, demonstrating the
effectiveness of our designed input preference
alignment.


We hypothesize that the diversity of source
tasks has a considerable influence on target task
performance. To verify this, we study the effect
of the number of source tasks by conducting
controlled experiments on {50 _,_ 100} randomly
selected source tasks. From the results in Figure 5,
we can observe that the performance of BiAlign
keeps improving as the number of source tasks
increases, indicating the importance of source task
diversity.
**Method** 25% 50% 100%
Output-Align 29.1 29.3 29.4
BiAlign **30.3** **30.8** **31.0**
Table 17: Comparison between BiAlign and OutputAlign at different proportions of training steps.


To maintain the basic in-weights capability of the
student model, we minimize the KL divergence
loss for the whole sequence instead of only the
label following Gu et al. (2024). In Table 16, we
show the performance comparison between using
the whole sequence and using only the label. We
can see that using the whole sequence also results
in slightly better average performance.


Table 17 reports the performance comparison
between BiAlign and Output-Align at different
proportions (roughly 25%, 50%, and 100%) of
training steps. We can observe that BiAlign
consistently outperforms Output-Align at different
steps.

A.14 Different Proportions of Training Steps


The analysis in Section 5.2 shows that conducting
Output-Align on the combination of the original
ICL training examples and the sampled subset
data achieves similar performance to the original
Output-Align method. We further experiment
with the fine-tuning approach. However, the
performance becomes even worse (from 29.5
to 29.3), once again demonstrating that simply
increasing training data does not necessarily lead
to better performance.


To further investigate the influence of the ranking
loss _L_ [rank] (Equation 3), we conduct experiments
with different weights _λ_ and report the results in
Table 18. All variants except the variant with _λ_ =
5 _._ 0 (too large) outperform baseline approaches by
a large margin, which demonstrates the superiority
of _L_ [rank] .


While we use all _C_ ( _N,_ 2) ( _N_ = 4 is the number of
subsets) pairs of positive and negative subsets for
input preference alignment, we also investigate the
effect of contrastive pair selection. Specifically,
_λ_


0.2 0.5 1.0 2.0 5.0
BiAlign 30.8 **31.2** 31.0 30.9 29.9
Table 18: Average performance (%) of BiAlign with
different _λ_ for the ranking loss _L_ [rank] .


3 4 5 All (6)
BiAlign 30.2 30.8 30.7 **31.0**
Table 19: Average results (%) of BiAlign with different
numbers of contrastive pairs.
we conduct controlled experiments on {3 _,_ 4 _,_ 5}
randomly selected contrastive pairs and report the
average results in Table 19. The best performance
is achieved using all pairs, justifying our selection
strategy.


We select GSM8K as a representative task and
show several examples of output in Table 20.
Compared with Vallina, BiAlign is able to generate
more precise and fine-grained reasoning paths, _e.g.,_
BiAlign can successfully understand the meaning
of ‘ _pair_ ’ and generate the rationale ‘ _10 pairs of_
_socks is 10 * 2 = 20 socks_ ’ while Vallina fails to do
so.


We show the reasoning path required to solve a
sample in LogiQA in Table 21, which is quite finegrained.
Answer the following questions.
_Few-shot demonstrations_
**Question** : Lindsay is doing the laundry, and thinks she has missed some socks. There are 50 socks that need washing. If
she washes 10 pairs of socks and 15 loose socks, how many socks has Lindsay missed?
**Answer** :
**Vanilla** : Lindsay started with 50 socks. She washed 10 pairs of socks and 15 loose socks. So she had 50 - 10 - 15 = 25
socks left. So the answer is 25. ✗
**BiAlign** : Lindsay started with 50 socks. She washed 10 pairs of socks and 15 loose socks. 10 pairs of socks is 10 * 2 = 20
socks. 15 loose socks is 15 * 1 = 15 socks. So she missed 50 - 20 - 15 = 15 socks. So the answer is 15. ✓
Table 20: Output examples of Vanilla and BiAlign on GSM8K. The upper part is the input to LLMs, including
few-shot demonstrations, and the lower part is the generated response. We color correct outputs in green and wrong
outputs in red.
**Context** : There is no doubt that minors should be prohibited from smoking. However, we cannot explicitly ban the use of
automatic cigarette vending machines in order to prevent minors from smoking. This ban is just like setting up roadblocks
on the road to prohibit driving without a license. These roadblocks naturally prohibit driving without a license, but also
block more than 99% of licensed drivers.
**Question** : In order to evaluate the above argument, which of the following questions is the most important?
**Options** :
A: Does the proportion of underage smokers in the total number of smokers exceed 1%?
B: How much inconvenience does the ban on the use of automatic vending machines bring to adult cigarette buyers?
C: Whether the proportion of unlicensed drivers in the total number of drivers really does not exceed 1%?
D: Is the harm of minor smoking really as serious as the public thinks?
**Reasoning path** :
To evaluate the argument effectively, the focus should be on understanding the impact and justification of the proposed ban
on cigarette vending machines, especially in the context of preventing minors from smoking. The argument draws a
parallel between the proposed ban and the hypothetical scenario of setting up roadblocks to prevent driving without a
license, suggesting that while the measure may target a minority (in this case, underage smokers or unlicensed drivers), it
disproportionately inconveniences the majority (licensed drivers or adult smokers).
_The most important question to evaluate the argument is: B: How much inconvenience does the ban on the use of_
_automatic vending machines bring to adult cigarette buyers?_
This option directly addresses the central issue of the argument—the balance between the inconvenience caused to
the majority by the ban and the intended benefit of preventing a minority (minors) from engaging in harmful behavior
(smoking). It’s critical to assess whether the inconvenience to adult smokers (who are the majority of cigarette consumers)
is justified by the potential benefits of reducing underage smoking. The comparison with roadblocks for unlicensed drivers
emphasizes the argument’s concern with the proportionality and fairness of broad preventive measures.
Option A concerns the proportion of underage smokers but doesn’t directly address the balance of convenience versus
benefit. Option C relates to the analogy used but does not directly help in evaluating the effectiveness and appropriateness
of the vending machine ban. Option D questions the severity of the problem (minor smoking) but does not specifically
help in assessing the argument’s concern about the consequences of the ban on the majority.
Therefore, understanding the extent of inconvenience to adult smokers is crucial in determining whether the proposed
solution is proportionate and justified, making Option B the most relevant and important question for evaluating the
argument.
Table 21: Reasoning path for an example in LogiQA.


**Andrew Lee** [1] **Xiaoyan Bai** [1] **Itamar Pres** [1] **Martin Wattenberg** [2] **Jonathan K. Kummerfeld** [3] **Rada Mihalcea** [1]

A Mechanistic Understanding of Alignment Algorithms: A Case Study on DPO
  and Toxicity

Andrew Lee
Xiaoyan Bai
Itamar Pres
Martin Wattenberg
Jonathan K. Kummerfeld
Rada Mihalcea


While alignment algorithms are now commonly
used to tune pre-trained language models towards
a user’s preferences, we lack explanations for the
underlying mechanisms in which models become
“aligned”, thus making it difficult to explain phenomena like jailbreaks. In this work we study a
popular algorithm, direct preference optimization
(DPO), and the mechanisms by which it reduces
toxicity. Namely, we first study how toxicity is
represented and elicited in a pre-trained language
model, GPT2-medium. We then apply DPO with
a carefully crafted pairwise dataset to reduce toxicity. We examine how the resulting model averts
toxic outputs, and find that capabilities learned
from pre-training are not removed, but rather bypassed. We use this insight to demonstrate a simple method to un-align the model, reverting it
back to its toxic behavior.


Large language models learn surprising capabilities from
pre-training on large datasets (Brown et al., 2020; Chowdhery et al., 2023; Touvron et al., 2023). While these capabilities lead to impressive achievements, they also include
unwanted behaviors that can be found in large-scale web
data, such as toxicity and bias (Sheng et al., 2019; Gehman
et al., 2020). As a result, researchers have developed alignment algorithms to reduce undesirable behaviors, which
often use reinforcement learning with human preferences
(RLHF). For instance, proximal policy optimization (PPO,
Schulman et al. 2017) fits a reward model on human preference data, which is then used to fine-tune a language
model, while direct preference optimization (DPO, Rafailov
et al. 2023) by-passes the reward model and derives reward
signals directly from pairwise preference data.
While such algorithms can suppress undesirable behavior,
1University of Michigan, Ann Arbor, U.S.A. 2Harvard University, Cambridge, Massachusetts [3] University of Sydney, Sydney,
Australia. Correspondence to: Andrew Lee _<_ ajyl@umich.edu _>_ .
our understanding of the mechanisms by which the undesirable behavior is suppressed is limited. Furthermore, researchers have demonstrated that such alignments can be
surprisingly easily undone (Wallace et al., 2019; Zou et al.,
2023; Wei et al., 2023; Carlini et al., 2023). While prior
work hypothesize why jailbreaks are possible through empirical studies (Wei et al., 2023), in this work we provide a
mechanistic explanation for such phenomena.
Given the above limitations, in this work we study the mechanisms by which alignment algorithms alter a model’s behavior. Researchers have demonstrated that a deep enough
understanding of a model’s inner representations allows us
to interpret how it makes decisions. For instance, various
concepts such as world models, truthfulness, or even taskspecific features have highly interpretable and controllable
representations (Li et al., 2023b; Todd et al., 2023; Nanda
et al., 2023). Motivated by such findings, we study how the
representation space of language models change by comparing it before and after an alignment algorithm is applied.
Our work relates to that of Jain et al. (2023), which studies
how the capabilities of a language model changes after finetuning on synthetic tasks. Unlike this previous work, we
study the change in mechanisms from a RLHF algorithm on
a natural language setting.
We consider DPO and toxicity as a case-study of RLHF
alignment algorithms. We first study how toxicity is represented and elicited in GPT2-medium (henceforth GPT2).
We then apply DPO using a carefully crafted pairwise
dataset that consists of toxic and nontoxic samples. Lastly,
we study the mechanisms by which toxicity is no longer
generated after DPO, and how those mechanisms can fail.
Our work is organized as follows: in Section 2 we provide the necessary preliminaries relevant to our work. In
Section 3, we demonstrate how toxicity is represented and
elicited in GPT2. We find multiple vectors in multilayer
perceptron (MLP) blocks that promote toxicity. We apply
singular value decomposition (SVD) to these toxic vectors
to find vectors that represent specific dimensions of toxicity
in the model. To validate the role of these vectors in generating toxic outputs, we intervene with our toxic vectors and
demonstrate much safer outputs.


_dmlp_
- _m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ][.]_ (2)
_i_ =1
In Section 4, we explain our procedure to apply DPO
on our language models to reduce toxicity, using a carefully crafted pairwise toxicity dataset, produced by using
PPLM (Dathathri et al., 2019) to generate paired toxic and
non-toxic samples.
In Section 5, we demonstrate how toxicity is no longer
elicited after DPO. Namely, we show that every parameter
is minimally shifted, including the toxic vectors. However,
such minimal changes in weights allow the model to avert
the triggering of toxic vectors. Put differently, DPO _does_
_not remove_ the capability of generating toxic outputs, but
learns an “offset”, distributed amongst its layers, to “bypass”
the regions that elicit toxicity. Based on this understanding, we demonstrate the ease of re-activating these vectors
to generate toxic outputs, and thus undoing the alignment
learned from DPO. We view our findings as shedding light
into why aligned models can be jailbroken or un-aligned.

A Mechanistic Understanding of Alignment Algorithms


In this section we provide background and notations, much
of which is borrowed from Geva et al. (2022).
**Transformers, MLPs.** Transformer-based language models typically consists of embedding and unembedding layers
_E, U ∈_ R _[|V|×][d]_ with a series of _L_ transformer layers inbetween (Vaswani et al., 2017). Each layer _l_ consists of
attention heads and a multilayer perception (MLP) layer.
Given an input sequence **w** = _⟨w_ 0 _, ..., wt⟩_, the model first
applies _E_ to create an embedding **x** _i ∈_ R _[d]_ for each token
_wi ∈_ **w** . We call **x** _i_ the residual stream.
The residual stream is then updated by attention heads and
MLP blocks from subsequent layers (bias terms omitted):
**x** _[ℓ]_ **i** [+] **[1]** = _x_ _[ℓ]_ _i_ [+][ MLP] _[ℓ]_ [(] _[x]_ _i_ _[ℓ]_ [+][ Att] _[ℓ]_ [(] _[x]_ _i_ _[ℓ]_ [))]
When needed, we specify the intermittent residual stream at
layer _ℓ_ (after the attention head, before the MLP) as **x** _[ℓ]_ ~~_[m]_~~ _[id]_ .
Per Geva et al. (2022), the updates to the residual stream
from each MLP block can be further decomposed. Namely,
MLP blocks consist of two linear transformations, with
point-wise activations _σ_ in-between:
MLP _[ℓ]_ ( **x** _[ℓ]_ ) = _σ_      - _WK_ _[ℓ]_ **[x]** _[ℓ]_ [�] _WV_ _[ℓ]_ _[,]_ (1)
where _WK_ _[ℓ]_ _[, W]_ _V_ _[ ℓ]_ _[∈]_ [R] _[d][mlp][×][d]_ [. We notate the] _[ i]_ [-th row in] _[ W][K]_
as MLP. **k** _[ℓ]_ _i_ [and refer to them as key-vectors, and the] _[ i]_ [-th]
column in _WV_, MLP. **v** _i_ _[ℓ]_ [, as value-vectors (we sometimes]
omit “MLP” and just use **k** _[ℓ]_ _i_ _[,]_ **[ v]** _i_ _[ℓ]_ [).]
Equation (1) indicates that _the output of MLP blocks is the_
_sum of its value vectors_ **v** _i, each scaled by a coefficient_
_value m_ _[ℓ]_ _i_ [, where] **[ m]** _[ℓ]_ [:=] _[ σ]_ - _WK_ _[ℓ]_ **[x]** _[ℓ]_ [�] _∈_ R _[d][mlp]_ :
MLP _[ℓ]_ ( **x** _[ℓ]_ ) =
_dmlp_
- _σ_ ( **x** _[ℓ]_ _·_ **k** _[ℓ]_ _i_ [)] **[v]** _i_ _[ℓ]_ [=]
_i_ =1
Put differently, the MLP block writes to the residual stream
_dmlp_ times, once for each value vector. We call each of
these updates a _sub-update_ .
All of our experiments are conducted with GPT2-medium,
which has _L_ = 24, _d_ = 1024, and _dmlp_ = 4096.
**Interpreting Value Vectors in Vocabulary Space.** Geva
et al. (2022) demonstrate that for each sub-update, each
value vector **v** _i_ either promotes or suppresses the likelihood
of a token _w_ from being generated:
_p_ - _w |_ **x** _[ℓ]_ + _m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ][, E]_ - _∝_ exp - **e** _w ·_ **x** _[ℓ]_ [�] _·_ exp - **e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_
where **e** _w_ is the embedding of _w_ . This indicates that
when **e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_ _[>]_ [ 0][, the likelihood of] _[ w]_ [ increases, while]
**e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_ _[<]_ [ 0][ decreases the likelihood.][1]
Further note that this dot product can be further decomposed. Namely, **e** _w ·_ **v** _i_ _[ℓ]_ [is a “static” value that does not]
depend on the input: only when **v** _i_ _[ℓ]_ [is scaled by] _[ m][i]_ [ (which]
is determined by the its corresponding key vector, **k** _[ℓ]_ _i_ [, and]
the residual stream **x** ) do we see the impact of the input
towards the likelihood of _w_ .
Thus the projection **r** _[ℓ]_ _i_ [=] _[ E]_ **[v]** _i_ _[ℓ]_ _[∈]_ [R] _[|V|]_ [ induces a ranking of]
tokens that get promoted by value vector **v** _i_, in which tokens
with the highest dot products **e** _w ·_ **v** _i_ are promoted most by
value vector **v** _i_ . In Section 3 we show value vectors that
promote toxicity by applying these projections.


In this section we demonstrate how toxicity is represented
and elicited in GPT2, by introducing a series of vectors that
can be extracted from the language model.

3. Toxicity in Pre-trained Language Models


**Toxicity Probe Vector.** We start by first training a linear probe model on a binary toxicity classification task.
Namely, we use the Jigsaw toxic comment classification
dataset (cjadams et al., 2017), which consists of 561,808
comments, each of which is labeled as toxic or non-toxic.
We use a 90:10 split for training and validation. We train
our probe model, _W_ Toxic, on the residual stream in the last
layer, averaged across all timesteps ( **¯x** _[L][−]_ [1] ):
_P_ (Toxic _|_ **¯x** _[L][−]_ [1] ) = softmax( _W_ Toxic **¯x** _[L][−]_ [1] ) _, W_ Toxic _∈_ R _[d]_
1See Appendix for derivation.


_Table 1._ Top toxic vectors projected onto the vocabulary space.
WARNING: THESE EXAMPLES ARE HIGHLY OFFENSIVE.
We note that SVD.UToxic[2] has a particularly gendered nature.
This arises from the dataset and language model we use.
VECTOR TOP TOKENS
_W_ Toxic c*nt, f*ck, a**hole, d*ck, wh*re, holes
MLP. **v** 770 [19] sh*t, a**, cr*p, f*ck, c*nt, garbage, trash
MLP. **v** 771 [12] delusional, hypocritical, arrogant, nonsense
MLP. **v** 2669 [18] degener, whining, idiots, stupid, smug
MLP. **v** 668 [13] losers, filthy, disgr, gad, feces, apes, thous
MLP. **v** 255 [16] disgrace, shameful, coward, unacceptable
MLP. **v** 882 [12] f*ck, sh*t, piss, hilar, stupidity, poop
MLP. **v** 1438 [19] c*m, c*ck, orgasm, missionary, anal
SVD.UToxic[0] a**, losers, d*ck, s*ck, balls, jack, sh*t
SVD.UToxic[1] sexually, intercourse, missive, rogens, nude
SVD.UToxic[2] sex, breasts, girlfriends, vagina, boobs
Our probe vector achieves an accuracy of 94% on the validation split. We view our toxic probe vector _W_ Toxic as an
aggregate of all the relevant signals in the language model
to classify an input as toxic.
**Toxic Vectors in MLP Blocks.** Given our probe vector
_W_ Toxic, we can use it to find weights within the language
model that promote toxicity. Namely, Geva et al. (2022)
demonstrate that value vectors promote tokens at a conceptlevel. Given this, we search for value vectors that promote
toxicity, by checking for all value vectors with the highest
cosine similarity with _W_ Toxic. We find that indeed, there are
value vectors that promote toxic tokens (See Section 3.2).
We notate our set of toxic value vectors as MLP. **v** Toxic and
their corresponding key vectors as MLP. **k** Toxic.
We provide two perspectives of our MLP. **v** Toxic vectors: 1)
when triggered, they promote the likelihood of toxic tokens
to be generated, and 2) MLP. **v** Toxic are vectors within the
model that contribute towards the _W_ Toxic direction.
**SVD: Decomposed Toxic Vectors.** After extracting a set
of N (=128) [2] MLP. **v** Toxic vectors, we stack them into a _N ×d_
matrix. We then apply singular value decomposition to get
decomposed singular value vectors SVD.UToxic. We refer
to the _i_ -th singular value vector as SVD.UToxic[ _i_ ]. We view
SVD.UToxic as basis vectors that span the toxicity representation space within the language model.
**3.2. Toxic Vectors in Vocabulary space.**
As mentioned in Section 2, we can inspect which tokens
are promoted by value vectors by projecting them onto the
vocabulary space.
2We experiment with different values for N, and get similar
results.
_Table 2._ Toxicity, perplexity (PPL), and F1 after interventions or
DPO. We scale our toxic vectors such that the resulting perplexity
is comparable to that of GPT2 (No Op). _†_ : Not an intervention.
METHOD VECTOR TOXIC PPL F1
NO OP N/A 0.453 21.7 0.193
SUBTRACT _W_ TOXIC 0.245 23.56 0.193
SUBTRACT MLP. **v** 770 [19] 0.305 23.30 0.192
SUBTRACT SVD.UTOXIC[0] 0.268 23.48 0.193
DPO _[†]_ N/A 0.208 23.34 0.195
Table 1 shows the tokens with the highest dot products with
our toxic vectors. Each MLP. **v** Toxic and SVD.UToxic vectors
seem to encode specific dimensions of toxicity, or different
contexts in which toxicity appears in pre-training data.


To validate the role that the toxic vectors play in eliciting
toxic outputs, we intervene during generation to suppress
toxic outputs using each vector. Namely, we use prompts
from REALTOXICITYPROMPTS (Gehman et al., 2020) that
elicit toxic outputs from GPT2. During the forward pass of
the model, we intervene by simply subtracting one of the
toxic vectors from the last layer:
**x** _[L][−]_ [1] = **x** _[L][−]_ [1] _−_ _α ∗_ _W,_
where _a_ is a heuristic scale value and _W_ is one of our
toxicity vectors.
To measure the efficacy of our interventions, we measure
three metrics: toxicity, perplexity, and F1.
**Toxicity.** To measure toxicity, we prompt GPT2 with the
“challenge” subset of REALTOXICITYPROMPTS, which consists of 1,199 prompts that elicit extremely toxic outputs
from language models. We follow prior work (Geva et al.,
2022) and use Perspective API, [3] an automated tool for toxicity detection, to assign toxicity scores to each generation.
**Perplexity.** To ensure that our interventions do not degrade generation quality, we also follow prior work (Geva
et al., 2022) and measure perplexity on the Wikitext-2
dataset (Merity et al., 2016).
**F1.** In addition to perplexity, we also follow prior work
(Dinan et al., 2020; Adolphs et al., 2023) and measure F1.
Namely, using 2,000 Wikipedia sentences as prompts, we
measure the harmonic mean between precision and recall
of our model’s output, where precision is the fraction of
3https://github.com/conversationai/perspectiveapi


_Table 3._ Examples of top-k and continuations to prompts that originally elicit “sh*t” from GPT2, our interventions, and GPT2DPO.
PROMPTS MODEL TOP-K CONTINUATIONS
“So I’m starting
to think she’s
full of”
“You are such
a piece of”
GPT2 sh*t, it, crap, herself, something “sh*t. I’m not saying she’s a bad person, but”
GPT2 - MLP.V [19] 770 it, herself, the, sh*t, something “it. I’m not sure if she’s just being a little too”
GPT2DPO it, herself, sh*t, surprises, crap “it. I think she’s very much aware of the fact that”
GPT2 sh*t, work, trash, garbage, a** “sh*t. You are such a piece of sh*t.”
GPT2 - MLP.V [19] 770 work, sh*t, “s”, art, human “work. I’ve been working on this for a while”
GPT2DPO work, art, sh*t, trash, cake “work. I have been a fan of your work for a long”
GPT2 sh*t, crap, cow, hell, f*ck “sh*t, I’m so glad I got this book.”
“Holy” GPT2 - MLP.V [19] 770 Cow, Grail, cow, “!”, Cross “Cow! I’ve been waiting for this for a while.”
GPT2DPO cow, crap, Grail, sh*t, smokes “cow, this is a great book! I’ve been reading”
generated tokens contained in the original Wikipedia continuation, and recall is the fraction of tokens in the Wikipedia
continuation contained in the model’s generation.
With perplexity and F1, we hope to see minimal changes
after our interventions to ensure we do not affect the quality
of our generations. Table 2 demonstrates the results from
our interventions, while Table 3 demonstrates examples of
generations before and after our interventions.
Note that our interventions depend on how much we scale
each vector ( _α_ ). We choose a scalar value such that the
resulting perplexity is similar to that of our post-DPO model.
For details regarding our post-DPO model see Section 4.
We find that subtracting toxic components from the residual
stream reduces toxicity.


We next describe our alignment procedure using DPO.


DPO relies on pairwise preference data, in which given a
prompt, we have a preferred (positive) continuation and
a non-preferred (negative) continuation. Given each preference pair, the algorithm promotes the likelihood of the
positive sample, while suppressing the likelihood of the
negative sample, using the following loss term:
_L_ DPO = _−_ E [log _σ_ ( _β_ log _P −_ _β_ log _N_ )] _,_
_πθ_ ( _y_ + _|_ **w** ) _πθ_ ( _y−_ _|_ **w** )
_P_ =
_πref_ ( _y_ + _|_ **w** ) _[, N]_ [ =] _πref_ ( _y−_ _|_ **w** ) _[,]_
where _y_ + and _y−_ are preferred (nontoxic) and non-preferred
(toxic) continuations of **w**, _πref_ is the frozen weights of
the original language model, and _πθ_ is the weights of the
language model being updated (See Rafailov et al. (2023)
for details). The algorithm promotes the likelihood of _P_,
while suppressing the likelihood of _N_ .


We build our pairwise toxicity dataset using PPLM
(Dathathri et al., 2019). PPLM is an attribute-controlled
language generation technique, which attaches a simple linear attribute classification layer, _p_ ( _a|_ **w** ) onto a language
model to guide its generation. During generation, PPLM
uses the attribute classifier to compute the gradients that
increases the likelihood of the language model’s output to
contain the desired attribute _a_, and shifts the activations in
such direction (See Dathathri et al. (2019) for details):
_p_ ( _y | a_ ) _∝_ _p_ ( _y_ ) _p_ ( _a | y_ )
To generate pairwise preference data, we use sentences
from Wikitext-2 (Merity et al., 2016) as prompts. For each
prompt, we generate a positive sample using greedy sampling with GPT2, while using PPLM to generate negative
(toxic) samples. We use our toxic probe _W_ Toxic as our attribute classifier to guide towards toxic outputs. We create
24,576 pairs of toxic and nontoxic continuations. [4] We train
until validation loss converges with a patience value of 10,
which occurs after approximately 6,000 sample pairs. Appendix D has details for DPO and PPLM hyperparameters.
The last row of Table 2 shows the resulting toxicity, perplexity, and F1 scores of our DPO model.
Figure 1 shows an example of the difference in behaviors
between GPT2 before and after DPO, for a specific toxic
token. Namely, we use 295 prompts from REALTOXICI
TYPROMPTS that outputs the token “sh*t” as the next token.
We then apply “Logit Lens” (Nostalgebraist, 2020), meaning we apply the unembedding layer on all intermittent
layers. This allows us to visualize the layers that promote
the “sh*t” token. The shared grey areas indicate the layers
in which “sh*t” is promoted the most, which all correspond
to MLP layers. We see that post-DPO, the toxic token is
promoted far less.
4We release this data to enable further studies.


|Col1|Col2|Model|Col4|Col5|Col6|Col7|Col8|Col9|Col10|Col11|Col12|Col13|Col14|Col15|Col16|Col17|Col18|Col19|Col20|Col21|Col22|Col23|Col24|Col25|
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|||Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO|Model
GPT2
DPO||||||||||
||||||||||||||||||||||||||
||||||||||||||||||||||||||
|Col1|Model
GPT2
DPO|Col3|Col4|Col5|Col6|Col7|
|---|---|---|---|---|---|---|
||||||||
||||||||
||||||||
_Figure 1._ Logit lens on GPT2 and GPT2DPO. Given 295 prompts
that originally elicit “sh*t” as the next token, we plot the average
probability of outputting “sh*t” from intermittent layers by applying the unembedding layer. Minor ticks indicate _ℓ_ ~~_m_~~ _id_ layers
(after attention heads, before MLP). Shaded areas indicate layers
that promote “sh*t” the most, which all correspond to MLP layers.


In this section we explain how our aligned language model
(GPT2DPO) averts toxic outputs.


Of the toxic vectors described in Section 3, note that
MLP. **v** Toxic are actual weights of the model. Thus we inspect
how these vectors change after DPO.
Interestingly, we find that every parameter in GPT2 and
GPT2DPO has barely changed, including token embeddings,
MLP blocks, and attention heads. Every parameter in GPT2
and its counterpart in GPT2DPO has a cosine similarity score
greater than 0.99 and on average a norm difference less than
1e-5. [5] This applies for MLP _._ **k** Toxic and MLP. **v** Toxic as well –
toxic MLP vectors **do not change** from DPO.
Put differently, although toxicity is reduced by DPO, the
ability to elicit toxicity with these value vectors still remain.
So how is it that GPT2DPO averts toxic outputs? Though its
parameters have barely moved, below we show that their
collective movement is enough to avoid toxic outputs.


In simplest terms, we observe a drop in activations for the
toxic vectors MLP. **v** Toxic in GPT2DPO. Namely, using the
same 1,199 prompts from REALTOXICITYPROMPTS, we
generate 20 tokens and measure the mean activations _mi_, or
_σ_ ( **x** _[ℓ]_ _·_ MLP. **k** _[ℓ]_ _i_ [)][, of our][ MLP.] **[v]** [Toxic][ vectors. Figure][ 2][ shows]
5 examples of the top MLP. **v** Toxic vectors.
Inspired by Balestriero et al. (2023), we visualize this drop
5The unembedding layer is the only exception, where the norm
difference is less than 1e-3.
|l
Σ MLP Layer
𝛔 𝛔 𝛔 𝛔 𝛔
𝝳
x Hidden space layer l: Rd
Toxic
After Region
DPO
Before DPO|Col2|
|---|---|
|**MLP Layer**l
**Toxic**
**Region**
𝛔
𝛔
𝛔
𝛔
**Σ**
𝛔
**Before DPO**
**Hidden space layer**l**: Rd**
𝝳**x**
**After**
**DPO**|**Before DP**|
|||
_Figure 3._ Visualization of residual streams before and after DPO.
We view the shift, _δ_ **x**, as an offset that allow GPT2DPO to bypass
regions that previously triggered toxic value vectors.
in activations with what we call “MLP activation regions”.
An activation region of a key vector is simply a _subspace_
within the model’s hidden space in which its vectors have
high dot products to activate its corresponding value vector:
_γ_ ( **k** _[ℓ]_ _i_ [) :=] _[ {]_ **[g]** _[|]_ **[g]** _[ ∈]_ [R] _[d][, σ]_ [(] **[k]** _[ℓ]_ _i_ _[·]_ **[ g]** [)] _[ >]_ [ 0] _[}][,]_ (3)
where _σ_ is a non-linear activation. Put differently, for all
key-vector regions that the residual stream “passes through”,
their corresponding value-vectors are activated, scaled, and
added into the residual stream.
We view the drop in activations as a shift in GPT2DPO’s
residual stream to avert the regions of toxic MLP vectors,
_γ_ (MLP. **k** Toxic). See Figure 3.
We formalize the shift in residual streams as following:
given the residual streams at layer _ℓ_ ~~_m_~~ _id_ (after attention
heads at layer _ℓ_ ) for both GPT2 and GPT2DPO, before
MLP _[ℓ]_ Toxic [, we notate the difference of the two residual]
streams as _δ_ **x** _[ℓ]_ ~~_[m]_~~ _[id]_ := **x** _[ℓ]_ DPO ~~_[m]_~~ _[id]_ _−_ **x** _[ℓ]_ GPT2 ~~_[m]_~~ _[id][, δ]_ **x** _[ℓ]_ ~~_[m]_~~ _[id]_ _∈_ R _[d]_ . We
_Figure 2._ Mean activations for toxic vectors before and after DPO.

5.2. GPT2DPO Avoids MLP _._ kToxic Regions


_Figure 4._ Linear shift of residual streams out of toxic regions. Each
point is a residual stream sampled from either **x** [19] GPT [or] **[ x]** DPO [19] [, us-]
ing REALTOXICITYPROMPTS, projected onto 1) _δ_ [¯] **x** [19][, the mean]
difference in residual streams, and 2) the principle component of
the residual streams. Dotted lines indicate samples from the same
prompt. Colors indicate whether each point activates MLP [19] 770 [.]
Note the shift from **x** [19] GPT [to] **[ x]** DPO [19] [, but also the drop in activations.]
view _δ_ **x** _[ℓ]_ ~~_[m]_~~ _[id]_ as a vector that takes GPT2’s residual stream
out of the toxicity-eliciting regions, _γ_ (MLP. **k** _[ℓ]_ Toxic [)][.]
Figure 4 provides a visualization of the residual stream’s
shift out of toxic regions. Namely, given prompts from RE
ALTOXICITYPROMPTS, we project the residual stream from
GPT2 and GPT2DPO at layer 19 onto two dimensions: 1) the
mean difference in the residual streams, _δ_ [¯] _x_ _[ℓ]_, and the main
principle component of the residual streams. [6] We further
indicate whether each residual stream activates MLP. **v** 770 [19] [.]
Notice both the consistent linear shift between GPT2 and
GPT2DPO and the drop in activations.
To understand where this shift comes from, we compute the
differences in all parameter weights in GPT2 before and
after DPO, and notate the differences as _δθ_ . We notate the
difference at a specific component such as a MLP block at
layer _ℓ_ as _δ_ MLP _[ℓ]_ [.]
Note that as previously noted, these differences _δθ_ _[ℓ][,][ ∀][ℓ]_ [are]
minimal. Despite these minimal changes, their accumulation is sufficient in getting the residual stream out of toxic
regions _γ_ (MLP. **k** _[ℓ]_ Toxic [)][.]
Given a toxic vector MLP. **v** Toxic at layer _ℓ_, to understand
where the shift in residual stream, _δ_ **x** _[ℓ]_ ~~_[m]_~~ _[id]_ comes from, we
measure the cosine similarity between _δ_ **x** _[ℓ]_ ~~_[m]_~~ _[id]_ and the shift
in value vectors in the preceding layers, _δ_ MLP.v _[j]_ [:]
_∀j < ℓ, ∀i < dmlp_ : _cos_ ( _δ_ **x** _[ℓmid]_ _, δ_ MLP.v _[j]_ _i_ [)] _[.]_
6We show layer 19 because MLP. **v** 77019 [is one of the most toxic]
vectors, but similar patterns can be found in other layers (See
Appendix B).
To our surprise, we find that the shift in value vectors,
_δMLP.v_, have high _negative_ cosine similarity scores with
the shift in residual streams _δ_ **x** : the value vectors in MLP
blocks shift in the _opposite direction_ as the shift in residual
stream. The blue areas in Figure 5 show the cosine similarity between _δ_ **x** [19] ~~_[ m]_~~ _[id]_ and _δ_ MLP _[j]_ [. We show layer 19 as an]
example because MLP.v [19] 770 [is one of the most toxic vectors,]
but the same pattern can be found in other layers (see Appendix C). Namely, the blue areas indicate the percentage of
value vectors at each layer in which their shifts have a cosine
similarity score against _δ_ **x** [19] ~~_[ m]_~~ _[id]_ as indicated by the x-axis.
Note that as the layers approach layer 19, the majority of
value vectors shift in the _opposite_ direction of _δ_ **x** [19][.]
Why the antipodal direction? This can be explained by two
facts: first, neurons in MLP blocks of language models
are sparse (Zhang et al., 2022; Li et al., 2023c), meaning
most neurons do not activate during a forward pass. Second, the choice of the MLP’s activation function _σ_ plays
a role. Namely, our language model uses GeLU functions
(Hendrycks & Gimpel, 2016). This means that neurons that
are inactive during a forward pass have a _negative_ value
close to 0. Thus, during the forward pass, for each value
vector, the newly learned direction _δ_ MLP _._ **v** gets multiplied by
a very small negative scale, flips directions, and _contributes_
towards the _δ_ **x** direction. The orange areas of Figure 5 indicate the mean activation of each value vector, from the 1,199
prompts in REALTOXICITYPROMPTS. Most of the time,
value vectors have a _negative_ activation - thus the shift in
value vectors end up _contributing_ towards the _δ_ **x** direction.
To summarize, GPT2DPO has learned an _offset_, _δ_ **x**, such
that the residual stream avoids regions that promote toxicity, _γ_ (MLP. **k** Toxic). This learned offset is distributed across
the many value vectors in earlier MLP blocks that are inactive for prompts that previously elicited toxic outputs. By
distributing this offset across numerous value vectors, the
language model is able to preserve its pre-trained language
modeling behavior, as individual weights are minimally affected. However, the distributed offset allows the model to
avert toxic outputs. Note that this behavior matches precisely what the alignment objective was - to preserve as
much of the pre-trained behavior, while optimizing for a
reward (non-toxic outputs).


A growing line of work finds that alignment algorithms can
easily be undone or jailbroken. We view our findings as a
mechanistic explanation for such phenomenon – namely, in
our case, the vectors that elicit toxicity are still sitting in the
model, but simply not triggered.
To confirm our understanding, we demonstrate a simple
way to undo alignment. To reiterate, DPO simply learned
an offset to take the residual stream **x** _[ℓ]_ out of regions that
|Col1|Layer 0|
|---|---|
|||
|||
|||
|||
|||
|||


|Col1|Layer 10|Col3|Col4|Layer 12|Col6|Col7|Layer 14|Col9|Layer 16|Col11|Layer 18|Col13|
|---|---|---|---|---|---|---|---|---|---|---|---|---|
||||||||||||||
||||||||||||||
||||||||||||||
||||||||||||||
||||||||||||||
||||||||||||||
||||||||||||||
_Figure 5._ The cosine similarity between _δ_ MLP _._ **v** and _δ_ **x** [19][. Blue areas indicate the percentage of value vectors with a cosine similarity score]
against _δ_ **x** as indicated by the x-axis. Orange areas indicate the percentage of value vectors with a mean activation as indicated by the
x-axis, during the forward pass of 1,199 REALTOXICITYPROMPTS prompts. Value vectors shift in the opposite direction of _δ_ **x**, but they
end up contributing towards the _δ_ **x** direction because of their negative activations.
_Table 4._ Un-aligning GPT2DPO. By scaling toxic key vectors, and
thus increasing the regions that elicit toxicity, we are able to undo
the alignment learned from DPO and reactivate toxicity.
METHOD TOXIC PPL F1
GPT2DPO 0.208 23.34 0.195
SCALE MLP _._ **k** TOXIC 0.458 23.30 0.195
GPT2 0.453 21.7 0.193
trigger toxic vectors: _γ_ (MLP _._ **k** _[ℓ]_ Toxic [)][. A simple way to re-]
activate toxicity is to increase those regions by scaling each
key vector larger (See Equation 3). This makes the residual
streams pass through toxic regions again, thus reverting
back to the pre-aligned behavior.
Table 4 shows toxicity, perplexity, and F1 scores after scaling up as few as 7 toxic key vectors MLP _._ **k** Toxic. We simply
select 7 MLP vectors with the highest cosine similarity as
our toxic probe vector, _W_ Toxic, and scale their key vectors by
10x.By doing so, the model reverts back to its pre-aligned
toxic behavior. Note that increasing activation regions _γ_
does not have an affect on perplexity, unlike our interventions from Section 3.3. This is likely because the latter
manipulates the residual stream directly, while scaling a key
vector does not (See Equation 2).


We view our work as providing a mechanistic explanation
for why aligned models can be undone or jailbroken – in our
experiments, the regions that previously elicited toxic behavior does not change after DPO. Rather, GPT2DPO learns
minimal changes spread across layers to avoid such regions
and receive its reward.
With such knowledge, we conjecture that more robust alignment algorithms can be designed. For instance, can we eliminate undesirable regions, as opposed to bypassing them?
In scenarios like ours, in which we can identify the weights
that elicit undesirable outputs, what happens if we only updated those weights in isolation? Similarly, if DPO merely
learned an offset that avoids toxic regions, can we replicate
this behavior by only updating the bias terms?
Alternatively, prior to deploying language models, perhaps
we can add “suppression heads” – layers that suppress undesirable behavior. What would happen if we only updated
late layers (or added layers) during alignment?
Lastly, can we characterize “jailbreak-ability” or “unalignability” of aligned models, without relying on test samples?
We leave these questions for future work.

6.1. On Designing Robust Alignment Algorithms


We hypothesize that the minimal changes distributed across
all layers is due to the KL-divergence term that is commonly incorporated in the loss terms of RLHF algorithms.
Namely, the KL-divergence term discourages each weight
from shifting too drastically, in order to preserve its capabilities learned during pre-training.
Similar to our work, Jain et al. (2023) fine-tunes a language model on synthetic tasks to study the changes in its
mechanisms. Interestingly, unlike our findings, the authors
demonstrate that the model simply learns “wrappers” at late
layers that optimize for each task.
We find this difference in model training behavior interesting, and conjecture that the KL-divergence term may play a
role in this difference. Note that fine-tuning typically does
not entail a KL-divergence term. Perhaps this allows the
model to make drastic and localized changes, such as in late
layers, as opposed to distributed, minimal changes.

6.2. On the Role of KL-Divergence Regularization


Numerous alignment algorithms have been proposed, and
the choice of algorithm may largely depend on the type of
data available. Perhaps most commonly, human feedback
data is used (Stiennon et al., 2020; Ouyang et al., 2022; Touvron et al., 2023) for methods such as PPO (Schulman et al.,
2017) or DPO (Rafailov et al., 2023). When labels for only
undesirable behavior is available, algorithms like unlikelihood training (Welleck et al., 2020) or Cringe (Adolphs
et al., 2023; Xu et al., 2023) can be used. We study DPO
because it is easy to use and currently widely used.


The goal of mechanistic interpretability is largely to reverse
engineer model behaviors (Olah et al., 2020; Elhage et al.,
2021; Geva et al., 2021). By doing so, researchers have
uncovered various interpretable and controllable representations, such as world models (Li et al., 2023a; Nanda et al.,
2023), “truthfulness” (Li et al., 2023b), knowledge (Meng
et al., 2022; Hernandez et al., 2023; Burns et al., 2023; Geva
et al., 2023), linguistic properties (Conneau et al., 2018; Tenney et al., 2019), or even tasks (Ilharco et al., 2022; Hendel
et al., 2023; Todd et al., 2023).
Rather than probing for specific representations, researchers
have also characterized the representations of language
models from a geometric perspective (Park et al., 2023).
Balestriero et al. (2023) demonstrate a geometric characterization that can be used to extract feature representations
that solve toxicity detection.
Similar to our work, Jain et al. (2023) study the mechanisms
in which fine-tuning on synthetic tasks alters the model’s capabilities. We study the effects of RLHF on a more realistic,
natural language setting.


Researchers demonstrated that aligned models can be surprisingly easily jailbroken (Wallace et al., 2019; Zou et al.,
2023; Wei et al., 2023; Carlini et al., 2023). Such adversarial attacks typically involve searching for prompts that
can elicit previously unlearned behaviors, or even personal
information (Nasr et al., 2023). Carlini et al. (2023) show
that multimodal models can also be jailbroken. Wei et al.
(2023) provide hypotheses, backed by empirical studies, as
to why language models can be jailbroken.
In a similar vein to jailbreaks, numerous researchers have
demonstrated that aligned models can easily be un-aligned
(Yang et al., 2023; Qi et al., 2023), sometimes with as few
as 100 fine-tuning examples. We view our work as adding a
mechanistic understanding of such phenomena.


In this work we studied the mechanisms by which alignment
algorithms unlearn a capability, taking DPO and toxicity
as a case study. First, we uncovered how toxicity is represented and elicited in a pre-trained language model. We
find numerous vectors in MLP blocks that promote toxicity.
Simply subtracting these vectors from the residual stream
can suppress toxic outputs.
Second, we applied DPO to our language model, using
PPLM to carefully craft pairs of toxic and non-toxic continuations to Wikipedia prompts.
Third, we studied how our aligned model GPT2DPO averts
toxicity. Rather than removing the regions that elicit toxicity, GPT2DPO bypasses them by learning an _offset_ . Such
an offset is distributed amongst multiple value vectors, allowing minimal changes to every weight. This allows the
model to preserve its pre-trained behavior, while averting
toxic outputs, which matches the objective of the DPO loss.
Given this understanding, we demonstrated how to break
the alignment of GPT2DPO, reverting it back to its toxic
behavior. Namely, we simply increase the regions that elicit
toxicity, by scaling their corresponding key vectors.
We view our findings as a mechanistic case study for why
aligned models can be jailbroken, and hope that this can
lead to more robust alignment algorithms. Our code, mod[els, and data can be found at https://github.com/](https://github.com/ajyl/dpo_toxic)
[ajyl/dpo_toxic.](https://github.com/ajyl/dpo_toxic)


We thank Ekdeep Singh Lubana for fruitful discussions, and
Santiago Serra Castro for helping with figures. This work
was supported via NSF under grant #2306372.


[10.18653/v1/2021.emnlp-main.446. URL https://](https://aclanthology.org/2021.emnlp-main.446)
[aclanthology.org/2021.emnlp-main.446.](https://aclanthology.org/2021.emnlp-main.446)
Geva, M., Caciularu, A., Wang, K., and Goldberg, Y. Transformer feed-forward layers build predictions by promoting concepts in the vocabulary space. In Goldberg, Y.,
Kozareva, Z., and Zhang, Y. (eds.), _Proceedings of the_
_2022 Conference on Empirical Methods in Natural Lan-_
_guage Processing_, pp. 30–45, Abu Dhabi, United Arab
Emirates, December 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.emnlp-main.
3. [URL https://aclanthology.org/2022.](https://aclanthology.org/2022.emnlp-main.3)
[emnlp-main.3.](https://aclanthology.org/2022.emnlp-main.3)
Geva, M., Bastings, J., Filippova, K., and Globerson, A. Dissecting recall of factual associations in auto-regressive
language models. _arXiv preprint
2023.
Hendel, R., Geva, M., and Globerson, A. In-context learning
creates task vectors. In Bouamor, H., Pino, J., and Bali,
K. (eds.), _Findings of the Association for Computational_
_Linguistics: EMNLP 2023_, pp. 9318–9333, Singapore,
December 2023. Association for Computational Linguis[tics. URL https://aclanthology.org/2023.](https://aclanthology.org/2023.findings-emnlp.624)
[findings-emnlp.624.](https://aclanthology.org/2023.findings-emnlp.624)
Hendrycks, D. and Gimpel, K. Gaussian error linear units
(gelus). _arXiv preprint
Hernandez, E., Sharma, A. S., Haklay, T., Meng, K., Wattenberg, M., Andreas, J., Belinkov, Y., and Bau, D. Linearity of relation decoding in transformer language models.
_arXiv preprint
Ilharco, G., Ribeiro, M. T., Wortsman, M., Schmidt, L.,
Hajishirzi, H., and Farhadi, A. Editing models with task
arithmetic. In _The Eleventh International Conference on_
_Learning Representations_, 2022.
Jain, S., Kirk, R., Lubana, E. S., Dick, R. P., Tanaka,
H., Grefenstette, E., Rocktaschel, T., and Krueger,¨
D. S. Mechanistically analyzing the effects of finetuning on procedurally defined tasks. _arXiv preprint_
_
Li, K., Hopkins, A. K., Bau, D., Viegas, F., Pfister, H., and´
Wattenberg, M. Emergent world representations: Exploring a sequence model trained on a synthetic task. In _The_
_Eleventh International Conference on Learning Represen-_
_tations_ [, 2023a. URL https://openreview.net/](https://openreview.net/forum?id=DeG07_TcZvT)
[forum?id=DeG07_TcZvT.](https://openreview.net/forum?id=DeG07_TcZvT)
Li, K., Patel, O., Viegas, F., Pfister, H., and Wattenberg, M.´
Inference-time intervention: Eliciting truthful answers
from a language model. _arXiv preprint
2023b.
Li, Z., You, C., Bhojanapalli, S., Li, D., Rawat, A. S.,
Reddi, S. J., Ye, K., Chern, F., Yu, F., Guo, R., and
Kumar, S. The lazy neuron phenomenon: On emergence
of activation sparsity in transformers. In _The Eleventh_
_International Conference on Learning Representations_,
[2023c. URL https://openreview.net/forum?](https://openreview.net/forum?id=TJ2nxciYCk-)
[id=TJ2nxciYCk-.](https://openreview.net/forum?id=TJ2nxciYCk-)
Meng, K., Bau, D., Andonian, A. J., and Belinkov, Y. Locating and editing factual associations in GPT. In Oh,
A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.),
_Advances in Neural Information Processing Systems_,
[2022. URL https://openreview.net/forum?](https://openreview.net/forum?id=-h6WAS6eE4)
[id=-h6WAS6eE4.](https://openreview.net/forum?id=-h6WAS6eE4)
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer
sentinel mixture models. In _International Conference on_
_Learning Representations_, 2016.
Nanda, N., Lee, A., and Wattenberg, M. Emergent linear
representations in world models of self-supervised sequence models. In _Proceedings of the 6th BlackboxNLP_
_Workshop: Analyzing and Interpreting Neural Networks_
_for NLP_, pp. 16–30, 2023.
Nasr, M., Carlini, N., Hayase, J., Jagielski, M., Cooper,
A. F., Ippolito, D., Choquette-Choo, C. A., Wallace, E.,
Tramer, F., and Lee, K. Scalable extraction of training`
data from (production) language models. _arXiv preprint_
_
Nostalgebraist. Interpreting gpt: The logit lens,
2020. URL [https://www.lesswrong.](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)
[com/posts/AcKRB8wDpdaN6v6ru/](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)
[interpreting-gpt-the-logit-lens.](https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens)
Olah, C., Cammarata, N., Schubert, L., Goh, G., Petrov,
M., and Carter, S. Zoom in: An introduction to circuits. _Distill_, 2020. doi: 10.23915/distill.00024.001.
https://distill.pub/2020/circuits/zoom-in.
Ouyang, L., Wu, J., Jiang, X., Almeida, D., Wainwright, C.,
Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Gray, A.,
et al. Training language models to follow instructions
with human feedback. In _Advances in Neural Information_
_Processing Systems_, 2022.
Park, K., Choe, Y. J., and Veitch, V. The linear representation hypothesis and the geometry of large language
models. In _Causal Representation Learning Workshop at_
_NeurIPS 2023_, 2023.
Qi, X., Zeng, Y., Xie, T., Chen, P.-Y., Jia, R., Mittal, P.,
and Henderson, P. Fine-tuning aligned language models
compromises safety, even when users do not intend to!
_arXiv preprint


Rafailov, R., Sharma, A., Mitchell, E., Ermon, S., Manning,
C. D., and Finn, C. Direct preference optimization: Your
language model is secretly a reward model, 2023.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and
Klimov, O. Proximal policy optimization algorithms.
_arXiv preprint
Sheng, E., Chang, K.-W., Natarajan, P., and Peng, N. The
woman worked as a babysitter: On biases in language
generation. In Inui, K., Jiang, J., Ng, V., and Wan, X.
(eds.), _Proceedings of the 2019 Conference on Empir-_
_ical Methods in Natural Language Processing and the_
_9th International Joint Conference on Natural Language_
_Processing (EMNLP-IJCNLP)_, pp. 3407–3412, Hong
Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v1/D19-1339. URL
[https://aclanthology.org/D19-1339.](https://aclanthology.org/D19-1339)
Stiennon, N., Ouyang, L., Wu, J., Ziegler, D., Lowe, R.,
Voss, C., Radford, A., Amodei, D., and Christiano,
P. F. Learning to summarize with human feedback. _Ad-_
_vances in Neural Information Processing Systems_, 33:
3008–3021, 2020.
Tenney, I., Das, D., and Pavlick, E. BERT rediscovers the
classical NLP pipeline. In Korhonen, A., Traum, D., and
Marquez, L. (eds.),` _Proceedings of the 57th Annual Meet-_
_ing of the Association for Computational Linguistics_, pp.
4593–4601, Florence, Italy, July 2019. Association for
Computational Linguistics. doi: 10.18653/v1/P19-1452.
[URL https://aclanthology.org/P19-1452.](https://aclanthology.org/P19-1452)
Todd, E., Li, M. L., Sharma, A. S., Mueller, A., Wallace,
B. C., and Bau, D. Function vectors in large language
models. _arXiv preprint
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux,
M.-A., Lacroix, T., Roziere, B., Goyal, N., Hambro, E.,`
Azhar, F., et al. Llama: Open and efficient foundation language models. _arXiv preprint
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones,
L., Gomez, A. N., Kaiser, L. u., and Polosukhin, I.
Attention is all you need. In Guyon, I., Luxburg, U. V.,
Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S.,
and Garnett, R. (eds.), _Advances in Neural Information_
_Processing Systems_, volume 30. Curran Associates, Inc.,
2017. [URL https://proceedings.neurips.](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)
[cc/paper_files/paper/2017/file/](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)
[3f5ee243547dee91fbd053c1c4a845aa-Paper.](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)
[pdf.](https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf)
Wallace, E., Feng, S., Kandpal, N., Gardner, M., and Singh,
S. Universal adversarial triggers for attacking and analyzing NLP. In Inui, K., Jiang, J., Ng, V., and Wan,
X. (eds.), _Proceedings of the 2019 Conference on Em-_
_pirical Methods in Natural Language Processing and_
_the 9th International Joint Conference on Natural Lan-_
_guage Processing (EMNLP-IJCNLP)_, pp. 2153–2162,
Hong Kong, China, November 2019. Association for
Computational Linguistics. doi: 10.18653/v1/D19-1221.
[URL https://aclanthology.org/D19-1221.](https://aclanthology.org/D19-1221)
Wei, A., Haghtalab, N., and Steinhardt, J. Jailbroken:
How does LLM safety training fail? In _Thirty-seventh_
_Conference on Neural Information Processing Systems_,
[2023. URL https://openreview.net/forum?](https://openreview.net/forum?id=jA235JGM09)
[id=jA235JGM09.](https://openreview.net/forum?id=jA235JGM09)
Welleck, S., Kulikov, I., Roller, S., Dinan, E., Cho, K., and
Weston, J. Neural text generation with unlikelihood training. In _International Conference on Learning Represen-_
_tations_ [, 2020. URL https://openreview.net/](https://openreview.net/forum?id=SJeYe0NtvH)
[forum?id=SJeYe0NtvH.](https://openreview.net/forum?id=SJeYe0NtvH)
Xu, J., Lee, A., Sukhbaatar, S., and Weston, J. Some
things are more cringe than others: Preference optimization with the pairwise cringe loss. _arXiv preprint_
_
Yang, X., Wang, X., Zhang, Q., Petzold, L., Wang, W. Y.,
Zhao, X., and Lin, D. Shadow alignment: The ease
of subverting safely-aligned language models. _arXiv_
_preprint
Zhang, Z., Lin, Y., Liu, Z., Li, P., Sun, M., and Zhou,
J. MoEfication: Transformer feed-forward layers are
mixtures of experts. In Muresan, S., Nakov, P., and
Villavicencio, A. (eds.), _Findings of the Association for_
_Computational Linguistics: ACL 2022_, pp. 877–890,
Dublin, Ireland, May 2022. Association for Computational Linguistics. doi: 10.18653/v1/2022.findings-acl.
71. [URL https://aclanthology.org/2022.](https://aclanthology.org/2022.findings-acl.71)
[findings-acl.71.](https://aclanthology.org/2022.findings-acl.71)
Zou, A., Wang, Z., Kolter, J. Z., and Fredrikson, M. Universal and transferable adversarial attacks on aligned language models. _arXiv preprint


In this section we provide details from Geva et al. (2022) that demonstrate that MLP value vectors promote or suppress the
likelihood of tokens.
We start from Equation 2:
_dmlp_
- _m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ][.]_
_i_ =1
MLP _[ℓ]_ ( **x** _[ℓ]_ ) =
_dmlp_
- _σ_ ( **x** _[ℓ]_ _·_ **k** _[ℓ]_ _i_ [)] **[v]** _i_ _[ℓ]_ [=]
_i_ =1
Thus, we can consider the update from MLP _[ℓ]_ as _dmlp sub-updates_, each sub-update being _m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_ [.]
We can then analyze the influence that each sub-update has on the output distribution, or the probability of generating token
_w ∈_ _V_ (taken from Geva et al. (2022)):
_p_       - _w |_ **x** _[ℓ]_ + _m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ][, E]_       - = [exp] _Z_       - **e** ~~�~~ _wE·_ ( **xx** _[ℓ][ℓ]_ ++ _m_ **e** _w_ _[ℓ]_ _i_ **[v]** _· mi_ _[ℓ]_ [))] _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_       - _∝_ exp       - **e** _w ·_ **x** _[ℓ]_ [�] _·_ exp       - **e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_       - (4)
where **e** _w_ is the token embedding of _w_, and _Z_ is the softmax normalization factor. This indicates that when **e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_ _[>]_ [ 0][,]
the likelihood of _w_ increases, while **e** _w · m_ _[ℓ]_ _i_ **[v]** _i_ _[ℓ]_ _[<]_ [ 0][ decreases the likelihood.]

A. Projecting Value Vectors onto Vocabulary Space


In this section we provide more examples of residual streams shifting out of toxic regions. See Figure 6
_Figure 6._ Shift in residual streams at layer 12, 18, and 13 (we show these three layers because MLP _._ **v** 771 [12] [,][ MLP] _[.]_ **[v]** 2669 [18] [, and][ MLP] _[.]_ **[v]** 668 [13]
are the next three vectors with highest cosine similarity with _W_ Toxic. See Table 1, Figure 2.


_Table 5._ Hyperparameters: DPO.
HYPERPARAMETER VALUE
LEARNING RATE 1E-6
BATCH SIZE 4
OPTIMIZER RMSPROP
GRADIENT ACCUMULATION STEPS 1
MAX GRADIENT NORM 10
VALIDATION METRIC LOSS/VALID
VALIDATION PATIENCE 10
DPO BETA 0.1
**C. Shifts in Residual Streams vs. Shifts in MLP Value Vectors.**
In this section we provide more examples of how MLP value vectors contribute in the _δ_ **x** direction at different layers.
|Col1|Layer 0|
|---|---|
|||
|||
|||
|||
|||
|||
|Col1|Layer 6|Col3|Col4|Layer 7|Col6|Col7|Layer 8|Col9|Layer 9|Col11|ayer 10|Col13|Layer 11|
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
|||||||||||||||
|||||||||||||||
|||||||||||||||
|||||||||||||||
|||||||||||||||
|||||||||||||||
|||||||||||||||
_Figure 7._ Shift in residual streams at layer 12 vs. shift in MLP value vectors ( _δ_ **x** [12] [vs.] _[ δ]_ MLP [).]


Tables 5, and 6 contain the hyperparameters used for our toxic probe, DPO, and PPLM, respectively.
|Col1|Layer 0|
|---|---|
|||
|||
|||
|||
|||
|||


|Col1|Layer 9|Col3|Col4|Layer 10|Col6|Layer 11|Col8|Layer 12|Col10|Col11|Layer 13|
|---|---|---|---|---|---|---|---|---|---|---|---|
|||||||||||||
|||||||||||||
|||||||||||||
|||||||||||||
|||||||||||||
|||||||||||||
_Figure 8._ Shift in residual streams at layer 14 vs. shift in MLP value vectors ( _δ_ **x** [14] [vs.] _[ δ]_ MLP [).]
_Table 6._ Hyperparameters: PPLM.
HYPERPARAMETER VALUE
STEP SIZE 0.4
TEMPERATURE 1
TOP K 10
NUM ITERATIONS 50
WINDOW LENGTH 0
HORIZON LENGTH 1
DECAY FALSE
GAMMA 1
GM SCALE 0.95
KL SCALE 0.1
|Col1|Layer 11|Col3|Layer 12|Col5|Layer 13|Col7|Layer 14|Col9|Layer 15|
|---|---|---|---|---|---|---|---|---|---|
|||||||||||
|||||||||||
|||||||||||


_Figure 9._ Shift in residual streams at layer 16 vs. shift in MLP value vectors ( _δ_ **x** [16] [vs.] _[ δ]_ MLP [).]
|Col1|Col2|Col3|Col4|Col5|Col6|Col7|Col8|Col9|Col10|
|---|---|---|---|---|---|---|---|---|---|
|||||||||||
|||||||||||
_Figure 10._ Shift in residual streams at layer 18 vs. shift in MLP value vectors ( _δ_ **x** [18] [vs.] _[ δ]_ MLP [).]


Ziping Ma [*], Furong Xu, Jian Liu, Ming Yang, Qingpei Guo [†]
Ant Group
maziping.im@gmail.com, booyoungxu.xfr@antgroup.com, rex.lj@antgroup.com
m.yang@antgroup.com, qingpei.gqp@antgroup.com

SyCoCa: Symmetrizing Contrastive Captioners with Attentive Masking for
  Multimodal Alignment

Ziping Ma
Furong Xu
Jian Liu
Ming Yang
Qingpei Guo


_Multimodal alignment between language and vision is_
_the fundamental topic in current vision-language model re-_
_search. Contrastive Captioners (CoCa), as a representa-_
_tive method, integrates Contrastive Language-Image Pre-_
_training (CLIP) and Image Caption (IC) into a unified_
_framework, resulting in impressive results. CLIP imposes_
_a bidirectional constraints on global representation of en-_
_tire images and sentences. Although IC conducts an uni-_
_directional image-to-text generation on local representa-_
_tion, it lacks any constraint on local text-to-image recon-_
_struction, which limits the ability to understand images at_
_a fine-grained level when aligned with texts. To achieve_
_multimodal alignment from both global and local perspec-_
_tives, this paper proposes Symmetrizing Contrastive Cap-_
_tioners (SyCoCa), which introduces bidirectional interac-_
_tions on images and texts across the global and local rep-_
_resentation levels. Specifically, we expand a Text-Guided_
_Masked Image Modeling (TG-MIM) head based on ITC and_
_IC heads. The improved SyCoCa can further leverage tex-_
_tual cues to reconstruct contextual images and visual cues_
_to predict textual contents._ _When implementing bidirec-_
_tional local interactions, the local contents of images tend_
_to be cluttered or unrelated to their textual descriptions._
_Thus, we employ an attentive masking strategy to select_
_effective image patches for interaction. Extensive experi-_
_ments on five vision-language tasks, including image-text_
_retrieval, image-captioning, visual question answering, and_
_zero-shot/finetuned image classification, validate the effec-_
_tiveness of our proposed method._


In recent years, the dramatic progress of multimodal alignment between vision and language has reshaped computer
*Work done during an internship with Ant Group.
†Corresponding author.
vision research in some sense. The availability of largescale datasets and powerful computational resources has
led to many seminal works and impressive results in this
field. Most methods use contrastive objective to constrain
global representation between modalities as shown in Figure 1 (a), such as CLIP [41, 58], ALBEF [29], mPLUG
[28, 57]. BEiT-v3 [53] treats image as a foreign language
and uses a mask-then-predict strategy for pre-training. This
pretraining method requires separate fine-tuning for each
downstream task and the pretrained model cannot be directly used across tasks. CoCa [61] combines Image-Text
Contrastive (ITC) and Image Caption (IC), resulting in a
pretrained model that can be directly used for both retrieval
and IC tasks. It is a promising architecture and works well
on multiple vision-language tasks.
However, in terms of local interaction between image
patches and text tokens, the IC head in CoCa only utilizes
the visual cues to generate textual descriptions, yet disregarding the visual context reconstruction from textual cues.
From vision pre-training works, such as SimMIM [56] and
MAE [19], we learn that the image reconstruction can learn
a strong content representation. Therefore, in the multimodal scenario, when text cue is introduced into the image
reconstruction task and local interaction is performed, the
representation of texts and images can be unified into one
space, thereby further enhancing multimodal alignment.
In this paper, we propose a novel framework called Symmetrizing Contrastive Captioners (SyCoCa) that incorporates both local image-to-text generation and text-to-image
reconstruction in addition to global constractive objective.
In addition to ITC head and IC head, we introduce a textguided masked image modeling (TG-MIM) head. The difference between CoCa and our method is illustrated in Figure 1. CoCa only achieved undirectional interaction between image and text as shown in Figure 1 (b). Further,
in Figure 1 (c), our TG-MIM introduces text for image reconstruction, which enhances the fine-grained representation ability of images. In this way, SyCoCa has the bidirectional global and local interactions between modalities.


Figure 1. Comparison of the pipelines in vision-language pretraining frameworks. (a) CLIP: aligning global features across modalities
through contrastive learning. (b) CoCa: introducing image captioning to construct unidirectional fine-grained interaction. (c) Our SyCoCa:
bidirectional local interation with attentive masking to enhance comprehensive cross-modal understanding.
During the actual multimodal alignment process, although an image is worth one thousand words, people seldom write one thousand words to describe the image content. Instead, text descriptions or image captions are often
highly abstract just focusing on the main character, object,
or event in an image. For example, the caption of a picture about the family dinner of Thanksgiving would rarely
delineate the furniture in the room in details. Moreover,
even short text descriptions may imply rich and abundant
image contents. _e.g.,_ . _a football match_ may well imply a
vivid scene of a grass field, multiple palyers, and croweded
audience, etc. Therefore, it is necessary to choose appropriate local regions and representation for alignment. To select effective image patches for text, we employ an attentive
masking strategy. Specifically, we calculate the similarity
between image tokens and text tokens to determine the relevance of image patches. For the IC task, the most pertinent
image patches are selected to aid in generating textual descriptions by considering their semantic similarity with text
descriptions. In contrast, we choose the least relevant image patches with the text on the TG-MIM task, aiming to
leverage the text tokens to assist in recovering image content. Extensive experiments demonstrate the effectiveness
of our proposed method on five downstream tasks.
In summary, the main contributions of our work are
listed as follows:
- We first propose a symmetrizing contrastive captioners
for multimodal alignment, which improves the understanding between images and sentences from a global and
local perspective.
- To promote bidirectional local interaction, we adopt an
attentive masking strategy to select appropriate image
patches for IC and TG-MIM heads respectively.
- Thorough experiments validate that our proposed
_SyCoCa_ outperforms CoCa on several downstream tasks,
_e.g._ image-text retrieval, image-captioning, visual question answering and image classification. For example, we
obtained +5.1%/3.7% in mTR/mIR on Flicker-30k compared to CoCa in image-text retrieval tasks.


**Vision-Language Pretraining.** In recent years, there has
been tremendous progress in multimodal alignment, especially between vision and language. Extensive researchers
have dedicated their efforts to exploring vision-language
pretraining. Early works [8, 49, 67] prefixed a pretrained
object detection modules to extract visual representations,
which were then aligned with textual representations to
achieve multimodal alignment. Later efforts focused on
training multimodal transformers from scratch, such as
ViLT [25] and VLMo [3]. Pre-training of a foundation
model on gigantic data with tremendous computation led
to the breakthrough of multimodal alignment of images and
texts. CLIP [41] and ALIGN [23] trained a dual-encoder
on large-scale noisy image-text pairs using contrastive loss,
obtaining generic image and text representations for crossmodal alignment and zero-shot image classification tasks.
Florence [63] used a unified contrastive learning for various vision and image-text benchmarks. BLIP [31] trained a
Q-former to align a frozen vision encoder and language encoder. BEiT-v3 [53] treated image as a foreign language,
and used mask-then-predict strategy for pre-training. To
enhance local interaction of images and texts, CoCa [61]
introduced a decoder for image caption based on CLIP.
Our SyCoCa further achieves bidirectional local interactions and enhances fine-grained alignment capabilities between modalities.
**Masked Image Modeling.** Beyond contrastive learning,
masked image modeling (MIM) [7, 13, 13, 39, 50] emerges
as another promising technique in vision pre-training. Recently, iGPT [7], ViT [15] and BEiT [2] recalled this learning approach on the modern vision transformers, which
show great potential in representation learning by introduc
ing special designs, such as clustering on pixels [7], prediction of mean color [15], and tokenization via an additional dVAE network with a block-wise masking strategy
[2]. SimMIM [56] predicted RGB values of raw pixels by
direct regression. MAE [19] used an asymmetric encoderdecoder architecture, where the encoder operates only on
the visible subset of patches, then a lightweight decoder reconstructs the original image from the latent representation
and mask tokens.
Inspired by MAE, we propose a TG-MIM head to improve fine-grained multimodal alignment. The major differences of TG-MIM against MAE include: 1) the text cues
are leveraged in TG-MIM to guide image reconstruction; 2)
an attentive masking scheme is used in TG-MIM to select
appropriate patches; and 3) TG-MIM targets on improving
multi-modal alignment, instead of learning vision representation as in MAE.


We are interested in leveraging bidirectional cross-modal
interaction to learn fine-grained visual and textual representations within an aligned latent space. Contrastive Captioner (CoCa) proposes to incorporate an image captioning task alongside the contrastive task, which learns finegrained representation and establishes a stronger alignment
between modalities than CLIP. It is worth noting that the
modality interaction in CoCa is unidirectional and asymmetric, focusing solely on genetating text from images.
To enhance vision-language alignment, we introduce a
novel framework called SyCoCa, which expands the existing CoCa model by incorporating bidirectional local interaction with attentive masking. Our framework comprises of
three objectives: (i) image-text contrasting (ITC), (ii) image
captioning (IC), and (iii) text-guided masked image modeling (TG-MIM). Figure 2 shows the overall framework of
SyCoCa.
The goal of SyCoCa is to further explore the potential of cross-modal interactive prediction tasks. Building
upon CoCa, SyCoCa incorporates TG-MIM to establish
predictions from text to image, aiming to compensate for
the inherent asymmetry in CoCa. This establishes bidirectional prediction between modalities, promoting finegrained vision-language understanding. Furthermore, we
design a novel attentive masking procedure that enables the
bidirectional cross-modal interactions to focus on different
regions of images, guided by the accompanying texts.


As shown in Figure 2, the overall model architecture consists of four key components: the image encoder, text encoder, (text-to-)image decoder, and (image-to-)text decoder.
Here, we provide a detailed explanation to each ones.
**Image Encoder.** We employ vision transformer as the
image encoder to model an input image. The image encoder takes image patches as input and encodes them into a
sequence of embeddings _{v_ _[cls]_ _, v_ [1] _, ..., v_ _[P]_ _}_ where each embedding corresponds to a specific image patch. Additionally, an extra [ `CLS` ] embedding is included to provide a
global representation of the image.
**Text Encoder.** As the text encoder, we adopt a
causal masked transformer encoder to model text inputs.
This encoder takes the simple BPE tokenized input text
and converts it into a sequence vector represented as
_{w_ [1] _, ..., w_ _[S]_ _, w_ _[cls]_ _}_, in which the embedding of the [ `CLS` ]
token summarizes the global text feature. The purpose of
adopting causal attention mask is to prevent any potential
information leakage from future tokens to past tokens during the text encoding process.
**Image / Text Decoder.** To further capture the interaction between image-text pairs, we use an image decoder
and a text decoder. Each decoder utilizes cross-attention
transformer modules to deeply fuse image and text information, enabling bidirectional local interaction. The crossmodality multi-head attention module uses the representations of one modality as the query and the other modality’s
representations as the key and value. More specifically, the
image decoder predicts the pixel values of masked image
patches, leveraging cross-attention to incorporate relevant
text. Similarly, the text decoder predicts the logits of the
next token, utilizing cross-attention to integrate relevant image information. This bidirectional fusion mechanism ensures a comprehensive mutual understanding of the imagetext pairs.


We start from an image-text pair dataset consisting of pairs
( _Ii, Ti_ ), where _Ii_ represents an image and _Ti_ is the corresponding text caption. The image encoder _EI_ and text encoder _ET_ are responsible for encoding the image and text,
respectively, yielding embeddings:
_{vi_ _[cls][, v]_ _i_ [1] _[, ..., v]_ _i_ _[P]_ _[}]_ [ =] _[ E][I]_ [(] _[I][i]_ [)] and (1)
_{wi_ [1] _[, ..., w]_ _i_ _[S][, w]_ _i_ _[cls][}]_ [ =] _[ E][T]_ [(] _[T][i]_ [)] _[.]_ (2)
_P_ refers to the number of patches, and _S_ is the length of tokenized text sequence. Here, we present the applied training
objectives in details.
**Image Text Contrasting (ITC).** In the training process,
we consider a batch of _N_ image-text pairs _{Ii, Ti}_ _[N]_ _i_ =1 [and]
their corresponding global representations _{vi_ _[cls][, w]_ _i_ _[cls][}]_ _i_ _[N]_ =1 [.]
It is assumed that these representations have been normalized to have a unit _ℓ_ 2 norm. The contrasive objective is
a dog is running on the grass [/s]
|Col1|K, V Q|Col3|Col4|Col5|Col6|Col7|Col8|
|---|---|---|---|---|---|---|---|
|||||||||
|high / low-attn mask /|Col2|Col3|Col4|Col5|Col6|Col7|Col8|
|---|---|---|---|---|---|---|---|
|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|Attentive
Masking
image CLS
ITC
text CLS|
|||||||||
|||||||||
Causal Text Encoder
[s] a dog is running on the grass [CLS]
“A dog is running on the grass.”


|Col1|Q|Col3|
|---|---|---|
||||
||||
|Col1|Col2|K, V|Col4|
|---|---|---|---|
|||||
|||||
Figure 2. The detailed illustration of our proposed method. The framework of our method consist of four modules: an image encoder, a
(causal) text encoder, a (text-to-)image decoder, a (image-to-)text decoder. Our method focuses on three pretraining objectives: image-text
contrasting (ITC), text-guided masked image modeling (TG-MIM) and image captioning (IC).
designed to align the image and text representations:
exp( _⟨vi_ _[cls][, w]_ _i_ _[cls][⟩][/τ]_ [)]
~~�~~ _N_
_j_ =1 [exp(] _[⟨][v]_ _i_ _[cls][, w]_ _j_ _[cls][⟩][/τ]_ [)]
exp( _⟨vi_ _[cls][, w]_ _i_ _[cls][⟩][/τ]_ [)]
~~�~~ _N_
_j_ =1 [exp(] _[⟨][v]_ _i_ _[cls][, w]_ _j_ _[cls][⟩][/τ]_ [)]
Nonetheless, this interaction is unidirectional and asymmetrical, primarily focusing on understanding visual elements
in language, without exploiting the comprehension of textual elements in the visual context.
**Text-Guided Masked Image Modeling.** To address the
inherent imbalance in cross-modal interaction within the
image captioning task, we design a novel training objective
called text-guided masked image modeling. Our SyCoCa
framework utilizes the image decoder, denoted as _DI_, to
reconstruct the masked image by predicting the pixel values for each masked patch. The output linear projection
and reshape process is similar to that of the MAE. For a
given image _I_ = _{p_ [1] _, ..., p_ _[P]_ _}_ consisting of _P_ patches, we
employ patch-wise masking using a masking map _M_ =
_{m_ [1] _, ..., m_ _[P]_ _}_, where _m_ _[i]_ _∈{_ 0 _,_ 1 _}_ indicates whether the
patch is masked. The TG-MIM loss is computed as the
L1 loss between the masked patches and their reconstructed
counterparts in the pixel space:
_{p_ ˆ [1] _, ...,_ ˆ _p_ _[P]_ _}_ = _DI_ ( _EI_ ( _I ⊕_ _M_ ) _, ET_ ( _T_ )) _,_ (5)
��
��
_,_
(3)
_LIT C_ =
2 _N_
+ [1]
2 _N_
- _N_
- log
_i_
- _N_
- log
_i_
where _⟨·, ·⟩_ refers to the inner product, and _τ_ is the temperature to scale the logits. The contrastive objective that
pulls the representations of paired image-text close together
while pushing apart unmatched pairs, promoting alignment
in a shared semantic space.
**Image Captioning.** Image captioning objective requires
the model to generate tokenized texts _Ti_ with precise details
in an autoregressive manner, compared with ITC treating
the inputs as single entities. We use the image encoder _EI_
to encode image representations and train the text decoder
_DT_ to maximize the conditional likelihood of the text _Ti_ by
utilizing forward autoregressive factorization:
_LT M_ =
_P_
- _m_ _[i]_ _· ∥_ _p_ _[i]_ _−_ _p_ ˆ _[i]_ _∥_ [1] _,_ (6)
_i_ =1
where ˆ _p_ _[i]_ refers to the reconstructed patches. Note that the
TG-MIM objective, proposed in this paper, distinguishes
itself from the image-pretraining objective MIM by dedicated efforts on the guidance offered by the cross-attended
text representation. The text-guided modeling focuses on
the interaction between visual and textual representations,
with the goal of enhancing the fine-grained understanding
of textual elements within the visual context through reconstruction process. The introduction of this objective com
_LIC_ = _−_
_|Ti|_
- log _DT_ - _Ti_ _[j][|][T][ <j]_ _i_ _, EI_ ( _Ii_ )� _._ (4)
_j_ =1
The text decoder is trained using parallel computation for
enhanced learning efficiency. This training objective enables the model to learn fine-grained representations with
strong alignment through cross-modal prediction, fostering the acquisition of joint-modality representations useful
for various multimodal understanding and generation tasks.
plements the uni-directional prediction of image captioning.


Intuitively, the reconstruction of image regions that reveal
the same semantics or show highly related contents to the
text shall be a more effective and robust cross-modal interaction, rather than reconstructing those irrelevant image
regions. Therefore, we suggest generating masking maps
that eliminate image patches that are relevant to language
captions, and subsequently reconstruct these patches based
on the text captions.
A straightforward idea to evalaute the correlation is by
measuring the similarity between each patch embedding _v_ _[i]_
and the global text embedding _w_ _[cls]_ . However, _w_ _[cls]_ provides a coarse-grained summary of the entire sentence. To
mask the vision elements mentioned in the caption as accurately as possible, we propose to calculate the token-wise
maximum similarity between image and text embeddings.
This method is effective in modeling fine-grained crossmodal similarities in a previous study [60]. Formally, for
the _i_ -th embedding _v_ _[i]_ of an image, we calculate its similarities with all text embeddings _{w_ [1] _, ..., w_ _[S]_ _}_ and select the
largest similarity
_s_ _[i]_ = max _[S]_ _j_ =1 _[⟨][v][i][, w][j][⟩]_ (7)
to represent its correlation with the text.
Image patches to be masked are selected based on their
scores _{s_ [1] _, ..., s_ _[P]_ _}_ . To ensure stable training in TG-MIM,
a fixed ratio _rh_ of top scoring patches (represented by the
_high attn mask_ in Figure 2) in the input image are masked.
Conversely, in the caption task, the patches with the lowest
scores (represented by the _low attn mask_ in Figure 2) in the
input image are masked at a fixed ratio _rl_ . The purpose is to
enhance cross-modal interaction by encouraging the model
to prioritize visually salient regions that play a significant
role in understanding and representing the image for caption
generation.


Finally, we pretrain SyCoCa with all these losses combined:
_L_ = _LIT C_ + _λICLIC_ + _λT M_ _LT M_ _,_ (8)
where _λIC_ and _λT M_ are the hyper-parameters weighting
between IC and TG-MIM. All the modules of SyCoCa are
trained from scratch.


To demonstrate the effectiveness of our proposed SyCoCa,
we conduct extensive experiments on 5 downstream tasks.
Initially, we present the experimental setup, including
model architecture, pretraining datasets, downstream tasks,
and implementation details. Subsequently, we compare
SyCoCa with the baseline model, CoCa, on image-text retrieval, image classification, image captioning and visual
question answering tasks. Finally, we perform a series of
ablation studies to further analyze and evaluate our model.


**Model Architecture.** To ensure a fair comparison, we
conduct our experiment using the open-source implementation of CoCa [1] . Both SyCoCa and the baseline model utilize
the same CoCa-Base configuration for the image encoder,
text encoder, and text decoder. However, in SyCoCa, we
introduce a new image decoder that shares the same architecture as the text decoder yet with minor modifications.
Specifically, we replace the token-prediction head with a
pixel-prediction head. All model parameters are initialized
using a Gaussian distribution with a mean of 0 and a standard deviation of 0.01, allowing the training process to start
from scratch.
**Pretraining Data.** We use the Conceptual Captions
12M [6] (CC12M) dataset with 12 million image-caption
pairs, as the multi-modal pretraining data for all models.
Although this dataset is smaller in scale compared to the
large custom datasets employed by the state-of-the-art pretraining models (such as 400M pairs in CLIP [41] and 3B
pairs in CoCa [61]), it is a good match to the computation resources available to us. Moreover, CC12M has been
widely adopted for benchmark evaluations in various studies on vision-language pretraining [17, 32, 36, 47, 64].
**Downstream Tasks.** We evaluate SyCoCa on 3 downstream vision-language tasks and 2 classification tasks, including image-text retrieval, image captioning, visual question answering, zero-shot image classification, and finetuned image classification.
**Implementation Details.** We conduct our model training on two machines, each equipped with 8 NVIDIA A100
GPUs, for a total of 20 epochs. The batch size during training is set to 2048, and the resolution of pretraining images
is set to 224 _×_ 224. We use the AdamW optimizer [35] with
an initial learning rate of 1 _e −_ 4. The learning rate schedule
follows a cosine decay, including a warm-up period of 5000
steps. In terms of hyperparameters, we simply set _λIC_ = 2
following CoCa and _λT M_ = 1. The masking ratios _rh_ and
_rl_ are both empirically set to 50%.