evie-kf

EvieKF is AdamW with a Kronecker-factored preconditioner built from gradient noise. It estimates the covariance of mini-batch gradient noise from a split-batch difference, factors it per layer as B ⊗ A, and shrinks the AdamW direction along the high-noise eigendirections:

y  = P^-1/2 m̂                              # the AdamW direction (whitened)
δ  = (g_A − g_B) / 2                        # noise sample from two half-batches
A, B ← EMA of (P^-1/2 δ)ᵀ(P^-1/2 δ), (P^-1/2 δ)(P^-1/2 δ)ᵀ
z  = (I + γ (B ⊗ A) / tr A)^-1/2 y          # via eigh of A and B, never forming B ⊗ A
z ← z · ‖y‖ / ‖z‖                           # global rescale: same step size as AdamW
W ← W (1 − lr·wd) − lr · P^-1/2 z

With γ = 0 it is exactly AdamW. Biases and norm parameters (0-/1-D) use a diagonal noise estimate.

Install

pip install evie-kf

Requires PyTorch ≥ 1.13.

Usage

EvieKF needs two gradients from disjoint halves of each batch. Use split_backward in place of loss.backward(). It uses the same examples and roughly the same FLOPs as one full-batch backward.

import torch
import torch.nn.functional as F
from evie_kf import EvieKF, split_backward

model = ...
optimizer = EvieKF(model.parameters(), lr=1e-3, weight_decay=1e-4, gamma=100.0)

def loss_fn(x, y):
    return F.cross_entropy(model(x), y)

for x, y in loader:
    optimizer.zero_grad()
    loss = split_backward(optimizer, loss_fn, x, y, max_grad_norm=5.0)
    optimizer.step()

max_grad_norm clips the gradient and scales the noise sample by the same coefficient. If you clip, do it through this argument, not separately.

If you already compute two half-batch gradients yourself, set p.grad = (g_A + g_B) / 2 and call optimizer.set_noise([(p, (g_A - g_B) / 2), ...]) before each step().

If step() runs past warmup without a noise sample, the step is plain AdamW and a RuntimeWarning is raised once.

Models with BatchNorm

Each half-batch goes through its own forward pass, so BatchNorm statistics are computed over half the batch. For a fair comparison with a baseline optimizer, run the baseline with the same split forward passes as well.

Hyperparameters

Ablation switches: diag_only=True uses a diagonal Σ throughout (EvieDiag), and centered=False uses uncentred second moments of the raw gradient (EvieKFu). diag_relative, kron_mean_norm and norm_preserve are diagnostic switches, and their defaults are the main method.

Weight matrices are treated as (out, in). Conv kernels are flattened to (out, in·kh·kw).

Cost

Per step, EvieKF runs two half-batch backward passes instead of one full-batch pass. Every refresh steps it also computes one eigendecomposition per Kronecker factor, in float64, for each factor of size ≤ maxf. The memory cost is the two factors per weight matrix plus one copy of the half-batch gradient.

Diagnostics

After each step(), these values are available for monitoring: optimizer.last_cos (cosine between the EvieKF and AdamW directions), tr_sigma, eff_rank, gain_lo, gain_hi and n_active_steps. evie_kf.self_check() compares the factored operator against a dense construction.

License

MIT