# Interpretability at Scale:
Identifying Causal Mechanisms in Alpaca

Zhengxuan Wu,∗ Atticus Geiger,∗ Christopher Potts, and Noah D. Goodman
Stanford University
{wuzhengx, atticusg, cgpotts<EMAIL_ADDRESS>Equal contribution.

###### Abstract

Obtaining human-interpretable explanations of large, general-purpose language
models is an urgent goal for AI safety. However, it is just as important that
our interpretability methods are faithful to the causal dynamics underlying
model behavior and able to robustly generalize to unseen inputs. Distributed
Alignment Search (DAS) [23] is a powerful gradient descent method grounded in
a theory of causal abstraction that uncovered perfect alignments between
interpretable symbolic algorithms and small deep learning models fine-tuned
for specific tasks. In the present paper, we scale DAS significantly by
replacing the remaining brute-force search steps with learned parameters – an
approach we call Boundless DAS. This enables us to efficiently search for
interpretable causal structure in large language models while they follow
instructions. We apply Boundless DAS to the Alpaca model (7B parameters),
which, off the shelf, solves a simple numerical reasoning problem. With
Boundless DAS, we discover that Alpaca does this by implementing a causal
model with two interpretable boolean variables. Furthermore, we find that the
alignment of neural representations with these variables is robust to changes
in inputs and instructions. These findings mark a first step toward deeply
understanding the inner-workings of our largest and most widely deployed
language models.

##  1 Introduction

Present-day large language models (LLMs) display remarkable behaviors: they
appear to solve coding tasks, translate between languages, engage in open-
ended dialogue, and much more. As a result, their societal impact is rapidly
growing, as they make their way into products, services, and people’s own
daily tasks. In this context, it is vital that we move beyond behavioral
evaluation to deeply explain, in human-interpretable terms, the internal
processes of these models; an initial step in auditing them for safety,
trustworthiness, and pernicious social biases.

The theory of causal abstraction [5; 22] provides a generic framework for
representing interpretability methods that faithfully assess the degree to
which a complex causal system (e.g., a neural network) implements an
interpretable causal system (e.g., a symbolic algorithm). Where the answer is
positive, we move closer to having guarantees about how the model will behave.
However, thus far, such interpretability methods have been applied only to
small models fine-tuned for specific tasks [19; 20; 28; 7] and this is
arguably not an accident: the space of alignments between the variables in the
hypothesized causal model and the representations in the neural network
becomes exponentially larger as models increase in size. When a good alignment
is found, one has specific formal guarantees. Where no alignment is found, it
could easily be a failure of the alignment search algorithm.

Distributed Alignment Search (DAS) [23] marks real progress on this problem.
DAS opens the door to (1) discovering structure spread across neurons and (2)
using gradient descent to learn an alignment between distributed neural
representations and causal variables. However, DAS still requires a brute-
force search over the dimensionality of neural representations, hindering its
use at scale.

In this paper, we introduce Boundless DAS which replaces the remaining brute-
force aspect of DAS with learned parameters, truly enabling explainability at
scale. We use Boundless DAS to study how Alpaca (7B) [46], an off-the-shelf
instruct-tuned LLaMA model, follows basic instructions in a simple numerical
reasoning task. Figure 1 summarizes the approach. We find that Alpaca achieves
near-perfect task performance _because_ it implements a simple algorithm with
interpretable variables. In further experiments, we show that Alpaca uses this
simple algorithm across a wide range of contexts and variations on the task.
These findings mark a first step toward deeply understanding the inner-
workings of our largest and most widely deployed language models.111Our code
is available publicly at https://github.com/frankaging/align-transformers

Figure 1: Our pipeline for scaling causal explainability to LLMs with billions
of parameters.

##  2 Related Work

#### Interpretability

Many methods have been developed in an attempt to explain and understand deep
learning models. These methods include analyzing learned weights [12; 1; 13],
gradient-based methods [44; 43; 60; 45; 51; 6], probing [14; 47; 24; 41; 12;
32; 8; 39], self-generated model explanations [27], and training external
explainers based on model behaviors [38; 31]. However, these methods rely only
on observational measurements of behavior and internal neural representations.
Such explanations are not guaranteed to be faithful to the underlying causal
mechanisms of the target models [29; 20; 49; 53].

#### Causal Abstraction

The theory of causal abstraction [40; 4; 5] offers a unifying mathematical
framework for interpretability methods aiming to uncover interpretable causal
mechanisms in deep learning models [19; 28; 20; 58; 54; 23] and training
methods for inducing such interpretable mechanisms [21; 59; 58; 25]. Causal
abstraction [22] can represent many existing interpretablity methods,
including interative nullspace projection [37; 15; 30], causal mediation
analysis [50; 33], and causal effect estimation [18; 16; 2]. In particular,
causal abstraction grounds the research program of _mechanistic
interpretability_ , which aims to reverse engineer deep learning models by
determining the algorithm or computation underlying their intelligent behavior
[34; 17; 35; 7; 54]. To the best of our knowledge, there is no prior work that
scales these methods to large, general-purpose LLMs.

#### Training LLMs to Follow Instructions

Instruction-based fine-tuning of LLMs can greatly enhance their capacity to
follow natural language instructions [10; 36]. In parallel, this ability can
also be induced into the model by fine-tuning base models with hundreds of
specific tasks [56; 11]. Recently, Wang et al. [55] show that the process of
creating fine-tuning data for instruction following models can partly be done
by the target LLM itself (“self-instruct”). Such datasets have led to many
recent successes in lightweight fine-tuning of ChatGPT-like instruction
following models such as Alpaca [46], instruct-tuned LLaMA [48], StableLM [3],
Vicuna [9], and GPT-J [52]. Our goal is to scale methods from causal
abstraction to understand _how_ these models follow a particular instruction.

##  3 Methods

###  3.1 Background on Causal Abstraction

#### Causal Models

We represent black box networks and interpretable algorithms using causal
models that consist of variables that take on values according to causal
mechanisms. An intervention is an operation that edits some of the causal
mechanisms in a model. We denote the output of a causal model $\mathcal{M}$
provided an input $\mathbf{x}$ as $\mathcal{M}(\mathbf{x})$.

#### Interchange Intervention

Begin with a causal model $\mathcal{M}$ with input variables $\mathbf{S}$,
base input $\mathbf{b}$, source input settings $\\{\mathbf{s}_{j}\\}_{1}^{k}$,
and disjoint sets of target variables $\\{\mathbf{Z}_{j}\\}_{1}^{k}$. An
interchange intervention yields a new model
$\textsc{IntInv}(\mathcal{M},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{Z}_{j}\\}_{1}^{k})$
which is identical to $\mathcal{M}$ except the values for each set of target
variables $\mathbf{Z}_{j}$ are fixed to be the value they would have taken for
source input $\mathbf{s}_{j}$. We then denote the output for the base input
with
$\textsc{IntInv}(\mathcal{M},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{Z}_{j}\\}_{1}^{k})(\mathbf{b})$.

#### Distributed Interchange Intervention

Let $\mathbf{N}$ be a subset of variables in $\mathcal{M}$, the _target
variables_. Let $\mathbf{Y}$ be a vector space with orthogonal subspaces
$\\{\mathbf{Y}_{j}\\}_{1}^{k}$. Let $\mathbf{R}$ be an invertible function
$\mathbf{R}:\mathbf{N}\to\mathbf{Y}$. Write $\mathsf{Proj}_{\mathbf{Z}}$ for
the orthogonal projection operator of a vector in $\mathbf{Y}$ onto subspace
$\mathbf{Z}$. A distributed interchange intervention yields a new model
$\textsc{DII}(\mathcal{M},\mathbf{R},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{Y}_{j}\\}_{1}^{k})$
which is identical to $\mathcal{M}$ except the causal mechanisms have been
rewritten such that each subspace $\mathbf{Y}_{j}$ is fixed to be the value it
would take for source input $\mathbf{s}_{j}$.222Previous work often focusses
on zeroing-out representations or substituting with mean values [33; 54],
rather than the more general interchange operation. Specifically, the causal
mechanism for $\mathbf{N}$ is set to be

$F^{*}_{\mathbf{N}}(\mathbf{v})=\mathbf{R}^{-1}\bigg{(}\mathsf{Proj}_{\mathbf{Y}_{0}}\Big{(}\mathbf{R}\big{(}F_{\mathbf{N}}(\mathbf{v})\big{)}\Big{)}\\\
+\sum_{j=1}^{k}\mathsf{Proj}_{\mathbf{Y}_{j}}\Big{(}\mathbf{R}\big{(}F_{\mathbf{N}}(\mathcal{M}(\mathbf{s_{j}}))\big{)}\Big{)}\bigg{)}$
(1)

where $\mathbf{Y}_{0}=\mathbf{Y}\setminus\bigoplus_{j=1}^{k}\mathbf{Y}_{j}$.

We then denote the output for the base input as
$\textsc{DII}(\mathcal{M},\mathbf{R},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{Y}_{j}\\}_{0}^{k})(\mathbf{b})$.

#### Causal Abstraction and Alignment

We are licensed to claim an algorithm $\mathcal{A}$ is a faithful
interpretation of the network $\mathcal{N}$ if the causal mechanisms of the
variables in the algorithm abstract the causal mechanisms of neural
representations relative to a particular alignment. For our purposes, each
variable of a high-level model is aligned with a linear subspace in the vector
space formed by rotating a neural representation with an orthogonal matrix. An
alignment and task input–output behavior together define a partial function
$\tau$ that translates low-level variable settings to high-level variable
settings.

#### Approximate Causal Abstraction

Interchange intervention accuracy (IIA) is a graded measure of abstraction
that computes the proportion of aligned interchange interventions on the
algorithm and neural network that have the same output. The IIA for an
alignment $(\tau,\Pi)$ of high-level variables $X_{j}$ to orthogonal subspace
$\mathbf{Y}_{j}$ between an algorithm $\mathcal{A}$ and network $\mathcal{N}$
is

$\frac{1}{|\mathbf{Inputs}|^{k+1}}\sum_{\mathbf{b},\mathbf{s}_{1},\dots,\mathbf{s}_{k}\in\mathbf{Inputs}_{L}}\biggl{[}\textsc{DII}(\mathcal{N},\mathbf{R}^{\theta},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{Y}_{j}\\}^{k}_{1})(\mathbf{b})=\\\\[-12.91663pt]
\textsc{IntInv}(\mathcal{A},\\{\tau(\mathbf{s}_{j})\\}^{k}_{1},\\{{X}_{j}\\}^{k}_{1}))(\tau(\mathbf{b}))\biggl{]}$
(2)

###  3.2 Boundless Distributed Alignment Search

Distributed alignment search (DAS) is a method for learning an alignment
between interpretable causal variables of a model $\mathcal{C}$ and fixed
dimensionality linear subspaces of neural representations in a network
$\mathcal{N}$ using gradient descent [23] as shown in Figure 3. Specifically,
an orthogonal matrix $\mathbf{R}:\mathbf{N}\to\mathbf{Y}$ is trained to
maximize interchange intervention accuracy under an alignment from each
variable $X_{j}$ to fixed dimensionality linear subspaces $\mathbf{Y}_{j}$ of
the rotated vector space.

Boundless DAS is our extension of DAS that learns the dimensionality of the
orthogonal linear subspaces in an $d$ dimensional vector space $\mathbf{Y}$
using a method inspired by work in neural PDE [57]. Specifically, for each
high-level variable $X_{j}$, we introduce a learnable continuous boundary
index parameter $b_{j}$ that can take on a value between $0$ and $d$, where
$b_{0}=0$, $b_{j}<b_{j+1}$, and $b_{j}<d$ for all $j$. The boundary mask
$\mathbf{M}_{j}$ for the source is a vector with $d$ values between $0$ and
$1$ where the $k$-th element of the array is defined to be

$(\mathbf{M}_{j})_{k}=\texttt{sigmoid}\left(\frac{k-b_{j}}{\beta}\right)*\texttt{sigmoid}\left(\frac{b_{j+1}-k}{\beta}\right)$
(3)

where $\beta$ is a temperature that we anneal through training. As $\beta$
approaches 0, the masks $\mathbf{M}_{j}$ converge to binary-valued vectors
that together encode an orthogonal decomposition of $\mathbf{Y}$. For
simplicity, we restrict each $b_{j}$ to be a multiple of $b_{j-1}$ with a
factor of 2.

#### Weighted Distributed Interchange Intervention

Let $\mathbf{N}$ be a subset of variables in $\mathcal{M}$, the _target
variables_. Let $\mathbf{Y}$ be a vector space with $d$ dimensions and let
$\\{\mathbf{M}_{j}\\}_{1}^{k}$ vectors in $[0,1]^{d}$. Let
$\mathbf{R}:\mathbf{N}\to\mathbf{Y}$ be an invertible transformation. A
weighted distributed interchange intervention yields a new model
$\textsc{SoftDII}(\mathcal{M},\mathbf{R},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{M}_{j}\\}_{1}^{k})$
which is identical to $\mathcal{M}$ except the causal mechanisms have been
rewritten such that each source input $\mathbf{s}_{j}$ contributes to the
setting of $\mathbf{Y}$ in proportion to its mask $\mathbf{M}_{j}$.
Specifically, the causal mechanism for $\mathbf{N}$ is set to be

$F^{*}_{\mathbf{N}}(\mathbf{v})=\mathbf{R}^{-1}\bigg{(}(\mathbf{1}-\sum_{j=1}^{k}\mathbf{M}_{j})\circ\mathbf{R}\big{(}F_{\mathbf{N}}(\mathbf{v})\big{)}\\\
+\sum_{j=1}^{k}\Big{(}\mathbf{M}_{j}\circ\mathbf{R}\big{(}F_{\mathbf{N}}(\mathcal{M}(\mathbf{s_{j}})\big{)}\Big{)}\bigg{)}$
(4)

where $\circ$ is element wise multiplication and $\mathbf{1}$ is a $d$
dimensional vector where each element is $1$. We then denote the output for
the base input as
$\textsc{SoftDII}(\mathcal{M},\mathbf{R},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{M}_{j}\\}_{1}^{k})(\mathbf{b})$.

#### Boundless DAS

Given a base input $\mathbf{b}$ and source inputs
$\\{\mathbf{s}_{j}\\}_{1}^{k}$, we minimize the following objective to learn a
rotation matrix $\mathbf{R}^{\theta}$ and masks
$\\{\mathbf{M}^{\theta}_{j}\\}_{1}^{k}$

$\sum_{\mathbf{b},\mathbf{s}_{1},\dots,\mathbf{s}_{k}\in\mathbf{Inputs}_{L}}\mathsf{CE}\biggl{(}\textsc{SoftDII}(\mathcal{N},\mathbf{R}^{\theta},\\{\mathbf{s}_{j}\\}_{1}^{k},\\{\mathbf{M}^{\theta}_{j}\\}^{k}_{1})(\mathbf{b}),\textsc{IntInv}(\mathcal{A},\\{\tau(\mathbf{s}_{j})\\}^{k}_{1},\\{{X}_{j}\\}^{k}_{1}))(\tau(\mathbf{b}))\biggl{)}$
(5)

where $\mathsf{CE}$ is the cross entropy loss. We anneal $\beta$ throughout
training and our weighted interchange interventions become more and more
similar to unweighted interchange interventions. During evaluation we snap the
masks to be binary-valued to create an orthogonal decomposition of
$\mathbf{Y}$ where each high-level variable $X_{j}$ is aligned with a linear
subspace of $\mathbf{Y}$ picked out by the mask $\mathbf{M}_{j}$, with the
residual (unaligned) subspace being picked out by the mask
$(\mathbf{1}-\sum_{j=1}^{k}\mathbf{M}_{j})$.

##  4 Experiment

###  4.1 Price Tagging

We follow the approach in Fig. 1 by first assessing the ability of Alpaca to
execute specific actions based on the instructions provided in input.
Formally, the input to the model $\mathcal{M}$ is given an instruction $t_{i}$
(e.g., “correct the spelling of the word:”) followed by a test query input
$x_{i}$ (e.g., “aplpe”). We use $\mathcal{M}(t_{i},x_{i})$ to depict the model
generation $y_{p}$ given the instruction and the test query input. We can
evaluate model performance by comparing $y_{p}$ with the gold label $y$. We
focus on tasks with high model performance, to ensure that we have a known
behavioral pattern to explain.

The instruction prompt of the Price Tagging game follows the publicly released
template of the Alpaca (7B) model. The core instruction contains an the
English sentence:

Please say yes only if it costs between [X.XX] and [X.XX] dollars, otherwise
no.

followed by an input dollar amount [X.XX], where [X.XX] are random continuous
real numbers drawn with a uniform distribution from [0.00, 9.99]. The output
is a single token ‘Yes’ or ‘No’.

For instance, if the core instruction says Please say yes only if it costs
between [1.30] and [8.55] dollars, otherwise no., the answer would be “Yes” if
the input amount is “3.50 dollars” and “No” if the input is “9.50 dollars”. We
restrict the absolute difference between the lower bound and the upper bound
to be [2.50, 7.50] due to model errors outside these values – again we need
behavior to explain.

Figure 2: Four proposed high-level causal models for how Alpaca solves the
price tagging task. Intermediate variables are in red. All these models
perfectly sole the task. Figure 3: Aligned distributed interchange
interventions performed on the Alpaca model that is instructed to solve our
Price Tagging game. It aligns the boolean variable representing whether the
input amount is higher than the lower bound in the causal model. To train
Boundless DAS, we sample two training examples and then swap the intermediate
boolean value between them to produce a counterfactual output using our causal
model. In parallel, we swap the aligned dimensions of the neural
representations in rotated space. Lastly, we update our rotation matrix such
that our neural network has a more similar counterfactual behavior to the
causal model.

###  4.2 Hypothesized Causal Models

As shown in Figure 2, we have identified a set of human-interpretable high-
level causal models, with alignable intermediate causal variables, that would
solve this task with 100% performance:

  * •

Left Boundary: This model has one high-level boolean variable representing
whether the input amount is higher than the lower bound, and an output node
incorporating whether the input amount is also lower than the high bound.

  * •

Left and Right Boundary: The previous model is sub-optimal in only abstracting
one of the boundaries. In this model, we have two high-level boolean variables
representing whether the input amount is higher than the lower bound and lower
than the higher bound. We take a conjunction of these boolean variables to
predict the output.

  * •

Mid-point Distance: We calculate the mid-point of the lower and upper bounds
(e.g., the mid-point of “3.50” and “8.50” is “6.00), and then we take the
absolute distance between the input dollar amount and the mid-point as $a$. We
then calculate one-half of the bounding bracket length (e.g., the bracket
length for “3.50” and “8.50” is “5.00”) as $b$. We predict output “Yes” if
$a\leq b$, otherwise “No”. We align only with the mid-point variable.

  * •

Bracket Identity: This model represents the lower and upper bound in a single
interval variable and passes this information to the output node. We predict
the output as “Yes” if the input amount within the interval, otherwise “No”.

#### Model Architecture

Our target model is the Alpaca (7B) [46], an off-the-shelf instruct-tuned
LLaMA model. It is a Transformer-based decoder-only autoregressive trained
language model with 32 layers and 32 attention heads. It has a hidden
dimension in size of 4096, which is also the dimension of our rotation matrix
which is applied for each token representation. In total, the rotation matrix
contains 16.8M parameters and has size 4096 $\times$ 4096.

#### Alignment Process

We separately train Boundless DAS with test query input token representations
(i.e., starting from the token of the first input digit till the last token in
the prompt) in a set of selected 7 layers: $\\{0,5,10,15,20,25,30\\}$. We also
train Boundless DAS on the token before the first digit as a control condition
where nothing should be expected to be aligned. Instead of interchanging with
multiple source input settings $\\{\mathbf{s}_{j}\\}_{1}^{k}$ at the same
time, we interchange with a single source at a time for simplicity, while
allowing multiple causal variables to be aligned across examples. We run each
experiment with three distinct random seeds. Since the global optimum of the
Boundless DAS objective corresponds to the best attainable alignment, but SGD
may be trapped by local optima, we report the best-performing seed in terms of
IIA. Figure 3 provides an overview of these analyses.

#### Evaluation Metric

To evaluate models, we use Interchange Intervention Accuracy (IIA) as defined
in Eqn. 2. IIA is bounded between 0.0 and 1.0. For most of the experiments,
the lower bound of IIA is at the chance (0.5) given the distribution of our
output labels. IIA can occasionally go above the model’s task performance,
when the interchange interventions puts the model in a better state, but for
the most part IIA is constrained by task accuracy.

Figure 4: Interchange Intervention Accuracy (IIA) for four different alignment
proposals. The Alpaca model achieves 85% task accuracy. The higher the number
is, the more faithful the alignment is. We color each cell by scaling IIA
using the model’s task performance as the upper bound and a dummy classifier
(predicting the most frequent label) as the lower bound. These results
indicate that the top two are highly accurate hypotheses about how Alpaca
solves the task, whereas the bottom two are inaccurate in this sense.

###  4.3 Boundless DAS Results

Figure 4 shows our main results, given in terms of IIA across our four
hypothesized algorithmic models (Figure 2). The results show very clearly that
‘Left Boundary’ and ‘Left and Right Right Boundary’ (top panels) are highly
accurate hypotheses about how Alpaca solves the task. For them, IIA is at or
above task performance (0.84), and intermediate variable representations are
localized in systematically arranged positions. By contrast, ‘Mid-point
Distance’ and ‘Bracket Identity’ (bottom panels) are inaccurate hypotheses
about Alpaca’s processing, with IIA peaking at around 0.72.

These findings suggest that, when solving our reasoning task, Alpaca
internally follows our first two high-level models by representing causal
variables that align with boundary checks for both left and right boundaries.
Interestingly, heatmaps on the top two rows also show a pattern of higher
scores around the bottom left and upper right and close to zero scores in
other positions. This is significantly different from the other two alignments
where, although some positions are highlighted (e.g., the representations for
the last token), all the positions receive non-zero scores. In short, the
accurate hypotheses correspond to highly structured IIA patterns, and the
inaccurate ones do not.

Additionally, alignments are better when the model post-processes the query
input with 1–2 additional steps: accuracy is higher at position 75 compared to
all previous positions. Surprisingly, there exist _bridging tokens_ (positions
76–79) where accuracy suddenly drops compared to earlier tokens, which
suggests that these representations have weaker causal effects on the model
output. In other words, boundary-check variables are fully extracted around
position 75, level 10, and are later copied into activations for the final
token before responding. By comparing the heatmaps on the top row, our
findings suggest that aligning multiple variables at the same time poses a
harder alignment process, in that it lowers scores slightly across multiple
positions.

###  4.4 Interchange Interventions with (In-)Correct Inputs

IIA is highly constrained by task performance, and thus we expect it to be
much lower for inputs that the model gets wrong. To verify this, we
constructed an evaluation set containing 1K examples that the model gets wrong
and evaluated our ‘Left and Right Boundary’ hypothesis on this subset. Table 1
reports these results in terms of max IIA and the correlation of the IIA
values with those obtained in our main experiments. As expected, IIA drops
significantly. However, two things stand out: IIA is far above task
performance (which is 0.0 by design now), and the correlation with the
original IIA map remains high. These findings suggest that the model is using
the same internal mechanisms to process these cases, and so it is likely that
is is narrowly missing the correct output predictions.

We also expect IIA to be higher if we focus only on cases that the model gets
correct. Table 1 confirms this expectationm using ‘Left and Right Boundary’ as
representative. IIA${}_{\text{max}}$ is improved from 0.86 to 0.88, and the
correlation with the main results is essentially perfect.

Full heatmaps for these analyses are given in Appendix A.4.

Experiment | Task Acc. | IIA${}_{\textbf{max}}$ | Correlation |  Variance ($\times 10^{-2}$)
---|---|---|---|---
Left Boundary ($\clubsuit$) | 0.85 | 0.90 | 1.00 | 2.01
Left and Right Boundary ($\varheartsuit$) | 0.85 | 0.86 | 1.00 | 1.14
Mid-point Distance | 0.85 | 0.70 | 1.00 | 0.04
Bracket Identity | 0.85 | 0.72 | 1.00 | 0.08
Correct Only | 1.00† | 0.88 | 0.99 ($\varheartsuit$) | 1.47
Incorrect Only | 0.00† | 0.71 | 0.84 ($\varheartsuit$) | 0.36
New Bracket (Seen) | 0.94 | 0.94 | 0.97 ($\clubsuit$) | 3.02
New Bracket (Unseen) | 0.95 | 0.95 | 0.94 ($\clubsuit$) | 1.66
Irrelevant Context (1) | 0.84 | 0.83 | 0.97 ($\varheartsuit$) | 0.94
Irrelevant Context (2) | 0.84 | 0.85 | 0.98 ($\varheartsuit$) | 1.04
Sibling Instructions | 0.84 | 0.83 | 0.87 ($\varheartsuit$) | 0.93
\+ _exclude_ top right | 0.84 | 0.83 | 0.92 ($\varheartsuit$) | 0.95

Table 1: Summary results for all experiments with task performance as accuracy
(range $[0,1]$), maximal interchange intervention accuracy (IIA) (range
$[0,1]$) across all positions and layers, Pearson correlations of IIA between
two distributions (compared to $\clubsuit$ or $\varheartsuit$; range
$[-1,1]$), and variance of IIA within a single experiment across all positions
and layers. †This is empirical task performance on the evaluation dataset for
this experiment.

###  4.5 Do Alignments Robustly Generalize to Unseen Instructions and Inputs?

One might worry that our positive results are highly dependent on the specific
input–output pairs we have chosen. We now seek to address this concern by
asking whether the causal roles (i.e., alignments) found using Boundless DAS
in one setting are preserved in new settings. This is crucial, as it tells how
robustly the causal model is realized in the neural network.

#### Generalizing Across Two Different Instructions

Here, we assess whether the learned alignments for the ‘Left Boundary’ causal
model transfer between different specific price brackets in the instruction.
To do this, we retrain Boundless DAS for the high-level model with a fixed
instruction that says “between 5.49 dollars and 8.49 dollars”. Then, we fix
the learned rotation matrix and evaluate with another instruction that says
“between 2.51 dollars and 5.51 dollars”. For both, Alpaca is successful at the
task, with around 94% accuracy. Our hypothesis is that if the found alignment
of the high-level variable is robust, it should transfer between these two
settings, as the aligning variable is a boolean-type variable which is
potentially agnostic to the specific comparison price.

Table 1 gives our findings in the ‘New Bracket’ rows. Boundless DAS is able to
find a good alignment for the training bracket with an IIA${}_{\text{max}}$
that is about the same as the task performance at 94%. For our unseen bracket,
the alignments also hold up extremely well, with no drop in
IIA${}_{\text{max}}$. For both cases, the found alignments also highly
correlate with the counterpart of our main experiment.

#### Generalizing with Inserted Context

Recent work has shown that language models are sensitive to irrelevant context
[42]. To address this concern, we add prefix strings to the input instructions
and evaluate how ‘Left and Right Boundary’ alignment transfers. Our
‘Irrelevant Context’ prefixes are the short string “Price Tagging game!” and
the longer string “Fruitarian Frogs May Be Doing Flowers a Favor” (from New
York Times April 28, 2023, the first post title in the ‘America’ column). For
both cases, the model achieves 84% task performance, which is slightly lower
than the average task performance. Nonetheless, the results in Table 1 suggest
the found alignments transfer surprisingly well here, with only 1% drop in
IIA${}_{\text{max}}$. Meanwhile, the IIA distribution across positions and
layers is highly correlated with our base experiment. Thus, our method seems
to identify causal structure that is robust to changes in irrelevant task
details _and_ position in the input string.

#### Generalizing to Modified Outputs

We further test whether the alignments found for instruction with a template
“Say yes …, otherwise no” can generalize to a new instruction with a template
“Say True …, otherwise False”. If there are indeed latent representations of
our aligning causal variables, the found alignments should persist and should
be agnostic to the output format.

Table 1 shows that this holds: the learned alignments actually transfer across
these two sibling instructions, with a minimum 1% drop in
IIA${}_{\text{max}}$. In addition, the correlation increases 6% when excluding
the top right corner (top 3 layers in the last position), as seen in the final
row of the table. This result is further substantiated in Appendix A.4, where
we see that IIA drops near the top right. This indicates a late fusion of
output classes and working representations, leading to different computations
for the new instruction only close to the output.

Figure 5: Learned boundary width for intervention site and in-training
evaluation interchange intervention accuracy (IIA) for two groups of data: (1)
aligned group where the boundary does not shrink to 0 at the end of the
training; (2) unaligned group where the boundary does shrink to 0 at the end
of the training. 1 on the y-axis means either 100% accuracy for IIA, or the
variable is occupying half of the hidden representation for the boundary
width.

###  4.6 Boundary Learning Dynamics

Boundless DAS automatically learns the intervention site boundaries (Section
3.2). It is important to confirm that we do not give up alignment accuracy by
optimizing boundaries, and to explore the dimensionality actually needed to
represent the abstract variables. We sample 100 experiment runs from our
experiment pool and create two groups: (1) the aligned group where the
boundary does not shrink to 0 at the end of the training; (2) the unaligned
group where the boundary does shrink to 0 at the end of the training. The
unaligned group is for those instances where there is no significant causal
role of representations found by our method. We plot how boundary width and
in-training evaluation IIA vary throughout training. Figure 5 shows that the
width for aligned cases shrink to 10–20% of one-half of the representation
space for a single aligning variable. IIA performance converges quickly for
aligned cases and, crucially, maintains stable performance despite shrinking
dimensionality. This suggests that only a small fraction of the whole
representation space is needed for aligning a causal variable, and that these
representations can be efficiently identified by Boundless DAS.

##  5 Analytic Strengths and Limitations

Explanation methods for AI should be judged by the degree to which they can,
in principle, identify models’ true internal causal mechanisms. If we reach
100% IIA for a high-level model using Boundless DAS we can assert that we have
positively identified a correct causal structure in the model (though perhaps
not the only correct abstraction). This follows directly from the formal
results of Geiger et al. [22]. On the other hand, failure to find causal
structure with Boundless DAS is not a proof that such structure is missing,
for two reasons. First, Boundless DAS explores a very flexible and large
hypothesis space, but it is not completely exhaustive. For instance a set of
variables represented in a highly non-linear way might be missed. Second, we
are limited by the space of causal models we think to test ourselves.

For models where task accuracy is below 100%, it is still possible for IIA to
reach 100% if the high-level causal model _explains errors_ of the low-level
model. However, in cases like those studied here where our high-level model
matches the idealized task but our language model does not completely do so,
we expect Boundless DAS to find only partial support for causal structure even
if we have identified the ideal set of hypotheses and searched optimally. This
too follows from the results of Geiger et al. [22]. However, as we have seen
in the results above, Boundless DAS can find structure even where the model’s
task performance is low, since task performance can be shaped by factors like
a suboptimal generation method or the rigidity of the assessment metric.
Future work must tighten this connection by modeling errors of the language
model in more detail.

##  6 Conclusion

We introduce Boundless DAS, a novel and effective method for scaling alignment
search of causal structure in LLMs to billions of parameters. Using Boundless
DAS, we find that Alpaca, off-the-shelf, solves a simple numerical reasoning
problem in a human-interpretable way. Additionally, we address one of the main
concerns around interpretability tools developed for LLMs – whether found
alignments generalize across different settings. We rigorously study this by
evaluating found alignments under several changes to inputs and instructions.
Our findings indicate robust and interpretable algorithmic structure. Our
framework is generic for any LLMs and is released to the public. We hope that
this marks a step forward in terms of understanding the internal causal
mechanisms behind the massive LLMs that are at the center of so much work in
AI.

##  7 Acknowledgements

We thank Alex Tamkin, John Hewitt, Jing Huang, and other members of the
Stanford NLP group for useful discussion. This research is supported in part
by faculty research awards to CP from Amazon, Meta AI, and the Stanford
Institute for Human-Centered Artificial Intelligence (HAI).

## References

  * Abnar and Zuidema [2020] Samira Abnar and Willem Zuidema.  Quantifying attention flow in transformers.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 4190–4197, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.385.  URL https://aclanthology.org/2020.acl-main.385.
  * Abraham et al. [2022] Eldar D Abraham, Karel D’Oosterlinck, Amir Feder, Yair Gat, Atticus Geiger, Christopher Potts, Roi Reichart, and Zhengxuan Wu.  Cebab: Estimating the causal effects of real-world concepts on nlp model behavior.  _Advances in Neural Information Processing Systems_ , 35:17582–17596, 2022.
  * Andonian et al. [2021] Alex Andonian, Quentin Anthony, Stella Biderman, Sid Black, Preetham Gali, Leo Gao, Eric Hallahan, Josh Levy-Kramer, Connor Leahy, Lucas Nestler, Kip Parker, Michael Pieler, Shivanshu Purohit, Tri Songz, Wang Phil, and Samuel Weinbach.  GPT-NeoX: Large Scale Autoregressive Language Modeling in PyTorch, 8 2021.  URL https://www.github.com/eleutherai/gpt-neox.
  * Beckers and Halpern [2019] Sander Beckers and Joseph Halpern.  Abstracting causal models.  In _AAAI Conference on Artificial Intelligence_ , 2019.
  * Beckers et al. [2019] Sander Beckers, Frederick Eberhardt, and Joseph Y. Halpern.  Approximate causal abstractions.  In _Proceedings of The 35th Uncertainty in Artificial Intelligence Conference_ , 2019.
  * Belinkov and Glass [2019] Yonatan Belinkov and James R. Glass.  Analysis methods in neural language processing: A survey.  _Trans. Assoc. Comput. Linguistics_ , 7:49–72, 2019.  doi: 10.1162/tacl\\_a\\_00254.  URL https://doi.org/10.1162/tacl_a_00254.
  * Chan et al. [2022] Lawrence Chan, Adrià Garriga-Alonso, Nicholas Goldowsky-Dill, Ryan Greenblatt, Jenny Nitishinskaya, Ansh Radhakrishnan, Buck Shlegeris, and Nate Thomas.  Causal scrubbing: a method for rigorously testing interpretability hypotheses, 2022.
  * Chi et al. [2020] Ethan A. Chi, John Hewitt, and Christopher D. Manning.  Finding universal grammatical relations in multilingual BERT.  In _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics_ , pages 5564–5577, Online, July 2020. Association for Computational Linguistics.  doi: 10.18653/v1/2020.acl-main.493.  URL https://www.aclweb.org/anthology/2020.acl-main.493.
  * Chiang et al. [2023] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, Ion Stoica, and Eric P. Xing.  Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023.  URL https://vicuna.lmsys.org.
  * Christiano et al. [2017] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei.  Deep reinforcement learning from human preferences.  _Advances in neural information processing systems_ , 30, 2017.
  * Chung et al. [2022] Hyung Won Chung, Le Hou, Shayne Longpre, Barret Zoph, Yi Tay, William Fedus, Eric Li, Xuezhi Wang, Mostafa Dehghani, Siddhartha Brahma, et al.  Scaling instruction-finetuned language models.  _arXiv preprint arXiv:2210.11416_ , 2022.
  * Clark et al. [2019] Kevin Clark, Urvashi Khandelwal, Omer Levy, and Christopher D. Manning.  What does BERT look at? an analysis of BERT’s attention.  In _Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_ , pages 276–286, Florence, Italy, August 2019. Association for Computational Linguistics.  doi: 10.18653/v1/W19-4828.  URL https://aclanthology.org/W19-4828.
  * Coenen et al. [2019] Andy Coenen, Emily Reif, Ann Yuan, Been Kim, Adam Pearce, Fernanda Viégas, and Martin Wattenberg.  Visualizing and measuring the geometry of bert, 2019.
  * Conneau et al. [2018] Alexis Conneau, German Kruszewski, Guillaume Lample, Loïc Barrault, and Marco Baroni.  What you can cram into a single $&!#* vector: Probing sentence embeddings for linguistic properties.  In _Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)_ , pages 2126–2136, Melbourne, Australia, July 2018. Association for Computational Linguistics.  doi: 10.18653/v1/P18-1198.  URL https://www.aclweb.org/anthology/P18-1198.
  * Elazar et al. [2020] Yanai Elazar, Shauli Ravfogel, Alon Jacovi, and Yoav Goldberg.  Amnesic probing: Behavioral explanation with amnesic counterfactuals.  In _Proceedings of the 2020 EMNLP Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP_. Association for Computational Linguistics, November 2020.  doi: 10.18653/v1/W18-5426.
  * Elazar et al. [2022] Yanai Elazar, Nora Kassner, Shauli Ravfogel, Amir Feder, Abhilasha Ravichander, Marius Mosbach, Yonatan Belinkov, Hinrich Schütze, and Yoav Goldberg.  Measuring causal effects of data statistics on language model’s ’factual’ predictions.  _CoRR_ , abs/2207.14251, 2022.  doi: 10.48550/arXiv.2207.14251.  URL https://doi.org/10.48550/arXiv.2207.14251.
  * Elhage et al. [2021] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  A mathematical framework for transformer circuits.  _Transformer Circuits Thread_ , 2021.  https://transformer-circuits.pub/2021/framework/index.html.
  * Feder et al. [2021] Amir Feder, Nadav Oved, Uri Shalit, and Roi Reichart.  CausaLM: Causal Model Explanation Through Counterfactual Language Models.  _Computational Linguistics_ , pages 1–54, 05 2021.  ISSN 0891-2017.  doi: 10.1162/coli_a_00404.  URL https://doi.org/10.1162/coli_a_00404.
  * Geiger et al. [2020] Atticus Geiger, Kyle Richardson, and Christopher Potts.  Neural natural language inference models partially embed theories of lexical entailment and negation.  In _Proceedings of the Third BlackboxNLP Workshop on Analyzing and Interpreting Neural Networks for NLP_ , pages 163–173, Online, November 2020\. Association for Computational Linguistics.  doi: 10.18653/v1/2020.blackboxnlp-1.16.  URL https://www.aclweb.org/anthology/2020.blackboxnlp-1.16.
  * Geiger et al. [2021] Atticus Geiger, Hanson Lu, Thomas Icard, and Christopher Potts.  Causal abstractions of neural networks.  In _Advances in Neural Information Processing Systems_ , volume 34, pages 9574–9586, 2021.  URL https://papers.nips.cc/paper/2021/hash/4f5c422f4d49a5a807eda27434231040-Abstract.html.
  * Geiger et al. [2022] Atticus Geiger, Zhengxuan Wu, Hanson Lu, Josh Rozner, Elisa Kreiss, Thomas Icard, Noah Goodman, and Christopher Potts.  Inducing causal structure for interpretable neural networks.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pages 7324–7338. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/geiger22a.html.
  * Geiger et al. [2023a] Atticus Geiger, Christopher Potts, and Thomas Icard.  Causal abstraction for faithful model interpretation.  Ms., Stanford University, 2023a.  URL https://arxiv.org/abs/2301.04709.
  * Geiger et al. [2023b] Atticus Geiger, Zhengxuan Wu, Christopher Potts, Thomas Icard, and Noah D. Goodman.  Finding alignments between interpretable causal variables and distributed neural representations.  Ms., Stanford University, 2023b.  URL https://arxiv.org/abs/2303.02536.
  * Hewitt and Manning [2019] John Hewitt and Christopher D Manning.  A structural probe for finding syntax in word representations.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 4129–4138, 2019.
  * Huang et al. [2022] Jing Huang, Zhengxuan Wu, Kyle Mahowald, and Christopher Potts.  Inducing character-level structure in subword-based language models with Type-level Interchange Intervention Training.  Ms., Stanford University and UT Austin, 2022.  URL https://arxiv.org/abs/2212.09897.
  * Kingma and Ba [2015] Diederick P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  In _International Conference on Learning Representations (ICLR)_ , 2015.
  * Kojima et al. [2022] Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa.  Large language models are zero-shot reasoners, 2022.
  * Li et al. [2021] Belinda Z. Li, Maxwell Nye, and Jacob Andreas.  Implicit representations of meaning in neural language models.  In _Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , pages 1813–1827, Online, August 2021. Association for Computational Linguistics.  doi: 10.18653/v1/2021.acl-long.143.  URL https://aclanthology.org/2021.acl-long.143.
  * Lipton [2018] Zachary C Lipton.  The mythos of model interpretability: In machine learning, the concept of interpretability is both important and slippery.  _Queue_ , 16(3):31–57, 2018.
  * Lovering and Pavlick [2022] Charles Lovering and Ellie Pavlick.  Unit testing for concepts in neural networks.  _Transactions of the Association for Computational Linguistics_ , 10:1193–1208, 2022.  doi: 10.1162/tacl_a_00514.  URL https://aclanthology.org/2022.tacl-1.69.
  * Lundberg and Lee [2017] Scott Lundberg and Su-In Lee.  A unified approach to interpreting model predictions, 2017.
  * Manning et al. [2020] Christopher D. Manning, Kevin Clark, John Hewitt, Urvashi Khandelwal, and Omer Levy.  Emergent linguistic structure in artificial neural networks trained by self-supervision.  _Proceedings of the National Academy of Sciences_ , 117(48):30046–30054, 2020.  ISSN 0027-8424.  doi: 10.1073/pnas.1907367117.  URL https://www.pnas.org/content/117/48/30046.
  * Meng et al. [2022] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov.  Locating and editing factual associations in GPT.  _Advances in Neural Information Processing Systems_ , 36, 2022.
  * Olah et al. [2020] Chris Olah, Nick Cammarata, Ludwig Schubert, Gabriel Goh, Michael Petrov, and Shan Carter.  Zoom in: An introduction to circuits.  _Distill_ , 2020.  doi: 10.23915/distill.00024.001.  https://distill.pub/2020/circuits/zoom-in.
  * Olsson et al. [2022] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah.  In-context learning and induction heads.  _Transformer Circuits Thread_ , 2022.  https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html.
  * Ouyang et al. [2022] Long Ouyang, Jeff Wu, Xu Jiang, Diogo Almeida, Carroll L. Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, John Schulman, Jacob Hilton, Fraser Kelton, Luke Miller, Maddie Simens, Amanda Askell, Peter Welinder, Paul Christiano, Jan Leike, and Ryan Lowe.  Training language models to follow instructions with human feedback.  _Advances in Neural Information Processing Systems_ , 36, 2022.
  * Ravfogel et al. [2020] Shauli Ravfogel, Yanai Elazar, Hila Gonen, Michael Twiton, and Yoav Goldberg.  Null it out: Guarding protected attributes by iterative nullspace projection.  In Dan Jurafsky, Joyce Chai, Natalie Schluter, and Joel R. Tetreault, editors, _Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, ACL 2020, Online, July 5-10, 2020_ , pages 7237–7256. Association for Computational Linguistics, 2020.  doi: 10.18653/v1/2020.acl-main.647.  URL https://doi.org/10.18653/v1/2020.acl-main.647.
  * Ribeiro et al. [2016] Marco Tulio Ribeiro, Sameer Singh, and Carlos Guestrin.  "why should i trust you?": Explaining the predictions of any classifier.  In _Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’16, page 1135–1144, New York, NY, USA, 2016. Association for Computing Machinery.  ISBN 9781450342322.  doi: 10.1145/2939672.2939778.  URL https://doi.org/10.1145/2939672.2939778.
  * Rogers et al. [2020] Anna Rogers, Olga Kovaleva, and Anna Rumshisky.  A primer in BERTology: What we know about how BERT works.  _Transactions of the Association for Computational Linguistics_ , 8:842–866, 2020.  doi: 10.1162/tacl_a_00349.  URL https://aclanthology.org/2020.tacl-1.54.
  * Rubenstein et al. [2017] Paul K. Rubenstein, Sebastian Weichwald, Stephan Bongers, Joris M. Mooij, Dominik Janzing, Moritz Grosse-Wentrup, and Bernhard Schölkopf.  Causal consistency of structural equation models.  In _Proceedings of the 33rd Conference on Uncertainty in Artificial Intelligence (UAI)_ , 2017.
  * Saphra and Lopez [2019] Naomi Saphra and Adam Lopez.  Understanding learning dynamics of language models with SVCCA.  In _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers)_ , pages 3257–3267, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.  doi: 10.18653/v1/N19-1329.  URL https://www.aclweb.org/anthology/N19-1329.
  * Shi et al. [2023] Freda Shi, Xinyun Chen, Kanishka Misra, Nathan Scales, David Dohan, Ed Chi, Nathanael Schärli, and Denny Zhou.  Large language models can be easily distracted by irrelevant context, 2023\.
  * Shrikumar et al. [2017] Avanti Shrikumar, Peyton Greenside, and Anshul Kundaje.  Learning important features through propagating activation differences.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pages 3145–3153. JMLR. org, 2017.
  * Simonyan et al. [2014] K Simonyan, A Vedaldi, and A Zisserman.  Deep inside convolutional networks: visualising image classification models and saliency maps.  In _Proceedings of the International Conference on Learning Representations (ICLR)_. ICLR, 2014.
  * Sundararajan et al. [2017] Mukund Sundararajan, Ankur Taly, and Qiqi Yan.  Axiomatic attribution for deep networks.  In _Proceedings of the 34th International Conference on Machine Learning-Volume 70_ , pages 3319–3328. JMLR. org, 2017.
  * Taori et al. [2023] Rohan Taori, Ishaan Gulrajani, Tianyi Zhang, Yann Dubois, Xuechen Li, Carlos Guestrin, Percy Liang, and Tatsunori B. Hashimoto.  Stanford alpaca: An instruction-following llama model.  https://github.com/tatsu-lab/stanford_alpaca, 2023.
  * Tenney et al. [2019] Ian Tenney, Dipanjan Das, and Ellie Pavlick.  BERT rediscovers the classical NLP pipeline.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 4593–4601, Florence, Italy, July 2019. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/P19-1452.
  * Touvron et al. [2023] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Uesato et al. [2022] Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins.  Solving math word problems with process- and outcome-based feedback, 2022\.
  * Vig et al. [2020] Jesse Vig, Sebastian Gehrmann, Yonatan Belinkov, Sharon Qian, Daniel Nevo, Yaron Singer, and Stuart Shieber.  Investigating gender bias in language models using causal mediation analysis.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 12388–12401. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/92650b2e92217715fe312e6fa7b90d82-Paper.pdf.
  * Voita et al. [2019] Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov.  Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned.  In _Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics_ , pages 5797–5808, Florence, Italy, July 2019. Association for Computational Linguistics.  doi: 10.18653/v1/P19-1580.  URL https://aclanthology.org/P19-1580.
  * Wang and Komatsuzaki [2021] Ben Wang and Aran Komatsuzaki.  GPT-J-6B: A 6 Billion Parameter Autoregressive Language Model.  https://github.com/kingoflolz/mesh-transformer-jax, May 2021.
  * Wang et al. [2022a] Boshi Wang, Sewon Min, Xiang Deng, Jiaming Shen, You Wu, Luke Zettlemoyer, and Huan Sun.  Towards understanding chain-of-thought prompting: An empirical study of what matters.  _arXiv preprint arXiv:2212.10001_ , 2022a.
  * Wang et al. [2023] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt.  Interpretability in the wild: a circuit for indirect object identification in gpt-2 small.  In _Proceedings of the International Conference on Learning Representations (ICLR)_. ICLR, 2023.
  * Wang et al. [2022b] Yizhong Wang, Yeganeh Kordi, Swaroop Mishra, Alisa Liu, Noah A Smith, Daniel Khashabi, and Hannaneh Hajishirzi.  Self-instruct: Aligning language model with self generated instructions.  _arXiv preprint arXiv:2212.10560_ , 2022b.
  * Wei et al. [2022] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le.  Finetuned language models are zero-shot learners.  In _International Conference on Learning Representations_ , 2022.
  * Wu et al. [2022a] Tailin Wu, Takashi Maruyama, and Jure Leskovec.  Learning to accelerate partial differential equations via latent global evolution.  _Advances in Neural Information Processing Systems_ , 36, 2022a.
  * Wu et al. [2022b] Zhengxuan Wu, Karel D’Oosterlinck, Atticus Geiger, Amir Zur, and Christopher Potts.  Causal Proxy Models for concept-based model explanations.  arXiv:2209.14279, 2022b.  URL https://arxiv.org/abs/2209.14279.
  * Wu et al. [2022c] Zhengxuan Wu, Atticus Geiger, Joshua Rozner, Elisa Kreiss, Hanson Lu, Thomas Icard, Christopher Potts, and Noah Goodman.  Causal distillation for language models.  In _Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies_ , pages 4288–4295, 2022c.
  * Zeiler and Fergus [2014] Matthew D Zeiler and Rob Fergus.  Visualizing and understanding convolutional networks.  In _European conference on computer vision_ , pages 818–833. Springer, 2014.

##  Appendix A Appendix

###  A.1 Price Tagging Game Experiment

Figure 6: Instruction template for the Price Tagging game. We have three
variables inside the prompt: lower bound amount, high bound amount as well as
the input amount where each number has two decimal points but bounded with
[0.00, 9.99].

Figure 6 shows our templates for our Price Tagging game Experiment. This
format is enforced by the instruction tuning template of the Alpaca model.

###  A.2 Training Details

#### Training Setups

During training for Boundless DAS, model weights are frozen except for the
rotation matrix. To train the rotation matrix, we use Adam Kingma and Ba
[2015] as our optimizer with an effective batch size of 64 and a learning rate
of $10^{-3}$. We use a different learning rate $10^{-2}$ for our boundary
indices for quicker convergence. We train our rotation matrix for 3 epochs and
evaluate with a separate evaluation set for each 200 training steps, and we
save the best rotation matrix evaluated on the evaluation set during training
and test on another hold-out testing set. Our training set has 20K examples.
Our in-training evaluation is limited to 200 examples for a quicker training
time, and the hold-out testing dataset has 1K examples. These datasets are
generated on the fly for different causal models guarded by our random seeds.
The rotation matrix is trained by using the orthogonalized parameterization
provided by the torch
library333pytorch.org/docs/stable/generated/torch.nn.utils.parametrizations.orthogonal.html.
For boundary learning, we anneal temperature from 50.0 to 0.10 with a step
number equal to the total training step, and temperature steps with gradient
backpropagation. Each alignment experiment is enabled with bfloat16 and can
fit within a single A100 GPU. Each alignment experiment takes about 1.5 hrs.
To finish all of our experiments, it takes roughly 2 weeks of run-time with 2
$\times$ A100 nodes with 8 GPUs each. We run experiments with 3 distinct
random seeds and report the best result using IIA as our metrics.

#### Vanilla Task Performance Remarks

To ensure the model we are aligning behaviorally achieves good task
performance, we evaluate the model task performance with 1K sampled triples of
values more than 300 times. The Alpaca (7B) model achieves 85% for randomly
drawn triples of numbers on average. We want to emphasize that it is really
hard to find a good reasoning task with intermediate steps (e.g., tasks like
“Return yes of it costs more than 2.50 dollars” do not require an intermediate
step) that 7B-30B LLMs can solve. We test different reasoning tasks (e.g.,
math, logic, natural QA, code, etc..) with instruct-tuned LLaMA Touvron et al.
[2023], Dolly444huggingface.co/databricks/dolly-v2-12b, StableLM Andonian et
al. [2021], Vicuna Chiang et al. [2023], GPT-J Wang and Komatsuzaki [2021] and
mostly struggled to find tasks they could do well.

###  A.3 Counterfactual Dataset Generation

To train Boundless DAS, we need to generate counterfactual datasets where each
example contains a pair of inputs as base and source, an intervention
indicator (i.e., what variable we are intervening), and a counterfactual
label. For instance, if we are training Boundless DAS for our Boundary Check
(Left Only) high-level causal model, we may sample the base example with the
bracket as [2.50, 7.50] and the input amount of 1.50 dollars, and the source
example with the bracket as [3.50, 8.50] and the input amount of 9.50 dollars.
If we intervene on the variable whether the input variable is higher than the
lower bound, and we swap the variable from the second example into the first
example, we will have a counterfactual label as “Yes”. For each high-level
causal model, we follow the same steps to construct the counterfactual
dataset.

###  A.4 Interchange Intervention Accuracy

Figure 7 to Figure 10 show breakdown of results presented in Table 1.

###  A.5 Pseudocode for Boundless DAS

Boundless DAS is a generic framework and can be implemented for models with
different architectures. Here we provide a pseudocode snippet for a decoder-
only model.

⬇

1def sigmoid_boundary_mask(population, boundary_x, boundary_y, temperature):

2 return torch.sigmoid((population - boundary_x) / temperature) * \

3 torch.sigmoid((boundary_y - population) / temperature)

4

5class AlignableDecoderLLM(DecoderLLM):

6 # overriding existing forward function

7 def forward(hidden_states, source_hiddne_states):

8 for idx, decoder_layer in enumerate(self.layers):

9 if idx == self.aligning_layer_idx:

10 # select aligning token reprs

11 aligning_hidden_states = hidden_states[:,start_idx:end_idx]

12 # rotate

13 rotated_hidden_states = rotate(aligning_hidden_states)

14 # interchange intervention with learned boundaries

15 boundary_mask = sigmoid_boundary_mask(

16 torch.arange(0, self.hidden_size),

17 self.learned_boundaries,

18 self.temperature

19 )

20 rotated_hidden_states = (1. - boundary_mask) * rotated_hidden_states +

21 boundary_mask * source_hidden_states

22 # unrotate

23 hidden_states[:,start_idx:end_idx] = unrotate(rotated_hidden_states)

24 hidden_states = self.layers[idx](hidden_states)

25 return hidden_states

A simplified version of Boundless DAS with pseudocode for a generic decoder-
only model.

###  A.6 Common Questions

In this section, we answer common questions that may be raised while reading
this report.

> _Can Boundless DAS work with 175B models?_

Yes, it does. Taking the recently released Bloom-176B as an example, the
hidden dimension is 14,336, which results in a rotation matrix containing 205M
parameters. This is possible if the model is sharded and the rotation matrix
is put into a single GPU for training. This is only possible given the fact
that Boundless DAS is now implemented as a token-level alignment tool. If the
alignment is for a whole layer (i.e., token representations are concatenated),
it is currently impossible. Worth to note that training the rotation matrix
consumes more resources than training a simple weight matrix, as the
orthogonalization process in torch requires matrix exponential requiring to
save multiple copies of the matrix.

> _Current paper limits to a set of tasks from the same family. How well will
> the findings in the paper generalize to other scenarios?_

Existing public models in 7B-30B regime are still not ineffective in solving
reasoning puzzles that require multiple steps. We test different reasoning
tasks (e.g., math, logic, natural QA, code, etc..) with instruct-tuned LLaMA
Touvron et al. [2023], Dolly555https://huggingface.co/databricks/dolly-v2-12b,
StableLM Andonian et al. [2021], Vicuna Chiang et al. [2023], GPT-J Wang and
Komatsuzaki [2021] models that have 6B-30B parameters and is hard to find
qualified reasoning tasks. Once larger LLMs are released and evaluated as
having stronger reasoning abilities, Boundless DAS will be ready as an
analysis tool. Although our Price Tagging game task is simple, we test with
different variants of the task and show the efficacy of our method.

> _Is it possible to apply Boundless DAS over the whole layer representation?_

It is possible for smaller models if we are learning a full-rank square
rotation matrix. It is impossible for Alpaca now. Let’s say we align 10 token
representations together, we will have a total dimension size of 40,960. This
leads to a rotation matrix containing 1.6B parameters. Given the current
orthogonalization process in torch, it will require a GPU with memory larger
than the current A100 GPU just to fit in a trainable rotation matrix.

> _Only limited high-level models are tested with Boundless DAS in the
> experiment._

There is a large number of high-level models to solve the task, and we only
test a portion of them. We think it is sufficient to demonstrate the
advantages of our pipeline. It would be interesting if humans can be in the
loop of testing different hypotheses of the high-level causal model in future
work. On the other hand, not all the high-level models are interesting to
analyze. For instance, if two high-level models are equivalent (e.g., swapping
variable names), it is not in our interest to find alignments.

> _Can Boundless DAS be applied to find circuits in a model?_

Yes, it can. Circuit finding or inductive head searching in any Transformer
model often involves using interventions at head level to find the mapping
between activations generated on certain heads and high-level concepts.
Previous work in this field usually rely on the assumption of localist
representation (i.e., there exists a one-to-one mapping between a group of
neurons and a high-level concept). This is often too ideal to be real in
practice, as multiple concepts can easily be aligned with the same group of
neurons or a group of neurons may map to multiple concepts. Boundless DAS
actually breaks this assumption by specifically focusing on distributed
alignment search. More importantly, Boundless DAS can be easily adapted into a
head-wise alignment search where we add a shared rotation matrix on top of
head representations of all the tokens.

Figure 7: Interchange Intervention Accuracy (IIA) evaluated with correct input
examples only as well as incorrect input examples only for our Left and Right
Boundary causal model. The higher the number is, the more faithful the
alignment is. We color each cell by scaling IIA using the model’s task
performance as the upper bound and a dummy classifier (predicting the most
frequent label) as the lower bound. Figure 8: Interchange Intervention
Accuracy (IIA) evaluated with different settings of brackets in the
instruction. The seen setting is for training, and the unseen setting is for
evaluation. The higher the number is, the more faithful the alignment is. We
color each cell by scaling IIA using the model’s task performance as the upper
bound and a dummy classifier (predicting the most frequent label) as the lower
bound. Figure 9: Interchange Intervention Accuracy (IIA) evaluated with
different irrelevant contexts inserted as prefixes. The higher the number is,
the more faithful the alignment is. We color each cell by scaling IIA using
the model’s task performance as the upper bound and a dummy classifier
(predicting the most frequent label) as the lower bound. Figure 10:
Interchange Intervention Accuracy (IIA) evaluated with different output
formats. The seen setting is for training, and the unseen setting is for
evaluation. The higher the number is, the more faithful the alignment is. We
color each cell by scaling IIA using the model’s task performance as the upper
bound and a dummy classifier (predicting the most frequent label) as the lower
bound.