sage.architecture.custom_losses.gradient_balancer

Generic gradient-norm budget balancer for multi-loss training.

This is the mechanism for combining a single primary loss (typically BCE) with n auxiliary losses by their gradient norms — it has nothing to do with any particular model and is reusable by any job that optimises more than one loss. It was extracted verbatim from the consistency training loop.

Every balance_every steps (after a balance_settle warm-up that skips the huge step-0 gradient) we measure ‖g_bce‖ and each auxiliary loss’s ‖g_i‖, then set per-aux weights so that (a) the aux gradient norms are EQUALISED among themselves and (b) their combined norm is <= balance_target * BCE's gradient scale. The budget mostly tracks BCE (aux saturates with it), with two small guards:

• budget floor B_ref = max(EMA(‖g_bce‖), floor_frac * g_bce_char) where g_bce_char is the settled scale captured at the first calibration — keeps aux from fully vanishing if BCE flattens, without the brittle peak-tracking;

• per-aux denom floor g_i_eff = max(‖g_i‖, denom_floor * EMA(‖g_i‖)) — a converged aux’s weight saturates at a finite ceiling instead of 1/‖g_i‖ -> infinity.

w_i = (balance_target / n) * B_ref / g_i_eff.

The caller supplies a recompute closure — recompute(model_output) -> (bce_term, [aux_term, ...]) — so the balancer stays agnostic about which losses it is balancing.

Classes

GradientNormBalancer

Combine bce + sum(w_i * aux_i) with gradient-norm-balanced weights.

Module Contents

class GradientNormBalancer(n_aux, balance_target=0.33, balance_every=250, balance_decay=0.7, balance_floor_frac=0.1, balance_denom_floor=0.1, balance_settle=500, aux_weights=None, autocast=False, aux_names=None)

Combine bce + sum(w_i * aux_i) with gradient-norm-balanced weights.

Parameters:

• n_aux (int) – Number of auxiliary losses being balanced against the primary (BCE).

• balance_target (float) – Budget cap: combined aux gradient norm <= balance_target * BCE scale. <= 0 is not used here (a job with no balancing simply passes no balancer); see the trainer’s no-balancer path.

• balance_every (int) – Re-calibrate the weights every this many steps.

• balance_decay (float) – EMA decay applied to the measured gradient norms per calibration.

• balance_floor_frac (float) – Budget floor as a fraction of the settled BCE-gradient scale.

• balance_denom_floor (float) – Per-aux denominator floor (mu): caps a converged aux’s weight.

• balance_settle (int) – Skip calibration for this many initial steps (the step-0 transient).

• aux_weights (sequence of float or None) – If given, FIXED-weight mode: no live gradient calibration (compile-safe production path). None -> live calibration (used to derive these weights offline).

• autocast (bool) – Whether the calibration forward runs under fp16 autocast (match training).

• aux_names (sequence of str or None) – Optional labels for the aux losses (inspection/logging only).

balance_target

balance_every = 250

balance_decay

balance_floor_frac

balance_denom_floor

balance_settle = 500

autocast = False

n_aux

aux_names

combine(bce, aux_terms, model, net_input, recompute, warmup=1.0)

Return bce + sum(w_i * aux_i), recalibrating weights on schedule.

Parameters:

• bce (torch.Tensor) – The primary (reference) loss term to backpropagate.

• aux_terms (list[torch.Tensor]) – The n_aux auxiliary loss terms (raw, pre-weight).

• model – The (possibly compiled) model and the current preprocessed input — reused for the clean eager calibration forward.

• net_input – The (possibly compiled) model and the current preprocessed input — reused for the clean eager calibration forward.

• recompute (callable) – recompute(model_output) -> (bce_term, [aux_term, ...]) rebuilding the same loss terms from a fresh forward, for gradient measurement.

• warmup (float) – Multiplier applied to all aux weights this step (e.g. ramp 0->1 over the first epoch so BCE/PE settle before the aux losses fully engage).