evie-optim

Evie: a risk-sensitive modulation of Adam-style optimizers for stochastic optimization, based on a risk-sensitive (Jacobson–Whittle) control formulation of neural network training. Evie multiplies the usual adaptive step by a noise-dependent, shrink-only gate: coordinates with higher gradient-noise variance get shrunk more.

For a variant that also allows bounded amplification on low-variance coordinates, see the companion package evie-amplify-optim.

Install

pip install evie-optim

Usage

import torch
from evie_optim import EvieOptimizer

model = torch.nn.Linear(10, 1)
optimizer = EvieOptimizer(model.parameters(), lr=1e-3, target_wp=0.85)

for step, (x, y) in enumerate(dataloader):
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(x), y)
    loss.backward()
    optimizer.step()

    # Calibrate gamma0 once, early in training (paper uses step 100-150).
    # Before this call, Evie behaves like plain AdamW.
    if step == 125:
        optimizer.calibrate_gamma0()

How it works

Evie treats mini-batch gradient noise as a stochastic disturbance in a control formulation of training, and derives a closed-form gain from the corresponding risk-sensitive Riccati equation. The practical update is:

sigma_t^2 = max(m2_hat - m1_hat^2, 0)              # gradient-variance estimate
gamma_t   = gamma0 / (1 + gamma0 * sigma_t^2)       # adaptive risk-aversion
well_pos  = 1 - gamma_t * sigma_t^2                 # in (0, 1]
update    = -lr * m1_hat * sqrt(well_pos) / sqrt(m2_hat)

gamma0 is calibrated once, from a separate fast-EMA variance estimate, so that well_pos starts near target_wp at the calibration step.

Because well_pos <= 1 always, Evie can only shrink the corresponding AdamW step — it never amplifies. See evie-amplify-optim for a variant that removes this restriction.

API

EvieOptimizer(params, lr=1e-3, gamma0=1.0, target_wp=0.85, beta1=0.9, beta2=0.999, eps=1e-8, weight_decay=0.0, fast_beta=0.9)

• lr: learning rate.
• gamma0: initial risk-aversion ceiling; overwritten by calibrate_gamma0().
• target_wp: target well-posedness factor used during calibration.
• beta1, beta2: Adam moment decay rates.
• weight_decay: decoupled weight decay (AdamW-style).
• fast_beta: decay rate for the separate fast-EMA variance estimate used only for calibration.

optimizer.calibrate_gamma0()

Call once, partway through training, to set gamma0 from the observed gradient variance. Safe to call only after at least one step().

License

MIT