Abstract:  Continuing from Wednesday, we take the local quadratic approximation seriously and ask how well it holds during LLM pretraining. Taylor expanding real training runs around their iterates, we find — somewhat surprisingly — that  linearized models track the true loss for a meaningful window and that this continues to hold deep into training. Using Lanczos quadrature, we then compute the full spectral density of 150M-parameter models, finding power-law spectra whose curvature is dominated by identifiable parts of the network. Together, these suggest Taylor's theorem is a more powerful lens on deep learning than its simplicity implies.

Abstract:  Neural scaling laws underlie many of the recent advances in deep learning, yet their theoretical understanding remains largely confined to linear models. In this work, we present a systematic analysis of scaling laws for shallow neural networks in the feature learning regime. Leveraging connections with matrix compressed sensing and LASSO, we derive a detailed phase diagram for the scaling exponents of the excess risk as a function of sample complexity and weight decay. This analysis uncovers crossovers between distinct scaling regimes and plateau behaviours, mirroring phenomena widely reported in the empirical neural scaling literature. Furthermore, we establish a precise link between these regimes and the spectral properties of the trained network weights, which we characterize in detail. Consequently, we provide a theoretical validation of recent empirical observations connecting the emergence of power-law tails in the weight spectrum with network generalization performance, yielding an interpretation from first principles.   

Abstract:  In this work, we consider a class of Muon-like updates, where we replace the update M with UΣᵖVᵀ for some parameter p. We call this a “spectral-shaping” operation, and develop a theory of how to pick p which depends on (a) local curvature of the loss function, (b) noise stemming from stochastic gradients and label noise, and (c) training stage. Our theory and experimentation reveal a previously overlooked behavior: positive p helps early by emphasizing high-curvature directions and accelerating signal contraction, while mildly negative p helps later by reallocating update strength toward low-curvature directions that still contain useful training signals. Building on the insight, we propose DynMuon, an efficient dynamic spectral shaping method that schedules p from positive to mildly negative over training. Extensive experiments across model sizes, architectures, and training settings show that DynMuon consistently achieves lower validation loss than Muon, while requiring 10.6–26.5% fewer steps to reach the same target loss. Our code is available at https://github.com/fzwark/DynMuon. 

Abstract:  Spectral optimizers such as Muon have recently shown strong empirical performance in large-scale language model training, but the source and extent of their advantage remain poorly understood. We study this question through the linear associative memory problem, a tractable model for factual recall in transformer-based models. In particular, we go beyond orthogonal embeddings and consider Gaussian inputs and outputs, which allows the number of stored associations to greatly exceed the embedding dimension. Our main result sharply characterizes the recovery rates of one step of Muon, SGD, and Newton's method on the logistic regression loss under a power law frequency distribution. We show that the storage capacity of Muon significantly exceeds that of SGD, and even matches Newton's method while only using first-order information. Moreover, Muon saturates at a larger critical batch size. We further analyze the multi-step dynamics under a thresholded gradient approximation and show that Muon achieves a substantially faster initial recovery rate than SGD, while both methods eventually converge to the information-theoretic limit at comparable speeds. Experiments on synthetic tasks validate the predicted scaling laws. Our analysis provides a quantitative understanding of the signal amplification of spectral preconditioners and lays the groundwork for establishing scaling laws across more practical language modeling tasks and optimizers. 

Abstract:  Large language models are increasingly trained in continual or open-ended settings, where the total training horizon is not known in advance. Despite this, most existing pretraining recipes are not anytime: they rely on horizon-dependent learning rate schedules and extensive tuning under a fixed compute budget. In this work, we provide a theoretical analysis demonstrating the existence of anytime learning schedules for overparameterized linear regression, and we highlight the central role of weight averaging—also known as model merging—in achieving the minimax convergence rates of stochastic gradient descent. We show that these anytime schedules polynomially decay with time, with the decay rate determined by the source and capacity conditions of the problem. Empirically, we evaluate 150M and 300M parameter language models trained at 1–32× Chinchilla scale, comparing constant learning rates with weight averaging and $1/\sqrt{t}$ schedules with weight averaging against a well-tuned cosine schedule. Across the full training range, the anytime schedules achieve comparable final loss to cosine decay. Taken together, our results suggest that weight averaging combined with simple, horizon-free step sizes offers a practical and effective anytime alternative to cosine learning rate schedules for large language model pretraining. 

Abstract: In this talk, we discuss stable scaling of deep learning models by identifying stable, feature-learning infinite width and depth limits of neural networks. The asymptotic description of randomly initialized networks in this regime will take the form of a dynamical mean field theory (DMFT). We will discuss how adoption of scaling strategies that admit such limits yields better hyperparameter transfer, where optimal hyperparameters in small models remain optimal in large models. We will provide examples of these results for multi-layer perceptrons, convolutional networks, self-attention blocks, and mixture-of-experts transformers. These exact limits assume parameters are much larger than the training data, and thus fail to capture the behavior of models in the scaling law regime where models of different sizes achieve different performance. To address this, we will introduce simplified, analytically tractable models which enable analysis of training dynamics far from the infinite limit. We show that early time deviations in model performance are universal, while late time deviations are architecture and data dependent. We use one of these toy models to analyze hyperparameter transfer across training horizons T through an optimal control paradigm, showing that the optimal strategy is highly dependent on feature structure and SGD noise statistics.  

Abstract:  Recent frontier large language models predominantly rely on Mixture-of-Experts (MoE) architectures. Despite empirical progress, there is still no principled understanding of how hyperparameters should scale with network width N, expert width Nₑ, number of experts M, sparsity K, and depth L to ensure both stability and optimal performance at scale. We take a principled step toward resolving this gap by analyzing three different scaling regimes: (I) co-scaling N ≈ Nₑ, (II) co-scaling N ≈ M ≈ K, and (III) full proportional scaling of N, Nₑ, M, and K. For each regime, we develop a novel Dynamical Mean Field Theory (DMFT) description of the limiting training dynamics of MoEs that provides a formal foundation for our analysis. Within this framework, we derive the unique parameterization for SGD and Adam satisfying all maximal-update (μ) desiderata. We then show that the resulting μP prescription does not reliably induce monotonic improvement with scale or robust learning-rate transfer. We trace these pathologies to scale-dependent observables in the aggregation dynamics, which motivates a refined set of desiderata that we term maximal scale stability. Guided by this principle, we derive a Maximally Scale-Stable Parameterization (MSSP) for both SGD and Adam in all three scaling regimes, and characterize the corresponding limiting dynamics — qualitatively distinct from the μP limit — through a separate DMFT analysis. Experiments verify that MSSP robustly recovers learning rate transfer and monotonic improvement with scale across regimes. Combined with existing depth-scaling theory, these results provide a complete scaling prescription for MoE architectures as a function of width, depth, expert width, and number of experts.

Abstract:  Stochastic momentum methods such as heavy ball (HB), Nesterov momentum, and variants of Accelerated SGD (ASGD) [Kidambi et al., 2018] are widely used in modern training, but their stochastic benefits depend on two distinct quantities: serial runtime, the number of iterations needed to reach a target accuracy, and compute efficiency (CE), the inverse total gradient-query or FLOP cost. Larger batches reduce serial runtime without hurting CE only when the contraction gap grows linearly with batch size. We study stochastic HB and ASGD for consistent linear regression with Gaussian covariates and prove finite-dimensional, discrete-time lower bounds on their batch-size tradeoffs. Our first result shows that HB does not improve the CE frontier over SGD for arbitrary spectra; rather, it preserves SGD-level CE over a larger batch-size window, allowing larger batches to reduce serial runtime until HB reaches its deterministic accelerated scale. This window can be a factor $\sqrt{\kappa}$ larger than the SGD critical batch size. For ASGD, the picture is more spectrum-dependent: for rapidly decaying power-law spectra, ASGD improves small-batch CE over HB/SGD, but as batch size grows it trades this CE advantage for improved serial runtime. Synthetic linear-regression experiments verify these qualitative regimes, including near-overlap of ASGD and HB for slowly decaying spectra and the predicted CE--serial tradeoff for rapidly decaying spectra.

Abstract: Large Language Models (LLMs) exhibit intriguing emergent phenomena in their high-dimensional activation spaces, most notably "attention sinks" (a few tokens receiving disproportionately large attention logits) and "residual sinks" (a few fixed dimensions with persistently large magnitudes). While often treated as anomalies or quantization hurdles, we propose a unified functional view: these outliers are essential mechanisms for high-dimensional rescaling.

Specifically, we demonstrate that a small fraction of outliers, coupled with normalization layers (e.g., Softmax, RMSNorm), effectively rescale the remaining non-outlier components—a phenomenon we term outlier-driven rescaling. Through interventional studies, we provide evidence that outliers function jointly with normalization and act primarily as rescaling factors rather than direct semantic contributors. Removing normalization eliminates the corresponding outliers but degrades training stability, whereas clipping outliers while retaining normalization similarly harms performance, indicating that outlier-driven rescaling is crucial for stabilizing high-dimensional optimization. Furthermore, this perspective provides explanations for the success of certain architectural modifications. Building on these insights, we show that outliers can be effectively absorbed into learnable parameters or mitigated via explicit gated rescaling, leading to improved training performance and enhanced quantization robustness.