Modular addition without black-boxes: Compressing explanations of MLPs that compute numerical integration

Thank you for reviewing our paper, and for providing constructive feedback. Per your general feedback, we've posted an updated manuscript: updated section 1 with more problem motivation, comparison with related work in section 2, and the remaining paper with notation changes and additional sentences and labels to make it easier to follow.

Is it possible to extend the proposed framework to MLP layers in a deeper transformers?

Yes. We can restrict the analysis to an MLP layer within a deeper transformer and get some analytic expression for what the layer is doing. However, piecing together what the network does as whole would be more difficult, since error bounds can explode over multiple layers of the network. Please refer to the general response (the section on applications of our methods).

The study seems only focuses on understanding MLPs within a constrained setting where the MLP approximates trigonometric functions, is it possible to extend to non-trigonometric functions?

Yes. As long as we have an analytic expression for the MLP weights, we can potentially use the idea of numerical integration to analyse the layer. Please refer to the general response (the section on applications of our methods).

I am still confused on the definition of “complete interpretation”, are there any evaluation metrics?

We have removed the framing of “complete interpretation” and instead rewritten section 3 to provide a clearer connection with [1], upon which we build our validation procedure. The evaluation metric is the computational complexity needed to check the interpretation. The lower the complexity, the better we've compressed the model, and hence the better our interpretation. We believe that a computation linear in the number of parameters is the best one can achieve here, since we would have to look through all the weights only once.

Please let us know if there are any other clarifications we can provide or improvements we can make!

[1] Gross et al. Compact Proofs of Model Performance via Mechanistic Interpretability. (2024)

Thank you for your response and for addressing some of my concerns. I think the revisions improve the clarity of your paper. The clarification on the challenges in extending the framework to deeper transformers, and the explainations on potential applications are beneficial. However, I feel these updates do not fully address the concern regarding the scalabilty and practical applications in more complex settings. Therefore, I will maintain my original score.

Thank you for your response. We want to note that we're the first paper to develop an approach to interpreting densely-connected MLP layers. While we cannot make changes right now, we plan on adding a section clearly detailing the structure of the infinite-width argument that we make in this paper, which is applicable to other settings.

Consider a nonlinear layer $\text{MLP}\left(a\right) = W_{out} \text{ReLU} \big (W_{in} a \big)$, where activation $n_i$ of neuron $i$ is given by $\text{ReLU}\big ((W_{in})_{i} ~ a \big)$. Our goal is to find an approximation $F(a)$ of each output $\text{MLP}_j\left(a\right)$ such that we can bound $|{\text{MLP}_j\left(a\right) - F(a)}| < \varepsilon$ without computing $\text{MLP}_j\left(a\right)$ for each $a$. While $\text{MLP}\left(a\right)$ is a discrete sum over finitely many neurons, we are trying to find an analytic approximation $F(a)$.

The infinite-width lens argument is constructed by an approximation of the discrete sum of the MLP output as follows: $\sum_i w_i f(a; \xi_i) \approx Z \int{\xi_0}^{\xi_n} f(a; \xi) ,\mathrm{d}\xi = F(a)$.

where

• $\xi_i$ is a ``locally one-dimensional" neuron-indexed variable of integration
• $w_i$ is approximately linear in $(\xi_{i+1} - \xi_{i-1})/2$
• $f(a; \xi_i) - f(a; \xi_{i-1})$ is uniformly small over all values of $a$, i.e. $f$ ``continuous" in $\xi$

If we are able to find values satisfying the constraints, then we can analytically evaluate the integral $\int_{\xi_0}^{\xi_n} f(a; \xi) ,\mathrm{d}\xi = F(a)$ independent of $a$, allowing us to bound the error of the numerical approximation of $f$ at each $\xi_i$, independent of $a$, for example via lipschitz constant $\times$ size of box.

Each mechanistic interpretability case study is a fair bit of human labor, and deriving these values will not be trivial. Moreover, without empirically checking, we cannot say how good a given approximation is. However, our approaches to error-bounding, and constructing the infinite-width / numerical integration approximation are sound, and general purpose solutions to interpreting MLPs.

Thanks for taking the time to review our paper. We've posted an updated manuscript incorporating feedback.

We want to start by addressing your concern that our contribution is not significant. We appreciate this note, and have modified sections 1, 2, 3, 7 to frame and describe our contribution more clearly.

1. We build upon [1] which is a fully formal approach to mechanistic explanation of models. In [1], the authors only interpret one layer transformers with no MLP layers, trained on a simpler task of computing the max of k integers. Our work applies the formalization to a new, harder case study.
2. Mechanistic interpretability is generally intensive labor, and formal mechanistic interpretability is even more intensive!
3. To our knowledge, we are the first to interpret MLP layers rigorously. We think this is a significant direction of exploration. By interpreting feature-maps, we are able to uncover a completely novel algorithmic description of the model. This is important for applications of mech interp geared at high stakes applications like anomaly detection and guarantees, but also for lower-stakes applications like model distillation and model steering. We consider our work as framing that feature-maps are important to interpret, and then concretely showing that interpreting them adds insight even to the models that are considered best understood in mech interp.

Have the authors investigated the effect of weight decay tuning on the clarity of their experimental results?

We use large weight decay when training the network so that the weights have smaller L2-norm, which gives rise to a more compact model. This is demonstrated in previous papers about grokking. However, our paper falls in the paradigm od post-hoc mechanistic interpretability, which we have clarified in section 3. Thus, we take as given that the networks are trained so that they perform perfectly on the dataset.

What factors influence the emergence of secondary frequencies in the model?

As stated in section 5, secondary frequencies improve the effectiveness of the model by improving its performance when cos(k/2 (a-b)) \approx 0. It can be interesting in future work to investigate the development of secondary frequencies over the training phase, and also the optimal strength of the secondary frequency with respect to the loss function.

Let us know if you have any more questions or suggestions! We're happy to incorporate any feedback you have on our work.

[1] Gross et al. Compact Proofs of Model Performance via Mechanistic Interpretability. (2024)

Given that the discussion period ends in approximately 16 hours, we would greatly appreciate the reviewer's thoughts on our updates and replies. We welcome any additional questions or comments, and remain happy to address them to the best of our ability within the remaining time.

Thank you very much for your time and feedback! We have uploaded an updated manuscript incorporating recommended changes.

The authors only consider Transformers with constant attention.

We’re following the terminology of Zhong et al, since we’re extending their interpretation. However, we've made the scope of our work clearer, and emphasized that we are studying the ReLU feature-map only.

Interpretability for MLPs trained on modular addition has already been extensively studied in literature.

We've modified section 2 to start with a description of prior work on analyzing modular addition models. We've clearly framed that our goal is to compress feature-maps, which prior work does not do. Moreover, we clarify that we worked on the modular addition model since they are some of the best understood models in literature. Thus, by finding more insight from analyzing the MLP, we demonstrate that the problem we're interested in is important.

Are all the experiments conducted with constant head Transformers? I couldn't find this detail in the Colab notebook since the checkpoints are directly imported for analysis.

Yes. (i.e. there is no attention that is dependent on the inputs.) We trained the models with a separate codebase, so it is not shown in the Colab notebook.

According to [3], networks can learn either "clock" or "pizza" algorithm, depending on the details of training. How do the authors ensure that the learned algorithm is always "pizza"?

Our preliminary analysis of the clock model (which is not included in the write-up) suggest that the primary difference between the "clock" and the "pizza" model is where the transition from $\cos a + \cos b \to \cos a \cos b$ happens. In the pizza model, we have (by construction) addition of embeds pre-ReLU, thus resulting in pre-activations which have significant coefficients in the $\cos a \cos b$ and $\sin a \sin b$ Fourier components. By contrast, the "clock" models seem to do some of this multiplication in the attention layer (as recognized by Nanda et al), but no neuron's pre-activations contain Fourier components of both $\cos a$ (or $\sin a$) and $\cos b$ (or $\sin b$). These components are then added post-ReLU to get a dependence on $a+b$. Thus the "clock" algorithm (using attention to get $\cos a \cos b$ or $\sin a \sin b$ components for unembed, and also forwarding only either a or b to the MLP) is impossible to implement in the architecture we use (constant attention). Regarding the consistency of learning the "pizza" algorithm, the results from 150 random seeds (Appendix C in our paper), most of the models agree with the formula for the weights that we suggested. One might consider this to define the “pizza” algorithm.

Are Figures 2,5 empirical results or just theoretical visualizations?

They are empirical results coming from the trained models and our calculations as detailed in section 4.1.

Lines 299-303: The omission of x/2 term seems rather hand-wavy. How does one see that computing the bounds after ignoring this term is a reasonable thing to do? It is unclear to me even after referring to Appendix K.

This is a quirk with our error metric. We are considering relative error, i.e. $\frac{error bound}{theoretical value}$. For the x/2 term, the theoretical value of the integral is 0, hence relative error would not make sense. Thus, we focused on the $|x|/2$ term (with non-zero theoretical value) instead, to demonstrate how our method works. It is possible to compute bounds for ReLU as a whole (see Appendix I). Also, the $x/2$ term is in a sense “simpler” than the $|x|/2 term$, because it is just a linear layer. Indeed, we can find a matrix $A \in \mathbb{R}^{p \times p}$ such that the logit contribution from this part of the network is $A[a,:] + A[b,:]$, for inputs $(a, b)$. Thus, this part of the network can be seen as already interpreted in time linear in the number of parameters.

Given that the discussion period ends in approximately 16 hours, we would greatly appreciate the reviewer's thoughts on our updates and replies. We welcome any additional questions or comments, and remain happy to address them to the best of our ability within the remaining time.

I thank the authors for the detailed responses to my questions and concerns.

The answers and the manuscript-edits address some of my concerns about clarity and soundness. However, I continue to have reservations about the significance of this work towards understanding Transformers trained on modular arithmetic and the field of mechanistic interpretability.

Consequently, I will increase my score but remain below acceptance.

Thank you for your support of our paper, we appreciate it!

Limited scope

We have clarified the contribution of this paper to be a prototype / first case study in compressing feature-maps. We believe that this is an important problem, and have updated our introduction.

Bounds are crude

The standard we adopt for mechanistic explanation is formal proof. Prior work has only empirically developed this on transformers with no MLP layers. Thus, despite the crudeness of bounds, we think that getting non-vacuous bounds itself in a step forward.

You suggest that this approach could be useful for practitioners in other domains. What properties of the model/problem would make it more likely that this approach would yield a good explanation?

When the model has a clear functional/mathematical form, this approach of numerical integration would work better, as we can use similar methods to guess the functional form implied by the network and bound the errors. As such, this can be useful when we look at neural networks used to solve mathematical problems (e.g. solving PDEs). However, it can also be used in other domains: if we spot some analytical expression in certain parts of the network, we could potentially convert this to a statement about the input / output space of the problem. Please refer to the general response for more details.

You leave deriving tighter error bounds to future work. What avenues for proving these bounds are you considering?

One can use properties of quadrature schemes to derive tighter bounds. Currently, we are using the naive estimate of transforming a weighted sum to rectangles under a curve. There may be better methods from numerical analysis (Gaussian quadratures etc.) which give better bounds.

We hope we've been able to address your concerns about the paper. Please let us know if there's any more information we can provide, or changes we can make!