Dynamical versus Bayesian Phase Transitions in a Toy Model of Superposition

It seems that the way we have written the introduction may have led to a misunderstanding. While it is listed as the first contribution of the paper, we do not believe that providing a classification of the critical points for the high-sparsity TMS potential is the primary contribution of the paper. Indeed, it's unclear how much generalisable insight such a classification would provide, since TMS is a very toy system and the classification is quite laborious even in this case.

The primary contribution of the paper lies in proving empirically that Bayesian phase transitions occur, that they are well-described by SLT, and that there are interesting (but not yet completely established) links between these Bayesian transitions and the dynamical transitions that many in the deep learning community have noted.

To respond to some of your individual points:

• You highlighted the importance of detailing the classification of critical points within the main text rather than relegating it to the appendices. Please note the definition of $k$-gon critical points is already presented in Section 3.2 of the main text. Although Figures A.1 to A.3 in Appendix A enhance the visualization of these $k$-gons, we regret that space constraints prevent their inclusion in the main text.
• It is true that changing the system would invalidate the current classification. However, as remarked above, the classification is not the primary contribution of the paper: it is one piece of the puzzle, that allows us to establish the other contributions listed above as the actually important contributions (about phase transitions). So the toy model we study is interesting not because the classification of critical points generalises, but because the Bayesian theory of phase transitions may generalise to larger models. This expectation is very reasonable theoretically, and the experiments in this paper are the first step towards verifying this theory empirically.
• You state that "I don't think the authors provide an adequate quantitative study of the phase transitions of this model beyond some general remarks". This paper contains, by far, the most detailed study so far in the literature of phases and phase transitions in any neural network and goes far beyond general remarks, providing a classification (both theoretically and empirically backed) of phases and phase transitions.
• Re: the role of the analytical results in the dynamical transition picture. There are ample evidence that properties of the critical points other than their energy level, play an important role. Analytical results about the geometrical properties like the learning coefficients can gives us a more detailed picture. For example, it leads to the curious observation that as loss drops during SGD training, there is a corresponding increase in complexity.

Regarding your questions:

• Re: whether other critical points are local minima. We can calculate Hessians at each of them and see that they contain no negative eigenvalues, but it is more difficult to prove they are local minima and this has not been done in all cases. As the paper is already 60 pages we do not include all these calculations.
• Re: phases with zero proportion. The question seems to be about phases that aren't present in Figure 2, such as $4^{----}$. Theoretically, the posterior has non-zero mass at all phases for any $n$. The posterior mass of each phase $W_\alpha$ is roughly given by $e^{-F_n(W_\alpha)} = e^{-n L_n(W_\alpha) + \lambda_\alpha \log n + c_\alpha}$. Thus, it is clear that high energy phases are exponentially suppressed and won't be dominant in the regime of $n$ we are probing. Figure F.1 contains the occupancy plots for more of the classified critical points, these are omitted in the main text for simplicity.
• Re: number of SGD steps. This is not a typo. We run 4500 epoch in our SGD experiments, each with 50 SGD steps (n / batch size = 1000 / 20) per epoch.
• Re: SGD escaping plateaus. We have not studied this, it does not seem to be common in the trajectories we studied.
• Re: Figure 3 and SGD initialisation. Different initialisations discover different phases with different probabilities. Studying this distribution was not the aim of the paper. Our aim was to find as many examples of phases and phase transitions as possible. The posterior distribution at low $n$ serves this purpose well. It has more probability mass near critical points of the energy function but is also diffused. When we initialise instead using a Gaussian at the origin, the plateaus in Figure 3 look roughly similar, but less phases are discovered (e.g. $5^+, 6, 4^{---}$ are not discovered). Our current view is that this is not very fundamental, and we would see similar results with either choice if enough SGD trajectories were sampled, but it would be interesting to study further.

Regarding the link between the small-scale learning task studied in our paper and larger-scale and more practical problems, our response is as follows. The paper makes several main claims, including:

• (A) Bayesian phase transitions are theoretically predicted and actually occur in TMS
• (B) Dynamical transitions occur in TMS, and
• (C) There is some relation between Bayesian and dynamical transitions
• (D) The local learning coefficient is a powerful tool for reasoning about phase transitions

In both (A) and (B) we don't see the scope of our results as really being limited to small-scale systems. Since (A) is a general phenomena of Bayesian learning once the general principles are seen to apply in TMS, it is reasonable to expect such transitions to occur in larger systems, since the larger the system the more we expect its behaviour to be governed by general laws. Dynamical transitions have already been observed in larger scale systems, including language models with billions of parameters in the work on Induction Heads and In-Context Learning from Anthropic.

We agree that in (C) the current evidence linking Bayesian and dynamical transitions is restricted to one toy system (TMS). It will require follow-up work to assess whether this can be extended to a general principle for working with larger-scale systems.

However in the case of (D) we believe our work is already informing people's plans about how to think about phase transitions in real-world systems and so we feel confident asserting that our approach will have a substantial impact. For example there has been some preliminary work done already looking at modular addition and grokking using the learning coefficient.

Regarding the context for the paper, the original work by Elhage et al motivates the Toy Model of Superposition as a simple model in which to study the phenomena of superposition, which they view as a key obstacle to mechanistic interpretability. Our motivations are different: we are interested in phase transitions, both in the context of Bayesian statistics (where the phenomena is understudied and lacks good examples) and in deep learning (where somewhat mysterious "phase transitions" over training, what we call dynamical transitions, are attracting increasing interest).

Our theoretical and experimental results give strong reason to believe the Toy Model of Superposition is also a good Toy Model of Phase Transitions, in both of these senses (Bayesian and dynamical).

Regarding what is known about the critical points in TMS. As far as we know nothing has been published on this that goes beyond the qualitative results in the original paper.

Re: making the theoretical results more formal. This is a good question, but we view it as out of scope for the present work. The aim of this paper is to show that Bayesian phase transitions exist, and to explain the theoretical context in which they make sense. To prove that the technical hypotheses of SLT apply to this example, and to further clarify the exact mathematical nature of Bayesian phase transitions, is quite difficult and may be the subject of future work. We agree it is important to clarify the status of hypotheses such as "relative finite variance" (necessary to apply the results of Watanabe's 2018 book) in neural networks.

Re: the BAH. The Bayesian Antecedent Hypothesis is, currently, just that - a hypothesis. This paper provides preliminary empirical evidence supporting the BAH. We are actively working on a rigorous theoretical foundation for the BAH, but it is beyond the scope of this paper.

Re: whether the conditions of Watanabe 2009 apply to the Gaussian prior considered in experiments. We don't quite understand the specific role of the prior in this question, could you please elaborate?

Thanks for your comments. Regarding the originality of results in Section 4 on Bayesian phase transitions: this is the first time this kind of treatment of Bayesian phase transitions with respect to the number of samples, with a division of the parameter space into phases, has appeared in the literature. We cite Watanabe's work because he has written about other kinds of phase transitions in Bayesian learning (where a hyperparameter changes in the prior or the true distribution). As far as we are aware the experiments in this paper are the first time clear phase transitions of the kind we are talking about have been shown to exist.

Re: Weaknesses. We have uploaded a revision where the caption of Figure 1 is expanded to include all the details necessary to read it, and Section 3.3 has been expanded slightly to include some more context on the local learning coefficient. We couldn't find a way to add in an introduction to SLT without pushing important content, currently in the main text, into Appendices, so while we agree an introduction would be valuable we've settled for the status quo of referring the reader to other papers.

Re: Question about momentum. Yes this is a natural question. I think you have in mind that if one uses an optimiser with momentum, it will tend to spend less time stuck near critical points and therefore may collapse the "plateaus" that we are displaying in the context of dynamical transitions. I expect this is correct but the plateaus will still be present. We're running the experiments now and will add a comment when we know more.