Unified Neural Network Scaling Laws and Scale-time Equivalence

We thank the reviewer for their insightful comments and for identifying areas where our manuscript can be improved. We address the concerns below.

Generality of the Theoretical Analysis

Please refer to our general response, where we clarify the scope of our theoretical model and its relevance to neural networks. While our analysis uses gradient flow (akin to full-batch gradient descent), we believe that the insights gained are valuable for understanding fundamental aspects of neural network training dynamics.

We acknowledge that extending the analysis to stochastic gradient descent (SGD) and larger models is important. Although our experiments use mini-batch SGD, prior work has shown that certain theoretical findings under gradient flow can provide useful approximations for SGD behavior, especially when the noise introduced by mini-batches averages out over time.

Scaling Laws and Applicability to Nonlinear Models

In section 5, we first present existing analyses of scaling laws with respect to training time in the context of linear models. These scaling laws have already been shown to be empirically effective in neural network models.

Then, by scale-time equivalence, we apply these analyses to scaling with respect to model scale. Importantly, scale-time equivalence does not impose a linearity assumption; rather, the existing theoretical analysis uses a linear model.

Related Work

We appreciate the reviewer bringing to our attention recent works on neural scaling laws, such as Lin et al. (2023) and Paquette et al. (2023). We have added these papers in the revised manuscript.

Unlike Lin et al. (2023), who study scaling laws in linear regression trained via one-pass SGD, our work focuses on the scale-time equivalence in the context of neural networks trained with gradient-based methods. We believe our findings complement this study by providing a theoretical analysis in more general, nonlinear models.

Limitations of the Work

We have added a dedicated section in the discussion to address the limitations of our work. We acknowledge that:

• Our experiments are conducted on vision tasks with relatively small models and datasets. Extending the analysis to larger models and diverse tasks is an important direction for future research.
• We currently lack a complete theoretical understanding of the cube root scaling observed in our experiments.

By explicitly stating these limitations, we aim to provide a balanced perspective on our contributions.

Additional Questions

• Figure 2 Batch Size: In Figure 2, the batch size is fixed at 32, which is a standard choice in neural network training. We used a consistent batch size across experiments to isolate the effects of model scale and training time.
• Derivation of Equation (10): Equation (10) is derived from Equation (9) simply by substituting $t$ with $pt$ to reflect the scale-time equivalence. We are unfortunately unsure what $M$ you are referring to but are happy to provide further clarification.

We appreciate the reviewer's thoughtful feedback and the opportunity to address the concerns raised. Below, we provide clarifications and additional details to strengthen our manuscript.

Strong Claims and Scope of the Paper

We acknowledge the reviewer's concern regarding the strength of our claims. Our intention is to present a theoretical and empirical analysis that provides insights into the relationship between model scale and training time in neural networks. While our theoretical results are derived under certain assumptions, we believe they offer valuable perspectives that are applicable beyond the specific settings studied.

We have added a discussion of the limitations in the revised manuscript, specifically acknowledging that our experiments are conducted on small-scale vision tasks and that we currently lack a full theoretical understanding of the cube root scaling. We are committed to refining our claims to accurately reflect the scope and applicability of our work.

Applicability of the Random Subspace Model to Neural Network

Please refer to our general response, where we clarify the rationale behind our theoretical model. Our random subspace model is designed to capture key aspects of neural network training, including nonlinearity and overparameterization. By considering a general differentiable loss function and parameterizing the model in terms of a random subspace, we aim to abstract the effect of increasing model size in practical neural networks.

Clarity of Section 3.1

We appreciate the reviewer's request for additional explanations in Section 3.1. We refer the reviewer to our general response; if these explanations are satisfactory, we are happy to integrate them into our revised Section 3.1. We would appreciate if the reviewer could point out any specific parts that need further clarification.

Specific Line Comments

• Line 130: We acknowledge that in practice, the projection $K$ may evolve during training. In our theoretical model, we assume $K$ is fixed to simplify the analysis and focus on the relationship between model scale and training time. This assumption allows us to derive analytical results.
• Line 139: Selecting a random subspace of parameters is conceptually similar to selecting a random subset of parameters to update. In fact, up to rotations in the parameter space, selecting a random subset and selecting a random linear subspace are equivalent.
• Line 152: We use a fixed learning rate in our theoretical analysis to isolate the effect of model scale and training time. We agree that in practice, learning rates can vary with model size, as explored in the μP literature. However, controlling hyperparameters like learning rate is also common in much theoretical literature on scaling laws.
• Line 209: By "minimum amount of training time required to achieve non-zero generalization," we refer to the number of epochs needed for the model to perform better than random chance on the test set.
• Line 244: The cube root argument is detailed in Appendix B. We recognize the importance of this scaling and we are happy to summarize the key points in more detail if it would be helpful. Nevertheless, we view this as not a central contribution to the main focus of our paper.
• Line 252: We used a fixed learning rate to maintain consistency across experiments. This approach is common in studies focusing on specific factors like model size, as it helps isolate their effects. Nevertheless, we acknowledge that tuning learning rates can impact performance.

Section 4 Details

We appreciate the feedback on Section 4 and are happy to include any specific details you suggest to improve clarity.

We thank the reviewer for their constructive feedback and for pointing out areas where our manuscript could be enhanced. We address the concerns below.

Comments on the Random Subspace Model

Please refer to our general response, where we clarify the motivation and applicability of our theoretical model. Our random subspace model is designed to abstract key aspects of neural network training, such as overparameterization and the effects of increasing model size. By considering a general differentiable loss function and a parameterization that reflects the addition of parameters in random directions, we aim to capture the essence of neural network behavior.

Potential Vacuity of the Bound in Theorem 1

The reviewer raises a concern regarding the interpretation of the error bound in Equation (4). The bound contains an exponential term $e^{\eta p t h \| K K^\top \|}$, which suggests that the error could grow exponentially with $p$ and $t$. However, it is important to note that $pt$ represents the product of model scale and training time, which we interpret as the total training effort or progress. Our theorem shows that as $p$ increases for a fixed $pt$, the error bound decreases, indicating that the scale-time equivalence becomes more precise with larger $p$. The exponential term reflects the sensitivity of the training trajectory to initial conditions over longer training times, which is a common characteristic in dynamical systems, and not reflective of the degree of scale-time equivalence.

Clarification on Equation (36)

Equation (36) in the appendix follows from substituting $pt$ for $t$ in Equation (3)​. We are happy to provide further clarification if we have misunderstood the reviewer's question.

Experimental Details and Figure Clarity

Regarding the plots with graded colors, the specific choices of hyperparameters used to generate each line are detailed in Appendix C. We are happy to include these hyperparameters in the caption or change the format of the plot if another visualization would be more helpful.

Figure 3 Column 2

In Figure 3 Column 2 shows the predicted performance of models trained for many epochs based on training larger models for a limited number of epochs. The dashed lines represent the predictions, and they are truncated because, in practice, we are limited by computational resources and cannot train arbitrarily large models.

Comparison with Advani and Saxe

We appreciate the reviewer noting the similarity to the work of Advani and Saxe (2017). In fact, in Section 5, we take a similar analysis of scaling with respect to training time in linear models, and then, using scale-time equivalence, apply it to make novel predictions about scaling with respect to model size. We are happy to further elaborate on the relationship with this work.

Questions on Figure 4

Figure 4 is a schematic illustration intended to conceptually demonstrate how different noise profiles can affect loss trajectories and potentially lead to double descent behavior. It is not derived from empirical data but serves as a visual aid to support the discussion.

Concerns on Scale-Time Equivalence

The reviewer raises an interesting question about whether scale-time equivalence holds when the model size becomes very small, potentially making the model underparameterized. First, ur theoretical analysis indicates that the equivalence may break down for small model scale $p$, as reflected in the error bound of Theorem 1, which becomes looser when $p$ is small.

However, assuming scale-time equivalence does hold, then if a large model is able to interpolate the data in finite time, so should a small model. On the other hand, it seems surprising that a small model would be able to fit the training data. How can we explain this seeming contradiction?

In regimes where scale-time equivalence holds, the appropriate measure of capacity is the product $pt$, representing the total training effort, rather than model scale $p$. Therefore, a smaller model trained for a proportionally longer time has enough capacity and can achieve similar performance to a larger model trained for a shorter time, provided that $pt$ remains sufficient for learning the data.

We thank the reviewer for their thoughtful comments and for highlighting areas where our manuscript could be improved. We address the specific concerns below.

Concerns on Theoretical Setup

Please see our general response above, where we clarify the theoretical model and its applicability to neural networks. Our model is designed to capture essential characteristics of neural network training, including nonlinearity and overparameterization. By abstracting the network's output dependency on certain parameter directions through $\alpha = K\beta$, we aim to provide a general framework that is relevant to deep neural networks.

Minimal Capacity to Fit the Data

The reviewer raises an important point regarding the model's capacity to fit the data. In our theoretical result (Theorem 1), the error bound contains a term inversely proportional to the square root of the model scale $p$ (i.e., $\frac{1}{\sqrt{p}}$​). This indicates that the scale-time equivalence holds more precisely for sufficiently large models. As $p$ becomes small, the error bound becomes looser, and the equivalence may break down.

We emphasize that our analysis does not assume that the model can fit the data at any scale; we do not assume that models are large enough to interpolate the training data. In fact, the maximum model size at which scale-time equivalence breaks down may be completely unrelated to its ability to fit the training data.

Cube Root Relationship Between Parameters and Model Scale

The reviewer correctly points out that our definition of model scale as the cube root of the number of parameters is heuristic. We based this choice on empirical observations where the cube root scaling provided consistent results across different architectures and datasets in our experiments.

This scaling reflects the idea that the effective capacity of neural networks increases with both width and depth, and their combined effect on parameter count is nonlinear. We agree that the relationship between parameters and model scale can vary depending on the architecture and data; the cube root scaling is just what we observed with our standard architectures on standard vision benchmarks.

Concerns on Figure 5

Regarding the observations in Figure 5, we note that in our experiments, both test and train error decrease monotonically with the amount of data, even as model size varies. This behavior aligns with prior work, such as Nakkiran et al. (2019), where increasing data volume generally leads to improved performance on CIFAR-10. Thus, we actually do not expect to see double-descent like behavior on this plot, consistent with the finding that it does not always appear.

Comments on Learning Rate

In our experiments, we used a fixed learning rate across different model scales to maintain consistency and to focus on the effects of model size and training time.

Our theoretical analysis assumes a constant learning rate, and varying it would introduce additional variables that could confound the relationship we are studying. We acknowledge that in practice, adjusting the learning rate with model size can be beneficial. However, our findings on scale-time equivalence are based on the assumption of a fixed learning rate, which is a common practice in theoretical studies to isolate specific effects.

Related Work

We appreciate the references to additional related work, including the papers on training time as a form of regularization. We have included these references in the revised manuscript.

Typos

Thank you for pointing out the missing articles in lines 267-269. We have corrected these typos in the revised manuscript.

I sincerely thank the authors for the rebuttal. Indeed, a few concerns (e.g. the learning rate issue) are resolved

I think the main concern with the connection of the proposed model and neural networks, shared with other reviewers in many forms, still persists. In fact, I find it a bit vague that the random subspace is a correct abstraction of the model's depth or width, in the sense these quantities introduce fundamental challenges to optimization and generalization in neural networks. Also, my questions regarding the comment [...] the model is linear. Also, why are $r, p, P$ all different? Why couldn’t you only have a large dimensional ambient space and a constant-sized low-rank subspace? still persist. In my view, I would suggest better justifying the model, and potentially considering studying a random feature model is very close and better tied to the literature. In the current form, I believe that the proposed model and heuristics on the relationship between model scale and the number of parameters are used to simplify the problem, rather than finding good abstraction to model neural networks. Also, as they are not introduced earlier in the literature, they should be more thoroughly justified.