Similar to a previous comment, tensor-transformers are a performant alternative,[1] which are more amenable to analytical tools (eg you can use linear algebra on tensors).
rather than running the network many times and seeing what it does, we read off behavioral properties of the network directly from the weights.
This just screams out tensor networks. They may make an easy test case when you generalize to non-random-init models.
I'm also aware of forthcoming work that can compute when two tensors are similar from the weights alone, with similarity being equivalent to "functional similarity on guassian inputs". I'm quite free next week if any of y'all would want to book a call.
A bilinear MLP is both more performant & similar to SOTA archs than a ReLU MLP
Indeed, tensor network diagrams show up in our algorithm (see Appendix A of the paper). We've also been thinking about mechanistic estimation for tensor network contractions as a problem in their own right, partly because they appear to be needed for harder MLP cases.
I just went on vacation, but this public comment is to shame myself to try to reproduce this with random complex matrixes instead of real ones (they scale more nicely with e.g. conditioning and tails of eigenvalues).
Do you think that research on cumulant propagation for LPE on randomly initialized networks (or even LPE in general) would be a good candidate for automation, like AlphaEvolve, Anthropic's automated weak-to-strong researcher, etc.?
The idea would be that the agent proposes code for a tuple (E, C) where E is the expected value vector for the activations and C is a representation of the propagated object (in principle, a factored cumulant tensor or something like it) and a layerwise update function f: (E0, C0), W, b -> (E1, C1). The score would depend on FLOPs and MSE.
Even LPE on trained models could be automated this way with sufficiently capable LLMs, although it would be harder if you needed to also discover the algorithm E that tracks structure across training and would require a good set of training datasets for the estimated networks, and well, another problem is that these AlphaEvolve-like algorithms are very expensive to run. But at least the checker could be optimized to run fast.
However, I don't know whether algorithms without theoretical guarantees or basis, even if they're short and human-readable, would be useful even if they scored well on the evaluation. For example, if some algorithm said "the expected output of the random MLP is just the sum of the last layer biases" and that outperformed random sampling on the eval, it would score well but wouldn't fit into the structure and randomness idea.
Yes, the fact that we have concrete metrics for progress is part of why we are excited about this line of work. We suspect that if we were given algorithms that achieve very low MSE for a given FLOP budget that we'd be able to extract useful insights from them, although it remains to be seen how well this will pan out in practice. We're planning to launch a contest soon to test out this idea, with LLM usage encouraged. (It will be slightly less prescriptive than the setup you suggested, code just has to take in network weights and produce expectations.)
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. ↩︎