A Comparative Study of deep neural network Scaling Methods And When Without-Replacement SGD Isn't Faster

A Comparative Study of deep neural network Scaling Methods And When Without-Replacement SGD Isn’t Faster Harshul Gupta1* 1*Applied Mathematics, Columbia University, 500 W 120th Street, New York, 10027, NY, USA. Corresponding author(s). E-mail(s): [email protected]; Abstract This project investigates two fundamental aspects of deep learning: deep neural network Scaling Methods and the Efficiency of Stochastic Gradient Descent (SGD) with and without Data Replacement. First, we conduct a comparative study of deep neural network parameterizations - NTP, NTP-LR, MuP, MuPLR, and MFP across different model architectures (MLPs, CNNs) and datasets (MNIST, FMNIST, CIFAR-10). Our results demonstrate that scaling rules like MuP and MFP not only achieve a lower test error but also improve learning rate transferability across models widths, enabling safer and more effective training of wide networks. Second, we evaluate the commonly held belief that withoutreplacement SGD leads to faster convergence. Empirical analysis reveals that without-replacement methods outperform with-replacement SGD only in regimes with extremely large iteration budgets relative to the problem’s condition number. These findings highlight critical nuances in the design of training protocols and offer practical guidelines for scaling and optimizing deep learning models. Keywords: deep neural network Scaling, MuP, Stochastic Gradient Descent, Without-Replacement Sampling 1 Introduction In recent years, deep learning has achieved remarkable success across a variety of domains, including computer vision, natural language processing, and scientific computing. Much of this progress has been driven by the increasing scale of deep neural 1 networks, both in terms of depth and width. However, scaling deep neural networks reliably is a non-trivial task, especially when it comes to preserving stable training dynamics and generalization performance. As networks grow wider, small mis-configurations in initialization or learning rates can lead to suboptimal or even unstable training. Understanding the mathematical principles behind these challenges is essential for developing robust and scalable deep learning systems. This project investigates two central questions in modern deep learning practice. First, how can deep neural network parameterizations be designed to allow smooth scaling across architectures of varying width. We explore and compare several prominent scaling methods: standard training with and without learning rate tuning (NTP, NTP-LR), µ-parameterization (MuP), its tuned variant (MuP-LR), and Maximal Feature Propagation (MFP). These methods differ in how they initialize weights and scale learning rates as model width increases. Our study evaluates their performance across multiple neural architectures (MLPs and CNNs) and benchmark datasets (MNIST, FMNIST, and CIFAR-10), focusing on test accuracy and the stability of hyper-parameter transfer. The second part of the project examines the common belief that withoutreplacement sampling in Stochastic Gradient Descent (SGD) always leads to faster convergence. Without-replacement SGD: where the dataset is shuffled and iterated through without repeating points in an epoch: is often assumed to be more dataefficient than with-replacement SGD. However, recent theoretical insights suggest that this advantage only materializes under specific conditions related to the optimization landscape, particularly its condition number. Through a series of controlled experiments, we explore when, and to what extent, without-replacement SGD truly outperforms its with-replacement counterpart. By bringing together empirical results and theoretical perspectives, this project aims to provide a more nuanced understanding of deep neural network scaling and SGD behavior. The findings serve as practical guidelines for training large-scale models and for selecting appropriate optimization strategies under different problem regimes. 2 Comparison of Neural Parametrizations Across Datasets Large deep neural networks have demonstrated superior performance on a wide range of tasks; however, the underlying principles that govern this performance remain only partially understood. A key determinant of success in training such models lies in how their hyper-parameters are configured, particularly the learning rate and the initialization scale. These factors influence whether the model operates in a lazy learning regime, where internal representations change minimally and resemble kernel methods, or in an active learning regime, where the network dynamically learns rich and expressive features during training. While lazy learning is stable and predictable, it often limits the model’s ability to fully capture complex patterns in data. Conversely, active learning enables greater representational power but introduces increased sensitivity to hyper-parameter choices 2 and optimization challenges. Navigating the trade-off between these regimes is essential for training deep and wide models effectively. The richness scale offers a useful conceptual framework for understanding this spectrum of behavior, linking model dynamics to choices in initialization and learning rate. Traditional parameterizations, such as the standard Neural Tangent Parameterization (NTP), often fail to maintain stability or effective performance as model width increases. These methods typically do not scale hyper-parameters in a way that preserves learning dynamics, which can result in inefficient training and poor generalization in wider architectures. The objective of this study is to investigate and compare following deep neural network scaling strategies - NTP, NTP-LR, MuP, MuP-LR, and MFP - across multiple architectures (MLP and CNN) and benchmark datasets (MNIST, Fashion-MNIST, and CIFAR-10). By evaluating their test-time performance, training stability, and the ability to transfer hyper-parameters across widths, we aim to provide actionable insights into the design of scalable and robust training protocols for modern deep learning systems. 3 Technical Comparison of Parametrizations To understand how different deep neural network scaling strategies behave, it is essential to compare their initialization and learning rate rules. These parameters directly affect how the model trains as width increases and play a crucial role in determining whether the network lies in the lazy or active learning regime. We define the activity of a model as follows: activity = (√width,for MuP, MFP, and MuP-LR, 1,for NTP and NTP-LR. Let q= activity2. For readout layers, the effective input dimension is scaled by 1/din, reflecting their role in final output computation. Each parametrization method—NTP, NTP-LR, MuP, MuP-LR, and MFP—adjusts the weight initialization scale and learning rate differently to accommodate changes in width. The initialization scale determines how weights are sampled at the start of training, while the learning rate scaling adjusts the step size of gradient updates in a width-aware fashion. Here is a detailed comparison of each approach: •NTP (Neural Tangent Parameterization) uses a fixed initialization and learning rate. However, as width increases, this fixed scheme often becomes unstable or inefficient, since it does not account for width-related variance in gradients. •NTP-LR (Neural Tangent Parameterization Learning Rate) improves on NTP by tuning the learning rate using a scale-aware adjustment. This yields better training dynamics, but hyper-parameters do not generalize well across widths. •MuP (µ-Parameterization) addresses scaling explicitly. By dividing the initialization by the activity factor and scaling the learning rate appropriately, MuP ensures consistent training behavior across different model widths. It is grounded 3 in the infinite-width limit and provides a stable transition from narrow to wide networks. •MuP-LR (µ-Parameterization Learning Rate) combines the MuP scaling rules with a tuned learning rate for improved optimization performance while preserving MuP’s transferability properties. •MFP (Maximal Feature Propagation) further enhances MuP by encouraging maximal signal propagation through layers, particularly during initialization and early training. It retains similar scaling formulas to MuP but focuses more on preserving rich feature dynamics. These differences are summarized in Table 1, showing how each method sets initialization scales, multipliers, and learning rate rules: Table 1 Comparison of initialization and learning rate scaling across parameterizations. Parametrization Init Scale Multiplier gLearning Rate Scaling NTP 1 √q/activity activity2 NTP-LR √q/activity 1 q MuP 1/activity √q1 MuP-LR √q/activity 1 q MFP 1 √q/activity activity2 This comparison highlights that MuP and MFP offer the most stable and transferable configurations for training across a wide range of model widths, while NTP-based approaches may require width-specific tuning to remain effective. 4 Experimental Setup To evaluate the efficacy of different deep neural network parametrizations, we conduct experiments across multiple architectures and datasets, focusing on their ability to scale effectively while maintaining training stability and performance. 4.1 Architectures and Datasets We consider two common deep neural network architectures: multi-layer perceptrons (MLPs) and convolutional deep neural networks (CNNs). The models are trained and evaluated on three benchmark datasets: MNIST, Fashion-MNIST (FMNIST), and CIFAR-10. Performance is measured primarily using test mean squared error (MSE). 4.2 MLP Model Configuration The MLP architecture used in our experiments consists of an input layer with dimension din = 784, corresponding to flattened 28×28 grayscale images, and an output layer with dout = 10 units, representing class probabilities. The network contains two hidden layers, each with 512 units. The width parameter utilized for scaling corresponds to 512. 4 The output scale of the network is set to approximately 0.0407. Weight initialization scales vary across layers as follows: 0.0357 for the first fully connected layer (fc1), 0.0441 for the second fully connected layer (fc2), and 0.0019 for the output layer (fc3). Corresponding layerwise learning rates are 0.6531 for fc1, 1.0000 for fc2, and 0.0020 for fc3. 4.3 CNN Model Configuration For the CNN experiments, the input dimension is din = 3, reflecting RGB image channels, and the output dimension remains dout = 10. The network architecture comprises four convolutional layers, each with 512 filters, which is also the scaling width. The network’s output scale is approximately 0.0364. Weight initialization scales for convolutional layers conv1 through conv4 range between 0.0442 and 0.0444, with the readout layer initialized at 0.0439. All convolutional layers and the readout layer utilize a uniform learning rate of 1.0000. These configurations provide a controlled setting to systematically compare the behavior of various parametrizations across different model scales and architectures. 5 Results We evaluate the performance of different deep neural network parametrizations—NTP, NTP-LR, MuP, MuP-LR, and MFP—across multiple datasets and architectures. Our focus is on their test error rates and the stability of learning rate transfer across model widths. 5.1 MLP Performance on MNIST and Fashion-MNIST Figures 1and 2illustrate the test mean squared error (MSE) for MLPs trained on MNIST and Fashion-MNIST datasets, respectively. Across both datasets, training methods employing MuP, MuP-LR, and MFP consistently outperform the baseline NTP and NTP-LR approaches. These advanced parametrizations yield lower test errors, demonstrating improved learning dynamics and more effective optimization. The results highlight the benefits of scaling rules that adjust initialization and learning rates to maintain stable training as model width increases. 5.2 CNN Performance on CIFAR-10 Figure 3presents test error results for CNN models trained on CIFAR-10. Similar to the MLP experiments, MuP, MuP-LR, and MFP methods achieve superior performance relative to NTP variants. The results emphasize the generality of these scaling methods across architectures, with better convergence and lower test errors observed for models trained using these advanced parametrizations. 5.3 Learning Rate Transfer Across Model Widths A critical aspect of scalable neural training is the ability to transfer hyper-parameters, such as the learning rate, across different model widths without loss of performance. 5 Fig. 1 Test MSE of MLP models on MNIST dataset using different parametrizations. MuP, MuPLR, and MFP outperform NTP variants. Fig. 2 Test MSE of MLP models on Fashion-MNIST dataset. MuP-based methods achieve lower errors consistently. Figure 4compares the transferability of learning rate settings between NTP and MuP training protocols. The left panel shows that under NTP, the optimal learning rate shifts noticeably as the model width increases. This makes it challenging to select a single learning rate that works well across all scales. Moreover, wider models trained with NTP are prone to instability at higher learning rates, necessitating careful hyper-parameter tuning for each width. 6 Fig. 3 Test MSE of CNN models on CIFAR-10 dataset. MuP, MuP-LR, and MFP maintain superior performance across widths. Fig. 4 Learning rate transferability across model widths. Left: NTP shows varying optimal learning rates and instability at larger widths. Right: MuP maintains stable learning rates across widths, enabling better hyper-parameter transfer. In contrast, the right panel demonstrates that MuP training maintains a nearly constant optimal learning rate across various widths. The learning curves for different widths align closely, indicating that hyper-parameters tuned on smaller models generalize well to larger ones. This stability simplifies training large-scale networks by reducing the need for exhaustive hyper-parameter search. 7 5.4 Summary of Findings Our results clearly indicate that MuP, MuP-LR, and MFP training methods enable more effective learning and better generalization compared to traditional NTP-based approaches. These methods improve both model performance and the robustness of hyper-parameter selection across widths. Notably, MuP parametrization facilitates reliable learning rate transfer, making it easier and safer to train large deep neural networks. 6 Background and Intuition: Stochastic Gradient Descent (SGD) Stochastic Gradient Descent (SGD) is a fundamental optimization algorithm widely used for training machine learning models, particularly deep neural networks. Instead of computing gradients over the entire dataset, SGD updates model parameters iteratively using gradients from individual data points or small batches. This stochasticity often results in faster convergence and improved generalization compared to full-batch gradient descent. During SGD, there are two prevalent strategies for sampling data points: •With replacement: At each iteration, a data point is randomly selected from the dataset, allowing the possibility of choosing the same point multiple times within an epoch. •Without replacement: The dataset is shuffled once at the start of each epoch, and data points are processed sequentially without repetition, ensuring each point is visited exactly once per epoch. A common intuition in the optimization community favors without-replacement SGD, as it supposedly converges faster by avoiding redundant updates and more efficiently utilizing the data. Since each data point is used exactly once per epoch, the algorithm is expected to make more consistent progress towards the optimum, reducing wasted computation from revisiting the same samples multiple times within an epoch. Understanding these sampling schemes and their implications is crucial for analyzing convergence properties and designing efficient training protocols in practice. 7 What the Research Shows: SGD With vs. Without Replacement Recent theoretical and empirical studies have examined the performance differences between stochastic gradient descent (SGD) with replacement and without replacement sampling schemes. Contrary to common intuition, the advantages of without-replacement SGD are nuanced and depend heavily on the conditioning of the optimization problem. In problems that are badly conditioned—that is, where the optimization landscape exhibits very different scales in different directions—the benefit of without-replacement 8 SGD over with-replacement SGD is not significant unless the algorithm performs many passes over the dataset. This phenomenon is closely linked to the condition number κof the problem, which measures the degree of skewness or stretching in the loss surface. Specifically, κis defined as the ratio of the largest to the smallest eigenvalue of the Hessian matrix associated with the objective function. A key theoretical insight established by recent works is that without-replacement SGD only begins to outperform with-replacement SGD after approximately κfull passes over the data. Since κcan be very large for many practical problems, reaching this regime can be computationally infeasible, limiting the practical advantages of without-replacement sampling in certain scenarios. Several convergence bounds in the literature formalize these conditions under which without-replacement SGD improves over with-replacement SGD. Table 2summarizes key results from prominent studies: Table 2 Convergence bounds and conditions for improvement of without-replacement SGD over with-replacement SGD. Here, nis the dataset size, kis the number of iterations, and λdenotes relevant spectral parameters related to the Hessian eigenvalues. Reference Condition for Improvement G¨urb¨uzbalaban et al. (2015b) k≳nwhen (λk)2dominates HaoChen and Sra (2018) k≳1/λ·max{1,pn/λ} Rajput et al. (2020) k≳1/λ2 Ahn et al. (2020) k≳1/λ2.5 These bounds demonstrate that the iteration count krequired for withoutreplacement SGD to outperform with-replacement SGD grows with the condition number κ(or related spectral quantities λ) and dataset size n. Consequently, for highly ill-conditioned problems, the practical speedup achievable by without-replacement sampling may be limited unless the training budget is sufficiently large to perform many epochs. In summary, while without-replacement SGD can theoretically offer faster convergence, the conditions to realize these benefits depend strongly on problem conditioning and computational budget. This nuanced understanding guides practitioners in choosing appropriate sampling strategies for SGD in large-scale and complex learning tasks. 8 SGD Setup: Finite-Sum Optimization We study finite-sum optimization problems of the form F(x) = 1 n n X i=1 fi(x), 9