Metadata-Version: 2.4
Name: evie-kf
Version: 0.1.0
Summary: EvieKF: AdamW with a Kronecker-factored gradient-noise preconditioner (PyTorch)
Author: Anonymous
License: MIT
Project-URL: Repository, https://github.com/rpextra2026-afk/evie-kf-optim
Keywords: optimizer,pytorch,deep-learning,adamw,kronecker,preconditioning
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=1.13
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Dynamic: license-file

# 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

```bash
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.

```python
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

| argument | default | notes |
|---|---|---|
| `lr` | `1e-3` | AdamW scale. Start from your tuned AdamW lr. |
| `gamma` | `100.0` | Noise-aversion strength, the one knob to tune. Search it on a log grid alongside lr, e.g. 3 … 3000, and extend the grid if the best value lands on an edge. |
| `betas`, `eps` | `(0.9, 0.999)`, `1e-8` | Adam moments |
| `weight_decay` | `0.0` | decoupled, `p *= 1 − lr·wd` |
| `beta_sig` | `0.95` | EMA of the noise factors |
| `warmup` | `20` | steps before the preconditioner engages |
| `refresh` | `20` | steps between eigendecompositions |
| `maxf` | `1024` | a Kronecker factor larger than this is kept diagonal |

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
