EmaTrainLossMonitor

class liesel.optim.EmaTrainLossMonitor(effective_window)

Bases: object

Configures EMA monitoring of the pre-update mini-batch training loss.

The pre-update mini-batch training loss is the scalar objective evaluated for an optimizer’s gradient at its supplied position, before that optimizer applies its update. The value and gradient share the same mini-batch, parameter position, PRNG key, and stochastic objective draw; monitoring is not an independent loss evaluation. The first active optimizer supplies this observation. If none is active, the engine evaluates the loss once at the unchanged position.

Let \(e\) be effective_window and \(N\) the configured number of full batches per epoch. The span \(W\) and smoothing coefficients are

\[W = \max(1, eN), \qquad \alpha = \frac{2}{W + 1}, \qquad \beta = 1 - \alpha.\]

Starting from \(m_0 = w_0 = 0\), mini-batch observation \(t\) updates

\[m_t = \beta m_{t-1} + \alpha L_t, \qquad w_t = \beta w_{t-1} + \alpha, \qquad \operatorname{EMA}_t = \frac{m_t}{w_t}.\]

Here \(L_t\) is the pre-update loss, \(m_t\) its unnormalized exponentially weighted sum, \(w_t\) the accumulated weight, \(\alpha\) the newest-loss weight, and \(\beta\) the decay applied to earlier losses. The index \(t\) continues across epoch boundaries. Since \(w_t = 1 - \beta^t\), the bias-correction factor \(1 / w_t\) automatically approaches one without changing the update rule.

Numerically, the engine stores the normalized EMA directly and uses the equivalent update

\[w_t = -\operatorname{expm1}(t\operatorname{log1p}(-\alpha)), \qquad \operatorname{EMA}_t = \operatorname{EMA}_{t-1} + \frac{\alpha}{w_t}(L_t - \operatorname{EMA}_{t-1}).\]

The update uses compensated summation to retain small increments. This avoids drift from repeatedly accumulating the normalization weight and reduces rounding error for long windows, including with float32 and JAX 64-bit mode disabled. Ordinary floating-point rounding still applies.

effective_window is an EMA span measured in epoch equivalents, not a hard inclusion window or half-life. With a typical multi-batch epoch, a span of one epoch equivalent has a half-life of roughly 0.35 epoch equivalents, so recent batches receive substantially more weight than early batches from the same epoch. The from_half_life() constructor uses the large-span approximation \(e = 2h / \log(2)\) for a requested half-life \(h\). The exact finite-step relationship also depends on the number of batches per epoch.

The EMA trails the current loss. Its weights have a mean age of \((W - 1) / 2\) observations, so once the bias correction has faded, a loss that changes at a steady rate is monitored with this delay. Each observation is taken before the updates of its batch, so relative to the position at the end of an epoch the delay is \((W + 1) / 2\) batches. For \(eN \geq 1\), this is \(e / 2 + 1 / (2N)\) epoch equivalents, roughly half of effective_window when epochs have many batches. The half-life is shorter than this delay because the weights decay exponentially.

Methods

from_half_life(half_life)

Constructs a monitor from an approximate half-life.

Attributes

effective_window

Effective averaging window of the exponential moving average.