Mechanistic estimation for wide random MLPs
Jacob_Hilton
Abstract
This post covers joint work with Wilson Wu, George Robinson, Mike Winer, Victor Lecomte and Paul Christiano. Thanks to Geoffrey Irving and Jess Riedel for comments on the post. In ARC's latest paper, we study the following problem: given a randomly initialized multilayer perceptron (MLP), produce an estimate for the expected output of the model under Gaussian input. The usual approach to this problem is to sample many possible inputs, run them all through the model, and take the average. Instead, we produce an estimate "mechanistically", without running the model even once. For wide models, our approach produces more accurate estimates, both in theory and in practice. Paper: Estimating the expected output of wide random MLPs more efficiently than sampling Code: mlp_cumulant_propagation GitHub repo We are excited about this result as an early step towards our goal of producing mechanistic estimates that outperform random sampling for any trained neural network. Drawing an analogy between this goal and a proof by induction, we see this result as (part of) the "base case": handling networks at initialization. We have a vision for the "inductive step", although we expect that to be much more difficult. Summary of results In our paper, we consider MLPs with weights , defined by where the activation function is applied coordinatewise, and is taken to be by default. Schematic of our ReLU MLP. An estimation algorithm takes in and a tolerance parameter , and aims to estimate to within an error of around . We evaluate estimation algorithms by checking their mean squared error over weights with randomly initialized entries drawn independently from .[1] Our baseline is Monte Carlo sampling, which draws samples, runs them through the model, and averages the results. As a function of the width and the tolerance (holding the depth constant), in the average case over initializations, this has mean squared error and runs in time . Our best-performing algorithm, on the oth