evie-amplify-optim

Evie-Amplify: a risk-sensitive modulation of Adam-style optimizers that allows bounded amplification on low-variance gradient coordinates, on top of the noise-dependent suppression inherited from evie-optim.

The base Evie update is structurally a shrink-only reweighting of AdamW — it can only ever damp the step. Evie-Amplify re-centers the same noise-dependent gate around a target well-posedness level so that sufficiently low-noise coordinates can receive an update larger than the corresponding AdamW step, while noisy coordinates are still strongly damped.

Install

pip install evie-amplify-optim

Usage

import torch
from evie_amplify_optim import EvieAmplifyOptimizer

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

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).
    if step == 125:
        optimizer.calibrate_gamma0()

How it works

sigma_t^2 = max(m2_hat - m1_hat^2, 0)                  # gradient-variance estimate
gamma_t   = gamma0 / (1 + gamma0 * sigma_t^2)           # adaptive risk-aversion
raw_wp    = 1 / (1 + gamma_t * sigma_t^2)               # in (0, 1], Evie's shrink-only gate
well_pos  = clamp(raw_wp / target_wp, eps, well_pos_max)
update    = -lr * m1_hat * sqrt(well_pos) / sqrt(m2_hat)

Dividing by target_wp re-centers the gate: when the raw gate exceeds target_wp (i.e. the coordinate is low-noise), well_pos > 1 and the update is amplified; otherwise it is damped, same as in base Evie. Amplification is bounded above by well_pos_max / target_wp; suppression is unbounded below as gradient variance grows.

gamma0 is calibrated once, from a separate fast-EMA variance estimate, the same way as in evie-optim.

API

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

• lr: learning rate.
• gamma0: initial risk-aversion ceiling; overwritten by calibrate_gamma0().
• target_wp: re-centering target for the well-posedness gate.
• beta1, beta2: Adam moment decay rates.
• weight_decay: decoupled weight decay (AdamW-style).
• well_pos_max: upper clamp on the well-posedness factor, bounding maximum amplification.
• 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