SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zhengxuan Wu*, Atticus Geiger*, Thomas Icard, Christopher Potts, and Noah D. Goodman
Stanford University  
{wuzhengx, atticusg, icard, cgpotts, ngoodman}@stanford.edu
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 has 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 faithfully understanding the inner-workings of our ever-growing and most widely deployed language models. Our tool is extensible to larger LLMs and is released publicly at https://github.com/frankaging/align-transformers.
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, as 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
Step 1
Q Identify a reasoning task that can be solved by a symbolic causal model.
Pick an off-the-shelf LLM.
Step 2
Step 3
Make sure the LLM passes the behavioral test (i.e., good task performance)
Step 4
Iteratively align the LLM with the causal model using Boundless DAS.
Propose a causal model to align.
Train with Boundless DAS for the proposed alignment.
Evaluate the alignment with our metrics.
Figure 1: Our pipeline for scaling causal explainability to LLMs with billions of parameters.
reprs
align reprs
Decoder-only Autoregressive LLMs
Instructions Representations
Test Query Representations
Instruction: Say yes only if the cost is between 2.00 and 3.00 dollars.
Input: 2.50 dollars [model output]
Test Query Representations
Causal Model
A  $\geq$  lower bound
Boundless DAS is for finding alignments between variables in the causal model and representations in the neural network.
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 interpretability at scale. We use Boundless DAS to study how Alpaca (7B) [47], 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 faithfully understanding the inner-workings of our largest and most widely deployed language models.
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 [45,44,60,46,52,6], probing [14,48,24,42,12,32,8,40], syntax-driven interventions [34], self-generated model explanations [27], and training external explainers based on model behaviors [39, 31]. However, these methods rely on observational measurements of behavior and internal neural representations. Such explanations are often not guaranteed to be faithful to the underlying causal mechanisms of the target models [29,20,50,53].
Causal Abstraction The theory of causal abstraction [41,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 interpretability methods, including interative nullspace projection [38, 15, 30], causal mediation analysis [51, 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 [35,17,36,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,37]. 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 [47], instruct-tuned LLaMA [49], StableLM [3], Vicuna [9]. Our goal is to scale methods from causal abstraction to understand how these models follow a particular instruction.
3.1 Background on Causal Abstraction
Causal Models We represent black box networks and interpretable algorithms using causal models that consist of variables taking on values according to causal mechanisms. We distinguish a set of input variables with possible values Inputs. An intervention is an operation that edits some 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 model  $\mathcal{M}$  (e.g., causal model or neural model) and a base input  $\mathbf{b}$  and source inputs  $\{\mathbf{s}_j\}_1^k$ , all elements of Inputs, together with disjoint sets of target variables  $\{\mathbf{Z}_j\}_1^k$  we are aligning. An interchange intervention yields a new model 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 intervened output (i.e., counterfactual output) for the base input with 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}\rightarrow \mathbf{Y}$ . A distributed interchange intervention yields a new model 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$ . Specifically, the causal mechanism for  $\mathbf{N}$  is,
$$
F _ {\mathbf {N}} ^ {*} (\mathbf {b}) = \mathbf {R} ^ {- 1} \left(\operatorname {P r o j} _ {\mathbf {Y} _ {0}} \left(\mathbf {R} \left(F _ {\mathbf {N}} (\mathbf {b})\right)\right) + \sum_ {j = 1} ^ {k} \operatorname {P r o j} _ {\mathbf {Y} _ {j}} \left(\mathbf {R} \left(F _ {\mathbf {N}} (\mathbf {s} _ {\mathbf {j}})\right)\right)\right) \tag {1}
$$
where  $\mathbf{b}$  is the base input,  $\mathbf{Y}_0 = \mathbf{Y} \setminus \bigoplus_{j=1}^k \mathbf{Y}_j$ , and  $\mathrm{Proj}_{\mathbf{Y}_0}$  or  $\mathrm{Proj}_{\mathbf{Y}_j}$  represents the orthogonal projection operator of the original rotated vector into the  $\mathbf{Y}_0$  or  $\mathbf{Y}_j$  subspace.  $F_{\mathbf{N}}(\cdot)$  represents the causal mechanisms of causal variables  $\mathbf{N}$  before any intervention, whereas  $F_{\mathbf{N}}^*(\cdot)$  represents the mechanisms after the distributed intervention. We then denote the intervened output (i.e., counterfactual output) for the base input as  $\mathrm{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 a causal model (i.e., an algorithm)  $\mathcal{A}$  is a faithful interpretation of the neural 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. We use  $\tau$  to represent the alignment mapping between a high-level causal variable and a neural representation.
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 of high-level variables  $Z_{j}$  to orthogonal subspace  $\mathbf{Y}_{j}$  between an algorithm  $\mathcal{A}$  and neural network  $\mathcal{N}$  is
$$
\frac {1}{| \mathbf {I n p u t s} | ^ {k + 1}} \sum_ {\mathbf {b}, \mathbf {s} _ {1}, \dots , \mathbf {s} _ {k} \in \mathbf {I n p u t s}} \left[ \tau \binom {\text {D I I} (\mathcal {N}, \mathbf {R}, \{\mathbf {s} _ {j} \} _ {1} ^ {k}, \{\mathbf {Y} _ {j} \} _ {1} ^ {k}) (\mathbf {b}))} {\text {I N T I N V} (\mathcal {A}, \{\tau (\mathbf {s} _ {j}) \} _ {1} ^ {k}, \{\mathbf {Z} _ {j} \} _ {1} ^ {k})) (\mathbf {b})} \right] \tag {2}
$$
where  $\tau$  translates from low-level causal variable values to high-level neural representation values.
${}^{2}$  Prior works focus on all-zero or mean value representation replacement [33,54] which is less general.
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}\rightarrow \mathbf{Y}$  is trained to maximize interchange intervention accuracy under an alignment from each variable  $Z_{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 a  $d$ -dimensional vector space  $\mathbf{Y}$  using a method inspired by work in neural PDE [57]. Specifically, for each high-level variable  $Z_{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
$$
\left(\mathbf {M} _ {j}\right) _ {k} = \operatorname {s i g m o i d} \left(\frac {k - b _ {j}}{\beta}\right) * \operatorname {s i g m o i d} \left(\frac {b _ {j + 1} - k}{\beta}\right) \tag {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}$ .
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}\rightarrow \mathbf{Y}$  be an invertible transformation. A weighted distributed interchange intervention yields a new model 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 {b}) = \mathbf {R} ^ {- 1} \left(\left(\mathbf {1} - \sum_ {j = 1} ^ {k} \mathbf {M} _ {j}\right) \circ \mathbf {R} \left(F _ {\mathbf {N}} (\mathbf {b})\right) + \sum_ {j = 1} ^ {k} \left(\mathbf {M} _ {j} \circ \mathbf {R} \left(F _ {\mathbf {N}} \left(\mathbf {s} _ {\mathbf {j}}\right)\right)\right)\right) \tag {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  $\mathrm{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}_j^\theta\}_{1}^k$
$$
\sum_ {\mathbf {b}, \mathbf {s} _ {1}, \dots , \mathbf {s} _ {k} \in \mathbf {I n p u t s}} \operatorname {C E} \left(\operatorname {S O F T D I I} \left(\mathcal {N}, \mathbf {R} ^ {\theta}, \left\{\mathbf {s} _ {j} \right\} _ {1} ^ {k}, \left\{\mathbf {M} _ {j} ^ {\theta} \right\} _ {1} ^ {k}\right) (\mathbf {b}), \operatorname {I N T I N V} \left(\mathcal {A}, \left\{\tau \left(\mathbf {s} _ {j}\right)\right\} _ {1} ^ {k}, \left\{\mathbf {Z} _ {j} \right\} _ {1} ^ {k}\right)\right) (\mathbf {b})\left. \right) \tag {5}
$$
where 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  $Z_{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  $(1 - \sum_{j = 1}^{k}\mathbf{M}_{j})$ .
Time Complexity Analysis DAS learns a rotation matrix but requires a manual search to determine how many neurons are needed to represent the aligning causal variable. Boundless DAS automatically learns boundaries (i.e., how many neurons are needed is determined by the "soft" boundary index via a boundary mask learning). For instance, given a representation with a dimension of 1024, DAS should in principle be run for all lengths  $k$  from 1 to 1024. In practice, this would be infeasible, so some subset of the lengths need to be chosen heuristically, which risks missing genuine structure. For Boundless DAS, we turn this search process into a mask learning process. DAS [23] is  $O(n \times m)$  where  $n$  is the number of total dimensions of the neural representation we are aligning and  $m$  is the number of causal variables, while Boundless DAS is  $O(m)$ .
4.1 Price Tagging
We follow the approach in Figure 1 by first assessing the ability of Alpaca to execute specific actions based on the instructions provided in the input. Formally, the input to the model  $\mathcal{M}$  is given as
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 solve the task.
Figure 3: Aligned distributed interchange interventions performed on the Alpaca model that is instructed to solve our Price Tagging game. To train Boundless DAS, we sample two training examples and then swap the intermediate boolean values 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.
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 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.
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 respectively. 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) [47], 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 train Boundless DAS with test query token representations (i.e., starting from the token of the first input digit till the last token in the prompt) in a selected set of 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. We interchange with a single source, while allowing multiple causal variables to be aligned across examples. For simplicity, we enforce the interval between two boundary indices to be the same (i.e., the same size of linear subspace for each aligning causal variable). 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. We measure the baseline IIA by replacing the learned rotation matrix with a random one. The lower bound of IIA for "Left Boundary" and "Left and Right Boundary" is about 0.50, and for "Mid-point Distance" or "Bracket Identity" it is about 0.60. The latter two baseline IIAs are higher than chance because they are conditioned on the distribution of our output labels (i.e., how many times the original label gets to be flipped due to the intervention). See Section 4.7 for more discussion of metric calibration. Additionally, IIA can occasionally go above the model's task performance, when the interchange interventions put the model in a better state, but IIA is constrained by task accuracy for the most part.
4.3 Boundless DAS Results
Figure 4 shows our main results, given in terms of IIA across our four hypothesized causal models (Figure 2). The results show very clearly that 'Left Boundary' and 'Left and Right Boundary' (top panels) are highly accurate hypotheses about how Alpaca solves the task. For them, IIA is at or above task performance (0.85), 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.