Reproducibility Study: Bilinear MLPs Enable Weight-Based Mechanistic Interpretability

This repository contains a reproduction and extension of
Pearce et al. (2024), Bilinear MLPs Enable Weight-Based Mechanistic Interpretability.

All experimental results corresponding to Claims 1–4 in the original paper are reproduced, as well as two additional extensions in a single executable notebook.

---

Main Results and Entry Point

full_report.ipynb

This notebook reproduces all results reported in the accompanying report, including:

Reproduced Claims

• Claim 1 (Image Classification)
  Eigenvector decomposition of bilinear MLP weights across image classification tasks reveals an interpretable low-rank structure.

• Claim 2 (Image Classification)
  Truncating lower-magnitude eigenvalue components has minimal impact on classification performance, indicating that most task-relevant information is captured by the leading eigenvectors.

• Claim 3 (Image Classification)
  The eigenvectors of bilinear MLPs are stable across random initializations and behave similarly across different model sizes.

• Claim 4 (Image Classification)
  Input noise regularization reduces overfitting and produces bilinear MLPs with eigenvector features that are more interpretable under eigenvalue analysis

Extensions

• Extension 1 (Adversarial Training on MNIST)
  Evaluation of FGSM (Goodfellow et al., 2014) and PGD (Madry et al., 2017) adversarial training on both robustness and eigenvector interpretability.

• Extension 2 (CIFAR-10 / CIFAR-100 / Tiny ImageNet)
  Testing whether low-rank structure, eigenvector stability, and truncation robustness generalize beyond MNIST.

All figures, tables, and metrics in the report are generated by running this notebook top-to-bottom.

---

Model Architectures

The core of this reproducibility study centers on the Bilinear MLP, an architecture designed to facilitate mechanistic interpretability directly from model weights. Unlike standard MLPs that use fixed non-linearities (e.g., ReLU or GeLU), Bilinear MLPs utilize a gated quadratic interaction.

1. The Bilinear MLP Layer

The fundamental transformation in these models replaces the standard activation function with an element-wise product of two linear projections. For an input vector $\mathbf{x}$, the output $\mathbf{y}$ is defined as:

$$\mathbf{y} = (W_1 \mathbf{x} \odot W_2 \mathbf{x})W_3$$

Where:

• $W_1, W_2 \in \mathbb{R}^{d_{hidden} \times d_{in}}$: These matrices project the input into a latent space where the bilinear interaction occurs.
• $W_3 \in \mathbb{R}^{d_{out} \times d_{hidden}}$: This matrix projects the gated features back to the output dimension.
• $\odot$: Denotes the Hadamard (element-wise) product.

2. Implementation Variants

We evaluate the following architecture:

• Bilinear Image Classifiers: Shallow networks used for Claims 1–4.
  • Configuration: Typically 2–3 layers deep with a consistent hidden dimension of 1024 units to ensure sufficient capacity for eigenvector decomposition.

3. Interpretability Tools

• Sparse Autoencoders (SAEs): To validate the interpretability of the weights, we train SAEs on the activations of the bilinear layer. This allows us to compare learned features (SAE encoder weights) against the leading eigenvectors of the $W$ matrices.
• Eigenvalue Analysis: We perform a low-rank approximation of the weight matrices to identify "interpretable directions" that correspond to specific semantic or visual features.

---

Datasets

---

Hardware Requirements

• GPU (Recommended): An NVIDIA GPU with at least 8GB of VRAM is recommended for training the image classification models and running the SAE analysis efficiently.
• CPU: A minimum of 16GB of System RAM is required for local execution of the full_report.ipynb notebook without a dedicated GPU.
• Storage: Ensure at least 5GB of free disk space to accommodate the datasets and saved model checkpoints.

Repository Structure

bilinear-decomposition_(cloned_from_authors)/
├── image/ # Image models, datasets, adversarial training code
├── language/ # Transformer language model with bilinear MLPs
├── results/ # Saved models, datasets, emissions logs, figures
├── sae/ # Sparse Autoencoder tools and eigenvector analysis
├── shared/ # Shared bilinear / linear components
├── toy/ # Author-provided toy model configuration 
├── tutorials/ # Author-provided reference notebooks
├── .gitignore # Author-provided gitignore
├── requirements.txt
├── readme.md # Author-provided readme of the original repository
environment-cpu.yml # The conda environment to run locally on your CPU
environment-gpu.yml # The conda environment to run on a GPU
full_report.ipynb # Main notebook for reproducing claims + extensions
.gitignore 

Environment Setup

All implementations use PyTorch.

Conda (recommended)

CPU:

conda env create -f environment-cpu.yml
conda activate fact-ai-cpu

GPU:

conda env create -f environment-gpu.yml
conda activate fact-ai-gpu