Bilinear MLPs enable weight-based mechanistic interpretability
Michael T Pearce Thomas Dooms Alice Rigg Jose Oramas Lee Sharkey
A mechanistic understanding of how MLPs do computation in deep neural net-
works 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 neverthe-
less 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 eigendecom-
position reveals interpretable low-rank structure across toy tasks, image classifi-
cation, 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.
interpretability, mechanistic interpretability, bilinear, feature extraction, weight-based, eigenvector, eigendecomposition, tensor network
The close-to-linear structure of bilinear MLPs enables weight-based analysis that reveals interpretable low rank structure across multiple modalities.