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.
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 other hand, has mean squared error and runs in time – a factor of faster.[2] This means that for any given depth, our algorithm is guaranteed to outperform Monte Carlo sampling at large enough width, although the dependence of our algorithms on depth is worse.[3]
Since this theoretical result provides no guarantee about the performance at any particular width, we empirically check the performance of our algorithm at realistic widths, looking at mean squared error versus the number of floating point operations (FLOPs) performed.[4] For ReLU MLPs with 4 hidden layers and width 256, our best algorithms outperform Monte Carlo sampling at FLOP budgets spanning 7 orders of magnitude, in some cases achieving the same mean squared error using fewer than as many FLOPs.[5]
Performance of our best cumulant propagation algorithms at estimating the mean of random ReLU MLPs with 4 hidden layers and width 256.
Perhaps more importantly, our algorithm outperforms Monte Carlo sampling in the tails of the distribution to an even greater extent than the body. For low probability estimation, Monte Carlo sampling is essentially useless for probabilities significantly below below , where is the number of samples used. On the other hand, using a similar number of FLOPs to Monte Carlo sampling with samples, our algorithm sometimes achieves a relative error of under 30% for probabilities 100 times lower than .[6]
Finally, it is straightforward to extend our method from the expected output of the network to bona fide loss functions. We demonstrate this for a simple distillation loss between two networks of different widths. Since our estimates are differentiable, we can train the student network using our estimates of this distillation loss, a process we refer to as mechanistic distillation. Although we do not outperform ordinary training, this serves as a proof of concept for using mechanistic estimation for training.
Significance of results
We are excited about these results not simply because of the performance improvements, but because our algorithms are "mechanistic": rather than running the network many times and seeing what it does, we read off behavioral properties of the network directly from the weights. If we could produce good mechanistic estimates for frontier models, we should be able to catch deceptive alignment at train time, even if the model's behavior looks benign on every training input.
This potentially offers us a completely different way to train models that avoids deceptive alignment. Instead of using stochastic gradient descent, we can apply gradient descent to a mechanistic estimate, a process we refer to as mechanistic training. Even if mechanistic training is about as efficient as ordinary training overall, it would produce models that allocate capacity in a very different way, because the kinds of errors that mechanistic estimates make are very different. These models would therefore generalize very differently.
For example, consider a loss function with a certain "dangerous" event that is very rare but has very high loss. Stochastic gradient descent may fail to sample the event even once during training, resulting in a model that does not account for the dangerous event at all. On the other hand, mechanistic training could notice how the rare event might transpire, resulting in a model that does a better job of avoiding it (potentially allocating capacity to this at the expense of other aspects of performance). This could be especially important if we are concerned about the dangerous event becoming significantly more likely under distribution shift.
Of course, we are a long way from being able to train frontier models using mechanistic training. Randomly initialized networks serve as the natural starting point for this goal, both because they are the simplest networks, and because every network starts out randomly initialized. We have produced mechanistic estimates for these networks that significantly outperform random sampling (for wide networks),[7] as well as providing demonstrations of low probability estimation and mechanistic training.
Note that the high-level algorithm in our paper is not new: it is a form of cumulant propagation, an algorithm we introduced in 2022 in Formalizing the presumption of indepdence (Appendix D). This works by propagating an approximate probability distribution through the model, without running the model on any particular input. However, a number of specific details about our algorithm are new, and these are essential for achieving low mean squared error.[8] It is only because we reoriented our research around outperforming random sampling that we ended up discovering these details and obtaining our results.
Extending to trained networks
Although we are able to outperform random sampling, our methods rely on the networks being randomly initialized. Roughly speaking, our methods start from Gaussian approximations to the activation distributions, and then track the lowest-order deviations from that. In a trained network, specific higher-order deviations can become much more important, and so our algorithm would need to be adapted to track these. We suspect that this can be made to work with some sort of auxiliary "advice" given to the estimation algorithm, which somehow points out which of these higher-order deviations to track, and which gets updated incrementally with each gradient step. In our analogy with a proof by induction, the algorithm for this incremental update would be our "inductive step".
One idea for what this could look like is for the advice to provide a different, more structured weight distribution such that the model's weights can be viewed as having been drawn randomly from this distribution. This structured distribution can be thought of as an efficient compression of the model's weights, an idea we discuss further in Compression as a possible MSP approach.[9] Regardless, being able to handle trained networks is clearly essential for our methods to be of practical utility, and we are pursuing this direction in further research.
Conclusion
Our paper achieves our goal of producing mechanistic estimates that outperform random sampling in the case of wide, randomly initialized MLPs. We are working on extending this to trained networks in order to produce new training methods that reduce the likelihood of deceptive alignment.
For , this is the standard He initialization, and ensures that the outputs of have variance around 1. ↩︎
This works providing that grows like for some . ↩︎
The depth scaling of our algorithms is discussed Appendix D of the paper. ↩︎
We also empirically validate that the scaling with width matches our theoretical predictions, as discussed in Section 6.3 of the paper. ↩︎
Our algorithms often underperform Monte Carlo sampling in wall-clock time, since we make no serious attempt to optimize their performance on hardware. We provide wall-clock time in Appendix I of the paper. ↩︎
We explain this result in Section 6.6 of the paper. ↩︎
It remains an open problem to match the performance of random sampling for random MLPs in other regimes, for example, with the depth growing linearly in the width. ↩︎
An ablation of some of these details is discussed in Section 6.4 of the paper. ↩︎
We don't think this approach works quite as we have stated it. For example, if the model starts by computing a one-way function of the weights, we may want an efficient compression of the output of this one-way function, rather than of the weights themselves. ↩︎
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
where the activation function
Schematic of our ReLU MLP.
An estimation algorithm takes in
to within an error of around
Our baseline is Monte Carlo sampling, which draws
Since this theoretical result provides no guarantee about the performance at any particular width, we empirically check the performance of our algorithm at realistic widths, looking at mean squared error versus the number of floating point operations (FLOPs) performed.[4] For ReLU MLPs with 4 hidden layers and width 256, our best algorithms outperform Monte Carlo sampling at FLOP budgets spanning 7 orders of magnitude, in some cases achieving the same mean squared error using fewer than
Performance of our best cumulant propagation algorithms at estimating the mean of random ReLU MLPs with 4 hidden layers and width 256.
Perhaps more importantly, our algorithm outperforms Monte Carlo sampling in the tails of the distribution to an even greater extent than the body. For low probability estimation, Monte Carlo sampling is essentially useless for probabilities significantly below below
Finally, it is straightforward to extend our method from the expected output of the network to bona fide loss functions. We demonstrate this for a simple distillation loss between two networks of different widths. Since our estimates are differentiable, we can train the student network using our estimates of this distillation loss, a process we refer to as mechanistic distillation. Although we do not outperform ordinary training, this serves as a proof of concept for using mechanistic estimation for training.
Significance of results
We are excited about these results not simply because of the performance improvements, but because our algorithms are "mechanistic": rather than running the network many times and seeing what it does, we read off behavioral properties of the network directly from the weights. If we could produce good mechanistic estimates for frontier models, we should be able to catch deceptive alignment at train time, even if the model's behavior looks benign on every training input.
This potentially offers us a completely different way to train models that avoids deceptive alignment. Instead of using stochastic gradient descent, we can apply gradient descent to a mechanistic estimate, a process we refer to as mechanistic training. Even if mechanistic training is about as efficient as ordinary training overall, it would produce models that allocate capacity in a very different way, because the kinds of errors that mechanistic estimates make are very different. These models would therefore generalize very differently.
For example, consider a loss function with a certain "dangerous" event that is very rare but has very high loss. Stochastic gradient descent may fail to sample the event even once during training, resulting in a model that does not account for the dangerous event at all. On the other hand, mechanistic training could notice how the rare event might transpire, resulting in a model that does a better job of avoiding it (potentially allocating capacity to this at the expense of other aspects of performance). This could be especially important if we are concerned about the dangerous event becoming significantly more likely under distribution shift.
Of course, we are a long way from being able to train frontier models using mechanistic training. Randomly initialized networks serve as the natural starting point for this goal, both because they are the simplest networks, and because every network starts out randomly initialized. We have produced mechanistic estimates for these networks that significantly outperform random sampling (for wide networks),[7] as well as providing demonstrations of low probability estimation and mechanistic training.
Note that the high-level algorithm in our paper is not new: it is a form of cumulant propagation, an algorithm we introduced in 2022 in Formalizing the presumption of indepdence (Appendix D). This works by propagating an approximate probability distribution through the model, without running the model on any particular input. However, a number of specific details about our algorithm are new, and these are essential for achieving low mean squared error.[8] It is only because we reoriented our research around outperforming random sampling that we ended up discovering these details and obtaining our results.
Extending to trained networks
Although we are able to outperform random sampling, our methods rely on the networks being randomly initialized. Roughly speaking, our methods start from Gaussian approximations to the activation distributions, and then track the lowest-order deviations from that. In a trained network, specific higher-order deviations can become much more important, and so our algorithm would need to be adapted to track these. We suspect that this can be made to work with some sort of auxiliary "advice" given to the estimation algorithm, which somehow points out which of these higher-order deviations to track, and which gets updated incrementally with each gradient step. In our analogy with a proof by induction, the algorithm for this incremental update would be our "inductive step".
One idea for what this could look like is for the advice to provide a different, more structured weight distribution such that the model's weights can be viewed as having been drawn randomly from this distribution. This structured distribution can be thought of as an efficient compression of the model's weights, an idea we discuss further in Compression as a possible MSP approach.[9] Regardless, being able to handle trained networks is clearly essential for our methods to be of practical utility, and we are pursuing this direction in further research.
Conclusion
Our paper achieves our goal of producing mechanistic estimates that outperform random sampling in the case of wide, randomly initialized MLPs. We are working on extending this to trained networks in order to produce new training methods that reduce the likelihood of deceptive alignment.
For
This works providing that
The depth scaling of our algorithms is discussed Appendix D of the paper. ↩︎
We also empirically validate that the scaling with width matches our theoretical predictions, as discussed in Section 6.3 of the paper. ↩︎
Our algorithms often underperform Monte Carlo sampling in wall-clock time, since we make no serious attempt to optimize their performance on hardware. We provide wall-clock time in Appendix I of the paper. ↩︎
We explain this result in Section 6.6 of the paper. ↩︎
It remains an open problem to match the performance of random sampling for random MLPs in other regimes, for example, with the depth growing linearly in the width. ↩︎
An ablation of some of these details is discussed in Section 6.4 of the paper. ↩︎
We don't think this approach works quite as we have stated it. For example, if the model starts by computing a one-way function of the weights, we may want an efficient compression of the output of this one-way function, rather than of the weights themselves. ↩︎