PaperlayerView exact revision on arXiv

Source-linked AI summary

SoftServe: A Scalable Quasi-Newton Method for Deep Learning

Joohwan Ko, Tetiana Parshakova, Diana Cai, Robert M. Gower

arXiv:2610.02182v1cs.LGcs.AI

TL;DR

Deep-learning quasi-Newton methods must accommodate non-convex, noisy objectives while remaining computationally feasible at modern model scales. SOFTSERVE uses structured, positive-definite curvature estimates; across several ill-conditioned tasks, it often outperforms Adam, Muon, and SOAP.

  • Problem

    Deep-learning quasi-Newton methods need curvature estimates that handle non-convex objectives and gradient noise efficiently, while dense curvature storage is prohibitively expensive for large models.

    Source: §1 · pp. 1–2

  • Method

    SOFTSERVE combines a soft secant constraint with structured inverse-Hessian approximations, including diagonal and Kronecker variants that preserve positive definiteness.

    Source: §1 · p. 2

  • Results

    SOFTSERVE often outperforms Adam, Muon, and SOAP on severely ill-conditioned tasks, including recurrent networks, autoencoders, and physics-informed models.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Experiments highlight the value of curvature on ill-conditioned tasks spanning recurrent networks, autoencoders, physics-informed models, and a 136M-parameter physics-informed diffusion model.

    Source: §7 · p. 9

    Convergence guarantees under realistic settings remain an open problem.

    Source: §7 · p. 9

Abstract

from arXiv · show

Quasi-Newton (QN) methods have long been among the most effective methods for large-scale unconstrained convex optimization. Two obstacles have limited their use in deep learning: non-convexity and enormous parameter sizes. We introduce SoftServe, a family of QN methods designed to overcome these obstacles without line searches or ad hoc curvature corrections. SoftServe derives positivedefinite curvature estimates from the variational objective of Berglund et al. (2025), even in the presence of negative curvature. We develop diagonal and Kroneckerfactored variants that preserve positive definiteness by construction and scale to massive neural networks. Finally, SoftServe relies on the stable coupled Newton-Schulz iteration for the required matrix operations, replacing costly matrix decompositions with GPU-friendly matrix multiplications. SoftServe excels on problems that are severely ill-conditioned, including tasks such as recurrent networks, deep autoencoders, physics-informed neural networks, and a 136M-parameter physics-informed diffusion model, often achieving lower losses than established baselines including Adam, Muon, and SOAP.

1 INTRODUCTION

Deep-learning quasi-Newton methods must handle non-convex, noisy objectives while keeping curvature estimates computationally feasible at modern model scales. SOFTSERVE combines structured positive-definite estimates with GPU-friendly updates, and often outperforms modern optimizers on ill-conditioned tasks.

1 INTRODUCTION

  • Non-convex objectives can give curvature estimates with negative eigenvalues, making second-order update directions ascent directions unless the estimate remains positive definite.Standard line searches are costly because they require access to the full-batch loss, which is infeasible in stochastic deep learning.Source: §1 · pp. 1–2
  • Storing ten parameter-and-gradient pairs for L-BFGS would increase model and gradient storage by 20-fold, making that approach infeasible beyond small models.This memory burden illustrates the scale challenge for curvature methods.Source: §1 · p. 2
  • Earlier soft-secant methods ensure positive-definite estimates under negative curvature, but their dense estimates are infeasible for deep learning.SOFTSERVE combines the soft-secant idea with structured approximations to address this limitation.Source: §1 · p. 2
  • SOFTSERVE uses diagonal and Kronecker-structured inverse-Hessian approximations that preserve positive definiteness, without full-batch evaluations or line searches.Its GPU-friendly numerical linear algebra includes variants of Newton-Schulz.Source: §1 · p. 2
  • SOFTSERVE often outperforms Adam, Muon, and SOAP on severely ill-conditioned tasks, including recurrent networks, autoencoders, physics-informed models, and a 136M-parameter diffusion model.The paper describes these as examples of tasks where curvature information is fundamental.Source: §1 · p. 2

2 QUASI-NEWTON PRELIMINARIES

This section reviews quasi-Newton curvature estimates and the BFGS secant condition, then introduces Soft QN as a positive-definite relaxation that can be restricted to structured matrices for deep learning.

2.1 THE BFGS ALGORITHM

  • The averaged Hessian along a step supplies the curvature relation that quasi-Newton methods use to estimate the inverse Hessian.The secant equation is underdetermined, so an additional criterion selects an estimate close to the previous one.Source: §2.1 · p. 3
  • BFGS chooses an inverse-Hessian estimate satisfying the secant equation while remaining close to the previous estimate, characterized using log-determinant divergence.This formulation applies when the curvature condition s_t^T y_t > 0 holds.Source: §2.1 · p. 3
  • When s_t^T y_t < 0, enforcing the secant equation forces s_t^T H y_t < 0, so H cannot be positive definite.In stochastic settings, the gradient difference uses the same random draw at both endpoints to reflect curvature for a single sampled loss.Source: §2.1 · p. 3

2.2 SOFT QUASI-NEWTON

  • Soft QN replaces the exact secant constraint with a least-squares penalty and balances its residual against proximity to the previous estimate using λ.The objective is scale invariant and strictly convex in H^-1.Source: §2.2 · p. 3
  • For every finite pair, Soft QN remains well-defined and positive definite even with negative curvature, while the dense solution costs O(D^2) operations and storage.The paper addresses the dense cost by restricting the estimate to structured matrices.Source: §2.2 · p. 3

3 SOFTSERVE: QUASI-NEWTON UPDATES OVER STRUCTURED FAMILIES

SOFTSERVE restricts the soft quasi-Newton objective to diagonal or Kronecker-structured inverse-Hessian families, producing efficient updates with reduced storage and computation.

3 SOFTSERVE: QUASI-NEWTON UPDATES OVER STRUCTURED FAMILIES

  • SOFTSERVE-Diag costs O(D), while SOFTSERVE-Kron costs O(n^2 + m^2) per n × m block, versus O(D^2) for dense Soft QN.These costs are comparable to Adam for the diagonal variant and to Shampoo and K-FAC for the Kronecker variant.Source: §3 · p. 4

3.1 DIAGONAL APPROXIMATION

  • SOFTSERVE-Diag uses one positive scale per coordinate and a unique closed-form scalar update, but ignores interactions between coordinates.Its inverse-Hessian estimate is diagonal.Source: §3.1 · p. 4

3.2 KRONECKER-STRUCTURED APPROXIMATION

  • The Kronecker estimate stores L and R using m^2 + n^2 entries and applies the inverse-Hessian estimate with two matrix products.Direct quadratic-equation solves can cost O(n^3 + m^3 + n^2m + nm^2), so the paper turns to GPU-friendly Newton-Schulz approximations.Source: §3.2 · pp. 4–5
  • SOFTSERVE-Kron alternates between optimizing L and R while holding the other factor fixed, with each conditional minimization yielding a unique positive-definite solution.The updates follow from reshaping the secant vectors into matrices and minimizing the soft quasi-Newton objective for each factor in turn.Source: §3.2 · pp. 4–5

4 PRACTICAL IMPLEMENTATION OF SOFTSERVE FOR DEEP LEARNING

SOFTSERVE’s Kronecker updates solve quadratic matrix equations with positive-definite solutions, using GPU-friendly approximations; its training procedure combines normalized steps with practical update and curvature-refresh practices.

4.1 EFFICIENT KRONECKER UPDATES

  • Both Kronecker factor updates reduce to quadratic matrix equations with unique positive-definite solutions.Direct evaluation uses matrix square roots and an inverse; the Gram formulation instead evaluates a square root of I + 4F^⊤VF and an inverse, avoiding separate square roots of V.Source: §4.1 · p. 5
  • Coupled Newton-Schulz iterations approximate matrix square roots using matrix multiplications and additions, while Newton-Schulz iterations also approximate inverses.The approach is used instead of GPU-unsuited SVD methods for efficient matrix operations.Source: §4.1 · p. 5

4.2 TRAINING WITH SOFTSERVE

  • The procedure uses momentum and Nesterov extrapolation or bias correction, and holds H_t fixed for K steps to amortize curvature-estimate refreshes.At each window’s start, it saves parameters, a stochastic gradient, and a random draw, then reevaluates the gradient with that draw at the end.Source: §4.2 · pp. 5–6
  • SOFTSERVE-Kron takes a normalized preconditioned step, evaluates gradient and iterate differences, then updates its Kronecker factors by solving quadratic matrix equations.SOFTSERVE-Diag follows the same procedure with the diagonal update.Source: §4.2 · p. 6

5 RELATED WORK

Related work includes stochastic and structured quasi-Newton approaches, matrix-aware preconditioners, and variational objectives connected to SOFTSERVE’s soft quasi-Newton objective.

Stochastic and structured Quasi-Newton

  • Stochastic quasi-Newton methods address unreliable curvature pairs through damping, sampled curvature information, batch strategies, diagonal rescaling, or secant penalties; K-BFGS uses layerwise Kronecker factors and damping.Classical quasi-Newton methods enforce the secant equation while changing the metric as little as possible.Source: §5 · p. 6

Matrix-aware methods

  • K-FAC, Shampoo, and SOAP use Kronecker factorizations for preconditioning, whereas Muon uses steepest descent in the matrix spectral norm.These methods specialize in matrix parameters of linear layers.Source: §5 · p. 6

Related variational objectives

  • With the mean held fixed, the covariance subproblem in batch-and-match variational inference has the same form as SOFTSERVE’s soft quasi-Newton objective.The soft quasi-Newton objective’s least-squares secant penalty has also been used to design preconditioned SGD variants called PSGD.Source: §5 · p. 6

6 EXPERIMENTS

Across recurrent, reconstruction, and physics-informed tasks, SOFTSERVE-Kron often achieves the lowest training loss, but it does not match Muon or SOAP on the language-model task. The experiments compare methods at equal gradient-evaluation budgets and include both deterministic and stochastic settings.

6.1 RNN ADDING

  • 6.6 × 10^-6 mean test MSE for SOFTSERVE-Kron beats Muon (1.5 × 10^-4) and SOFTSERVE-I (1.0 × 10^-2) on the ill-conditioned Adding problem.The full loss Hessian is nearly singular at initialization across four seeds, and SOFTSERVE-Kron outperforms the baselines on all three seeds.Source: §6.1 · p. 7

6.2 MNIST AUTOENCODER

  • SOFTSERVE-Kron reaches the lowest training objective with full-batch MNIST autoencoder training and is close to Muon with minibatches, improving on SOFTSERVE-I in both settings.The two settings use the same 2.8M-parameter network and compare deterministic full-batch with stochastic minibatch training.Source: §6.2 · pp. 7–8

6.3 PHYSICS-INFORMED NEURAL NETWORKS

  • SOFTSERVE-Kron attains the lowest final training objective in five of six small-PINN settings, including all three with resampled points.On deterministic Wave, Adam and Adam+L-BFGS finish lower, with Adam+L-BFGS lowest; Figure 2 distinguishes fixed-point and resampled-point settings.Source: §6.3 · p. 8, Figure 2 · p. 8
  • On PirateNet, SOFTSERVE-Kron is competitive with Muon on Allen–Cahn and reaches the lowest mean loss on KdV.The comparison uses the unweighted PDE and initial-condition MSE on fixed points; Figure 3 shows Allen–Cahn at left and KdV in the middle.Source: §6.3 · p. 8, Figure 3 · p. 9

6.4 PHYSICS-INFORMED DIFFUSION

  • At 100,000 gradient evaluations, SOFTSERVE-Kron reaches a lower mean training loss than Adam, Muon, and SOAP on the 136M-parameter physics-informed diffusion model.The model trains on 30,000 optimized structures; Figure 3 shows its training objective in the right panel.Source: §6.4 · p. 9, Figure 3 · p. 9

6.5 LANGUAGE MODEL

  • SOFTSERVE-Kron records mean validation NLL 4.661 on the language model, below AdamW (4.786) but above Muon (4.440) and SOAP (4.416).Thus, the gains on preceding tasks do not extend to matching Muon and SOAP on this language model.Source: §6.5 · p. 9

7 DISCUSSION AND LIMITATIONS

SOFTSERVE structures its quasi-Newton updates to retain positive definiteness, including under negative curvature, and derives diagonal and Kronecker variants from a soft quasi-Newton objective. The discussion identifies fixed λ and missing realistic-setting convergence guarantees as open limitations.

7 DISCUSSION AND LIMITATIONS

  • SOFTSERVE uses structured Hessian approximations that are always positive definite and GPU-friendly Newton–Schulz approximations, with experiments on ill-conditioned recurrent, autoencoder, and physics-informed tasks.The evaluated physics-informed models include a 136M-parameter PIDM.Source: §7 · p. 9
  • The method currently keeps λ fixed within each training run; adaptive λ and convergence guarantees under realistic settings remain open directions.Source: §7 · p. 9

A.1 BFGS

  • BFGS can be characterized by a log-determinant-divergence formulation, which provides the starting point for replacing its hard secant equation with a soft penalty.The formulation recovers the same BFGS update, according to the passage.Source: §A.1 · p. 14

A.2 SOFT QUASI-NEWTON

  • The soft quasi-Newton objective has a unique positive-definite solution, and BFGS is recovered as a limiting case under positive curvature.The positive-curvature condition is stated as s⊤y > 0.Source: §A.2 · p. 14

A.3 CONNECTIONS WITH VARIATIONAL INFERENCE

  • The objective is the one-score covariance objective of Cai et al. (2024) with the mean fixed, linking the covariance subproblems without equating variational inference and quasi-Newton optimization.Source: §A.3 · p. 14

B SOFTSERVE: STRUCTURED UPDATES

B.1 DIAGONAL UPDATE

  • The diagonal restriction has a unique minimizer regardless of the signs of the coordinatewise secant products.The update is separable across coordinates.Source: §B.1 · p. 15

B.2 KRONECKER UPDATE

The Kronecker update alternates between its two factors, solving positive-definite quadratic matrix equations. The implementation computes the resulting Gram-form updates with Newton–Schulz routines and matrix multiplications.

B.2 KRONECKER UPDATE

  • The Kronecker update alternates between L and R, holding the other factor fixed; each update has a unique closed-form solution, monotonically decreases the soft QN objective, and costs O(n^3 + m^3 + n^2m + nm^2).Each factor update cannot increase the objective because the current factor is feasible.Source: §B.2 · p. 16

B.3 QUADRATIC MATRIX EQUATIONS

  • Each Kronecker factor update reduces to a quadratic matrix equation with a unique positive-definite solution.The solution can be expressed using a Gram factorization U = F F⊤.Source: §B.3 · p. 17

C.1 EVALUATING THE UPDATES IN GRAM FORM

  • For the Gram solution, the implementation forms equation coefficients using matrix operations and reuses coupled Newton–Schulz outputs for square-root and inverse-square-root terms.Reusing both outputs avoids a separate Newton–Schulz call to approximate the inverse factor.Source: §C.1 · pp. 17–18

C.2 NEWTON-SCHULZ FOR THE SQUARE ROOT AND INVERSE

  • The Gram-form solution requires square roots and inverses, which the implementation approximates with Newton–Schulz iterations built from matrix multiplications and additions.The coupled iteration simultaneously approximates a matrix square root and its inverse; a separate routine approximates matrix inverses.Source: §C.2 · pp. 18–19

C.3 COMPUTING THE KRONECKER UPDATES VIA NEWTON-SCHULZ

  • The Newton–Schulz Gram update combines the square-root and inverse routines with matrix multiplications to approximate the quadratic-equation solution.Source: §C.3 · p. 19
  • The implementation uses 18 iterations for square-root routines and 10 for the inverse routine, symmetrizing outputs and recording normalized solve residuals.The residuals are recorded to check closeness to the exact quadratic matrix equation solution.Source: §C.3 · p. 20

D.1 TRAINING WITH SOFTSERVE

SOFTSERVE combines momentum-based parameter updates with periodically refreshed diagonal or Kronecker Hessian estimates. Its implementation uses constrained descent, routes parameter blocks by shape, and replays sampled-loss inputs when forming secant pairs.

D.1 TRAINING WITH SOFTSERVE

  • Constrained descent was used throughout the experiments after initial experiments found it produced better methods than regularized descent.The method normalizes the update and uses momentum; the inverse-Hessian estimate is refreshed periodically to amortize its cost.Source: §D.1 · pp. 20–21

D.2 ROUTING AND STATE

  • Matrix parameters use Kron blocks, other parameter shapes use Diag blocks or a first-order fallback, and the metric remains fixed between refreshes.Each refresh uses one raw interval pair; equal-shaped Kron factors can be stacked for batched matrix multiplication.Source: §D.1 · p. 20, §D.2 · p. 21

D.3 REPLAY AND GRADIENT ACCOUNTING

  • Replay keys reproduce sampled loss inputs at both secant endpoints, removing sampling mismatch from the gradient difference used for metric refreshes.Replayed gradients count toward the budget but do not update momentum, the learning-rate schedule, or parameters; checkpoints within an interval retain the replay key and metric state.Source: §D.3 · p. 22

E.1 EXPERIMENTAL SETTING

The experimental setup tunes learning rates and evaluates selected configurations on fresh seeds. The supplied RNN and MNIST passages describe their tasks, models, training budgets, and evaluation measures.

E.1 EXPERIMENTAL SETTING

  • Learning rates are tuned on seed 0 and selected configurations are evaluated on seeds 1–3; methods share initialization seeds and data within each task.Test errors are not used for selection, and failed settings appear as individual traces.Source: §E.1 · p. 22

E.2 RNN ADDING

  • The RNN Adding task uses 100-number sequences with two marked positions and trains a 16-unit tanh RNN to predict their sum from batches of 128 newly sampled sequences.The model has 321 parameters; validation and test sets each contain 8,192 fixed sequences per seed.Source: §E.2 · p. 23

E.3 MNIST AUTOENCODER

  • The MNIST task trains a 2.8M-parameter autoencoder with a 30-dimensional bottleneck; evaluation reports test reconstruction BCE and regularized training objectives.Full-batch runs use 4,000 gradient evaluations and minibatch runs use 6,000; the latter use batches of 1,000 images.Source: §E.3 · p. 25

E.4 SMALL PHYSICS-INFORMED NEURAL NETWORKS

The small PINN experiments solve three one-dimensional PDE problems with a fully connected network and compare deterministic fixed-point training with resampled-point training. They also repeat FP32-selected configurations in FP64 without retuning.

Model and objective

  • The PINN uses three width-200 tanh hidden layers and 81,201 parameters, with an objective summing PDE-residual, initial-condition, and boundary-condition errors.For Wave, the conditions include initial displacement and velocity and values at both spatial boundaries.Source: §E.4 · p. 27

Training and tuning

  • Training uses 200,000 gradient evaluations; deterministic runs keep 10,000 interior collocation points fixed, while stochastic runs resample them at every update.Both gradients in each secant pair use the same points.Source: §E.4 · p. 27

Evaluation

  • FP64 runs reuse FP32-selected configurations without retuning, and failed runs are shown individually; an L-BFGS zero-displacement line search is plotted as an unchanged solution.The supplied figure descriptions cover solution errors, training objectives, parameter updates, and wall-clock time.Source: §E.4 · p. 28, Figure E.8 · p. 29, Figure E.9 · p. 29

E.5 PIRATENET

The supplied passages describe evaluation setups for PirateNet, a physics-informed diffusion model, and a language model, alongside ablations of the secant penalty and QME solver. In the reported ablations, penalty choice affects both RNN accuracy and stability, while solver behavior depends on iteration count and numerical precision.

E.5 PIRATENET

  • PirateNet evaluates Allen–Cahn and KdV with a 0.73M-parameter residual PINN, comparing the unweighted PDE and initial-condition MSE on common fixed points.Training uses 32,768 collocation points per update; the supplied passages describe fixed gradient budgets and relative L2 error evaluation, but do not report comparative outcomes.Source: §E.5 · p. 32

E.6 PHYSICS-INFORMED DIFFUSION

  • The diffusion experiment uses a released 136M-parameter U-Net and combines reconstruction and physics-based terms for structural-design outputs.Validation uses EMA parameters on 128 fixed examples; runs end at 100,000 gradient evaluations, yielding about 90,900 updates for Kron versus 100,000 for baselines.Source: §E.6 · p. 34

E.7 LANGUAGE MODEL

  • The language-model experiment trains a 16.17M-parameter GPT-style model on 200M FineWeb tokens and reports next-token cross-entropy as NLL per token.Selected configurations are evaluated on fresh seeds 1–3, with mean NLL reported.Source: §E.7 · p. 36

F ABLATION STUDIES

  • The ablation studies vary the secant penalty, Gram QME solver accuracy, and refresh interval, selecting learning rates on seed 0 and evaluating on fresh seeds 1–3.Unless varied, λ = 999, K = 10, and solver counts are 18 root and 10 inverse iterations.Source: F ABLATION STUDIES · p. 36

F.1 SENSITIVITY TO THE SECANT PENALTY

  • At λ = 999, RNN validation error is lower than at λ = 0.01, and λ values from 0.1 to 999 outperform tuned Muon; λ = 9999 fails on all three evaluation seeds.The learning rate is retuned for each penalty value, so the best penalty in this experiment is not a task-independent default.Source: §F.1 · p. 36

F.2 ACCURACY OF THE QME SOLVER

  • At the default 18/10 solver iterations, median relative errors are approximately 2.3 × 10−6, 7.3 × 10−5, and 7.2 × 10−3 across the three tested κ values.Too few iterations can leave substantial error, while additional iterations do not uniformly reduce FP32 error once rounding limits accuracy.Source: §F.2 · p. 37
  • In RNN training, the four stable solver settings have similar learning curves, and counts above 18/10 do not consistently improve validation error.The 12/10 setting failed at all three tested seed-0 learning rates before the full budget.Source: §F.2 · p. 37

F.3 REFRESH FREQUENCY AND TRAINING COST

Refresh frequency trades off convergence against training time, with the best interval depending on the task. Refreshing every update slows convergence on both tasks.

F.3 REFRESH FREQUENCY AND TRAINING COST

  • More frequent refreshes require more secant replays, leaving fewer parameter updates within the same gradient budget.Refreshes at every update slow convergence on both tasks.Source: §F.3 · p. 38
  • K = 5 reaches the lowest mean RNN validation MSE, while larger intervals shorten MNIST training with similar final objectives.The comparison uses task-specific gradient budgets and plots both gradient evaluations and wall-clock time.Source: §F.3 · p. 38, Figure F.4 · p. 39
  • Figure F.4 shows three-seed mean learning curves with seed-range shading; wall-clock time includes evaluation and diagnostic overhead.All runs stop at the same task-specific gradient budget, including replay: 40,000 for RNN and 6,000 for MNIST.Source: Figure F.4 · p. 39

Related papers

Loading 2610.02182v1…