Modular addition without black-boxes: Compressing explanations of MLPs that compute numerical integration
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.
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}
}