SoftServe

A scalable quasi-Newton method
for deep learning.

Gradients describe the local slope of the loss. Changes in gradients also provide information about its curvature. SoftServe uses these changes to learn a structured inverse-Hessian approximation and precondition neural-network updates.

Joohwan KoTetiana ParshakovaDiana CaiRobert M. Gower

01Learn from
gradient changes
02Maintain a
positive-definite estimate
03Scale with
structured matrices

The idea in brief

Quasi-Newton methods estimate curvature from parameter and gradient changes, without forming a Hessian. For deep learning, we want that estimate to stay positive definite even when the loss has negative curvature—and small enough to store for a large network.

Soft quasi-Newton (Soft-QN) handles negative curvature by fitting the observed changes with a penalty rather than an exact constraint [1]. SoftServe makes this formulation scalable: we derive diagonal and Kronecker-factored updates from the same objective, evaluated with GPU-friendly matrix multiplications.

Our experiments span a recurrent network, an autoencoder, physics-informed neural networks (PINNs), and a 136M-parameter physics-informed diffusion model.

01 / Learning curvature

From gradient changes
to curvature.

When curvature varies across directions, a step size small enough for one direction can be unnecessarily small for another. A quasi-Newton method learns an inverse-Hessian approximation H. Multiplying by this matrix rescales and mixes the gradient components before the update:

x⁺ = x − ηHg

Here x is the parameter vector, g = ∇f(x) the gradient of the loss f, and η > 0 the learning rate. Applying H is called preconditioning.

We require H to be positive definite, so gᵀHg > 0 for every nonzero gradient. This makes −Hg a local descent direction. How can we learn such a matrix without computing the Hessian?

What a secant pair tells us

Evaluate the same loss at two parameter points. Their displacement s = x_new − x_old and gradient change y = g_new − g_old form a secant pair.

For a smooth loss and nearby points, y ≈ ∇²f(x_old)s, where ∇²f is the Hessian, the matrix of second derivatives. An inverse estimate should therefore map y back to s. Classical methods such as BFGS impose this as the secant equation:

H_new y = s.

But these two requirements can conflict. A positive-definite H_new satisfying the equation requires sᵀy > 0 when y ≠ 0. Nonconvex losses can violate this condition. In stochastic training, an exact fit also reflects only one sampled loss.

Replace the constraint with a penalty

Soft-QN [1] allows an imperfect fit. Starting from a positive-definite estimate H_old, it balances fitting the pair against changing that estimate:

The Soft-QN variational objective
H_new = argmin over H ≻ 0: ‖s − Hy‖²_H⁻¹ + D_ld(H_old, H) / λ

The first term penalizes the secant mismatch. The log-determinant divergence D_ld measures the change from the previous estimate. The parameter λ > 0 sets the tradeoff: small values favor the old estimate; large values emphasize the new pair.

Unlike the exact constraint, this objective has a unique positive-definite minimizer for every finite pair, including negative-curvature pairs.

What do the norm and divergence mean?

For a vector v, the weighted norm is ‖v‖²_H⁻¹ = vᵀH⁻¹v. For positive-definite d × d matrices, the log-determinant divergence is

D_ld = tr(H⁻¹H_old) − log det(H⁻¹H_old) − d.

It is nonnegative and vanishes only when the matrices agree. It equals twice the Kullback–Leibler divergence from the Gaussian N(0, H_old) to N(0, H). Related matrix equations arise in the variational-inference method Batch and Match [2].

02 / Explore the update

A two-dimensional
Soft-QN update.

Start from the identity estimate H_old = I. Rotate y relative to s, then vary λ. The orange arrow is H_new y; compare it with the green target s to see how closely the secant equation is fitted.

Dense Soft-QNIllustration · exact 2D update
zᵀH⁻¹z = 1H remains positive definite
How a soft secant pair changes a positive-definite metric An ellipse shows the new inverse-curvature estimate. Sliders control lambda and the angle between the secant vectors. Numerical readouts below summarize the result.
Previous metricUpdated metricTarget s
0.01 · stronger regularization100 · stronger secant penalty
AlignedOpposite
This pair admits a positive-definite exact secant fit.
Observed sᵀy0.819
Smallest eigenvalue of H
Relative secant residual

For a 2D coordinate z, each ellipse satisfies zᵀH⁻¹z = 1; it is not a loss contour. Both axes rescale together. This illustrates one dense Soft-QN update, not a SoftServe-Kron training run.

Try the negative pair, then increase λ. When sᵀy < 0, no positive-definite matrix can map y exactly to s. The residual stays nonzero, but the smallest eigenvalue remains positive: the fit is relaxed, not the definiteness requirement.

The calculation behind the animation

With slider angle θ, we use s = (1, 0)ᵀ, y = (cos θ, sin θ)ᵀ, and H_old = I. Let V = I + λssᵀ and b = yᵀVy. The exact dense solution is

H_new = V − 4λ (Vy)(Vy)ᵀ / (1 + √(1 + 4λb))².

It satisfies λHyyᵀH + H = V. The displayed residual is ‖s − Hy‖₂ / ‖s‖₂. These browser calculations are independent of the paper’s experimental measurements.

03 / Diag and Kron

Keep the objective.
Structure the matrix.

Soft-QN resolves the curvature-sign problem, but a dense estimate for d parameters still stores d² entries. SoftServe restricts the same objective to smaller matrix families and derives their updates directly.

SoftServe-Diag

The diagonal stores one positive scale h_i per parameter. Each scale has a closed-form update, since the objective separates across entries.

H = diag(h), hᵢ > 0

SoftServe-Kron

Two positive-definite factors act on the rows and columns of each weight-matrix block. This retains interactions between parameters without storing the full matrix:

H = L ⊗ R, Hg = vec(RGL)

Here ⊗ denotes the Kronecker product. For a block X ∈ ℝⁿˣᵐ, G is its gradient and g = vec(G) stacks its columns into a vector. The factors R and L are n × n and m × m. We compute RGL directly, storing n² + m² entries instead of (nm)².

What does structure save?

A 1,024 × 1,024 weight matrix

644,096
Dense curvaturen⁴ entries4 TiB
SoftServe-Diagn² entries4 MiB
SoftServe-Kron2n² entries8 MiB

FP32 curvature storage for one unblocked, square layer. This excludes parameters, momentum, secant snapshots, and temporary solver memory; it is not a total optimizer-memory measurement.

Update one factor at a time

Each factor depends on the other, so we update one while holding the other fixed, then switch. Each exact solve has a unique positive-definite solution and cannot increase the curvature-fitting objective. One alternating sweep does not jointly minimize both factors or guarantee a decrease in training loss.

04 / The training loop

How it fits into
a training loop.

1. Apply the current metric

At each step, form a momentum-filtered gradient m for each parameter block. Apply the current estimate H, then normalize the resulting direction:

d = Hm / √(mᵀHm), x⁺ = x − ηd.

The denominator gives d unit length in the H⁻¹ norm: dᵀH⁻¹d = 1. A zero momentum gives no update. Setting H = I gives our SoftServe-I control, which keeps normalized momentum but removes curvature learning.

2. Form a secant pair every K updates

We keep the estimate fixed for K updates to reduce refresh cost. To measure curvature rather than a change of minibatch, both secant gradients use the same batch and randomness. Intermediate training steps can still use fresh batches.

Default K = 1010 training gradients + 1 replay
12345678910↶

Save the starting parameters, gradient, batch, and random state. After ten updates, replay the saved loss at the new parameters to form the secant pair.

At K = 10, replay adds 10% more gradient evaluations; factor computations add separate runtime. Deterministic training can reuse its endpoint gradients without an extra evaluation.

3. Refresh the factors with matrix multiplications

Once we have a secant pair, refreshing a Kronecker factor means solving for a positive-definite matrix Q:

QUQ + Q = V, U ⪰ 0, V ≻ 0.

This is a quadratic matrix equation. The secant pair and current factors determine its coefficients: U is positive semidefinite and V is positive definite. These conditions give a unique positive-definite solution.

The solution involves matrix square roots and inverses. We approximate them with Newton–Schulz iterations, which use repeated matrix multiplications and additions. This makes the factor solves suitable for GPUs without requiring matrix decompositions.

Numerical solver settings

Our Gram formulation writes U = FFᵀ, with Gram factor F, and avoids a separate square root of V. Coupled Newton–Schulz jointly approximates square and inverse square roots; another Newton–Schulz iteration approximates inverses.

The default uses 18 coupled root iterations and 10 inverse iterations. These approximate the exact factor solves; the paper reports their numerical residuals and the effect of iteration count on training.

05 / Experiments

From recurrent networks
to physics-informed diffusion.

Does the learned curvature improve training? We tune each method’s learning rate separately, then compare the selected configurations on three fresh seeds. The plots match gradient evaluations, counting secant replays, line searches, and loss-balancing gradients as well as ordinary training gradients.

RNN Adding and MNIST autoencoding

321 → 2.8M parameters
RNN Adding and MNIST reconstruction. Left: a recurrent neural network with tanh activations retains two marked values in a length-100 sequence and predicts their sum, minimizing mean squared error (MSE). Middle and right: a deep autoencoder compresses MNIST images into a 30-dimensional code and reconstructs them, minimizing regularized binary cross-entropy.

On Adding, Kron’s mean test MSE is 6.6 × 10⁻⁶, compared with 1.5 × 10⁻⁴ for Muon and 1.0 × 10⁻² for the identity-metric control. On MNIST, Kron reaches the lowest full-batch training objective and is close to Muon with minibatches.

Lines are three-seed arithmetic means; bands span the seed minimum–maximum. Settings with failed seeds retain individual traces. The paper’s appendix reports solution and test errors, along with wall-clock comparisons.

06 / Use SoftServe

Use SoftServe
in PyTorch.

Pass your model to SoftServeKron. The optimizer routes matrix weights and convolution kernels to Kron, and uses your chosen fallback—Adam, AdamW, or SoftServe-Diag—for the remaining parameters.

PyTorch · model-based API
from softserve import SoftServeKron

optimizer = SoftServeKron(
    model,
    lr=1e-3,
    lam=99,
    K=10,
    fallback="adamw",
    fallback_lr=3e-4,
)

These are example hyperparameters, not universal defaults. For stochastic training, pass optimizer.step(closure) a function that recomputes the loss and gradients on its captured minibatch. The optimizer handles snapshots and replay.

See the complete training examples ↗

References and citation

  1. Berglund, Zhang, and Johansson. Soft quasi-Newton: guaranteed positive definiteness by relaxing the secant constraint. Optimization Methods and Software, 2025.
  2. Cai et al. Batch and Match: Black-Box Variational Inference with a Score-Based Divergence. ICML, 2024.
  3. Ko, Parshakova, Cai, and Gower. SoftServe: A Scalable Quasi-Newton Method for Deep Learning. 2026. See the paper for proofs, benchmark references, and complete experimental settings.
Cite this work BibTeX
arXiv:2610.02182
@misc{ko2026softserve,
  title = {{SoftServe}: A Scalable Quasi-Newton Method
           for Deep Learning},
  author = {Joohwan Ko and Tetiana Parshakova and
            Diana Cai and Robert M. Gower},
  year = {2026},
  eprint = {2610.02182},
  archivePrefix = {arXiv},
  primaryClass = {cs.LG},
  url = {https://arxiv.org/abs/2610.02182}
}

Scroll horizontally on a small screen to inspect the axes and legends.