English

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

Machine Learning 2024-12-06 v1 Artificial Intelligence

Abstract

The goal of mechanistic interpretability is discovering simpler, low-rank algorithms implemented by models. While we can compress activations into features, compressing nonlinear feature-maps -- like MLP layers -- is an open problem. In this work, we present the first case study in rigorously compressing nonlinear feature-maps, which are the leading asymptotic bottleneck to compressing small transformer models. We work in the classic setting of the modular addition models, and target a non-vacuous bound on the behaviour of the ReLU MLP in time linear in the parameter-count of the circuit. To study the ReLU MLP analytically, we use the infinite-width lens, which turns post-activation matrix multiplications into approximate integrals. We discover a novel interpretation of} the MLP layer in one-layer transformers implementing the ``pizza'' algorithm: the MLP can be understood as evaluating a quadrature scheme, where each neuron computes the area of a rectangle under the curve of a trigonometric integral identity. Our code is available at https://tinyurl.com/mod-add-integration.

Keywords

Cite

@article{arxiv.2412.03773,
  title  = {Modular addition without black-boxes: Compressing explanations of MLPs that compute numerical integration},
  author = {Chun Hei Yip and Rajashree Agrawal and Lawrence Chan and Jason Gross},
  journal= {arXiv preprint arXiv:2412.03773},
  year   = {2024}
}
R2 v1 2026-06-28T20:23:37.919Z