Bilinear MLPs enable weight-based mechanistic interpretability

We value the reviewer's positive feedback about the technique's potential and are motivated to clarify the work as best we can. We first clarify the reviewer’s questions and then list how we included all proposed changes. > I couldn't understand what Figure 2 is demonstrating. Figure 2 provides intuition on a given eigenvector’s activation behavior. In short, the activation is defined by the squared inner product of the vector with the input. If an input matches strongly with the red (positive) or blue (negative) part of the eigenvector, the activation will be high. If it matches both, the terms for the inner product will be canceled. The activation will depend on how well the input matches each shape. If the input has a shorter stripe “-”, the activation will be lower. Next, the eigenvectors are not positionally invariant (as we don’t use CNNs), so they only match the specific indicated region. Sometimes, the model creates more spatially independent edge detectors, such as the digit 5 eigenvector, which will match with top strokes at several locations. Does the reviewer have any proposals on how to clarify the figure? > Why is input noise an effective regularizer? Training with input noise forces the model to learn robust features that can be detected despite the noise. These features tend to activate more sparsely [1] and generalize better to the test set since spurious patterns due to overfitting are less robust to noise. > Is the bilinear MLP related to the gated linear network literature? While the idea is similar, replacing non-linearities with a linear gating mechanism, the operationalization is quite different. The two main differences are that we perform weight-based interpretability by utilizing the fact that bilinear MLPs can be rewritten in terms of tensors. The halfspace gating used in GLNs is non-linear and cannot be rewritten as such. The other is that our work aims to remain as close as possible (both in spirit and accuracy) to MLPs, while their approach is quite different. To address your remaining questions/suggestions, we: - We have added named dimensions when introducing W, V, and B in Sec. 2 - added A) and B) captions to Figure 2. - added a clarifying sentence (both in sections) about what we mean by eigenvectors (and which matrix it corresponds to). - changed the W on line 121 to be bold (nice catch!). - renamed variables on lines 138 and 189 (as they were not the same). We encourage further questions, remarks, or suggestions that would benefit this work. [1] Trenton Bricken, Rylan Schaeffer, Bruno Olshausen, and Gabriel Kreiman. Emergence of sparse representations from noise. In International Conference on Machine Learning, pp. 3148–3191. PMLR, 2023a.

I thank the authors for taking the time to respond and revise the manuscript. I now understand Figure 2 better but still have some doubts about the use of the term "edge detector". If the authors mean something that "only match the specific indicated region", perhaps they shouldn't call them edge detectors. The term edge detector is long-standing in image processing. It typically refers to algorithms like Canny edge detector [1] and Sobel edge detector, which detect edges of different shapes in all regions of an image. [1] Canny, John. A computational approach to edge detection. TPAMI (1986).

We thank the reviewer for their comments regarding the benefits of our interpretability techniques. We are pleased the reviewer appreciated the significance and presentation of our work, particularly the choice of experiments and examples. > I did not notice much analysis using SVD, perhaps I missed it. Given the symmetric nature of all the matrices we consider, we opted to use the eigendecomposition over the SVD (which are similar in nature and equivalent up to sign for symmetric matrices). The eigendecomposition has two appealing qualities over SVD: It splits positively and negatively contributing features (according to eigenvalues). The projection matrix is the same on each side, making interpretability easier. > Can you justify the use of the HOSVD as used? Two concerns: 1) The tensor represents a non-linear function. 2) If removed all but but the largest singular values of B, would the resulting tensor functionally approximate B? Also, what were the results of the SVD analysis, what section is that in? 1) The B tensor is a multilinear map, meaning it is linear in each dimension while holding the others constant. But when we use the same vector x as the input to both of its input dimensions, then the overall output is non-linear in x. 2) Yes, HOSVD and other approaches for truncating the B tensor will approximate B. They provide a good approximation even when different input vectors x1 and x2 are used instead of the same vector x as in our use case. 3) We did not sufficiently highlight the HOSVD approach, which we have added to appendix D. This approach combines both SVD (to find the shared output features) and the eigendecomposition (to find the principal interactions for those features). > Minor Comment: Giving some dimensions on W, V could help with the presentation. We have also added named dimensions when introducing new matrices in Sec. 2. Thanks for the suggestion. We look forward to hearing whether the changes are suitable and whether the reviewer has additional follow-ups, suggestions, or questions.

We thank the reviewer for their insightful comments. We are encouraged to hear that the paper was understandable and helpful for people less familiar with the area. >...hence the lack of robust demonstrations at a broader range of tasks, models, and circuits. Being able to showcase such range would significantly strengthen the influence of the work in the nearer future. Balancing the clarity of the presentation with the range and depth of demonstrations was challenging. We appreciate any suggestions on the types of demonstrations or tasks that would further strengthen our work. For image classification, we opted to do a “deep dive” to exhaustively demonstrate the utility and robustness of the features derived from model weights (e.g., truncating the model with minimal accuracy loss, assessing similarity across runs and model sizes, using features to reveal the impact of regularization, etc). For language modeling (likely the main use-case for our approach), our approach showed that SAE features for the MLP outputs can be computed using surprisingly few MLP input directions and highlighted an interesting example. Since our submission, we have trained a larger GPT-2 medium-sized model and will add results similar to those in Sec. 4 to the paper to demonstrate that our results are generalized. > Perhaps some of the existing problems can be solved with advances in discovering semantic “output directions” as part of the method. Regarding semantic output directions, one approach we did not sufficiently highlight is HOSVD. This allows us to find the most important output directions (in contrast to only picking one). Intuitively, this finds the most important “shared features”. We did not include these experiments as the sharing in these small models is limited and less immediately interpretable than the main results. However, we have added an additional appendix D and added references to it in Sec. 3.3. On this note, an advantage of bilinear MLPs is that they allow for flexible analysis, particularly using various tensor decompositions. In future work, we plan to develop approaches suitable for analyzing multiple MLP layers simultaneously. Please let us know if the changes are satisfactory and whether you have further suggestions for improving this work.

We added additional experiments on GPT2-level models trained on FineWeb (instead of TinyStories) in Section 5, specifically Figure 9. We find our results to generalize to these models that are an order of magnitude larger, strengthening our claims. We also added Appendix H, which discusses that high-quality SAEs are better explained using fewer eigenvectors. Additionally, in a previous revision, we added Appendix D, which discusses HOSVD to find more ‘semantic’ output directions and may interest the reviewer regarding their raised suggestions. If the reviewer feels we have not sufficiently addressed their concerns, we would kindly appreciate some targeted feedback. Otherwise, we would be grateful if the reviewer considered increasing their score.

We are delighted the reviewer shares our enthusiasm for the paper and the proposed approach. We are also happy to hear that the presentation was clear. > a more detailed explanation of the role of sparse autoencoders in Section 3.1 would enhance the clarity and comprehensiveness of the work We have expanded our explanation of the use of SAEs in Section 3.1. We clarified that the overall approach can be used for any dictionary of input and output features, but it is not specific to SAEs. With SAEs, however, the encoder directions can be used to transform the B tensor such that output feature activations can be inferred directly from input activations. We would strongly appreciate any further feedback (even nitpicks) to strengthen this work further.

With the discussion period closing, we wondered whether the reviewer was satisfied with the modifications to the paper. In short, we scaled our experiments on language models to GPT2-medium, a 10x scale increase, and showed they generalized; we also added two new appendices (D and H). If the reviewer feels the work is deserving, a score increase would significantly increase the paper’s visibility (and impact).