Improving Neuron-level Interpretability with White-box Language Models

## Abstract

Abstract Neurons in auto-regressive language models like GPT-2 can be interpreted by analyzing their activation patterns. Recent studies have shown that techniques such as dictionary learning, a form of post-hoc sparse coding, enhance this neuron-level interpretability.
In our research, we are driven by the goal to fundamentally improve neural network interpretability by embedding sparse coding directly within the model architecture, rather than applying it as an afterthought.
In our study, we introduce a white-box transformer-like architecture named Coding RAte TransformEr ( crate ), explicitly engineered to capture sparse, low-dimensional structures within data distributions.
Our comprehensive experiments showcase significant improvements (up to 103 103 103 103 % lorelative improvement) in neuron-level interpretability across a variety of evaluation metrics.
Detailed investigations confirm that this enhanced interpretability is steady across different layers irrespective of the model size, underlining crate ’s robust performance in enhancing neural network interpretability.
Further analysis shows that crate ’s increased interpretability comes from its enhanced ability to consistently and distinctively activate on relevant tokens. These findings point towards a promising direction for creating white-box foundation models that excel in neuron-level interpretation.

## 1 Introduction

Representation learning aims to learn a continuous mapping, to transform a random vector in a high dimensional space that is sampled from a dataset, to a feature vector in another (typically lower-dimensional) space []. Recently, deep learning has witnessed tremendous empirical success in modeling massive amounts of high-dimensional data, and the predominant practice has been to learn first a task-agnostic representation by pre-training a large neural network, which is commonly known as the foundation model [, ]. Among language foundation models, the transformers architecture [] with Generative Pre-Training [] (GPT) has recently demonstrated a strong capability of modeling sequential data and thus predicting subsequent tokens [, ].
Such strong capability has emerged significant success in downstream applications [, ], yet the large neural network is known to be black-box, where the representations in the model are not independently interpretable, introducing difficulty in designing effective paradigms for major known challenges of (visual) language models like hallucination [, ], bias [, ], and catastrophic forgetting [, ].

To interpret the functions of individual modules in the language models, mechanistic interpretation was proposed to reverse-engineer such models, through identifying meaningful patterns in the data representations and computational mechanisms of the model components [, ]. Recent studies on auto-regressive models like GPT-2 have delved into neuron-level interpretation, where the focus is on understanding the activations within the model’s MLP layers []. This approach helps to reveal the specific roles of individual neurons, which is crucial for precise model editing and control [, ]. Recent research has also proposed that sparse auto-encoder (SAE)-based dictionary learning effectively promotes mono-semanticity of neurons, thus enhancing neuron-level interpretability [].
However, as a post-hoc method, sparse auto-encoders introduce non-negative reconstruction loss, causing noise and reducing fidelity when interpreting neurons []. Recent studies also highlight their scalability challenges, as neuron decomposition becomes difficult in larger language models due to the vast number of directions [, , ].

Can we instead build sparsity directly into the language model? In this paper, we develop the crate language model, a GPT-2-size language model that builds sparse coding into the model with a mathematically principled way. crate handles the problems introduced by sparse encoders at scale: it (i) escapes the loss introduced in reconstructing the language model representations, enabling loss-free steering, and (ii) escapes the unsteady process of training a sparse auto-encoder. To avoid adding inconsistency into the language model, we develop on top of a mathematically principled white-box model framework, named crate [].(^1^11In the remaining parts of the paper, we use “crate” or “crate language model” to refer to our language model architecture, while “original crate” denotes the architecture framework described in the literature.) After encoding the text tokens into numbers, we apply language-domain-specific modifications to the original crate architecture and obtain the token representations. The final representations are then used to predict the next token, while the intermediate representations gets interpreted.

Figure: Figure 1: Instances are systematically identified where the interpretability of crate (ours, row 1) outperforms GPT-2 (row 2). For each neuron (rounded box), we show two top activated text excerpts (excerpt 1 and 2) and one randomly activated excerpt (excerpt 3). Results show that crate consistently activates on and only on semantically relevant text excerpts (first two excerpts), leading to more precise explanations predicted by agents like Mistral.
Refer to caption: x1.png

To this end, the main contribution of this work is to propose a causal language model architecture based on the crate model framework that builds sparsity inherently, that achieves significantly better neuron-level interpretability (up to 103% relative increase) than language models with the GPT-2 architecture under a similar configuration. crate forms a family of models from single-layer model up to a 12-layer configuration. Comparative qualitative analysis of neuron activations between crate and GPT-2 is provided in Figure 1, alongside extensive quantitative evaluations demonstrating that by explicitly integrating sparse coding into the language model, crate achieves markedly improved interpretability across layers compared to GPT-2, applicable across a wide range of model sizes under a variety of evaluation metrics.

## 2 Related Work

Neuron-level interpretation. Recent studies have provided insights into how auto-regressive models like GPT-2 work at the level of individual neurons, named neuron-level interpretation. These studies focus on analyzing activations, which are the outputs from the activation functions within the model’s multi-layer perceptron (MLP) blocks []. Neuron-level interpretation is crucial to understanding the mechanisms in a model, including what concepts are learned in the neurons of the network, whether specific neurons are learning particular concepts, and how localized/distributed and redundantly the knowledge is preserved within neurons of the network.

A higher neuron-level interpretability indicates that more neurons are interpretable or neurons are more interpretable []. As interpretations of the neurons can help localize the knowledge obtained in a neural network, neuron-level interpretation can be used for editing the knowledge in models [, ], model pruning and distillation [], adapting the model to different domains and steering the output [, ], and debugging model errors []. Improved neuron-level interpretability increases reliability and performance in the applications above.

Sparse auto-encoders. To enhance interpretability, post-hoc sparse coding methods like dictionary learning [] are used, but these techniques result in imperfect reconstructions and thus always introduces loss when steering the model [, ]. Literature also indicates that sparse-autoencoders are hard to scale, i.e., a dramatic drop in the number of interpretable features can be observed when the model being interpreted becomes deeper []. Additionally, tuning SAE models for larger $L$ values involves extensive hyperparameter tuning and time-consuming training, requiring multiple metrics (reconstruction rate, L1 loss, number of dead neurons) for reliable judgment, which can’t be easily optimized with automatic engineering tricks [].

Evaluation of neuron-level interpretability. Metrics now exist to evaluate neuron-level interpretability in language models, examing if neurons trigger on relevant tokens in given contexts [, ]. Recent works have demonstrated that a small number of circuits in language models are interpretable [, ], but comprehending each neuron, out of millions, is vital for thorough model safety audits. Given the prohibitive cost of human evaluation on such a scale, OpenAI introduced an automated metric using large language models for interpretability assessment [], which Anthropic later refined for sparse activations []. These methods align closely with human judgment and have gained broad acceptance within the research community [, , , ]. These metrics show that neuron-level interpretability in auto-regressive models is limited [], where the popular hypothesis is that neurons are superpositions of simpler semantics, which makes them fire (produce a high activation) at multiple semantically distinct sets of tokens []. Metrics mentioned above are what we use in this work.

White-box models and structured representation learning.
In the domain of structured representation learning, white-box models stand out for their ability to generate explicit, structured data representations that adhere to specific, desirable configurations such as sparsity and piece-wise linearity, as discussed by Gregor and LeCun [] and Chan et al. []. Within this framework, Yu et al. [] introduced an innovative approach to constructing deep networks based on unrolled optimization. Specifically, Yu et al. [] proposed the crate foundation model framework, utilizing an information-theoretic objective aimed at promoting the compression and sparsity of data towards a predefined statistical structure.
Recently, empirical experiments suggest that the white-box design of crate inherently develops segmentation capabilities from the data representations at both holistic and component levels with supervised training in the vision domain [], which directly motivates us to further explore the data representations within such architecture for language models. Recent work has also shown that the crate framework is scalable: it can be effectively scaled up to comparable performance as Vision Transformer (ViT) with careful engineering []. Furthermore, the fine-tuning performance of the pretrained crate model is also proven to be comparable in both the language domain [] and vision domain [].

## 3 Preliminaries

This section introduces the original crate architecture introduced in Yu et al. [].

Notations. In this paper, we denote the one-hot input tokens by $\bm{X}=[\bm{x}_{1},\dots,\bm{x}_{N}]\in\mathbb{R}^{V\times N}$, where $\bm{x}_{i}\in\mathbb{R}^{V\times 1}$ represents the $i$-th one-hot token, $N$ is the total number of input tokens, and $V$ is the vocabulary size.
We use $f\in\mathcal{F}:\mathbb{R}^{V\times N}\rightarrow\mathbb{R}^{d\times N}$ to denote the mapping induced by the model, which is a composition of $L+1$ operators (layers) $f=f^{L}\circ\cdots f^{\ell}\circ\cdots f^{1}\circ f^{\mathrm{pre}}$, where $f^{\ell}:\mathbb{R}^{d\times N}\rightarrow\mathbb{R}^{d\times N}(1\leq\ell\leq
L)$ represents the mapping of the $\ell$-th operator, and $f^{\mathrm{pre}}:\bm{X}\in\mathbb{R}^{V\times N}\rightarrow\bm{Z}^{1}\in
\mathbb{R}^{d\times N}$ represents the pre-processing layer that transforms the one-hot token representations $\bm{X}=[\bm{x}_{1},\dots,\bm{x}_{N}]$ to semantic embeddings $\bm{Z}^{1}=[\bm{z}_{1}^{1},\dots,\bm{z}_{N}^{1}]$.
We let $\bm{Z}^{\ell}$ denote the input token representations of the $\ell$-th operator $f^{\ell}$ for $1\leq\ell\leq L$, so that $\bm{z}_{i}^{\ell}\in\mathbb{R}^{d}$ denotes the representation of the $i$-th token $\bm{x}_{i}$ before the $\ell$-th layer. We denote $\bm{Z}=\bm{Z}^{L+1}$ as the output token representations of the last ($L$-th) layer.

Framework, objective, and optimization. The transformation of input data into parsimonious (piecewise linearized and compact) representations is accomplished by adopting a local signal model for the marginal distribution of the tokens $\bm{z}_{i}$. This statement suggests that the tokens can be approximately considered to occupy a union of several (identified as $K$) low-dimensional spaces, each with a dimension $p\ll d$. These spaces are characterized by orthonormal bases, represented as $\bm{U}_{[K]}=(\bm{U}_{k})_{k=1}^{K}$ where $\bm{U}_{k}\in\mathbb{R}^{d\times p}$.
Within the framework of this local signal model, crate aims to optimize the sparse rate raduction objective:

$$ $\max_{f\in\mathcal{F}}\mathbb{E}_{\bm{Z}}\big{[}\Delta R(\bm{Z}\mid\bm{U}_{[K] })-\lambda\|\bm{Z}\|_{0}\big{]}=\max_{f\in\mathcal{F}}\mathbb{E}_{\bm{Z}}\big{ [}R(\bm{Z})-\lambda\|\bm{Z}\|_{0}-{R}^{c}(\bm{Z};\bm{U}_{[K]})\big{]}.$ (1) $$

where $\lambda$ is the sparsification regularizer and $\bm{Z}=f(\bm{X})$. The coding rate $R(\bm{Z})$ serves as a close estimate (following Ma et al. []) for the average amount of bits necessary for encoding the tokens $\bm{z}_{i}$ to a precision level $\varepsilon$ using a Gaussian codebook. Additionally, $R^{c}(\bm{Z}\mid\bm{U}_{[K]})$ represents the theoretical maximum average amount of bits needed to encode the projection of the tokens onto each low dimensional subspace defined in the local signal model, specifically $\bm{U}_{k}^{*}\bm{z}_{i}$, to the same precision level $\varepsilon$ utilizing a Gaussian codebook, as outlined by Yu et al. []. If the subspaces are adequately incoherent from each other, the solutions that minimize the object function, viz. Equation 1, in terms of $\bm{Z}$, are associated with subspace configurations that are both incoherent and aligned with the axes, as pointed out by Yu et al. [].

A network aimed at optimizing the sparse coding rate reduction objective through unrolled optimization gradually shifts the distribution of $\bm{X}$ towards the intended canonical forms, where each iteration of the unrolled optimization process acts as a layer:

$$ $f\colon\bm{X}\xrightarrow{\hskip 2.84526ptf^{\mathrm{pre}}\hskip 2.84526pt}\bm {Z}^{1}\rightarrow\cdots\rightarrow\bm{Z}^{\ell}\xrightarrow{\hskip 2.84526ptf ^{\ell}\hskip 2.84526pt}\bm{Z}^{\ell+1}\rightarrow\cdots\to\bm{Z}^{L+1}=\bm{Z}.$ (2) $$

The iterative optimization framework incorporates multiple design choices, among which is a two-step alternating minimization approach grounded in robust theoretical principles []. This approach delineates two distinct blocks: the MSSA and the ISTA block, collectively defining a single crate layer:

$$ $\bm{Z}^{\ell+1/2}\doteq\bm{Z}^{\ell}+\texttt{MSSA}(\bm{Z}^{\ell}\mid\bm{U}_{[K ]}^{\ell}),\qquad f^{\ell}(\bm{Z}^{\ell})=\bm{Z}^{\ell+1}\doteq\texttt{ISTA}( \bm{Z}^{\ell+1/2}\mid\bm{D}^{\ell}).$ (3) $$

Compression operator: Multi-Head Subspace Self-Attention (MSSA). Given local models $\bm{U}_{[K]}^{\ell}$, to form the incremental transformation
$f^{\ell}$ optimizing Equation 1 at layer $\ell$, crate first compresses the token set $\bm{Z}^{\ell}$ against the subspaces by
minimizing the coding rate $R^{c}(\,\cdot\mid\bm{U}_{[K]}^{\ell})$.
As Yu et al. [] show, doing this minimization locally by performing a step of gradient descent on $R^{c}(\,\cdot\mid\bm{U}_{[K]}^{\ell})$ leads to the so-called multi-head subspace self-attention (MSSA) operation, defined as

$$ $\operatorname{\texttt{MSSA}}(\bm{Z}\mid\bm{U}_{[K]})\doteq\frac{p}{(N+1) \varepsilon^{2}}\begin{bmatrix}\bm{U}_{1},\dots,\bm{U}_{K}\end{bmatrix}\begin{ bmatrix}(\bm{U}_{1}^{\ast}\bm{Z})\operatorname{softmax}\left((\bm{U}_{1}^{\ast }\bm{Z})^{\ast}(\bm{U}_{1}^{\ast}\bm{Z})\right)\\ \vdots\\ (\bm{U}_{K}^{\ast}\bm{Z})\operatorname{softmax}\left((\bm{U}_{K}^{\ast}\bm{Z}) ^{\ast}(\bm{U}_{K}^{\ast}\bm{Z})\right)\end{bmatrix},$ (4) $$

In practice, the calculation of the intermediate representations $\bm{Z}^{\ell+1/2}$ with the output from the MSSA block is calculated in a weighted form:

$$ $\bm{Z}^{\ell+1/2}\approx\left(1-\kappa\cdot\frac{p}{(N+1)\varepsilon^{2}} \right)\bm{Z}^{\ell}+\kappa\cdot\frac{p}{(N+1)\varepsilon^{2}}\cdot \operatorname{\texttt{MSSA}}(\bm{Z}^{\ell}\mid\bm{U}_{[K]}),$ (5) $$

where $\kappa>0$ is a learning rate hyperparameter. This block resembles to the multi-head self-attention block of the GPT-2 model, but the query, key, and value projection matrices within a single head are all identical in the MSSA block.

Sparsification operator: Iterative Shrinkage-Thresholding Algorithm (ISTA). The remaining term to optimize in Equation 1 is the
difference of the global coding rate $R(\bm{Z})$ and the $\ell^{0}$ norm of the
tokens, which together encourage the representations to be both sparse and
non-collapsed. Yu et al. [] show that local minimization of this
objective in a neighborhood of the intermediate representations $\bm{Z}^{\ell+1/2}$ is approximately achieved by a LASSO problem with respect to a
sparsifying orthogonal dictionary $\bm{D}^{\ell}\in\mathbb{R}^{d\times h}$. Taking an iterative step towards solving this LASSO problem gives the iterative shrinkage-thresholding algorithm (ISTA) block [, ]. The ReLU non-linearity appearing in this block arises from an additional non-negativity constraint on the representations in the LASSO program, motivated by the goal of better separating distinct modes of variability in the token distribution:

$$ $\bm{Z}^{\ell+1}=f^{\ell}(\bm{Z}^{\ell})=\operatorname{ReLU}(\bm{Z}^{\ell+1/2}+ \eta\bm{D}^{\ell*}(\bm{Z}^{\ell+1/2}-\bm{D}^{\ell}\bm{Z}^{\ell+1/2})-\eta \lambda\bm{1})\doteq\operatorname{\texttt{ISTA}}(\bm{Z}^{\ell+1/2}\mid\bm{D}^{ \ell}).$ (6) $$

## 4 The CRATE Language Model

This section introduces the difference between our work and the original crate paper [], thus introducing what changes we made to the crate architecture.

First, we note that the task in this work is different: we apply the crate architecture to the next-token prediction task in the language domain, while the original crate paper applies the architecture to the image classification task in the vision domain. This difference immediately leads to differences when designing the model architecture: (i) we apply a causal mask to the original crate model to avoid the model seeing tokens after the current token, and (ii) we change the embedding layer and heads of the original crate model to fit the language vocabulary.

Second, we’re interested in interpreting the neurons within crate on the next-token prediction task and making direct comparisons to the GPT-2 model architecture. As neuron-level interpretation is commonly evaluated on the hidden states in the MLP block of the GPT-2 model [, ], (iii) we over-parameterize (increase the hidden dimension of) the ISTA block of the original crate model to align with the hidden dimension of the MLP block of the GPT-2 model.

Below we show specific definitions of the modifications we made. We illustrate the architecture in Figure 2, and show implementation details in Appendix A.

Figure: Figure 2: Block architecture for the crate language model, where $S_{\lambda}(x)=\operatorname{ReLU}(x-\eta\cdot\lambda\cdot 1)$. Differences from the original architecture mentioned in Yu et al. [] are marked bold: we (i) add a causal mask $\text{Mask}(\cdot)$ and (ii) over-parameterize the ISTA block.
Refer to caption: x2.png

Embedding and Head. In order to apply the crate architecture to the language domain, we define the pre-processing layer $f^{\mathrm{pre}}$ that transforms language tokens into position-aware semantic embeddings, and define post-processing head $f^{\mathrm{head}}(\bm{Z})$ that maps the final representations to output token distributions in the vocabulary space:

$$ $\displaystyle f^{\mathrm{pre}}(\bm{X})=\bm{E}_{\mathrm{sem}}(\bm{X})+\bm{E}_{ \mathrm{pos}},\qquad f^{\mathrm{head}}(\bm{Z})=\bm{W}^{\mathrm{head}}\bm{Z},$ (7) $$

where $\bm{E}_{\mathrm{sem}}$ is a semantic embedding matrix that maps input tokens $\bm{x}_{i}$ to embedding vectors in $\mathbb{R}^{d}$, $\bm{E}_{\mathrm{pos}}\in\mathbb{R}^{d\times N}$ is a positional embedding matrix, and $\bm{W}^{\mathrm{head}}\in\mathbb{R}^{V\times d}$ maps the (contextualized) token representations $\bm{Z}^{L+1}$ to the distribution of the next token. All parameters mentioned above are learnable.

MSSA Block. To align with the next word prediction task used in GPT-2 [], we replace the attention matrix in MSSA (Equation 4) with a causally masked self-attention, defined as

$$ $\displaystyle\begin{split}\operatorname{softmax}\left((\bm{U}_{k}^{\ast}\bm{Z} )^{\ast}(\bm{U}_{k}^{\ast}\bm{Z})\right)&\to\operatorname{softmax}\left( \texttt{CausalMask}((\bm{U}_{k}^{\ast}\bm{Z})^{\ast}(\bm{U}_{k}^{\ast}\bm{Z})) \right),\\ \text{where}\ \texttt{CausalMask}(\bm{M})_{ij}&=\begin{cases}\bm{M}_{ij},&i \leq j\\ -\infty,&i>j.\end{cases}\end{split}$ (8) $$

ISTA Block. To investigate the neuron interpretability of the activation matrix $\bm{A}\in\mathbb{R}^{h\times N}$, we design an overcomplete version of the original ISTA block (Equation 6) with $\bm{D}^{\ell}\in\mathbb{R}^{d\times h}$ where $h=nd$, and $n=4$ to keep a fair comparison to the GPT-2 [] and standard transformer architecture Vaswani et al. []):

$$ $\displaystyle\begin{split}&\bm{A}_{t}\doteq\operatorname{\texttt{ISTA}}(\bm{Z} ^{\ell+1/2}\mid\bm{D}^{\ell}),\\ &\bm{A}_{k}=\operatorname{ReLU}(\bm{A}_{k-1}-\eta(\bm{D}^{\ell})^{\ast}(\bm{D} ^{\ell}\bm{A}_{k-1}-\bm{Z})-\eta\cdot\lambda\cdot\bm{1}),\quad\bm{A}_{0}=\bm{0 },\,\,k\in[t],\\ &\bm{Z}^{\ell+1}=\bm{D}^{\ell}\bm{A}_{t},\end{split}$ (9) $$

where $\eta>0$ is the step size, $\lambda>0$ is the sparsification regularizer, and $t$ is the number of ISTA iterations. In practice, we set $t=2$ to keep computation efficient.
The ISTA block resembles the MLP block in the GPT-2 model, but with a relocated skip connection.

Token Processing. The desired token processing procedure is illustrated in Figure 3. The process starts with random token representations ($\bm{Z}^{1}$). Through successive layers, the representations ($\bm{Z}^{\ell}$) are compressed to align with the axis via the MSSA block, forming $\bm{Z}^{\ell+1/2}$ that are semantically more consistent among relevant tokens. This is then refined by sparse coding (the ISTA block) to produce the representations $\bm{Z}^{\ell+1}$ that align on incoherent axes, leading to semantically more specified token representations. Repeated across layers, this culminates in distinct token representations $Z^{L+1}$ aligned on unique semantic axes. More detailed explanation of this optimization process can be found in LABEL:app:learning-process.

Figure: Figure 3: crate iteratively compresses (MSSA block) and sparsifies (ISTA block) the token representations (colored points) across its layers from $1$ to $L$, transforming them into parsimonious representations aligned on axes (colored lines) with distinct semantic meanings.
Refer to caption: x3.png

## 5 Empirical Experiments

We examine the next token prediction performance and neuron interpretability of crate in this section. We detail the architecture, size, pre-training recipe (Section 5.1), performance (Section 5.2), and neuron-level interpretability (Section 5.3) of crate compared to the standard transformer architecture. In this section, we denote $K$ as the number of attention heads, $d$ as the dimension of the residual stream in the model, and $h$ as the hidden (inner) dimension of the ISTA or MLP module.

### 5.1 Setup

Model architecture and size. The crate model is designed with various sizes $L\in\{1,2,3,6,12\}$, where each size matches the GPT-2 configurations for direct comparisons, as shown in Table 1. Configurations for $L\in\{1,2,3\}$ adhere to GPT-2 models as per Bricken et al. [], while $L\in\{6,12\}$ follow configurations from Sanh et al. [] and Radford et al. [], respectively. Notably, crate maintains approximately $2/3$ the size of GPT-2 at scale. Both models utilize the Byte-level BPE tokenizer with a vocabulary size of 50,257, following Radford et al. []. We explain the parameter size difference between crate and GPT-2 in Section A.1.

**Table 1: Model configuration of crate and model size comparison to GPT-2.**
| Model Config | $d$ | $K$ | $L$ | $h$ | CRATE | GPT-2 |
| --- | --- | --- | --- | --- | --- | --- |
| 1L | 128 | 4 | 1 | 512 | 6.54M | 6.64M |
| 2L | 128 | 4 | 2 | 512 | 6.64M | 6.83M |
| 3L | 128 | 4 | 3 | 512 | 6.74M | 7.03M |
| S(mall) | 768 | 12 | 6 | 3,072 | 55.9M | 81.1M |
| B(ase) | 768 | 12 | 12 | 3,072 | 81.2M | 123.6M |

Datasets and optimization. We pre-train both models on the next token prediction task using the uncopyrighted Pile dataset [] and the Adam optimizer []. Following Bricken et al. [], we pre-train both crate and GPT-2 of smaller sizes ($L\in\{1,2,3\}$) on 100 billion tokens with a context window of 1,024 tokens.
Following the pre-training setup in Karpathy [] and scaling law in Touvron et al. [], we pre-train using 100 billion tokens for the Small models, and 160 billion tokens for the Base models.(^2^22Practicaly, we train with a batch size of 768 for 125,000 iterations for $L\in\{1,2,3,6\}$, and a batch size of 256 for 600,000 iterations for $L=12$.) It takes about 4 days to pre-train crate-Base on 160 billion tokens with 32 A5000 GPUs.

### 5.2 Performance

This section demonstrates that crate, despite not outperforming GPT-2, still generates reasonable predictions, as evidenced through quantitative and qualitative comparisons.

We observe that both training and validation loss curve of crate-Base on the Pile dataset converges well, as presented in Figure 4 (left). Although the convergence is slower than GPT-2, the loss curve of crate keeps decreasing after training on 160 billion tokens, while GPT-2 already tends to converge.

We also demonstrate the zero-shot validation loss curve of crate evaluated on OpenWebText as well as other datasets [] in Figure 4 (right). Results show that crate effectively learns transferable representations across a number of datasets, and achieves comparable performance to GPT-2 after full training on the 160 billion tokens.
We also demonstrate the scalability of the crate architecture by comparing the validation loss of crate and GPT-2 with respect to the model size in Figure 5 (left). Results show that the performance of crate is close to GPT-2 across all model sizes. However, we do recognize that forcing sparsification in a model potentially leads to a higher compute cost on the next-token-prediction objective, which aligns with observations in Bricken et al. [] that enabling monosemanticity might hurt model performance. Analysis on the sparsity of crate and GPT-2 can be found in LABEL:app:analysis-sparsity.

Qualitative examples from crate and GPT-2 are demonstrated in Figure 5 (right). We conclude that crate can make reasonable predictions, encouraging us to further look into its neuron-level interpretability.

Figure: Figure 4: Left: loss curve when pre-training crate-Base and GPT2-Base on the Pile dataset. Right: zero-shot validation loss of crate evaluated on a variety of datasets (Pile, LAMBADA, OpenWebText and WikiText).
Refer to caption: extracted/5943414/figures/crate_loss.png

Figure: Figure 5: Left: Validation loss of crate compared to GPT-2 on the Pile dataset, with respect to the model size. Right: Qualitative examples of predictions made by crate and GPT-2. The tokens in blue are considered good. We compare crate-Base to GPT2-Base on the next word prediction task.
Refer to caption: x5.png

### 5.3 Interpretability

In order to quantitatively evaluate the interpretability of the neuron activations, we adopt the large language model-based algorithm introduced in Bills et al. [] and Bricken et al. [], as demonstrated in Algorithm 1. We retrieve the sparse code $\bm{A}_{t}$ (activations after the $\operatorname{ReLU}$ unlinearlity in the ISTA block) of crate for interpretation, and compare to activations from the MLP block of GPT-2.

Figure: Algorithm 1 Interpretability Evaluation Algorithm []

In practice, we adopt three evaluation metrics: two from OpenAI (top-and-random and random-only) [] and one from Anthropic []. We adopt the official implementation from Wu et al. [], where details on the implementation are elaborated in LABEL:app:impl-interp-metric.
Note that the Anthropic metric has much shorter text excerpts than the OpenAI metrics, so it is biased to sparse activations. For all evaluations, we discard the last layer of crate and GPT-2, according to the empirical observation that the last layer of crate is biased to the pre-training task [].

CRATE achieves markedly improved and more steady neuron-level interpretability across layers compared to GPT-2, applicable across a wide range of model sizes. We show evaluation results of the interpretability of crate and GPT-2 averaged across layers in Table 2 (left).
We observe that the interpretability of crate comprehensively outperforms GPT-2 on all metrics for $L\in\{2,3,6,12\}$. When averaging the mean interpretability across all metrics, crate outperforms GPT-2 up to strikingly $103$% relative improvement under the OpenAI Random-only metric when $L=6$.
We also present the layer-wise interpretation scores in Figure 6, which demonstrates that crate achieves higher interpretability than GPT-2 on almost all layers. For detailed distributions of the layer-wise scores of crate-Base compared to GPT2-Base on different metrics, refer to LABEL:app:interp-distribution.

**Table 2: Mean and variance of the average interpretability across layers for different model sizes.**
|  | Mean ($\uparrow$, darker green means more interpretable) | Variance ($\downarrow$, darker red means less steady) |  |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | Top-and-Random | Random-only | Anthropic | Top-and-Random | Random-only | Anthropic |  |  |  |  |  |  |
|  | crate | GPT-2 | crate | GPT-2 | crate | GPT-2 | crate | GPT-2 | crate | GPT-2 | crate | GPT-2 |
| 1L | $3.9$ | $8.8$ | $4.8$ | $8.9$ | $10.1$ | $14.2$ | $0.0$ | $0.0$ | $0.0$ | $0.0$ | $0.0$ | $0.0$ |
| 2L | $8.05$ | $4.2$ | $6.95$ | $1.95$ | $11.35$ | $10.2$ | $0.06$ | $0.01$ | $1.1$ | $0.12$ | $0.0$ | $0.25$ |
| 3L | $9.1$ | $3.57$ | $8.43$ | $1.37$ | $11.23$ | $9.2$ | $0.26$ | $7.51$ | $1.2$ | $1.93$ | $1.14$ | $19.21$ |
| 6L | $7.96$ | $5.4$ | $6.36$ | $3.14$ | $10.4$ | $8.52$ | $2.29$ | $20.85$ | $1.87$ | $18.39$ | $2.01$ | $32.56$ |
| 12L | $6.8$ | $6.34$ | $5.12$ | $2.67$ | $8.88$ | $8.65$ | $7.09$ | $11.35$ | $2.83$ | $7.48$ | $18.3$ | $24.65$ |

Figure: Figure 6: Interpretation scores evaluated using the OpenAI Random-only metric, Top-and-Random metric, and Anthropic metric, respectively. Top: interpretation scores of crate and GPT-2 for $L=12$. Middle: interpretation scores of crate and GPT-2 for $L=6$. Bottom: interpretation scores of crate, GPT-2, and GPT-2 with sparse auto-encoder for $L\in\{1,2,3\}$. Variance bars are normalized to $1/10$ of its original size.
Refer to caption: x7.png

The variances of the average interpretation scores of crate and GPT-2 across layers are shown in Table 2 (right). From the results we draw a solid conclusion that the interpretability of crate is much more steady than GPT-2 across all model sizes. Figure 6 further demonstrates a clear pattern that, for all model sizes, crate maintains a higher interpretability than GPT-2 among almost all layers.

The built-in sparse coding approach introduces consistent and specific neuron-level behaviors. The strong interpretability of crate on the OpenAI Top-and-Random metric and the Anthropic metric, as shown in Figure 6, indicates its consistent behavior on relevant tokens. These two methods contain a large portion of top-activated text excerpts, so they are valid for measuring whether the activations are consistent with the summarized explanation [, ].
Additionally, the larger interpretability gap of crate and GPT-2 on the OpenAI Random-only metric versus the Top-and-Random metric highlights the specificity of crate in avoiding firing on irrelevant tokens. The random-only metric usually includes highly irrelevant text excerpts, so it effectively measures the capability of the language model to avoid activating on semantically irrelevant tokens [].

Qualitatively, we refer back to the interpretation examples shown in Figure 1. In this figure, we list three neurons from crate (row 1) and GPT-2 (row 2), respectively. For each neuron, we show two top-activated text excerpts and one random excerpt.
Results show that crate is able to consistently activate on sementically similar tokens within the most relevant text excerpts, and does not activate on random tokens that are semantically distinguished from the top tokens. This promotes a more precise explanation given by the explanation model (Mistral in the figure).
On the other hand, GPT-2 is much worse at distinguishing tokens from different contexts, because it also has high activations on random text excerpts where the semantic meanings deviate far from the top activations.

Comparing CRATE to GPT-2 with post-hoc sparse auto-encoders. We follow Bricken et al. [] and train SAEs for models with layers $L\in\{1,2,3\}$, using output activations from GPT-2 on the Pile dataset’s training split, leading to the GPT2-SAE model. Details on the SAEs’ architecture and training are in LABEL:app:sae.

The interpretability scores of GPT2-SAE compared to crate and GPT-2, as depicted in Figure 6, reveal that under the long-context OpenAI metrics, GPT2-SAE matches GPT-2 but falls short of crate. This is attributed to its neuron activations becoming nearly 99% sparse after sparse auto-encoding, diminishing interpretability in long contexts. Conversely, under the Anthropic metric, GPT2-SAE surpasses both GPT-2 and crate in interpretability, corroborating findings in Bricken et al. [] that post-hoc approaches enhance short-context interpretability, often a sign of mono-semanticity. However, the interpretability of GPT2-SAE on the Anthropic metric decreases significantly when $\ell$ increases, while crate remains steady, introducing concern in scalability. Further qualitative comparisons are can be found in LABEL:app:qual-neuron-interps.

Besides its good performance on the Anthropic metric, the post-hoc dictionary learning approach requires considerable manual effort. To get a taste, training a sparse auto-encoder for a single GPT-2 layer takes 4 hours when $h=512$ and a day when $h=3072$ on an A100 GPU.

Table: Table 3: Interpretability improvement of crate and GPT-2 after applying SAE. Results are obtained by subtracting the interpretation scores of the language model and the SAE model trained on that language model.

Does crate have more optimal representation than GPT-2 in terms of interpretability? Alternatively, we train SAE models upon the crate model, denoted crate-sae, and compare the interpretability improvement of crate-sae over crate to the interpretability improvement of GPT2-SAE over GPT-2. As shown in Table 3, the improvement of interpretability of crate-sae over crate is smaller than GPT2-SAE over GPT-2 under both OpenAI and Anthropic metrics. This suggests that crate has more optimal representations than GPT-2 in terms of interpretability. Experimental details can be found in LABEL:app:interp-crate-sae.

Steering the CRATE model. Following Bricken et al. [], we manually activate some neurons and observe the logit effects (changes of the token probability of the language model head). Some qualitative examples are shown in Figure 7. Compared to the lossy steering of the SAE models, crate can be steered without loss. Discussions on the lossy steering process can be found in LABEL:sec:steer-lm-or-sae.

Figure: Figure 7: Qualitative examples on logit effects of manually activating feature (i) in crate. Text shown on the right side are the most positive changes in token prediction probability. The logit effects align with feature interpretations.
Refer to caption: x16.png

### 5.4 Ablations: Disentangling Interpretability from Other Factors

Despite the superior interpretability of crate, we’re interested in whether this interpretability comes from some side factors. Specifically, we’re interested in (i) whether it’s due to the worse performance of crate and (ii) whether it’s because crate has less parameters.

Is the improved interpretability due to performance gap? We first investigate whether the interpretability improvement is due to worse performance by comparing the interpretability of two different checkpoints of crate (intermediate checkpoint and full checkpoint). Results in Table 4 show that the interpretability of the intermediate checkpoint is lower than the full checkpoint, which suggests that “sacrificing” performance does not necessarily introduce better interpretation scores.

**Table 4: Validation loss and interpretability of the crate model at different checkpoints. The interpretation scores are measured under the OpenAI TaR metric.**
| Checkpoint | Loss | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | Avg |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| 79B Tokens | 2.38 | 13.9 | 8.3 | 6.5 | 6.0 | 4.3 | 6.7 | 5.1 | 4.6 | 3.7 | 5.2 | 8.0 | 6.6 |
| 158B Tokens | 2.29 | 13.6 | 8.2 | 6.7 | 5.4 | 5.0 | 7.2 | 5.1 | 5.1 | 3.6 | 5.5 | 9.4 | 6.8 |

Another piece of evidence is that crate-2L has higher interpretability scores than crate-1L. As shown in Figure 5 (left), crate-2L has much better performance than crate-1L on the next-token prediction task. On the other hand, as shown in Table 2, the interpretability of crate-2L is also much higher than crate-1L. Thus, a lower performance does not necessarily introduce higher interpretability.

Is the interpretability gap due to number of parameters? We observe that crate-Base (81.2M) has a similar number of parameters as GPT2-Small (81.1M). However, results in Table 2 indicate that the interpretability of crate-Base is higher than GPT2-Small on all metrics, and their layer-wise interpretation scores are also different in Figure 6. This evidence suggests that two models with similar number of parameters does not necessarily have similar interpretability.

Another piece of evidence is that the interpretability of crate/GPT2-1L all the way up to crate/GPT2-12L does not have a consistent trend of increasing/decreasing interpretability, but their number of parameters both monotonously increases. This indicates that a model with larger number of parameters does not necessarily has better/worse interpretability.

## 6 Conclusion, Limitation, and Future Work

In this paper, we demonstrated that replacing the standard transformer architecture with the white-box model crate as a foundational architecture significantly improves the interpretability. Our empirical findings on the capability of crate to be consistent and distinctive on the neuron-level activations underscore the importance of the white-box design in developing better language foundation models, fostering optimism that the introduction of built-in sparse coding approaches will catalyze further advancements in neuron-level interpretations.

Despite these findings, we acknowledge that the performance of crate is not as good as GPT-2 on the next-token prediction task, which is potentially due to the introduction of the ISTA operator that introduces sparsity. This aligns with previous work suggesting that the performance might drop when explicitly introducing sparsity [].
Future work should investigate towards a better trade-off between performance and interpretability of language models with built-in sparsity. It would also be meaningful to research on more qualitative mechanisms in the white-box language model and how to use these mechanisms for downstream edits.

## Acknowledgement

We thank the anonymous reviewers for their helpful comments. Hao Bai would like to thank Yaodong Yu for valuable discussions on the idea and paper, Yuexiang Zhai for discussions on compute resources, Shengbang Tong, Tianzhe Chu, and Druv Pai for discussions on details of the crate architecture.

## Appendix A Details on the CRATE Architecture

### A.1 Parameter Size of crate and GPT-2

crate is smaller than GPT-2 because of the architecture difference. The vanilla GPT-2 architecture has two main parameterized blocks: Attention block and MLP block.

Parameter size of the MSSA Block. In crate, the MSSA block resembles the Attention block, but instead of K, Q, V matrices, we only have one matrix. Therefore, compared to standard transformers, crate uses $1/3$ of the parameters for the multi-head attention part.

Parameter size of the ISTA block. The MLP block in vanilla GPT-2 has one parametric matrix that transforms the input representations to the inner space (usually 4x larger), and another parametric matrix that transforms the inner representations back to the output space (as large as the input space). In crate, the MLP block is replaced by the overcomplete ISTA block, which transforms the input representation to the overcomplete basis ($4\times$ larger) and transforms back with the same parametric matrix. Therefore, compared to standard transformers, crate uses $1/2$ of the parameters for the MLP part.

### A.2 Causal MSSA Block

This process can be implemented by PyTorch-like code shown in LABEL:alg:torchcode-mssa.

Figure: Algorithm 2 PyTorch-Like Code for Causal MSSA Forward Pass

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: - How exactly is the model trained to predict next token? I don’t see any task related loss terms discussed.
- Why do you discard the last layer for interpretability evaluation? Seems a little unfair to me. Don’t we also want last layer to be interpretable?
- Algorithm 1 says you are measuring correlation. Why is the interpretation score not between [-1, 1]?
- How many neurons were used for evaluation in figure 4?
- Are the fig 4 error bars variance or standard deviation?
- Why is the variance zero for models with few layers and then explodes for bigger ones?
- Table 2 says variance of interpretability scores is much smaller for 12L crate models than 12L gpt models, however in fig 12 it looks like the opposite. Why is this the case?
- What is the sparsity i.e. average L0 of the trained SAE models in Table 8?
- What expansion factor mu was used for SAEs in fig 4?
- How does the training time compare to training a GPT-model of similar size?
- The methods mostly improves random-only scores. Could this be because the neurons are less sparse?
soundness: 1
strengths: Addresses an important problem. Some promising results, such as increased interpretability of individual neurons over standard transformers when measured with automated interpretability. Mostly well written. Overall I think the method has promise but is currently lacking evidence to really show whether it's a meaningful improvement over previous methods.
summary: This paper proposes CRATE language models, a new alternative architecture to transformers that aims to be more inherently interpretable by encouraging a sparse/disentangled representation. This architecture was introduced for vision models in a recent paper and this paper adapts it to the vision setting. They show it increases interpretability of individual neurons in automated evaluations.
weaknesses: Overall I feel like the paper is overclaiming a bit in many places:
- To begin with, I really don't like the name white-box, as that implies the model is easy to understand. While it is an improvement over standard transformer, we are far from fully understanding the model, and I think a more accurate name would be dark-grey-box or something. I think calling this an interpretable-by-design model instead of white-box would be better.
- I disagree with some of the authors claims regarding limitations of SAE models, in particular claims that they scale poorly are not supported by the references provided
    - Templeton et al. cited on line 78 actually shows how to successfully scales SAEs to frontier language models
    - citation in line 124 simply shows later layers of the same attention sae are less interpretable, which is not related to scaling. Also don’t think this is true for saes trained on MLP neurons or residual stream. 
    - Overall seems like sparse-autoencoders scale better than the proposed models, i.e. people have trained saes on much bigger models than gpt-2 without issues, and the computational cost of training an SAE is smaller than the cost of training a new language model.
- The paper claims to solve the issue of introducing reconstruction loss when training SAEs, but the introduced method has worse language modeling capabilities than standard models, and the scale of this difference seems similar to that introduced by SAEs


I'm not convinced these models are much more interpretable than previous methods:
- The shown examples show good explanation for the highly activating input the random activations seem not related to explanation at all
- Are the example neurons cherry-picked? I would like to examples of randomly selected neurons
- The paper would benefit for an additional interpretability study with human evaluators
- Doesn't look like this method is significantly more interpretable than original GPT neurons on larger model i.e. L=12. Only random-only score is improved, but it is very low to begin with so I'm not sure how meaningful and improvement there is. Most metrics seem to be within the error margin (why are error bars scaled down 10x? feels intentionally misleading). 

Some parts hard to follow, in particular section 3:
- I feel like section 3 does not describe the architecture in sufficient detail. For example, what is the coding rate R? I should be able to follow the main method without reading the references.

Minor:
- Eq. 15 should have \bar{a} instead of \bar{x]

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: The loss curve for CRATE in Figure 2 qualitatively looks like it has not plateaued. Do the authors expect that further training could cause CRATE to catch up to the transformer baseline in terms of perplexity? How do the loss curves look for model sizes other than Base? It would be quite compelling if, at a smaller model size trained on more tokens, CRATE was found to catch up with the transformer.
soundness: 3
strengths: The CRATE algorithm places the work on a strong mathematical basis. The simulation scoring evaluation shows that the CRATE architecture is quantifiably more interpretable than a naive baseline.
summary: In this study, the Coding RAte TransformEr (CRATE) architecture is adapted to the task of autoregressive language modeling for the first time. CRATE integrates sparse coding directly into the model architecture for better interpretability, moving beyond post-hoc methods. While it performs worse than a vanilla GPT-2 baseline in terms of language modeling perplexity, CRATE shows improved interpretability of neurons as compared to GPT-2.
weaknesses: The reported increase in interpretability for the largest, 12-layer model is fairly modest (16.3%). Mean interpretability scores seem low in absolute terms across the board (usually less than 10 out of 100), indicating that CRATE is only slightly less black box than a normal language model, rather than a genuinely "white box" architecture, as the title suggests. The examples in Figure 1 seem to be cherrypicked to have much higher than average interpretability scores (in the 36 to 50 range). Indeed, because Table 2 shows that CRATE has lower variance in interpretability scores than GPT-2, the examples shown in Fig. 1 must be several standard deviations more interpretable than the mean.

The abstract says "we introduce a white-box transformer-like architecture named Coding RAte TransformEr (CRATE)," but this is misleading since the body of the text recognizes that CRATE was previously introduced by Yu et al. (2023), and the main contribution of the present work is to build a language model on top of the CRATE representation learning framework. The abstract should be updated to more accurately reflect the novel contribution of this work.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.  Given that ISTA essentially operates like a ReLU with thresholded features, would a simpler sparsity constraint on MLP activations achieve similar results?
2. If I did not misunderstood, the CRATE model is trained in a similar approach to common GPT models. Although the paper derived all the components (ISTA and MSSA blocks) via the rate reduction principle, there is no absolute guarantee that the model acts according to this principle and may behave unexpectedly. In fact, I think the separation of compression and sparsification feels somewhat artificial when discussing deep models trained with backpropagation. Could you clarify why this architecture is still considered white-box?
soundness: 3
strengths: 1. The CRATE architecture, integrated with the sparse coding mechanism in the ISTA block, has better neuron-level interpretability than GPT. Individual neurons in CRATE respond more consistently to specific semantic concepts, making the correspondence between neuron activations to language features purer than those in GPT neurons.
2. The model got high scores in neuron-level probing and outperforms GPT in neuron interpretability metrics, suggesting that CRATE’s architectural design aligns well with interpretability goals.
summary: The paper introduces the CRATE architecture for language modeling, aimed at enhancing interpretability of language models, particularly at the neuron level. The architecture integrates a sparse coding mechanism in its ISTA block (an alternative to the MLP block in GPT models), which allows for improved sparsity and interpretability. However, the model’s language pretraining performance is suboptimal compared to GPT by a noticeable margin. The CRATE model is evaluated on interpretability benchmarks, showing improvements in neuron activation consistency on semantics. But the ablation on the architecture design choices are rather insufficient.
weaknesses: 1. While CRATE improves interpretability, it performs unfavorably in language pretraining tasks as shown in Fig 2, indicating a possible trade-off between architectural interpretability and language modeling performance.
2. When interpreted using SAE, CRATE model underperforms in OpenAI metrics, implying that CRATE’s interpretability gains are not consistent across all interpretability methods.
3. The ISTA block’s thresholded activation approach might be over-engineered. The ISTA block resembles a MLP block with two layers, but presented with more complex formulations . Even though the paper (and the CRATE paper it follows) suggest that this design is derived from theoretical principles, there is no ablations in this paper to discuss the trade-offs of using such design, compared to other alternatives. I am curious whether a simpler sparsity constraint on intermediate MLP activations could potentially achieve better balance between interpretability and performance.
4. The paper discussed that improved interpretability may not come from a performance gap. However, the more important question should be if the gap is unavoidable due to the sparsity imposed on the ISTA block, which is essential for the interpretability improvement. One can train the CRATE model really bad to achieve both bad interpretability and performance, but that does not prove that there is no negative correlation between interpretable architecture design and language pre-training performance.