PaperlayerView exact revision on arXiv

Source-linked AI summary

Bilinear MLPs enable weight-based mechanistic interpretability

Michael T. Pearce, Thomas Dooms, Alice Rigg, Jose M. Oramas, Lee Sharkey

arXiv:2410.08417v2cs.LGstat.ML

TL;DR

MLP weights are difficult to interpret because element-wise nonlinearities obscure how inputs and outputs relate, while activation-based analyses generally do not explain how features are formed. The paper analyzes bilinear MLPs through weight decompositions and finds interpretable low-rank structure across tasks, including language-model circuits. It concludes that weight-based interpretability is viable for large language models, subject to limits on available output directions and eigenvector interpretability.

  • Problem

    Element-wise nonlinearities obscure how MLP weights construct features, and activation-based approaches generally describe which features are present rather than how they form or affect downstream computation.

    Source: §1 · p. 1

  • Method

    The authors express bilinear MLPs with a third-order tensor and decompose the weights into eigenvectors that explain outputs along selected directions.

    Source: §1 · p. 2

  • Results

    Across image-classification tasks the decompositions reveal interpretable low-rank structure, and in language models they identify a negation-related sentiment circuit and find many output features correlated with low-rank approximations.

    Source: §1 · p. 2, §6 · p. 10

  • Takeaways & Limitations

    The results support weight-based interpretability as viable for large language models and bilinear MLPs as an interpretable drop-in replacement for ordinary MLPs.

    Source: §1 · p. 2, §6 · p. 10

    Methods typically need meaningful output directions, which in deeper models rely on dataset-dependent SAE features, and eigenvectors are not guaranteed to be monosemantic.

    Source: §6 · p. 10

Abstract

from arXiv · show

A mechanistic understanding of how MLPs do computation in deep neural networks remains elusive. Current interpretability work can extract features from hidden activations over an input dataset but generally cannot explain how MLP weights construct features. One challenge is that element-wise nonlinearities introduce higher-order interactions and make it difficult to trace computations through the MLP layer. In this paper, we analyze bilinear MLPs, a type of Gated Linear Unit (GLU) without any element-wise nonlinearity that nevertheless achieves competitive performance. Bilinear MLPs can be fully expressed in terms of linear operations using a third-order tensor, allowing flexible analysis of the weights. Analyzing the spectra of bilinear MLP weights using eigendecomposition reveals interpretable low-rank structure across toy tasks, image classification, and language modeling. We use this understanding to craft adversarial examples, uncover overfitting, and identify small language model circuits directly from the weights alone. Our results demonstrate that bilinear layers serve as an interpretable drop-in replacement for current activation functions and that weight-based interpretability is viable for understanding deep-learning models.

1 INTRODUCTION

The paper targets the gap between identifying features in activations and explaining how MLP weights construct them. It presents bilinear MLPs as an interpretable alternative and applies weight decompositions to image and language-model analyses.

1 INTRODUCTION

  • Element-wise nonlinearities obscure how MLP weights, inputs, and outputs relate, while activation-based methods generally identify features without explaining their formation or downstream effects.The paper argues that weight-grounded nonlinear feature interactions would help make circuit analysis more faithful.Source: §1 · p. 1
  • Bilinear layers express computations with a third-order tensor, enabling direct weight decompositions, and are proposed as an interpretable drop-in replacement for ordinary MLPs.The paper notes their language-modeling performance is above ReLU-based transformers and slightly below SwiGLU.Source: §1 · p. 2
  • Eigenvector analyses across image-classification tasks reveal interpretable low-rank structure, allow smaller eigenvalue terms to be truncated while preserving performance, and expose overfitting patterns useful for adversarial masks.The analyses also show how regularization reduces signs of overfitting in extracted features.Source: §1 · p. 2
  • In language models, the authors identify a small circuit that flips next-token sentiment when the current token is a negation, and find many output features correlated with low-rank approximations.They present these findings as evidence that weight-based interpretability can be viable in large language models.Source: §1 · p. 2

2 BACKGROUND

Bilinear MLPs are GLU variants without an element-wise nonlinear activation; their pairwise input interactions can be represented by matrices organized in a third-order tensor. The section also describes a symmetry simplification and a bias-handling transformation.

2 BACKGROUND

  • A bilinear layer removes the GLU’s nonlinear activation, beats ordinary ReLU MLPs on language modeling, and performs almost as well as SwiGLU.With training time held constant, bilinear layers achieve equal loss; with data held constant, their loss is marginally worse.Source: §2 · p. 2
  • Each output’s pairwise input interactions form an interaction matrix, and the matrices across outputs can be organized into a third-order bilinear tensor.The tensor can be contracted along its output axis to obtain an interaction matrix for a chosen output direction.Source: §2 · p. 3
  • Because each interaction matrix is evaluated on two copies of the same input, its antisymmetric part contributes zero and can be omitted, leaving a symmetric form with real eigenvalues and orthogonal eigenvectors.The paper states that this simplifies eigendecomposition.Source: §2 · p. 3
  • Biases can be incorporated by augmenting the weight matrices and input with a constant coordinate; experiments otherwise omit biases for simplicity without harming performance.The same framework is illustrated on a toy classification task with biases.Source: §2 · p. 3

3 ANALYSIS METHODS

Bilinear MLP weights can be analyzed with linear-algebra decompositions, with the choice of method depending on available input or output features. The section presents direct feature interactions, output-conditioned eigendecomposition, and higher-order SVD when features are unavailable.

3.1 INPUT / OUTPUT FEATURES →DIRECT INTERACTIONS

  • Given input and output feature dictionaries, the bilinear tensor can be transformed into their feature basis to study how output features are constructed from input features.For sparse activations, the authors study the statistics of active interactions; SAE output encoder directions can also be used in the transformation.Source: §3.1 · p. 3

3.2 OUTPUT FEATURES →EIGENDECOMPOSITION

  • For an output feature direction, contracting the tensor produces an interaction matrix whose eigenvectors identify important input directions and whose largest-magnitude eigenvalues give the best low-rank approximation.The eigenvectors diagonalize the matrix, so there are no cross-interactions between them; Figure 1 diagrams the tensor-to-matrix decomposition.Source: §3.2 · p. 4, Figure 1 · p. 4

3.3 NO FEATURES →HIGHER-ORDER SVD

  • When no prior features are available, the tensor can be reshaped by flattening its two input dimensions and applying a standard SVD; the resulting interaction matrices can then be decomposed into eigenvectors.The paper demonstrates this approach for an MNIST model in Appendix D.Source: §3.3 · p. 4

4 IMAGE CLASSIFICATION: INTERPRETING VISUAL FEATURES

Across image-classification tasks, bilinear-layer eigenvectors reveal interpretable, often low-rank features; their consistency, behavior on a ground-truth task, and use in adversarial masks illustrate what can be learned directly from weights.

4.1 QUALITATIVE ASSESSMENT: TOP EIGENVECTORS APPEAR INTERPRETABLE

  • MNIST eigenvectors emphasize class-specific curve segments or prototypical images, whereas Fashion-MNIST eigenvectors act as localized edge detectors.The quadratic activation is large when input aligns with either the positive or negative eigenvector regions, but not both.Source: §4.1 · p. 5, Figure 2 · p. 5
  • Only a small fraction of eigenvalues are substantial, and dense Gaussian input noise makes features more interpretable while revealing overfitting patterns in unregularized models.Increasing noise produces more digit-like eigenvectors and a lower-rank spectrum.Source: §4.1 · p. 6, Figure 4 · p. 6

4.2 QUANTITATIVE ASSESSMENT: EIGENVECTORS LEARN CONSISTENT PATTERNS

  • Across runs, top eigenvectors have cosine similarities of 0.8–0.9, and retaining a handful yields minimal accuracy drops (0.01%).The figure compares ordered eigenvectors and accuracy after truncating to the most important eigenvalues.Source: §4.2 · pp. 6–7, Figure 5 · p. 6

4.3 COMPARING WITH GROUND TRUTH: EIGENVECTORS FIND COMPUTATION

  • On the ground-truth MNIST task, one positive eigenvector computes the target similarity, while two remaining eigenvectors appear to overfit the data.The decomposition uses the True − False output difference; the negative eigenvector seems to penalize diagonal structures.Source: §4.3 · p. 7, Figure 6 · p. 7
  • Including biases benefited this task because the model needed a learned threshold for its binary decision.Without bias, the model sought quadratic invariances that did not generalize well.Source: §4.3 · p. 7

4.4 ADVERSARIAL MASKS: GENERAL ATTACKS FROM WEIGHTS

  • Eigenvector-derived adversarial masks reduced accuracy more than random-mask baselines and increased misclassification toward the targeted digit, even with Gaussian-noise regularization.For unregularized models, eigenvectors also exposed overfitting on rarely active outer-edge pixels, enabling masks to be constructed by hand.Source: §4.4 · p. 8, Figure 7 · p. 8

5 LANGUAGE: FINDING INTERACTIONS BETWEEN SAE FEATURES

The section uses weight eigendecomposition to examine interactions between sparse-autoencoder input and output features. It identifies a sentiment-negation circuit and finds that many output-feature activations are well approximated by low-rank decompositions.

5 LANGUAGE: FINDING INTERACTIONS BETWEEN SAE FEATURES

  • The analysis applies eigendecomposition to input and output feature dictionaries from sparse autoencoders trained on a six-layer TinyStories bilinear transformer.The authors focus on single-layer circuits and describe nonlinear interactions grounded in layer weights.Source: §5 · p. 8, §5.1 · p. 8

5.1 SENTIMENT NEGATION CIRCUIT

  • The layer-4 not-good and not-bad output features form an opposing linear subspace and respond to negation tokens alongside positive- or negative-sentiment tokens.Their decoder-vector cosine similarity is -0.975.Source: §5.1 · p. 9
  • Top interactions for the not-good feature resemble an AND gate: negative-sentiment and negation input features interact positively, while their self-interactions are negligible.The positive-sentiment input feature interacts with the opposite sign.Source: §5.1 · p. 9, Figure 8 · p. 9
  • The two largest eigenvalues, 0.62 and -0.66, capture the negation circuit’s geometry; the top-two-eigenvector approximation correlates 0.66 with active not-good activations and 0.76 at large activations.A sentiment feature alone activates both eigenvectors and cancels, while a negation feature strongly activates the positive eigenvector.Source: §5.1 · p. 9, Figure 8 · p. 9

5.2 LOW-RANK APPROXIMATIONS OF OUTPUT FEATURE ACTIVATIONS

  • Most output features have high active-only correlation with low-rank approximations, often above 0.75 using two eigenvectors.Average correlation starts around 0.65 with one eigenvector and increases with additional eigenvectors; approximation improves with longer SAE training.Source: §5.2 · p. 9, Figure 9 · p. 10

6 DISCUSSION

The discussion argues that bilinear layers support intrinsic, weight-based interpretability by exposing feature interactions and low-rank structure. It also notes limits on the available output directions and on the interpretability of eigenvectors.

6 DISCUSSION

  • The authors identify an example low-rank circuit that flips next-token sentiment when the current token is a negation, and say its behavior is easier to understand through top eigenvectors.They contrast this with finding a similar circuit in conventional MLPs.Source: §6 · p. 10
  • Bilinear layers can be transformed into explicit feature-interaction matrices and decomposed in a way equivalent to the original computations, supporting weight-based interpretability.The paper reports that many language-model SAE output features are well approximated by low-rank interaction matrices, particularly at large activations.Source: §6 · p. 10
  • The method typically needs meaningful output directions, which in deeper models come from dataset-dependent sparse-autoencoder features; eigenvectors are not guaranteed to be monosemantic.The authors expect orthogonality may limit interpretability for high-rank spectra.Source: §6 · p. 10

CONTRIBUTIONS

The authors describe their respective contributions to the MNIST and language-model work and say they wrote the paper in tandem.

CONTRIBUTIONS

  • Michael did most of the MNIST analysis and contributed insights across the presented topics.The passage assigns Thomas the language-model section and code infrastructure.Source: CONTRIBUTIONS · p. 11
  • The authors wrote the paper in tandem, each focusing on their respective section.Michael focused on MNIST; Thomas worked on language models and code infrastructure.Source: CONTRIBUTIONS · p. 11

ETHICS STATEMENT

The paper presents methods for analyzing model internals rather than advances in model capabilities; the authors identify dataset-detail leakage as a possible misuse and suggest safety evaluation as a way to avoid it.

ETHICS STATEMENT

  • The authors say the work does not advance state-of-the-art model capabilities and identify efficient leakage of learned dataset details as its envisioned misuse.They say this misuse can be avoided by using the technique during safety evaluation.Source: ETHICS STATEMENT · p. 11

REPRODUCIBILITY STATEMENT

The appendices document reproducibility provisions, illustrate digit-specific eigenvector features, and summarize how regularization and augmentation affect feature quality.

REPRODUCIBILITY STATEMENT

  • The authors describe reproducible training setups, figure-generation code, seed use, planned model sharing, and an interactive demo.The demo is intended to support independent analysis of figures in Appendices A, B, and O.Source: REPRODUCIBILITY STATEMENT · p. 11

A EIGENSPECTRA: SHOWING EIGENVECTORS ACROSS DIGITS

B REGULARIZATION & AUGMENTATION: ABLATIONS & OBSERVATIONS

  • Input noise sparsifies features, geometric transformations blur them, and dropout has little impact on features.The passage summarizes the appendix’s regularization and augmentation observations.Source: B REGULARIZATION & AUGMENTATION:… · p. 16

B.1 REGULARIZATION

Input noise and weight decay affect eigenfeatures differently, while several geometric augmentations smooth or broaden features without preventing overfitting.

B.1 REGULARIZATION

  • Dense Gaussian input noise gives the best reported trade-off between feature interpretability and accuracy; salt-and-pepper noise lowers both, and Perlin noise does not fix overfitting.Input noise has the largest feature impact among the explored techniques.Source: §B.1 · p. 16
  • Adding Gaussian noise to activations had no measurable impact in these experiments, possibly because the models were shallow.The authors qualify this result as potentially specific to their shallow models.Source: §B.1 · p. 16
  • Weight decay generally sparsifies eigenvalues without significantly changing eigenvectors, potentially eliminating an unimportant long tail and reducing analysis effort.This makes it easier to analyze a model without substantially altering the eigenvectors.Source: §B.1 · p. 16
  • Dropout did not appear to affect the models, and overfitting persisted even at rates above 0.5.The authors suggest this result may differ in larger or capacity-constrained models.Source: §B.1 · p. 16

B.2 AUGMENTATION

  • Translation smooths features but does not prevent overfitting; high translation can split the same pattern across locations and increase rank.The reported maximum shift is about 7 pixels in each direction.Source: §B.2 · p. 16, Figure 13 · p. 16
  • Rotation broadens features, while blur makes them somewhat smoother; neither alone prevents edge-pixel overfitting without noise.Combining blur and rotation produces smooth, curvy features.Source: §B.2 · p. 16, Figure 14 · p. 17

C EXPLAINABILITY: A SMALL CASE STUDY WITH EIGENVECTORS

The appendices demonstrate eigenvector-based explanations and shared-feature analysis, quantify sparsity effects, compare model sizes, and document experimental setups.

C EXPLAINABILITY: A SMALL CASE STUDY WITH EIGENVECTORS

  • A correctly classified five had similarly positive activations for classes 5, 6, and 8, but class 5 was least suppressed; a two was misclassified as 7 when its strokes activated 7’s eigenvector more strongly.The first case illustrates classification through suppression, while the second had relatively unimportant negative eigenvectors.Source: C EXPLAINABILITY: A SMALL CASE STUDY… · p. 18, Figure 15 · p. 18, Figure 16 · p. 18

D HOSVD: FINDING THE MOST IMPORTANT SHARED FEATURES

E SPARSITY: WEIGHT DECAY VERSUS INPUT NOISE

  • Input noise increases the number of important eigenvectors and affects eigenvector sparsity, whereas weight decay decreases their number but does not substantially affect that sparsity.The analysis treats eigenvalues as indicating effective rank and top eigenvectors as indicating effective pixel count.Source: E SPARSITY: WEIGHT DECAY VERSUS INPUT… · p. 20, Figure 19 · p. 20

F TRUNCATION & SIMILARITY: A COMPARISON ACROSS SIZES

  • Across model sizes, few eigenvectors generally recover full accuracy, narrower models tend to be more low-rank, and top features of similarly sized models are mostly similar.The figure also compares eigenvector similarity across all model sizes.Source: Figure 20 · p. 20

G.1 IMAGE CLASSIFICATION SETUP

  • The image-classification models comprise an embedding, a bilinear layer, and a head or unembedding; most listed training parameters have little effect in these small models.Input noise is an exception to the parameter-insensitivity statement.Source: §G.1 · p. 21, Table 1 · p. 21

G.2 LANGUAGE MODEL SETUP

  • For language modeling, a six-layer modern transformer (about 33 million parameters) replaces SwiGLU with bilinear MLPs and trains on simplified, cleaned TinyStories; separate FineWeb models use GPT-2-small/medium architectures.The FineWeb models have 12 and 16 layers, respectively, with 162M and 335M parameters; both use the Mixtral tokenizer.Source: §G.2 · pp. 21–22

G.3 SPARSE AUTOENCODER SETUP

  • The sparse autoencoders use TopK activations with k = 30, and the section’s dictionaries use a 4x expansion.The appendix points to the training hyperparameters and layerwise added loss in its table and figure.Source: §G.3 · p. 22, Figure 23 · p. 22, Table 4 · p. 23

H CORRELATION: ANALYZING THE IMPACT OF TRAINING TIME

Across these appendices, longer SAE training improves low-rank feature approximation, while bilinear layers are less data-efficient but equally compute-efficient in the reported comparison. The remaining sections describe tensor decomposition, bias terms, adversarial masks, and language-model feature analyses.

H CORRELATION: ANALYZING THE IMPACT OF TRAINING TIME

I BILINEAR TRANSFORMERS: A LOSS COMPARISON

  • 6% fewer epochs let SwiGLU reach the bilinear variant’s final loss, while their compute-efficiency differences vanish at constant time.With abundant data, the authors report little disadvantage to bilinear layers over the other variants.Source: Table 6 · p. 24, I BILINEAR TRANSFORMERS: A LOSS… · p. 24

J FINETUNING: YOUR TRANSFORMER IS SECRETLY BILINEAR

  • Fine-tuning TinyLlama toward β = 0 increased loss by about 0.05 relative to continued training, but the difference decreased quickly with more training.The tested schedule interpolated β to zero over the first 120M tokens, then trained for 280M tokens.Source: J FINETUNING: YOUR TRANSFORMER IS… · p. 24, Figure 25 · p. 24

K TENSOR DECOMPOSITIONS: EXTRACTING SHARED FEATURES

  • The tensor decomposition expresses each interaction matrix through eigenvectors and eigenvalues, and a full-rank set of output directions permits recovery using a pseudoinverse.Eigenvectors are orthogonal within an output direction but may overlap across different output directions.Source: K TENSOR DECOMPOSITIONS: EXTRACTING… · p. 25

L BILINEAR LAYERS: A PRACTICAL GUIDE

  • Bilinear layers model pairs of features; adding biases supplies linear and constant terms, though training often defaults to quadratic invariances instead of learned constants.The authors note that linear structure matters for some toy tasks.Source: L BILINEAR LAYERS: A PRACTICAL GUIDE · p. 25

M ADVERSARIAL MASKS: ADDITIONAL FIGURE

  • The additional adversarial-mask examples compare eigenvector-based masks for models with no regularization, Gaussian noise regularization at 0.15, and Gaussian noise regularization at 0.3.Source: Figure 26 · p. 25

N LANGUAGE: FURTHER DETAILS FOR FEATURE CIRCUITS

  • The language-model figure shows a histogram of self- and cross-interactions and an eigenvalue spectrum for the sentiment-negation feature, alongside a random-matrix reference spectrum.Source: Figure 27 · p. 26

O INPUT FEATURES OF THE NEGATION FEATURE

Loading 2410.08417v2…