Common Causes for Sudden Shifts: Linking Phase Transitions in Sinusoidal Networks

We are very grateful for the detailed review. We have the following comments:

For instance, I cannot validate the quality of the approximations...either additional numerical experiments or theoretical justifications would be highly appreciated.

In our original submission, we numerically evaluated some of our approximations in Section D of the supplementary materials (now called Section E in light of other changes). We should have pointed to this section more explicitly (we now do this at the end of Section 3.3). We have also added more numerical experiments to this section to help validate the quality of our approximations (ex gauging the cosine alignment between the true principal eigenvector and our approximation from Section 3.3).

Finally, I am not sure about the simultaneity of the order parameters since according to Figure 4, the confidence regions of MAG-Ma and many other order parameters hardly overlap. In this sense, MAG-Ma can be interpreted as a time-shifted order parameter.

Thank you for the feedback. We believe this was an error in how we were performing change detection for MAG-Ma. In the symmetric phase, the MAG-MA statistic is close to zero, and at the phase transition, it grows rapidly. To identify the critical point, as we did for CKA, AUC, and other parameters, we fit a sigmoid function to the trajectory of the order parameter. The inflection point $t^*$ of the sigmoid then corresponds to the critical point, with the slope $w$ providing a confidence interval $[t^*-2w, t^*+2w]$ (by analogy with the Gaussian erf, this corresponds to the region with $2\sigma$ certainty).

The issue with this, of course, is that, for MAG-Ma, that inflection points happens well after the statistic has grown past zero! To this end, it makes more sense to use the interval $[t^*-3w, t^*-w]$ for MAG-Ma. With this change, the confidence interval for MAG-Ma cleanly overlaps with the other order parameters, as seen in Figure 4.

From my perspective, additional numerical experiments for Figures 4-6 are required to disprove that the results are only valid for this specific configuration. Could you please perform these experiments on other datasets or in modified settings?

We have expanded the supplementary materials with additional figures to show that the findings indeed hold across a number of settings. Though the specific shapes and timings vary, the observations described in Section 4.1 and 4.2, pertaining to how the shapes and timing vary with hyperparams, hold across datasets.

We are very grateful for the detailed review. We have the following comments:

However, these analyses all appear to be speculative and do not explain why these phenomena occur.

Thank you for the feedback. We agree that a full theoretical explanation of why these phenomena occur is a complex undertaking and remains an open research question. However, as with previous ICLR submissions, we believe the novelty of our analysis and observations is sufficient to inspire further research, even though we cannot, at present, provide a detailed accounting of their origins.

The analysis is based on fitting isotropic Gaussians to key parameters which the paper itself states (in Sec 3.2) is not appropriate and disconnected from actual empirical training. Further, the paper’s arguments for symmetry breaking appear to contradict this assumption as well.

We believe we can state the main contribution of 3.1 more clearly. The key point of this section is to demonstrate how the experimentally-observed locality of the SIREN NTK (as visualized in Figure 2), when analysed in the continuum, leads naturally to the wavelike evolution of the residuals. The diffusion equation merely offers us some theoretical justification for our choice of orders parameters (the asymptotic value $C_{\infty}$ of the $C_{NTK}$, along with the correlation length scale $\xi_{corr}$).

We initially chose the uniform Gaussian approximation to simplify our presentation, as a means to build intuition about how strongly-local kernels result in diffusion equations. We have now updated our derivation in Section 3.1 to include a spatially-varying correlation lengthscale $\xi$ (the form of the diffusion equation is mostly unchanged). Note that it is possible to carry out the derivation with a anisotropic Gaussian as well. The main difference is that the simplification in equations 50-51 is not possible, but the resulting equation (55) would still be a diffusion equation, just an anisotropic one (the second order term would look like $\text{tr}(\Sigma(x)\nabla_x^2 r(x))$).

We believe this addition is unnecessary, however, as in Section 3.2, we refine the local approximation to encompass other spatial variations in the NTK. Likewise, our discussion on symmetry breaking does not refute the assumption of locality: we simply observe that, at the beginning of training, there is very little spatial variation in the local structure of the NTK, whereas at the phase transition, we observe significant increase in spatial variation.

However, please correct me if I’m wrong, but this is tautologically true – the NTK always changes when training starts to work...the argument comes down to “training works when the image edges start appearing”...

To clarify, it would be more correct to say that, as edges start appearing in the principal eigenvector $v_0(x)$, training begins to slow down. Per the argument in Section 3.3, $v_0(x)$ grows as $\mathcal{O}(||\nabla_{\theta}f(x)||^4)$, which in turn grows for data points that are noisy/influential in training. One candidate explanation is accumulating errors in modelling the edges leads to increased gradient confusion, which in turn slows down learning. However, at present, the exact causal mechanism falls outside the scope of the paper (though some preliminary analysis is found in Section C). What our paper does report is the surprising fact that these different dynamical characteristics exhibit sudden, simultaneous changes at a critical point during training. As far as we can tell, this observation is not obvious, nor has it been reported before. We would be very grateful if you could point us towards a reference.

We would like to thank the reviewer for their thoughtful review. As we understand it, the reviewer's primary concerns are:

1. Certain results in the paper seem to be obvious.
2. How applicable are our observations for the broader range of tasks that SIREN models are used for?

We believe that both of these concerns are the result of a mixup regarding the goals of our paper. For example, objection (1) seems to stem from the following:

...if there is a critical point in training, it is unclear to me why we would expect these characteristics (which are correlated with image edges) to not all change at the same time coinciding with the appearance of edges in the reconstruction and why this is a non-obvious statement.

We would like to clarify that the critical point does not concern the appearance of edges in the reconstruction (ie the model outputs $f_{\theta}(x)$), but in the principal eigenvector $v_0(x)$ of the NTK (ie the Gram matrix for the gradients $\nabla_{\theta}f_{\theta}(x)$). Whereas the former is directly trained to mimic a target signal, the latter is an emmergent property of the system, whose evolution during training is still an active area of investigation.

Here's why this isn't obvious: In CNNs, where pixels correspond to features, feature-based interpretability tools highlight how the early layers of the network learn edge detectors. In INRs, however, pixels correspond to individual datapoints. Unfortunately, the highly nonlinear nature of gradient descent obscures the relationship between the final learned parameters of a DNN, and the dataset it was trained on.

In the case of 2D image regression, we have some intuition about what datapoints are important, based on the wealth of Computer Vision literature. The question is, can we demonstrate that these intuitions are valid for SIREN models?

The NTK offers a simplified picture of dynamics under gradient descent, recasting evolution in terms of the residuals, instead of the model's parameters. There are, however, many gaps here, chief amongst them being the gap between static NTK approximations (which empirically describe SIRENs very poorly, see Figure 4 here) and the true dynamics of the network. In practice, the dynamics of the empirical NTK, and its principal eigenvectors, is a topic of active research, though many theoretical analyses focus on linear architecture, on a single task (classification) (Atanasov et al., (2021); Shan & Bordelon (2022); Baratin et al. (2021)).

Our work is a significant step forward in that regards, because SIRENs (the focus of this paper) are fully-nonlinear DNNs. In particular, we analyze:

• What does the principal eigenvector learn? Phrased another way, what is the NTK aligning with? The empirical observation that the SIREN NTK is dominated by its local structure allows us to relate (in Section 3.3) the principal eigenvector to the distribution of edges in an image.
• When does NTK alignment take place? Our theoretical and empirical studies demonstrate that NTK alignment, in SIREN models, happens suddenly, and around the same time that other phase transitions occur.

We have reorganized our introduction in order to more clearly emphasize our focus.

We are very grateful for the detailed review. We have the following comments:

It is unclear to what extent the results are specific to sinusoidal networks...A simple question to ask is the following: does it make sense for this paper to be rewritten using ReLU MLPs applied to low-dimensional regression tasks (which I suppose would look more like the references [23,24,25])?

We have added a new section to the supplementary materials (section D) which compares SIREN models to ReLU-MLPs with a positional encoder, as used in [44]. Our main empirical takeaway is as follows: while ReLU-PE models also exhibit Neural Tangent Kernel alignment, it is a much slower, non-local process, that does not coincide with loss-rate collapse or translational symmetry breaking. In contrast, in SIRENs, the NTK is dominated by its local structure, and NTK alignment takes place alongside loss rate collapse, translational symmetry breaking, and the onset of diffusion.

It is plausible that this is the underlying mechanism explaining the diffusion phase in the training dynamics, but it is not clear if this is a coincidence or not...{if} we use the gradients of the INR to parameterize the diffusion equation (15), would they yield similar results?

Thank you for the feedback. We think that the diffusion behaviour is a consequence of the locality of the NTK, as seen in Figure 2. Intuitively, due to locality, a residual at a point will only change significantly if it is has neighbours with high residuals. Over time, this disturbance propagates to the neighbours' neighbours, resulting in the formation of a propagating wavecrest. Our derivation of the diffusion equation merely expresses this insight mathematically.

Regarding the comparison between the solution of the diffusion equation and the true gradient dynamics, this is something we have thought a lot about. While using the gradients of the INR might be a tricky, one nice thing about our diffusion equation is that it is very easy to identify the steady-state solution:

This is just a Helmholtz equation, which would be possible to analyze using numerical methods. However, defining "similarity" in this context - that is, comparing the analytical solution to that obtained by gradient descent - requires some careful thinking. We would be very open to hear suggestions.

That said, we think this analysis may fall outside the current scope of the paper, which is more concerned with deriving, and understanding the relation between, different order parameters that summarize the phases of learning. In deriving the diffusion equation, our goal was to identify the order parameters that described the onset of the wave behsviour (which we found were the correlation lengthscale, and the asymptotic value of the $C_{NTK}$). The remainder of the paper then examines how the evolution of these parameters matches the evolution of other order parameters.

It would be nice, either via more experiments or an expanded discussion, to see how the proposed metrics could be used by practitioners to guide their design and use of INRs.

Our work reinforces previous observations about the critical roles played by gradient confusion [13] and correlation length-scale [12] in determining the dynamics, and generalization behaviour, of DNNs. Though we do not claim these as a contribution, in our paper, for SIREN models specifically, the locality of the NTK has allowed us to derive alternative ways to calculate proxies for these metrics, which are potentially easier to compute, and thus guide practitioners. The utility of metrics like MagMa and $AUC(v_0, \nabla I)$ towards this end would be interesting as future work, but at present we feel it is outside the scope of the paper.

Happily. On page 10, in the related works section, we describe some of the models that are commonly studied in the NTK alignment literature. A common strategy is to analytically derive insights from linear models or deep linear models where the evolution of the NTK is more tractable, and then empirically show that the insights translate to nonlinear architectures [23,24,35]. By contrast, the models we study, SIRENs, are fully nonlinear (sinusoidal activations), deep (3-7 layer) neural networks. What makes our analysis tractable is that SIRENs are trained on low dimensional inputs, in this case, single 2D images. Thus, it is much easier to identify how the individual datapoints, and the image gradient at those points, impact model learning. In particular, we focus on why the principal eigenvectors of the NTK align with these image gradients during training.

Thank you for the reply.

it contains important observations about learning in a more complex set of models than is usually studied in the NTK alignment literature, and as such represents a step towards practical advances in the field.

Could you elaborate on this point in detail?

Thank you for taking the time to read our paper. We appreciate your feedback and understand your perspective on theoretical work. While the work is presently specialized, we do believe it contains important observations about learning in a more complex set of models than is usually studied in the NTK alignment literature, and as such represents a step towards practical advances in the field.

We are very grateful for the detailed review. We have the following comments:

It is important to contrast phase transitions in CNNs vs INRs

Thank you for your feedback. It is crucial to note that training Implicit Neural Representations (INRs) and Convolutional Neural Networks (CNNs) are fundamentally different learning problems with distinct characteristics. In this case, we train an INR to represent a single image, where the inputs are pixel coordinates, and the outputs are grayscale values. The resulting model is a learned continuous representation of a single image.

By contrast, CNNs are trained on multiple image to create predictors for a task that generalise to new data. In an INR, each pixel coordinate/grayscale value is a distinct training datapoint, whereas in a CNN, each pixel grayscale value is a distinct input feature. Thus the impact of "edges" are very different in both models. In CNNs, models are encouraged to learn filters which are templates of edges. However, the relationship between edges and model weights is much more subtle in INRs. Our paper illustrates that insight can be gained from the phenomenon of Neural Tangent Kernel alignment. For example, we demonstrate that the principal eigenvector for the NTK of a SIREN is sensitive to data points where the gradient $\nabla_x f$ is high. In the case of a CNN classifier, the same eigenvectors would not detect edges in images, but rather, edges between image manifolds (e.g. an image that is midway between being a cat and a dog).

Given how different the use cases, the architectures, and the conclusions are for INRs and CNNs, we don't think it is reasonable to consider CNNs, and their phase transitions, in the main body of the text.

Computational experiments are insufficient

Our revised submission contains results from ten additional images (full set shown in Figure 9), which we do draw from ten different ImageNet classes. We believe performing our experimental sweep on each image of ImageNet (as you recommend) is impractical - with our variations in seeds, architecture, and hyperparams, this would necessitate us to train on the order of 100M models.

We do recognize that some interesting results could be found from such a study. From the preliminary analysis in Section C, we note that NTK alignment and loss rate collapse do not always take place, and there is some relationship between the occurrence of these phase transitions and the complexity of the image. While tackling this would certainly make for a thought-provoking paper, we think it is out of scope for the present work, but certainly worth looking at in the future!

Why not evaluate training dynamics over a substantive set, e.g., a set of models for 1000 different images

We have included more figures in Appendix G. As outlined above, each image produces one INR, so the order parameter trajectories can only be collected for one image at a time. The specific timings for the evolution of the order parameters vary immensely across different architectural choices and images. In Figure 11 of the Appendix, you can also see all runs plotted on a single histogram, comparing SIRENs with ReLU-Pe.

That said, the high variance of the order parameters between runs makes these sorts of plots very cluttered. Instead, we compare the relative timings between critical points of the order parameter trajectories (for example, looking at how certain critical points tend to cluster). An example is the right side of Figure 4, which examines the co-occurence of detected phase transitions, aggregated across all architectures and images studied.

Can the order parameters be computed efficiently during training to be used as stopping criteria or to adjust the learning rate?

It is important to emphasize that many of the order parameters considered in this work are themselves referenced from other works that investigated just this question. Indeed, our work reinforces previous observations about the critical roles played by gradient confusion [13] and correlation lengethscale [12] in determining the dynamics, and generalization behaviour, of DNNs.

For SIREN models specifically, the locality of the NTK has allowed us to derive alternative ways to calculate proxies for these metrics. In particular, our decomposition of the NTK (A.1) and the local structure (A.3.2) have a lower memory footprint than direct computation. However, we are dubious about extending this approximation to higher-dimensional datasets.

MAG-Ma is efficient to compute, and would be interesting to explore in future work. Presently, we feel it's outside the scope of the paper.

see Joseph Goodman's book

Thank you for pointing us towards this reference! We have added a citation to the Goodman book in the main text.