Neural Network Optimization

Ask about this topic

Overview

Neural network optimization asks how to adjust a model’s parameters so it learns useful patterns efficiently and reliably. Common approaches include stochastic gradient descent and Adam, matrix-aware and spectral optimizers such as Muon, curvature-based methods, learning-rate schedules, and techniques that reduce memory or computation. Researchers also analyze loss landscapes, scaling laws, and initialization, and explore alternatives such as local or forward-only learning and analytic solutions for restricted models. A central tension is that gains depend on architecture, data, and training budget, while faster loss reduction does not necessarily improve generalization; connecting practical performance to reliable theory remains an open challenge.

Latest papers

CardsList

332 papers37 papers in the last four weeks, up 28% on the four weeks before. 0.2% of all new papers.

Mar 30, 2026

EdgeDiT: Hardware-Aware Diffusion Transformers for Efficient On-Device Image Generation

Diffusion Transformers (DiT) have established a new state-of-the-art in high-fidelity image synthesis; however, their massive computational complexity and memory requirements hinder local deployment on resource-constrained edge devices. In this paper, we introduce EdgeDiT, a family of hardware-efficient generative transformers specifically engineered for mobile Neural Processing Units (NPUs), such as the Qualcomm Hexagon and Apple Neural Engine (ANE). By leveraging a hardware-aware optimization framework, we systematically identify and prune structural redundancies within the DiT backbone that are particularly taxing for mobile data-flows. Our approach yields a series of lightweight models that achieve a 20-30% reduction in parameters, a 36-46% decrease in FLOPs, and a 1.65-fold reduction in on-device latency without sacrificing the scaling advantages or the expressive capacity of the original transformer architecture. Extensive benchmarking demonstrates that EdgeDiT offers a superior Pareto-optimal trade-off between Frechet Inception Distance (FID) and inference latency compared to both optimized mobile U-Nets and vanilla DiT variants. By enabling responsive, private, and offline generative AI directly on-device, EdgeDiT provides a scalable blueprint for transitioning large-scale foundation models from high-end GPUs to the palm of the user.

Sravanth Kodavanti, Manjunath Arveti, Sowmya Vajrala +2

Samsung Research Institute Bangalore, India

Diffusion TransformerDiffusion-Based Image GenerationDiffusion Model Inference AccelerationOn-Device InferenceNeural Network Optimization
AskPDF
Mar 30, 2026

MuonEq: Balancing Before Orthogonalization with Lightweight Equilibration

Orthogonalized-update optimizers such as Muon improve training of matrix-valued parameters, but existing extensions typically either rescale updates after orthogonalization or use heavier whitening-based preconditioners before it. We introduce {\method}, a lightweight family of pre-orthogonalization equilibration schemes for Muon with three forms: two-sided row/column normalization (RC), row normalization (R), and column normalization (C). By rebalancing the momentum matrix before finite-step Newton--Schulz orthogonalization, {\method} improves the geometry seen by orthogonalization. We show that finite-step orthogonalization is governed by the input spectrum, especially stable rank and condition number, and that row/column normalization acts as a zeroth-order surrogate for whitening. For hidden matrix weights, R is the default variant. Theoretically, {\method} (R) retains the standard O~(T−1/4)\widetilde{\mathcal O}(T^{-1/4}) Muon-type nonconvex stationarity guarantee with decoupled weight decay and a horizon-free diminishing learning-rate schedule, and extends it to finite-step NS5 up to an explicit inexactness constant. In LLaMA2 pretraining on C4, {\method} (R) consistently outperforms Muon on 130M, 350M, and 1B models, with faster convergence and lower validation perplexity. The code is available at the MuonEq codebase.

Da Chang, Qiankun Shi, Lvgang Zhang +5

Shenzhen Institute of Advanced Technology, Chinese Academy of Sciences · Pengcheng Laboratory · University of Chinese Academy of Sciences +4

Neural Network OptimizationMuon OptimizerOrthogonal Optimization
AskPDF
Mar 28, 2026

Preconditioned Attention: Enhancing Efficiency in Transformers

Central to the success of Transformers is the attention block, which effectively models global dependencies among input tokens associated to a dataset. However, we theoretically demonstrate that standard attention mechanisms in transformers often produce ill-conditioned matrices with large condition numbers. This ill-conditioning is a well-known obstacle for gradient-based optimizers, leading to inefficient training. To address this issue, we introduce preconditioned attention, a novel approach that incorporates a conditioning matrix into each attention head. Our theoretical analysis shows that this method significantly reduces the condition number of attention matrices, resulting in better-conditioned matrices that improve optimization. Conditioned attention serves as a simple drop-in replacement for a wide variety of attention mechanisms in the literature. We validate the effectiveness of preconditioned attention across a diverse set of transformer applications, including image classification, object detection, instance segmentation, long sequence modeling and language modeling.

Hemanth Saratchandran

Australian Institute for Machine Learning (AIML), Adelaide University

Attention MechanismsNeural Network OptimizationNeural Network Training Dynamics
AskPDF
Mar 27, 2026

Sharp Capacity Scaling of Spectral Optimizers in Learning Associative Memory

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.

Juno Kim, Eshaan Nichani, Denny Wu +2

UC Berkeley · Princeton University · New York University +1

Spectral OptimizationNeural Network OptimizationMuon OptimizerSpectral MethodsAssociative Memory
AskPDF
Mar 26, 2026

Second-Order, First-Class: A Composable Stack for Curvature-Aware Training

Second-order methods promise improved stability and faster convergence, yet they remain underused due to implementation overhead, tuning brittleness, and the lack of composable APIs. We introduce Somax, a composable Optax-native stack that treats curvature-aware training as a single JIT-compiled step governed by a static plan. Somax exposes first-class modules -- curvature operators, estimators, linear solvers, preconditioners, and damping policies -- behind a single step interface and composes with Optax by applying standard gradient transformations (e.g., momentum, weight decay, schedules) to the computed direction. This design makes typically hidden choices explicit and swappable. Somax separates planning from execution: it derives a static plan (including cadences) from module requirements, then runs the step through a specialized execution path that reuses intermediate results across modules. We report system-oriented ablations showing that (i) composition choices materially affect scaling behavior and time-to-accuracy, and (ii) planning reduces per-step overhead relative to unplanned composition with redundant recomputation.

Mikalai Korbit, Mario Zanon

IMT School for Advanced Studies Lucca Lucca, Italy

Curvature-Aware OptimizationDeep Learning OptimizationNeural Network OptimizationSecond-Order Optimization
AskPDF
Mar 26, 2026

Improving Infinitely Deep Bayesian Neural Networks with Nesterov's Accelerated Gradient Method

As a representative continuous-depth neural network approach, stochastic differential equation (SDE)-based Bayesian neural networks (BNNs) have attracted considerable attention due to their solid theoretical foundations and strong potential for real-world applications. However, their reliance on numerical SDE solvers inevitably incurs a large number of function evaluations (NFEs), resulting in high computational cost and occasional convergence instability. To address these challenges, we propose a Nesterov-accelerated gradient (NAG) enhanced SDE-BNN model. By integrating NAG into the SDE-BNN framework along with an NFE-dependent residual skip connection, our method accelerates convergence and substantially reduces NFEs during both training and testing. Extensive empirical results show that our model consistently outperforms conventional SDE-BNNs across various tasks, including image classification and sequence modeling, achieving lower NFEs and improved predictive accuracy.

Chenxu Yu, Wenqi Fang

Shenzhen Institutes of Advanced Technology Chinese Academy of Sciences Shenzhen, P.R.China

Bayesian Neural NetworksNeural Network OptimizationAccelerated Gradient Methods
AskPDF
Mar 26, 2026

A Systematic Empirical Study of Grokking: Depth, Architecture, Activation, and Regularization

Grokking the delayed transition from memorization to generalization in neural networks remains poorly understood, in part because prior empirical studies confound the roles of architecture, optimization, and regularization. We present a controlled study that systematically disentangles these factors on modular addition (mod 97), with matched and carefully tuned training regimes across models. Our central finding is that grokking dynamics are not primarily determined by architecture, but by interactions between optimization stability and regularization. Specifically, we show: (1) depth has a non-monotonic effect, with depth-4 MLPs consistently failing to grok while depth-8 residual networks recover generalization, demonstrating that depth requires architectural stabilization; (2) the apparent gap between Transformers and MLPs largely disappears (1.11×\times delay) under matched hyperparameters, indicating that previously reported differences are largely due to optimizer and regularization confounds; (3) activation function effects are regime-dependent, with GELU up to 4.3×\times faster than ReLU only when regularization permits memorization; and (4) weight decay is the dominant control parameter, exhibiting a narrow ``Goldilocks'' regime in which grokking occurs, while too little or too much prevents generalization. Across 3--5 seeds per configuration, these results provide a unified empirical account of grokking as an interaction-driven phenomenon. Our findings challenge architecture-centric interpretations and clarify how optimization and regularization jointly govern delayed generalization.

Shalima Binta Manir, Anamika Paul Rupa

Howard University

GrokkingNeural Network Training DynamicsNeural Network GeneralizationNeural Network Optimization
AskPDF
Mar 25, 2026

Local learning for stable backpropagation-free neural network training towards physical learning

While backpropagation and automatic differentiation have driven deep learning's success, the physical limits of chip manufacturing and rising environmental costs of deep learning motivate alternative learning paradigms such as physical neural networks. However, most existing physical neural networks still rely on digital computing for training, largely because backpropagation and automatic differentiation are difficult to realize in physical systems. We introduce FFzero, a forward-only learning framework enabling stable neural network training without backpropagation or automatic differentiation. FFzero combines layer-wise local learning, prototype-based representations, and directional-derivative-based optimization through forward evaluations only. We show that local learning is effective under forward-only optimization, where backpropagation fails. FFzero generalizes to multilayer perceptron and convolutional neural networks across classification and regression. Using a simulated photonic neural network as an example, we demonstrate that FFzero provides a viable path toward backpropagation-free in-situ physical learning.

Yaqi Guo, Fabian Braun, Bastiaan Ketelaar +3

Department of Materials Science and Engineering, Delft University of Technology, Mekelweg 2, Delft, 2628CD, The Netherlands · Department of Precision and Microsystems Engineering, Delft University of Technology, Mekelweg 2, Delft, 2628CD, The Netherlands · Department of Intelligent Systems, Delft University of Technology, Van Mourik Broekmanweg 6, Delft 2628XE, The Netherlands

Physical Neural NetworksNeural Network OptimizationLocal LearningBackpropagationForward-Forward Learning
AskPDF
Mar 25, 2026

Stochastic Dimension Zeroth-Order Estimator: Stable and Memory-Efficient Training of PINNs

Physics-Informed Neural Networks (PINNs) for high-dimensional and high-order partial differential equations (PDEs) are primarily constrained by the O(dk)\mathcal{O}(d^k) spatial derivative complexity and the O(P)\mathcal{O}(P) memory overhead of backpropagation (BP). While randomized spatial estimators successfully reduce the spatial complexity to O(1)\mathcal{O}(1), their reliance on first-order optimization still leads to prohibitive memory consumption at scale. Zeroth-order (ZO) optimization offers a BP-free alternative; however, naively combining randomized spatial operators with ZO perturbations triggers a variance explosion of O(1/ε2)\mathcal{O}(1/\varepsilon^2), leading to numerical divergence. To address these challenges, we propose the Stochastic Dimension-free Zeroth-order Estimator (SDZE), a unified framework that achieves dimension-independent complexity in both space and memory. Specifically, SDZE leverages Common Random Numbers Synchronization (CRNS) to algebraically cancel the O(1/ε2)\mathcal{O}(1/\varepsilon^2) variance by locking spatial random seeds across perturbations. Furthermore, an implicit matrix-free subspace projection is introduced to reduce parameter exploration variance from O(P)\mathcal{O}(P) to O(r)\mathcal{O}(r) while maintaining an O(1)\mathcal{O}(1) optimizer memory footprint. Empirical results demonstrate that SDZE enables the training of 10-million-dimensional PINNs on a single NVIDIA A100 GPU, delivering significant improvements in speed and memory efficiency over state-of-the-art baselines.

Zhangyong Liang, Huanhuan Gao

National Center for Applied Mathematics, Tianjin University, Tianjin 300072, China. · School of Mechanical and Aerospace Engineering, Jilin University, Changchun 130025, China

Zeroth-Order OptimizationNeural Network OptimizationPartial Differential EquationsPhysics-Informed MLVariance Reduction
AskPDF
Mar 25, 2026

Unveiling Hidden Convexity in Deep Learning: a Sparse Signal Processing Perspective

Deep neural networks (DNNs), particularly those using Rectified Linear Unit (ReLU) activation functions, have achieved remarkable success across diverse machine learning tasks, including image recognition, audio processing, and language modeling. Despite this success, the non-convex nature of DNN loss functions complicates optimization and limits theoretical understanding. In this paper, we highlight how recently developed convex equivalences of ReLU NNs and their connections to sparse signal processing models can address the challenges of training and understanding NNs. Recent research has uncovered several hidden convexities in the loss landscapes of certain NN architectures, notably two-layer ReLU networks and other deeper or varied architectures. This paper seeks to provide an accessible and educational overview that bridges recent advances in the mathematics of deep learning with traditional signal processing, encouraging broader signal processing applications.

Emi Zeger, Mert Pilanci

Department of Electrical Engineering, Stanford University

ReLU Neural NetworksDeep Learning OptimizationNeural Network Optimization
AskPDF
Mar 22, 2026

On the Role of Batch Size in Stochastic Conditional Gradient Methods

We study the role of batch size in stochastic conditional gradient methods under a μμ-Kurdyka-Łojasiewicz (μμ-KL) condition. Focusing on momentum-based stochastic conditional gradient algorithms (e.g., Scion), we derive a new analysis that explicitly captures the interaction between stepsize, batch size, and stochastic noise. Our study reveals a regime-dependent behavior: increasing the batch size initially improves optimization accuracy but, beyond a critical threshold, the benefits saturate and can eventually degrade performance under a fixed token budget. Notably, the theory predicts the magnitude of the optimal stepsize and aligns well with empirical practices observed in large-scale training. Leveraging these insights, we derive principled guidelines for selecting the batch size and stepsize, and propose an adaptive strategy that increases batch size and sequence length during training while preserving convergence guarantees. Experiments on NanoGPT are consistent with the theoretical predictions and illustrate the emergence of the predicted scaling regimes. Overall, our results provide a theoretical framework for understanding batch size scaling in stochastic conditional gradient methods and offer guidance for designing efficient training schedules in large-scale optimization.

Rustem Islamov, Roman Machacek, Aurelien Lucchi +3

University of Basel, Switzerland · MBZUAI, UAE · University of Bern, Switzerland +2

Stochastic OptimizationNeural Network OptimizationStochastic Optimization Convergence
AskPDF
Mar 22, 2026

Rethinking Plasticity in Deep Reinforcement Learning

This paper investigates the fundamental mechanisms driving plasticity loss in deep reinforcement learning (RL), a critical challenge where neural networks lose their ability to adapt to non-stationary environments. While existing research often relies on descriptive metrics like dormant neurons or effective rank, these summaries fail to explain the underlying optimization dynamics. We propose the Optimization-Centric Plasticity (OCP) hypothesis, which posits that plasticity loss arises because optimal points from previous tasks become poor local optima for new tasks, trapping parameters during task transitions and hindering subsequent learning. We theoretically establish the equivalence between neuron dormancy and zero-gradient states, demonstrating that the absence of gradient signals is the primary driver of dormancy. Our experiments reveal that plasticity loss is highly task-specific; notably, networks with high dormancy rates in one task can achieve performance parity with randomly initialized networks when switched to a significantly different task, suggesting that the network's capacity remains intact but is inhibited by the specific optimization landscape. Furthermore, our hypothesis elucidates why parameter constraints mitigate plasticity loss by preventing deep entrenchment in local optima. Validated across diverse non-stationary scenarios, our findings provide a rigorous optimization-based framework for understanding and restoring network plasticity in complex RL domains.

Zhiqiang He

The University of Electro-Communications, Tokyo, Japan

Loss of PlasticityNon-Stationary RLNeural Network Training DynamicsNeural Network Optimization
AskPDF
Mar 20, 2026

RMNP: Row-Momentum Normalized Preconditioning for Scalable Matrix-Based Optimization

Preconditioned adaptive methods have gained significant attention for training deep neural networks, as they capture rich curvature information of the loss landscape. The central challenge in this field lies in balancing preconditioning effectiveness with computational efficiency of implementing the preconditioner. Among recent advances, Muon stands out by using Newton-Schulz iteration to obtain preconditioned updates without explicitly constructing the preconditioning matrix. Despite its advantages, the efficiency of Muon still leaves room for further improvement. In this paper, we introduce RMNP (Row Momentum Normalized Preconditioning), an optimizer that replaces Newton-Schulz iteration with a simple row-wise (dind_{\text{in}}) ℓ2\ell_2 normalization operation, motivated by the empirically observed diagonal block structure of the Transformer layerwise Hessian. We empirically verified that orthogonalization and row-wise (on input dim) ℓ2\ell_2 normalization are asymptotically equivalent in the case of the transformer. This substitution reduces the per-iteration computational complexity from O(mn⋅min⁡(m,n)){O}(mn\cdot\min(m,n)) to O(mn){O}(mn) for an m×nm\times n weight matrix while maintaining comparable optimization performance. Theoretically, we establish convergence guarantees for RMNP in the non-convex setting that match recent results for Muon optimizers, achieving the minimax optimal complexity. Extensive experiments on large language model pretraining show that RMNP delivers competitive optimization performance compared with Muon while substantially reducing preconditioning wall-clock time. Our code is available at https://github.com/Dominator-Index/RMNP.

Shenyang Deng, Zhuoli Ouyang, Tianyu Pang +4

Dartmouth College · International Computer Science Institute · University of California, Berkeley +1

Neural Network OptimizationMuon Optimizer
AskPDF
Mar 20, 2026

Does This Gradient Spark Joy?

Policy gradient computes a backward pass for every sample, even though the backward pass is expensive and most samples carry little learning value. The Delightful Policy Gradient (DG) provides a forward-pass signal of learning value: delight, the product of advantage and surprisal (negative log-probability). We introduce the Kondo gate, which compares delight against a compute price and pays for a backward pass only when the sample is worth it, thereby tracing a quality--cost Pareto frontier. In bandits, zero-price gating preserves useful gradient signal while removing perpendicular noise, and delight is a more reliable screening signal than additive combinations of value and surprise. On MNIST and transformer token reversal, the Kondo gate skips most backward passes while retaining nearly all of DG's learning quality, with gains that grow as problems get harder and backward passes become more expensive. Because the gate tolerates approximate delight, a cheap forward pass can screen samples before expensive backpropagation, suggesting a speculative-decoding-for-training paradigm.

Ian Osband

Google DeepMind

Bandit LearningPolicy GradientNeural Network OptimizationAdaptive Computation
AskPDF
Mar 17, 2026

Linearized Bregman Iterations for Sparse Spiking Neural Networks

Spiking Neural Networks (SNNs) offer an energy efficient alternative to conventional Artificial Neural Networks (ANNs) but typically still require a large number of parameters. This work introduces Linearized Bregman Iterations (LBI) as an optimizer for training SNNs, enforcing sparsity through iterative minimization of the Bregman distance and proximal soft thresholding updates. To improve convergence and generalization, we employ the AdaBreg optimizer, a momentum and bias corrected Bregman variant of Adam. Experiments on three established neuromorphic benchmarks, i.e. the Spiking Heidelberg Digits (SHD), the Spiking Speech Commands (SSC), and the Permuted Sequential MNIST (PSMNIST) datasets, show that LBI based optimization reduces the number of active parameters by about 50% while maintaining accuracy comparable to models trained with the Adam optimizer, demonstrating the potential of convex sparsity inducing methods for efficient neuromorphic learning.

Daniel Windhager, Bernhard A. Moser, Michael Lunglmayr

Silicon Austria Labs, Linz, Austria · Software Competence Center Hagenberg, Hagenberg, Austria · Institute of Signal Processing, Johannes Kepler University, Linz, Austria

Spiking Neural NetworksSparse Neural NetworksNeural Network OptimizationNeuromorphic Computing
AskPDF
Mar 17, 2026

Prior-Informed Neural Network Initialization: A Spectral Approach for Function Parameterizing Architectures

Neural network architectures designed for function parameterization, such as the Bag-of-Functions (BoF) framework, bridge the gap between the expressivity of deep learning and the interpretability of classical signal processing. However, these models are inherently sensitive to parameter initialization, as traditional data-agnostic schemes fail to capture the structural properties of the target signals, often leading to suboptimal convergence. In this work, we propose a prior-informed design strategy that leverages the intrinsic spectral and temporal structure of the data to guide both network initialization and architectural configuration. A principled methodology is introduced that uses the Fast Fourier Transform to extract dominant seasonal priors, informing model depth and initial states, and a residual-based regression approach to parameterize trend components. Crucially, this structural alignment enables a substantial reduction in encoder dimensionality without compromising reconstruction fidelity. A supporting theoretical analysis provides guidance on trend estimation under finite-sample regimes. Extensive experiments on synthetic and real-world benchmarks demonstrate that embedding data-driven priors significantly accelerates convergence, reduces performance variability across trials, and improves computational efficiency. Overall, the proposed framework enables more compact and interpretable architectures while outperforming standard initialization baselines, without altering the core training procedure.

David Orlando Salazar Torres, Diyar Altinses, Andreas Schwung

Department of Automation Technology and Learning Systems, South Westphalia University of Applied Sciences, Soest, Germany

Neural Network InitializationNeural Network Optimization
AskPDF
Mar 14, 2026

Align Forward, Adapt Backward: Closing the Discretization Gap in Logic Gate Networks

In neural network models, soft mixtures of fixed candidate components (e.g., logic gates and sub-networks) are often used during training for stable optimization, while hard selection is typically used at inference. This raises questions about training-inference mismatch. We analyze this gap by separating forward-pass computation (hard selection vs. soft mixture) from stochasticity (with vs. without Gumbel noise). Using logic gate networks as a testbed, we observe distinct behaviors across four methods: Hard-ST achieves zero selection gap by construction; Gumbel-ST achieves near-zero gap when training succeeds but suffers accuracy collapse at low temperatures; Soft-Mix achieves small gap only at low temperature via weight concentration; and Soft-Gumbel exhibits large gaps despite Gumbel noise, confirming that noise alone does not reduce the gap. We propose CAGE (Confidence-Adaptive Gradient Estimation) to maintain gradient flow while preserving forward alignment. On logic gate networks, Hard-ST with CAGE achieves over 98% accuracy on MNIST and over 58% on CIFAR-10, both with zero selection gap across all temperatures, while Gumbel-ST without CAGE suffers a 47-point accuracy collapse.

Youngsung Kim

Department of Artificial Intelligence and Department of Electrical and Computer Engineering, Inha University, Incheon, Korea.

Neural Network OptimizationDifferentiable Logic Gate Networks
AskPDF
Mar 13, 2026

Scaling Laws and Pathologies of Single-Layer PINNs: Network Width and PDE Nonlinearity

We establish empirical scaling laws for Single-Layer Physics-Informed Neural Networks on canonical nonlinear PDEs. We identify a dual optimization failure: (i) a baseline pathology, where the solution error fails to decrease with network width, even at fixed nonlinearity, falling short of theoretical approximation bounds, and (ii) a compounding pathology, where this failure is exacerbated by nonlinearity. We provide quantitative evidence that a simple separable power law is insufficient, and that the scaling behavior is governed by a more complex, non-separable relationship. This failure is consistent with the concept of spectral bias, where networks struggle to learn the high-frequency solution components that intensify with nonlinearity. We show that optimization, not approximation capacity, is the primary bottleneck, and propose a methodology to empirically measure these complex scaling effects.

Faris Chaudhry

Department of Computer Science, Imperial College London

Physics-Informed MLNeural Scaling LawsNeural Network OptimizationShallow Neural Networks
AskPDF
Mar 11, 2026

SCORE: Replacing Layer Stacking with Contractive Recurrent Depth

Residual connections are central to modern deep neural networks, enabling stable optimization and efficient information flow across depth. In this work, we propose SCORE (Skip-Connection ODE Recurrent Embedding), a discrete recurrent alternative to classical layer stacking. Instead of composing multiple independent layers, SCORE iteratively applies a single shared neural block using an ODE (Ordinary Differential Equation)-inspired contractive update: ht+1 = (1 - dt) * ht + dt * F(ht) This formulation can be interpreted as a depth-by-iteration refinement process, where the step size dt explicitly controls stability and update magnitude. Unlike continuous Neural ODE approaches, SCORE uses a fixed number of discrete iterations and standard backpropagation without requiring ODE solvers or adjoint methods. We evaluate SCORE across graph neural networks (ESOL molecular solubility), multilayer perceptrons, and Transformer-based language models (nanoGPT). Across architectures, SCORE generally improves convergence speed and often accelerates training. SCORE is reducing parameter count through shared weights. In practice, simple Euler integration provides the best trade-off between computational cost and performance, while higher-order integrators yield marginal gains at increased compute. These results suggest that controlled recurrent depth with contractive residual updates offers a lightweight and effective alternative to classical stacking in deep neural networks.

Guillaume Godin

1Osmo Labs PBC New York, USA

Residual LearningRecurrent Neural NetworksNeural Network Optimization
AskPDF
Mar 11, 2026

What do near-optimal learning rate schedules look like?

A basic unanswered question in neural network training is: what is the best learning rate schedule shape for a given workload? The choice of learning rate schedule is a key factor in the success or failure of the training process, but beyond having some kind of warmup and decay, there is no consensus on what makes a good schedule shape. To answer this question, we designed a search procedure to find the best shapes within a parameterized schedule family. Our approach factors out the schedule shape from the base learning rate, which otherwise would dominate cross-schedule comparisons. We applied our search procedure to a variety of schedule families on three workloads: linear regression, image classification on CIFAR-10, and small-scale language modeling on Wikitext103. We showed that our search procedure indeed generally found near-optimal schedules. We found that warmup and decay are robust features of good schedules, and that commonly used schedule families are not optimal on these workloads. Finally, we explored how the outputs of our shape search depend on other optimization hyperparameters, and found that weight decay can have a strong effect on the optimal schedule shape. To the best of our knowledge, our results represent the most comprehensive results on near-optimal schedule shapes for deep neural network training, to date.

Hiroki Naganuma, Atish Agarwala, Priya Kasimbeg +1

Mila · Université de Montréal · Google DeepMind

Learning Rate SchedulingNeural Network OptimizationHyperparameter Optimization
AskPDF
Mar 10, 2026

On the Width Scaling of Neural Optimizers Under Matrix Operator Norms I: Row/Column Normalization and Hyperparameter Transfer

A central question in modern deep learning is how to design optimizers whose behavior remains stable as the network width ww increases. We address this question by interpreting several widely used neural-network optimizers, including AdamW and Muon, as instances of steepest descent under matrix operator norms. This perspective links optimizer geometry with the Lipschitz structure of the network forward map, and enables width-independent control of both Lipschitz and smoothness constants. However, steepest-descent rules induced by standard p→qp \to q operator norms lack layerwise composability and therefore cannot provide width-independent bounds in deep architectures. We overcome this limitation by introducing a family of mean-normalized operator norms, denoted \pmean→\qmean\pmean \to \qmean, that admit layerwise composability, yield width-independent smoothness bounds, and give rise to practical optimizers such as rescaled AdamW, row normalization, and column normalization. The resulting learning rate width-aware scaling rules recover μμP scaling~\cite{yang2021tensor} as a special case and provide a principled mechanism for cross-width learning-rate transfer across a broad class of optimizers. We further show that Muon can suffer an O(w)\mathcal{O}(\sqrt{w}) worst-case growth in the smoothness constant, whereas a new family of row-normalized optimizers we propose achieves width-independent smoothness guarantees. Based on the observations, we propose MOGA (Matrix Operator Geometry Aware), a width-aware optimizer based only on row/column-wise normalization that enables stable learning-rate transfer across model widths. Large-scale pre-training on GPT-2 and LLaMA shows that MOGA, especially with row normalization, is competitive with Muon while being notably faster in large-token and low-loss regimes.

Ruihan Xu, Jiajin Li, Yiping Lu

University of Chicago · Sauder School of Business, University of British Columbia · Industrial Engineering & Management Science, Northwestern University

Neural Network OptimizationHyperparameter TransferMatrix Optimization
AskPDF
Mar 7, 2026

Combining Adam and its Inverse Counterpart to Enhance Generalization of Deep Learning Optimizers

In the training of neural networks, adaptive moment estimation (Adam) typically converges fast but exhibits suboptimal generalization performance. A widely accepted explanation for its defect in generalization is that it often tends to converge to sharp minima. To enhance its ability to find flat minima, we propose its new variant named inverse Adam (InvAdam). The key improvement of InvAdam lies in its parameter update mechanism, which is opposite to that of Adam. Specifically, it computes element-wise multiplication of the first-order and second-order moments, while Adam computes the element-wise division of these two moments. This modification aims to increase the step size of the parameter update when the elements in the second-order moments are large and vice versa, which helps the parameter escape sharp minima and stay at flat ones. However, InvAdam's update mechanism may face challenges in convergence. To address this challenge, we propose dual Adam (DualAdam), which integrates the update mechanisms of both Adam and InvAdam, ensuring convergence while enhancing generalization performance. Additionally, we introduce the diffusion theory to mathematically demonstrate InvAdam's ability to escape sharp minima. Extensive experiments are conducted on image classification tasks and large language model (LLM) fine-tuning. The results validate that DualAdam outperforms Adam and its state-of-the-art variants in terms of generalization performance. The code is publicly available at https://github.com/LongJin-lab/DualAdam.

Tao Shi, Liangming Chen, Long Jin +1

School of Information Science and Engineering, Lanzhou University, Lanzhou 730000, China · Helen and John C. Hartmann Department of Electrical and Computer Engineering, New Jersey Institute of Technology, Newark, NJ 07102, USA

Neural Network OptimizationAdaptive Gradient MethodsNeural Network Generalization
AskPDF
Mar 6, 2026

Weak-SIGReg: Covariance Regularization for Stable Deep Learning

Modern neural network optimization relies heavily on architectural priorssuch as Batch Normalization and Residual connectionsto stabilize training dynamics. Without these, or in low-data regimes with aggressive augmentation, low-bias architectures like Vision Transformers (ViTs) often suffer from optimization collapse. This work adopts Sketched Isotropic Gaussian Regularization (SIGReg), recently introduced in the LeJEPA self-supervised framework, and repurposes it as a general optimization stabilizer for supervised learning. While the original formulation targets the full characteristic function, a computationally efficient variant is derived, Weak-SIGReg, which targets the covariance matrix via random sketching. Inspired by interacting particle systems, representation collapse is viewed as stochastic drift; SIGReg constrains the representation density towards an isotropic Gaussian, mitigating this drift. Empirically, SIGReg recovers the training of a ViT on CIFAR-100 from a collapsed 20.73% to 72.02% accuracy without architectural hacks and significantly improves the convergence of deep vanilla MLPs trained with pure SGD. Code is available at github.com/kreasof-ai/sigreg.

Habibullah Akbar

Kreasof AI, Jakarta, Indonesia

Neural Network OptimizationRepresentation Collapse
AskPDF
Mar 5, 2026

Multilevel Training for Kolmogorov Arnold Networks

Algorithmic speedup of training common neural architectures is made difficult by the lack of structure guaranteed by the function compositions inherent to such networks. In contrast to multilayer perceptrons (MLPs), Kolmogorov-Arnold networks (KANs) provide more structure by expanding learned activations in a specified basis. This paper exploits this structure to develop practical algorithms and theoretical insights, yielding training speedup via multilevel training for KANs. To do so, we first establish an equivalence between KANs with spline basis functions and multichannel MLPs with power ReLU activations through a linear change of basis. We then analyze how this change of basis affects the geometry of gradient-based optimization with respect to spline knots. The KANs change-of-basis motivates a multilevel training approach, where we train a sequence of KANs naturally defined through a uniform refinement of spline knots with analytic geometric interpolation operators between models. The interpolation scheme enables a ``properly nested hierarchy'' of architectures, ensuring that interpolation to a fine model preserves the progress made on coarse models, while the compact support of spline basis functions ensures complementary optimization on subsequent levels. Numerical experiments demonstrate that our multilevel training approach can achieve orders of magnitude improvement in accuracy over conventional methods to train comparable KANs or MLPs, particularly for physics informed neural networks. Finally, this work demonstrates how principled design of neural networks can lead to exploitable structure, and in this case, multilevel algorithms that can dramatically improve training performance.

Ben S. Southworth, Jonas A. Actor, Graham Harper +1

Theoretical Division, Los Alamos National Laboratory, USA. · Center for Computing Research, Sandia National Laboratories, USA.

Kolmogorov-Arnold NetworksNeural Network Optimization
AskPDF
Mar 4, 2026

Exploiting Subgradient Sparsity in Max-Plus Neural Networks

Deep Neural Networks are powerful tools for solving machine learning problems, but their training often involves dense and costly parameter updates. In this work, we use a novel Max-Plus neural architecture in which classical addition and multiplication are replaced with maximum and summation operations respectively. This is a promising architecture in terms of interpretability, but its training is challenging. A particular feature is that this algebraic structure naturally induces sparsity in the subgradients, as only neurons that contribute to the maximum affect the loss. However, standard backpropagation fails to exploit this sparsity, leading to unnecessary computations. In this work, we focus on the minimization of the worst sample loss which transfers this sparsity to the optimization loss. To address this, we propose a sparse subgradient algorithm that explicitly exploits the algebraic sparsity. By tailoring the optimization procedure to the non-smooth nature of Max-Plus models, our method achieves more efficient updates while retaining theoretical guarantees. This highlights a principled path toward bridging algebraic structure and scalable learning.

Ikhlas Enaieh, Olivier Fercoq

LTCI, S2A · Image, Data, Signal Department (IDS), Telecom Paris, Institut Polytechnique de Paris, France. · S2A, LTCI

Neural Network OptimizationGradient Descent
AskPDF
Mar 3, 2026

Same Error, Different Function: The Optimizer as an Implicit Prior in Financial Time Series

Neural networks applied to financial time series operate in a regime of underspecification, where model predictors achieve indistinguishable out-of-sample error. Using large-scale volatility forecasting for S&\&P 500 stocks, we show that different model-training-pipeline pairs with identical test loss learn qualitatively different functions. Across architectures, predictive accuracy remains unchanged, yet optimizer choice reshapes non-linear response profiles and temporal dependence differently. These divergences have material consequences for decisions: volatility-ranked portfolios trace a near-vertical Sharpe-turnover frontier, with nearly 3×3\times turnover dispersion at comparable Sharpe ratios. We conclude that in underspecified settings, optimization acts as a consequential source of inductive bias, thus model evaluation should extend beyond scalar loss to encompass functional and decision-level implications.

Federico Vittorio Cortesi, Giuseppe Iannone, Giulia Crippa +2

Massachusetts Institute of Technology, Cambridge, MA, USA · Princeton University, Princeton, NJ, USA.

Neural Network OptimizationFinancial Time SeriesTime Series Forecasting
AskPDF
Feb 28, 2026

Geodesic Gradient Descent: A Generic and Learning-rate-free Optimizer on Objective Function-induced Manifolds

Euclidean gradient descent algorithms barely capture the geometry of objective function-induced hypersurfaces and risk driving update trajectories off the hypersurfaces. Riemannian gradient descent algorithms address these issues but fail to represent complex hypersurfaces via a single classic manifold. We propose geodesic gradient descent (GGD), a generic and learning-rate-free Riemannian gradient descent algorithm. At each iteration, GGD uses an n-dimensional sphere to approximate a local neighborhood on the objective function-induced hypersurface, adapting to arbitrarily complex geometries. A tangent vector derived from the Euclidean gradient is projected onto the sphere to form a geodesic, ensuring the update trajectory stays on the hypersurface. Parameter updates are performed using the endpoint of the geodesic. The maximum step size of the gradient in GGD is equal to a quarter of the arc length on the n-dimensional sphere, thus eliminating the need for a learning rate. Experimental results show that compared with the classic Adam algorithm, GGD achieves test MSE reductions ranging from 35.79% to 48.76% for fully connected networks on the Burgers' dataset, and cross-entropy loss reductions ranging from 3.14% to 11.59% for convolutional neural networks on the MNIST dataset.

Liwei Hu, Guangyao Li, Wenyong Wang +2

the School of Information Science and Engineering, Hebei University of Science and Technology, Shijiazhuang 050091, China · the School of Computer Science and Engineering, University of Electronic Science and Technology of China, Chengdu 611731, China · the International Institute of Next Generation Internet, Macau University of Science and Technology, Macau 519020

Neural Network OptimizationRiemannian Optimization
AskPDF
Feb 26, 2026

MSINO: Curvature-Aware Sobolev Optimization for Manifold Neural Networks

We introduce Manifold Sobolev Informed Neural Optimization (MSINO), a curvature aware training framework for neural networks defined on Riemannian manifolds. The method replaces standard Euclidean derivative supervision with a covariant Sobolev loss that aligns gradients using parallel transport and improves stability via a Laplace Beltrami smoothness regularization term. Building on classical results in Riemannian optimization and Sobolev theory on manifolds, we derive geometry dependent constants that yield (i) a Descent Lemma with a manifold Sobolev smoothness constant, (ii) a Sobolev Polyak Lojasiewicz inequality giving linear convergence guarantees for Riemannian gradient descent and stochastic gradient descent under explicit step size bounds, and (iii) a two step Newton Sobolev method with local quadratic contraction in curvature controlled neighborhoods. Unlike prior Sobolev training in Euclidean space, MSINO provides training time guarantees that explicitly track curvature and transported Jacobians. Applications include surface imaging, physics informed learning settings, and robotics on Lie groups such as SO(3) and SE(3). The framework unifies value and gradient based learning with curvature aware convergence guarantees for neural training on manifolds.

Suresan Pareth

Riemannian Neural NetworksNeural Network OptimizationCurvature-Aware OptimizationRiemannian Optimization
AskPDF
Feb 25, 2026

When Learning Hurts: Fixed-Pole RNN for Real-Time Online Training

Recurrent neural networks (RNNs) can be interpreted as discrete-time state-space models, where the state evolution corresponds to an infinite-impulse-response (IIR) filtering operation governed by both feedforward weights and recurrent poles. While, in principle, all parameters including pole locations can be optimized via backpropagation through time (BPTT), such joint learning incurs substantial computational overhead and is often impractical for applications with limited training data. Echo state networks (ESNs) mitigate this limitation by fixing the recurrent dynamics and training only a linear readout, enabling efficient and stable online adaptation. In this work, we analytically and empirically examine why learning recurrent poles does not provide tangible benefits in data-constrained, real-time learning scenarios. Our analysis shows that pole learning renders the weight optimization problem highly non-convex, requiring significantly more training samples and iterations for gradient-based methods to converge to meaningful solutions. Empirically, we observe that for complex-valued data, gradient descent frequently exhibits prolonged plateaus, and advanced optimizers offer limited improvement. In contrast, fixed-pole architectures induce stable and well-conditioned state representations even with limited training data. Numerical results demonstrate that fixed-pole networks achieve superior performance with lower training complexity, making them more suitable for online real-time tasks.

Alexander Morgan, Ummay Sumaya Khan, Lingjia Liu +1

dept. of ECE, Virginia Tech, Alexandria, VA, USA

Recurrent Neural NetworksNeural Network Training DynamicsNeural Network OptimizationOnline LearningLinear RNNs
AskPDF
Feb 24, 2026

Neural network optimization strategies and the topography of the loss landscape

Neural networks are trained by optimizing multi-dimensional sets of fitting parameters on non-convex loss landscapes. Low-loss regions of the landscapes correspond to the parameter sets that perform well on the training data. A key issue in machine learning is the performance of trained neural networks on previously unseen test data. Here, we investigate neural network training by stochastic gradient descent (SGD) - a non-convex global optimization algorithm which relies only on the gradient of the objective function. We contrast SGD solutions with those obtained via a non-stochastic quasi-Newton method, which utilizes curvature information to determine step direction and Golden Section Search to choose step size. We use several computational tools to investigate neural network parameters obtained by these two optimization methods, including kernel Principal Component Analysis and a novel, general-purpose algorithm for finding low-height paths between pairs of points on loss or energy landscapes, FourierPathFinder. We find that the choice of the optimizer profoundly affects the nature of the resulting solutions. SGD solutions tend to be separated by lower barriers than quasi-Newton solutions, even if both sets of solutions are regularized by early stopping to ensure adequate performance on test data. When allowed to fit extensively on the training data, quasi-Newton solutions occupy deeper minima on the loss landscapes that are not reached by SGD. These solutions are less generalizable to the test data however. Overall, SGD explores smooth basins of attraction, while quasi-Newton optimization is capable of finding deeper, more isolated minima that are more spread out in the parameter space. Our findings help understand both the topography of the loss landscapes and the fundamental role of landscape exploration strategies in creating robust, transferrable neural network models.

Jianneng Yu, Alexandre V. Morozov

Department of Physics & Astronomy, Rutgers, The State University of New Jersey, 136 Frelinghuysen Rd., Piscataway, NJ 08854, U.S.A.

Neural Network OptimizationStochastic Gradient DescentNeural Network GeneralizationQuasi-Newton Methods
AskPDF