Curvature-Aware Optimization

Follow this topic Ask about this topic

Latest papers

21 papers1 paper in the last four weeks, against 1 the four weeks before. 0.0% of all new papers.

Oct 14, 2026

One Size Does Not Fit All: Curvature-Aware Adaptive Perturbation for SAM

Poor generalisation in deep neural networks is frequently associated with convergence to sharp minima. Sharpness-Aware Minimisation (SAM) improves generalisation of deep neural networks by minimising a combination of the training loss and the local sharpness of the loss landscape, biasing convergence towards flat minima. However, SAM relies on a fixed perturbation radius

ρ

that ignores local curvature, leading to inefficient behaviour: the perturbation may be too small to escape poor basins defined by sharp local minima, or too large to allow convergence to solutions with good loss-sharpness tradeoff. We reinterpret

ρ

as the intensity of the perturbation signal and propose Curvature-Aware Adaptive Perturbation for SAM (CAP-SAM), a geometry-aware method that modulates the perturbation radius based on a sharpness ratio, enabling a more effective search for robust minima on the loss landscape. We prove that CAP-SAM preserves the same convergence rate as standard SAM in non-convex stochastic settings. Extensive experiments indicate that our approach consistently outperforms SAM and its most prominent variants across diverse network architectures on CIFAR-10, CIFAR-100, and ImageNet.

Alessia Rossi, Felipe Campelo, Gabriel L. Oliveira

School of Engineering Mathematics and Technology, University of Bristol, Bristol, Lovelace Building, Ada, Bristol, BS8 1TW, UK

Curvature-Aware Optimization Sharpness-Aware Minimization Neural Network Optimization Neural Network Generalization

Aug 28, 2026

Curvature-Conditioned Multiscale Momentum with Sphere Constraints for LLM Pretraining

Pretraining accounts for a large fraction of the total computational cost in LLM training. However, noise-dominant gradients and the highly ill-conditioned loss landscape bring severe challenges. Although modern adaptive optimizers such as AdamW and Muon have achieved great success in large-scale pretraining, their reliance on gradient normalization offers limited mitigation of the ill-conditioned curvature. The progress along flat directions (eigen-directions of small eigenvalues), which dominates the final loss reduction, remains relatively slow. To enhance training dynamics along flat directions, we propose a curvature-conditioned multiscale momentum method with sphere constraints, delivering steady acceleration in LLM pretraining. This multiscale momentum, applied only along flat directions, pairs a slow-decay component for noise reduction with a fast-decay component for rapid curvature adaptation, harnessing their complementary strengths. Crucially, we employ a sphere constraint technique to prevent parameter inflation and excessively rapid effective learning rate decay that would otherwise arise from a naive combination. Extensive experiments show that the proposed method significantly accelerates Muon across diverse architectures (dense, MoE) and model sizes (0.12B--2.3B parameters). Theoretically, we verify the acceleration effect and provide insight into the design principles underlying the flat-direction multiscale momentum.

Shuchen Zhu, Yuxin Fang, Mingze Wang +1

ByteDance Seed · Peking University

Curvature-Aware Optimization Momentum Methods Accelerated Gradient Methods Language Model Pretraining

Aug 12, 2026

Curvature-Aware Zeroth-Order Optimization for Memory-Efficient Test-Time Adaptation

Test-time adaptation (TTA) aims to enhance the cross-domain performance of pre-trained models by adapting to unlabeled test data. While most existing TTA methods rely on backpropagation (BP) for finetuning, BP-free methods such as zeroth-order (ZO) methods are more desired in practical on-device scenarios. ZO methods rely only on forward computation, which can largely reduce the complexity and memory overhead of on-device deployment. However, ZO methods suffer from much higher variance compared with first-order methods in estimating the gradient. To address this, we propose an improved ZO method to substantially boost the performance of ZO optimization based TTA. First, we provide an observation to reveal the persistent low-rank Hessian structure of the loss during the adaptation process. Based on this insight, we then propose a loss-landscape curvature-aware zeroth-order (CAZO) method, which leverages a sliding-average estimation of the diagonal Hessian to construct a covariance matrix for anisotropic perturbation sampling. CAZO operates by freezing pretrained weights and optimizing minimal adapter parameters via forward-only passes based gradient estimation, which can substantially reduce the memory overhead compared to BP-based methods. Extensive experiments demonstrate that CAZO significantly outperforms existing TTA methods, achieving state-of-the-art performance while maintaining an excellent balance between accuracy and memory efficiency. Code is available at https://github.com/Hollyming/CAZO.

Junming Zhang, Shuyu Yin, Peilin Liu +2

Shanghai Jiao Tong University, China

On-Device Training Test-Time Adaptation Memory-Efficient Fine-Tuning Zeroth-Order Optimization Curvature-Aware Optimization

Aug 5, 2026

MALT: Lightweight Curvature-Aware Muon via Diagonal Preconditioning

Muon has recently emerged as a promising alternative to AdamW for language model pretraining by orthogonalizing momentum matrices using Newton-Schulz iterations. Although Muon mitigates gradient anisotropy, it does not explicitly account for the curvature geometry of the loss landscape and may therefore remain sensitive to curvature anisotropy. We bridge this gap by proposing MALT (Muon Augmented by Lightweight Two-sided Preconditioning), which uses lightweight diagonal preconditioners to reduce the sensitivity of Muon to curvature anisotropy. Specifically, MALT uses two-sided diagonal preconditioners with low memory and computational overhead to approximately capture the curvature geometry of the loss landscape. It orthogonalizes the preconditioned momentum using Newton-Schulz iterations and maps the result back to define the update direction, while norm grafting controls the update magnitude. To improve the robustness of MALT to stochastic gradient noise, we further propose MALTER (MALT with Adaptive stEpsize Rescaling). Convergence guarantees are provided for MALT in the stochastic non-convex setting. Experiments on GPT-2 Small, Medium, and Large pretraining show that the proposed methods outperform Muon while maintaining nearly the same memory footprint and wall-clock time.

Tongle Wu, Huanyu Dong, Ying Sun +1

School of Electrical Engineering and Computer Science, The Pennsylvania State University, University Park, PA, USA · Department of Computer Science, City University of Hong Kong, Hong Kong SAR, China

Language Model Pretraining Muon Optimizer Curvature-Aware Optimization

Jul 4, 2026

Directional Curvature from Armijo Backtracking: A Low-Cost Sharpness Probe and a Calibration-Free Learning-Rate Safeguard for Adam

Local sharpness, defined by the largest Hessian eigenvalue

λ_1

, sets the maximum stable gradient update size, but its computation would usually require running Lanczos or Hessian-vector products. However, we notice that even a single Armijo backtracking line search already contains this information with just a few forward passes, as the accepted step

α

determines the directional curvature along the search direction up to the multiplicative band set by the backtracking factor. The correlation between

\logα

and

\logλ_1

on CIFAR-10, Fashion-MNIST and Imagenette reaches

-0.91

to

-0.95

in Pearson correlation, and even after removing the trend per run the correlation remains at

-0.60

to

-0.70

. This allows for a cheap online Edge-of-Stability estimate of the slow sharpness component. The employed probing mechanism searches along Adam's first-step update direction, at initialisation and nine times over the course of the first 50 optimiser steps. The learning-rate cap is set as twice the smallest observed step size, and in the studied learning-rate ranges (

10^{-3}

to

3.0

) and GPT-2 pretraining experiments all capped runs avoid divergence; in the general architecture analysis, one MLP architecture is still sensitive to the initial batch order. This probing protocol incurs an approximately one percent overhead, and using a non-binding cap means that the optimiser's state and first update are bit-identical. There is no fine-tuning of any of the protocol parameters to any specific architecture; this is the sense in which this is a calibration-free safeguard. It is meant as a way of avoiding divergence, not achieving accuracy. At GPT-2 scale, multiple measurements also illustrate why an initialisation-only cap is insufficient: directional curvature rises substantially in the first five optimiser steps, which motivates the short probationary window.

Ashmitha R, Jörg Frochte

Department of Artificial Intelligence and Data Science, Sri Ramakrishna Engineering College, Anna University, Coimbatore 641022, Tamil Nadu, India · Interdisciplinary Institute for Applied AI and Data Science Ruhr (AKIS), Department of Electrical Engineering and Computer Science, Bochum University of Applied Sciences, Am Hochschulcampus 1, 44801 Bochum, Germany

Gradient Descent Adaptive Gradient Methods Gradient Descent Dynamics Layer-Wise Learning Rate Adaptation Curvature-Aware Optimization Loss Landscape Sharpness

Jun 29, 2026

Curvature-Weighted Gradient Diversity: A Noise Measure for Geometry-Adaptive SGD Schedules

The standard convergence analysis of mini-batch stochastic gradient descent (SGD) models gradient noise using a single variance term that treats all parameter directions equally, ignoring the fact that noise in high-curvature directions has less impact because learning rates are already constrained there. We introduce Curvature-Weighted Gradient Diversity (CWGD), a geometry-aware measure that weights per-sample gradient diversity by the inverse square root of the Hessian, providing a tighter proxy for the effective optimization noise. For strongly convex quadratic objectives with diagonal Hessians and isotropic noise, we prove that a CWGD-modulated cosine learning-rate schedule can reduce the asymptotic optimization error floor by up to a factor of two compared with standard cosine annealing. We implement this idea as CWGD-Cosine using a Hutchinson-based diagonal Hessian estimator that is exact for quadratic objectives. Across a range of condition numbers, batch sizes, and noise structures, CWGD-Cosine consistently achieves approximately 20% lower final optimization error than standard cosine annealing while incurring negligible overhead in the quadratic setting. We also identify and correct a degenerate curvature estimator, analyze the robustness of the proposed estimator, and explicitly discuss the limitations of the method, including Hessian staleness in non-convex optimization. These results establish CWGD as a principled geometry-aware measure of optimization noise and motivate future extensions to more general learning problems.

Muhammad Hamza, Ayush Goel

Indian Institute of Technology Kharagpur · Indian Institute of Technology Kharagpur (IIT KGP)

Learning Rate Scheduling Stochastic Optimization Convergence Stochastic Gradient Descent Curvature-Aware Optimization

Jun 3, 2026

Why Muon Outperforms Adam: A Curvature Perspective

Muon improves training efficiency over Adam in large language-model training by about two times, but the local geometric source of this advantage remains unclear. Our work takes a first step toward demystifying Muon's superiority over Adam from a curvature perspective. First, we apply a second-order Taylor approximation to the training landscape and show that Muon achieves a larger one-step loss decrease than Adam at matched validation loss. The two optimizers have comparable first-order gains, but Muon consistently incurs a smaller second-order curvature penalty. Second, we decompose this curvature penalty into the squared update norm and Normalized Directional Sharpness (NDS). We find that Muon and Adam have comparable update norms, so Muon's smaller curvature penalty is driven by lower NDS, not update scale. Third, we study how training data and model structure shape Muon's NDS advantage. Using Zipf-Probabilistic Context-Free Grammar (PCFG) data with controlled imbalance, we show that data imbalance amplifies Muon's NDS advantage over Adam. A within-/cross-layer decomposition further shows that, in the middle and late stages of training, Muon's lower NDS is mainly sustained by smaller within-layer curvature. Beyond empirical evidence, we analyze stylized quadratic problems with heterogeneous curvature and gradient alignment toward high-curvature modes. We prove that Muon attains a smaller average NDS than GD by balancing update energy across curvature groups; when curvature heterogeneity is sufficiently strong, this also yields lower local quadratic loss after the same number of steps.

Shuche Wang, Fengzhuo Zhang, Jiaxiang Li +2

National University of Singapore · Yale University · University of Minnesota

Loss Landscape Sharpness Neural Network Optimization Muon Optimizer Adaptive Gradient Methods Loss Landscape Geometry Curvature-Aware Optimization

Jun 1, 2026

From Non-Convex to Strongly Convex: Curvature-Adaptive FTPL for Online Optimization

Curvature adaptivity is a classical theme in online optimization: for convex Lipschitz losses, adaptive methods interpolate between the optimal

O(\sqrt{T})

regret for general convex losses and

O(\log T)

regret under strong convexity. Recent work has shown that Follow-the-Perturbed-Leader (FTPL) achieves optimal

O(\sqrt{T})

regret even for online non-convex Lipschitz losses, assuming access to an approximate offline-optimization oracle, but these guarantees do not exploit curvature. We show that FTPL can be made curvature-adaptive in the non-convex setting, without knowing in advance how curvature will accumulate over time. Our algorithm replaces the fixed perturbation scale of standard FTPL with a time-varying scale chosen using only past information. We give a simple follow-the-leader tuning rule for this scale and show that it competes, up to constants, with the best choice in hindsight. The resulting method achieves

O(\sqrt{T})

regret for arbitrary non-convex Lipschitz losses and improves as cumulative curvature grows; with sufficiently accurate oracle calls, it achieves

O(\log T)

regret when cumulative curvature grows linearly, which includes the classical strongly convex regime. We complement these upper bounds with matching lower bounds for prescribed cumulative-curvature sequences, already for one-dimensional convex losses, showing that the tradeoff between worst-case non-convex regret and curvature-driven fast rates is intrinsic.

Moses Charikar, Chirag Pabbaraju, Ambuj Tewari

Stanford University. · University of Michigan, Ann Arbor.

Online Convex Optimization Follow-the-Perturbed-Leader Nonconvex Optimization Curvature-Aware Optimization

May 30, 2026

Exploiting weight-space symmetries for approximating curvature

Many machine learning techniques rely on approximating a loss function's curvature, but this is notoriously hard to do at the scale of modern deep networks. Surprisingly, no previous work has exploited the curvature constraints that arise from well known weight-space symmetries in loss landscapes. By analytically averaging over group actions that leave the loss invariant, we construct structured Hessian approximations from single gradients that can be tractably estimated, stored, and inverted. The choice of user-specified symmetry group directly governs the trade-off between approximation accuracy and computational cost. Moreover, our framework provides a unifying theoretical lens for viewing existing methods; in particular, a specific choice of symmetry group recovers Shampoo/Muon-like curvature estimates. We validate our method on a range of network architectures, and deploy it to second-order optimization benchmarks, including a small language model. Our curvature estimation framework might find applications in other machine learning problems such as uncertainty estimation, continual learning, compression/pruning, training data attribution, and more.

Artem Artemev, Rui Xia, Benjamin M. Boyd +4

MediaTek, Cambridge, UK · University of Cambridge, Cambridge, UK · Vector Institute, Toronto, Canada

Neural Network Optimization Second-Order Optimization Weight-Space Analysis Loss Landscape Geometry Curvature-Aware Optimization

May 29, 2026

Towards Efficient LLMs Annealing with Principled Sample Selection

The annealing phase is a pivotal convergence stage in LLM pre-training that ultimately determines final model quality. However, effectively selecting training data during this phase remains a key challenge. Current strategies rely on empirical heuristics, such as domain filtering or context extension, which lack a principled grounding in optimization theory. In this work, we characterize the annealing phase through the lens of the loss landscape's spectral geometry. We argue that optimal convergence requires gradient updates to satisfy heterogeneous constraints across different eigen-directions. Building on this insight, we formulate data selection as a problem of satisfying these directional constraints. To this end, we propose DiReCT (Directionally-Restrained Constrained Training), a novel framework that reformulates sample selection in the annealing stage as a constrained optimization problem. By imposing explicit directional constraints on per-sample gradients based on the spectral properties of the Hessian, DiReCT identifies samples that align with the optimal curvature-aware descent path. Extensive experiments across various model scales demonstrate that DiReCT consistently achieves state-of-the-art performance. For future research, code is available at https://github.com/xuyj233/Direct.

Yuanjian Xu, Jianing Hao, Wanbo Zhang +2

The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China · Microsoft Research Asia, Beijing, China · Fudan University, Shanghai, China

Language Model Pretraining Efficient Language Model Training Curvature-Aware Optimization Training Data Selection

May 15, 2026

Navigating Potholes with Geometry-Aware Sharpness Minimization

Sharpness-aware minimization (SAM) encourages flat minima by perturbing parameters along directions of high loss curvature, but treats all parameter directions uniformly, ignoring the underlying loss geometry. We introduce LLQR+SAM, which combines SAM with a learned preconditioner obtained from the recently proposed LLQR framework, a second-order method that recasts steepest descent as a layerwise linear-quadratic regulator problem. The preconditioner is updated sparsely and maintained as a slow exponential moving average, so it captures a smoothed, low-resolution picture of the loss landscape geometry. The SAM perturbation then operates on top of this learned geometry, probing curvature at a faster timescale. We show that this two-timescale structure is not merely a computational convenience: theoretically, the preconditioner amplifies the SAM escape signal in directions that are flat under the average geometry but locally sharp (potholes). Wide, flat basins, by contrast, remain stable. Empirically, LLQR+SAM gives consistent gains over both SAM and LLQR alone across standard vision and sequence modeling benchmarks, supporting the view that slow learned geometry and fast sharpness correction are genuinely complementary.

Simon Dufort-Labbé, Mehrab Hamidi, Razvan Pascanu +3

Mila, Université de Montréal · Samsung – SAIL Montreal

Sharpness-Aware Minimization Flat Minima Loss Landscape Geometry Curvature-Aware Optimization Loss Landscape Sharpness

May 15, 2026

Accelerated Gradient Descent for Faster Convergence with Minimal Overhead

In this paper, we present CT-AGD (Curvature-Tuned Accelerated Gradient Descent), an optimization method for non-convex optimization problems in deep learning training tasks. CT-AGD is a general boosting procedure that accelerates first-order methods by explicitly capturing the local curvature using finite-difference quotients, and the development of heuristics aimed at mitigating noise and bias introduced by stochastic mini-batch training. CT-AGD has a comparable storage and computational overhead as adaptive gradient methods such as Adam. Our extensive experiments demonstrate that CT-AGD achieves the same level of accuracy as the baseline first-order methods, yet reduces the required training epochs by 33% on average.

Manuel Graca, L. Miguel Silveira, Arlindo Oliveira +1

IST, ULisboa Lisbon, Portugal · Old Dominion University Norfolk, VA USA

Deep Learning Optimization Gradient Descent Nonconvex Stochastic Optimization Accelerated Gradient Methods Curvature-Aware Optimization

May 8, 2026

Curvature Beyond Positivity: Greedy Guarantees for Arbitrary Submodular Functions

Submodular functions -- functions exhibiting diminishing returns -- are central to machine learning. When the objective is monotone and non-negative, the greedy algorithm achieves a tight

63\%

approximation. But many practical objectives incorporate costs that make them negative on some inputs, and all existing multiplicative guarantees require non-negativity. Prior work handles negativity through additive bounds for the special class of decomposable functions and non-monotonicity through partial-monotonicity parameters, but these address each difficulty in isolation and neither extends the classical structural theory. We extend curvature -- a parameter measuring how far a function deviates from linearity -- to all submodular functions, handling both non-monotonicity and negativity through a single classical concept. A greedy algorithm with pruning achieves a curvature-controlled multiplicative ratio for any submodular function, including those taking negative values -- the first such guarantee beyond monotonicity and non-negativity. In the non-monotone regime

1 \le c_g < 2.2

, the bound strictly beats the best known uniform ratio of

0.401

(for non-negative

f

), and it recovers the classical

(1-e^{-c_g})/c_g

guarantee for monotone functions. A multilinear-extension variant extends the framework to general combinatorial constraints via multilinear relaxation. Experiments on cost-penalized experimental design, coverage, feature selection, and a curvature sweep on Multi-News passage selection support the theory.

Yixin Chen, Alan Kuhnle

Department of Computer Science & Engineering, Texas A&M University

Submodular Optimization Approximation Algorithms Curvature-Aware Optimization

May 5, 2026

Layerwise LQR for Geometry-Aware Optimization of Deep Networks

Geometry-aware optimizers such as Newton and natural gradient can improve conditioning in deep learning, but scalable variants such as K-FAC, Shampoo, and related preconditioners usually impose structural approximations early, often discarding cross-layer interactions induced by the network computation. We introduce Layerwise LQR (LLQR), a framework for learning structured inverse preconditioners under a global layerwise optimal-control objective. The starting point is an exact equivalence: the steepest-descent step under a broad class of divergence-induced quadratic models--including Newton, Gauss-Newton, Fisher/natural-gradient, and intermediate-layer metrics--can be written as a finite-horizon Linear Quadratic Regulator (LQR) problem. This formulation serves as a reference that exposes the layerwise dynamics and cost matrices encoding the original dense geometry. We then derive a scalable relaxation that learns diagonal, (E-)Kronecker-factored, or other structured inverse preconditioners by minimizing the LQR objective and reusing them across iterations. The resulting optimizer wraps standard methods while retaining a principled connection to second-order geometry, without forming or inverting the global curvature matrix. Experiments on ResNets and Transformers show that LLQR improves optimization dynamics and often translates these gains into improved final test performance, while adding only modest wall-clock overhead. It establishes LLQR as a practical framework for geometry-aware second-order methods and a reference for evaluating scalable approximations.

Simon Dufort-Labbé, Pierre-Luc Bacon, Razvan Pascanu +2

Mila, Université de Montréal · Samsung – SAIL Montreal

Deep Learning Optimization Second-Order Optimization Linear Quadratic Regulator Curvature-Aware Optimization

Apr 16, 2026

Lightweight Geometric Adaptation for Training Physics-Informed Neural Networks

Physics-Informed Neural Networks (PINNs) often suffer from slow convergence, training instability, and reduced accuracy on challenging partial differential equations due to the anisotropic and rapidly varying geometry of their loss landscapes. We propose a lightweight curvature-aware optimization framework that augments existing first-order optimizers with an adaptive predictive correction based on secant information. Consecutive gradient differences are used as a cheap proxy for local geometric change, together with a step-normalized secant curvature indicator to control the correction strength. The framework is plug-and-play, computationally efficient, and broadly compatible with existing optimizers, without explicitly forming second-order matrices. Experiments on diverse PDE benchmarks show consistent improvements in convergence speed, training stability, and solution accuracy over standard optimizers and strong baselines, including on the high-dimensional heat equation, Gray--Scott system, Belousov--Zhabotinsky system, and 2D Kuramoto--Sivashinsky system.

Kang An, Chenhao Si, Shiqian Ma +1

School of Data Science The Chinese University of Hong Kong, Shenzhen Shenzhen, China

Neural Network Optimization Second-Order Optimization Physics-Informed ML Curvature-Aware Optimization

Mar 31, 2026

Efficient Bilevel Optimization with KFAC-Based Hypergradients

Bilevel optimization (BO) is widely applicable to many machine learning problems. Scaling BO, however, requires repeatedly computing hypergradients, which involves solving inverse Hessian-vector products (IHVPs). In practice, these operations are often approximated using crude surrogates such as one-step gradient unrolling or identity/short Neumann expansions, which discard curvature information. We build on implicit function theorem-based algorithms and propose to incorporate Kronecker-factored approximate curvature (KFAC), yielding curvature-aware hypergradients with a better performance efficiency trade-off than Conjugate Gradient (CG) or Neumann methods and consistently outperforming unrolling. We evaluate this approach across diverse tasks, including meta-learning and AI safety problems. On models up to BERT, we show that curvature information is valuable at scale, and KFAC can provide it with only modest memory and runtime overhead. Our implementation is available at https://github.com/liaodisen/NeuralBo.

Disen Liao, Felix Dangel, Yaoliang Yu

University of Waterloo · Vector Institute

Kronecker-Factored Approximation Bilevel Optimization Hypergradient Optimization Differentiable Optimization Curvature-Aware Optimization Second-Order Optimization

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 Optimization Deep Learning Optimization Neural Network Optimization Second-Order Optimization

Mar 8, 2026

Vision Transformers that Never Stop Learning

Loss of plasticity refers to the progressive inability of a model to adapt to new tasks and poses a fundamental challenge for continual learning. While this phenomenon has been extensively studied in homogeneous neural architectures, such as multilayer perceptrons, its mechanisms in structurally heterogeneous, attention-based models such as Vision Transformers (ViTs) remain underexplored. In this work, we present a systematic investigation of loss of plasticity in ViTs, including a fine-grained diagnosis using local metrics that capture parameter diversity and utilization. Our analysis reveals that stacked attention modules exhibit increasing instability that exacerbates plasticity loss, while feed-forward network modules suffer even more pronounced degradation. Furthermore, we evaluate several approaches for mitigating plasticity loss. The results indicate that methods based on parameter re-initialization fail to recover plasticity in ViTs, whereas approaches that explicitly regulate the update process are more effective. Motivated by this insight, we propose ARROW, a geometry-aware optimizer that preserves plasticity by adaptively reshaping gradient directions using an online curvature estimate for the attention module. Extensive experiments show that ARROW effectively improves plasticity and maintains better performance on newly encountered tasks.

Caihao Sun, Mingqi Yuan, Shiyuan Wang +1

University of Hong Kong · The Hong Kong Polytechnic University · Technische Universität Dresden +1

Curvature-Aware Optimization Loss of Plasticity Vision Transformer

Mar 1, 2026

Curvature-Weighted Capacity Allocation: A Minimum Description Length Framework for Layer-Adaptive Large Language Model Optimization

Layer-wise capacity in large language models is highly non-uniform: some layers contribute disproportionately to loss reduction, whereas others are nearly redundant. Existing layer-scoring methods provide sensitivity estimates but do not give a principled rule for converting those estimates into allocation or pruning decisions under a global hardware budget. We introduce a curvature-aware, MDL-inspired framework built around the layer gain

ζ_k^2=g_k^\top\widetilde H_{kk}^{-1}g_k

. This quantity equals twice the maximal decrease predicted by the regularized layer-restricted quadratic model and incorporates inverse local curvature; it is therefore a local surrogate for reducible risk, not a universal dominance claim over gradient-norm scores. After normalizing the gains into scores

q_k

, we formulate two convex programs: one allocates expert slots under diminishing returns, and the other assigns layer-wise pruning ratios while protecting high-score layers. Both continuous programs have unique globally optimal solutions characterized by one dual variable and computable in

O(K\log(1/\varepsilon))

time by bisection. We also prove a quadratic transfer-regret bound: when source and target score vectors differ by at most

δ

, the target surrogate cost of the transferred decision is within

O(δ^2)

of the target optimum. Experiments on Mistral-7B and Gemma-7B show clear allocation gains in some settings and competitive, though mixed, pruning performance. The framework therefore replaces an empirical score-to-decision heuristic with a budget-feasible optimization procedure whose guarantees apply to the stated continuous surrogates. Code is available on github repo - TKAI-LAB-Mali/Curvature-Weighted-Capacity-Allocation

Theophilus Amaefuna, Hitesh Vaidya, Anshuman Chhabra +1

Bellini College of Artificial Intelligence, Cybersecurity and Computing, University of South Florida, Tampa, Florida, USA

LLM Pruning LLM Compression Curvature-Aware Optimization

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 Networks Neural Network Optimization Curvature-Aware Optimization Riemannian Optimization

Oct 17, 2025

Second-Order Optimization with Asynchronous Curvature

Second-order methods make fast progress per iteration, but forming and factorizing the Hessian often dominates their wall-clock time. We study Split-Client, a task-parallel design in which a curvature client factorizes Hessians in the background while a gradient client takes cubic-regularized steps with the most recently published matrix, so that curvature cost appears as a delay rather than as blocking time. For delayed exact Hessians, we prove a first-order stationarity bound that depends on the average age of the curvature rather than its maximum, using an overlap count and a weighted Hardy inequality that allows any bounded-error initial matrix; the analysis extends to subsampled and uniformly inexact curvature. Under a fixed-cost model, the resulting wall-clock bound matches the scaling of optimally tuned Lazy Hessian without choosing a reuse period. In experiments on two LIBSVM datasets, an implementation decreases the objective earlier in wall-clock time than Lazy and Vanilla cubic Newton, with the largest gains at intermediate accuracy.

El Mahdi Chayti, Martin Jaggi

Machine Learning and Optimization Laboratory EPFL

Second-Order Optimization Curvature-Aware Optimization Optimization Convergence Analysis