dicee.weight_averaging

Classes

AMWA	Adaptive Momentum Weight Averaging (AMWA).
ASWA	Adaptive stochastic weight averaging
SWA	Stochastic Weight Averaging callback.
SWAG	Stochastic Weight Averaging - Gaussian (SWAG).
EMA	Exponential Moving Average (EMA) callback.
TWA	Train with Weight Averaging (TWA) using subspace projection + averaging.

Module Contents

class dicee.weight_averaging.AMWA(amwa_start_epoch: int = 0, amwa_c_epochs: int = 1, path: str | None = None, monitor: str = 'MRR', maximize: bool = True, beta: float | None = None, beta_window: int = 10, beta_init: float = 1.0, beta_floor: float = 1e-08)

Bases: dicee.abstracts.AbstractCallback

Adaptive Momentum Weight Averaging (AMWA).

Paper: https://www.sciencedirect.com/science/article/pii/S0031320325009586

BaseNet is the ordinary model optimized during training, while StableNet is updated once per epoch via weight averaging,

w_stable <- m * w_stable + (1 - m) * w_base

where the adaptive momentum is computed from the validation-performance difference between StableNet and BaseNet,

m(Delta) = 1 / (1 + exp(-max(0, Delta) / beta_n))

with Delta defined so that Delta > 0 means StableNet performs better than BaseNet. This yields the exact behavior described in the paper: StableNet stays close to its current solution when it already validates better, and it takes a larger step toward BaseNet when BaseNet validates better. The paper sets beta_n to the standard deviation of Delta over the past 10 epochs; this implementation follows that default and also allows a fixed beta for the ablation setting.

param amwa_start_epoch:

Epoch at which AMWA starts tracking StableNet.

type amwa_start_epoch:

int

param amwa_c_epochs:

Update interval in epochs.

type amwa_c_epochs:

int

param path:

Directory used to persist amwa.pt and amwa_history.json.

type path:

str, optional

param monitor:

Validation metric to monitor, e.g. "MRR".

type monitor:

str

param maximize:

Whether larger values of monitor are better.

type maximize:

bool

param beta:

If provided, use a fixed beta instead of the adaptive beta_n from the paper. This matches the paper’s fixed-beta ablation.

type beta:

float, optional

param beta_window:

Number of past Delta values used to estimate beta_n.

type beta_window:

int

param beta_init:

Beta used before enough Delta history exists.

type beta_init:

float

param beta_floor:

Numerical floor used when the empirical standard deviation is zero.

type beta_floor:

float

amwa_start_epoch = 0

amwa_c_epochs = 1

path = None

monitor = 'MRR'

maximize = True

beta = None

beta_window

beta_init

beta_floor

current_epoch = -1

initial_eval_setting: str | None = None

stable_state_dict: Dict[str, torch.Tensor] | None = None

delta_history: List[float] = []

momentum_history: List[float] = []

beta_history: List[float] = []

base_history: List[float] = []

stable_history: List[float] = []

history: List[Dict[str, Any]] = []

on_train_epoch_start(trainer, model)

Called when the train epoch begins.

on_train_epoch_end(trainer, model)

Called at the end of each training epoch.

Parameters:

• trainer (AbstractTrainer or pl.Trainer) – The active trainer instance.

• model (BaseKGE) – The model being trained. model.loss_history contains the per-epoch average losses accumulated so far.

on_fit_end(trainer, model)

Called once after the final training epoch completes.

Override to perform post-training actions such as saving the final model state, computing evaluation metrics, or cleaning up resources.

class dicee.weight_averaging.ASWA(num_epochs, path)

Bases: dicee.abstracts.AbstractCallback

Adaptive stochastic weight averaging ASWE keeps track of the validation performance and update s the ensemble model accordingly.

path

num_epochs

initial_eval_setting = None

epoch_count = 0

alphas = []

val_aswa = -1

on_fit_end(trainer, model)

Called once after the final training epoch completes.

Override to perform post-training actions such as saving the final model state, computing evaluation metrics, or cleaning up resources.

static compute_mrr(trainer, model) → float

get_aswa_state_dict(model)

decide(running_model_state_dict, ensemble_state_dict, val_running_model, mrr_updated_ensemble_model)

Perform Hard Update, software or rejection

Parameters:

• running_model_state_dict

• ensemble_state_dict

• val_running_model

• mrr_updated_ensemble_model

on_train_epoch_end(trainer, model)

Called at the end of each training epoch.

Parameters:

• trainer (AbstractTrainer or pl.Trainer) – The active trainer instance.

• model (BaseKGE) – The model being trained. model.loss_history contains the per-epoch average losses accumulated so far.

class dicee.weight_averaging.SWA(swa_start_epoch, swa_c_epochs: int = 1, lr_init: float = 0.1, swa_lr: float = 0.05, max_epochs: int | None = None)

Bases: dicee.abstracts.AbstractCallback

Stochastic Weight Averaging callback.

Initialize SWA callback.

swa_start_epoch: int

The epoch at which to start SWA.

swa_c_epochs: int

The number of epochs to use for SWA.

lr_init: float

The initial learning rate.

swa_lr: float

The learning rate to use during SWA.

max_epochs: int

The maximum number of epochs. args.num_epochs

swa_start_epoch

swa_c_epochs = 1

swa_lr = 0.05

lr_init = 0.1

max_epochs = None

swa_model = None

swa_n = 0

current_epoch = -1

static moving_average(swa_model, running_model, alpha)

Update SWA model with moving average of current model. Math: # SWA update: # θ_swa ← (1 - alpha) * θ_swa + alpha * θ # alpha = 1 / (n + 1), where n = number of models already averaged # alpha is tracked via self.swa_n in code

on_train_epoch_start(trainer, model)

Update learning rate according to SWA schedule.

on_train_epoch_end(trainer, model)

Apply SWA averaging if conditions are met.

on_fit_end(trainer, model)

Replace main model with SWA model at the end of training.

class dicee.weight_averaging.SWAG(swa_start_epoch, swa_c_epochs: int = 1, lr_init: float = 0.1, swa_lr: float = 0.05, max_epochs: int | None = None, max_num_models: int = 20, var_clamp: float = 1e-30)

Bases: dicee.abstracts.AbstractCallback

Stochastic Weight Averaging - Gaussian (SWAG). Parameters

swa_start_epochint

Epoch at which to start collecting weights.

swa_c_epochsint

Interval of epochs between updates.

lr_initfloat

Initial LR.

swa_lrfloat

LR in SWA / GSWA phase.

max_epochsint

Total number of epochs.

max_num_modelsint

Number of models to keep for low-rank covariance approx.

var_clampfloat

Clamp low variance for stability.

swa_start_epoch

swa_c_epochs = 1

swa_lr = 0.05

lr_init = 0.1

max_epochs = None

max_num_models = 20

var_clamp = 1e-30

mean = None

sq_mean = None

deviations = []

gswa_n = 0

current_epoch = -1

get_mean_and_var()

Return mean + variance (diagonal part).

sample(base_model, scale=0.5)

Sample new model from SWAG posterior distribution.

Math: # From SWAG, posterior is approximated as: # θ ~ N(mean, Σ) # where Σ ≈ diag(var) + (1/(K-1)) * D D^T # - mean = running average of weights # - var = elementwise variance (sq_mean - mean^2) # - D = [dev_1, dev_2, …, dev_K], deviations from mean (low-rank approx) # - K = number of collected models

# Sampling step: # 1. θ_diag = mean + scale * std ⊙ ε, where ε ~ N(0, I) # 2. θ_lowrank = θ_diag + (D z) / sqrt(K-1), where z ~ N(0, I_K) # Final sample = θ_lowrank

on_train_epoch_start(trainer, model)

Update LR schedule (same as SWA).

on_train_epoch_end(trainer, model)

Collect Gaussian stats at the end of epochs after swa_start.

on_fit_end(trainer, model)

Set model weights to the collected SWAG mean at the end of training.

class dicee.weight_averaging.EMA(ema_start_epoch: int, decay: float = 0.999, max_epochs: int | None = None, ema_c_epochs: int = 1)

Bases: dicee.abstracts.AbstractCallback

Exponential Moving Average (EMA) callback.

Parameters:

• ema_start_epoch (int) – Epoch to start EMA.

• decay (float) – EMA decay rate (typical: 0.99 - 0.9999) Math: θ_ema <- decay * θ_ema + (1 - decay) * θ

• max_epochs (int) – Maximum number of epochs.

ema_start_epoch

decay = 0.999

max_epochs = None

ema_c_epochs = 1

ema_model = None

current_epoch = -1

static ema_update(ema_model, running_model, decay: float)

Update EMA model with exponential moving average of current model. Math: # EMA update: # θ_ema ← (1 - alpha) * θ_ema + alpha * θ # alpha = 1 - decay, where decay is the EMA smoothing factor (typical 0.99 - 0.999) # alpha controls how much of the current model θ contributes to the EMA # decay is fixed in code –> can be extended to scheduled

on_train_epoch_start(trainer, model)

Track current epoch.

on_train_epoch_end(trainer, model)

Update EMA if past start epoch.

on_fit_end(trainer, model)

Replace main model with EMA model at the end of training.

class dicee.weight_averaging.TWA(twa_start_epoch: int, lr_init: float, num_samples: int = 5, reg_lambda: float = 0.0, max_epochs: int | None = None, twa_c_epochs: int = 1)

Bases: dicee.abstracts.AbstractCallback

Train with Weight Averaging (TWA) using subspace projection + averaging.

Parameters

twa_start_epochint

Epoch to start TWA.

lr_initfloat

Learning rate used for β updates.

num_samplesint

Number of sampled weight snapshots to build projection subspace.

reg_lambdafloat

Regularization coefficient for β updates.

max_epochsint

Total number of training epochs.

twa_c_epochsint

Interval of epochs between TWA updates.

twa_start_epoch

num_samples = 5

reg_lambda = 0.0

max_epochs = None

lr_init

twa_c_epochs = 1

current_epoch = -1

weight_samples = []

twa_model = None

base_weights = None

P = None

beta = None

sample_weights(model)

Collect sampled weights from the current model and maintain rolling buffer.

build_projection(weight_samples, k=None)

Build projection subspace from collected weight samples. :param weight_samples: list of flat weight tensors [(D,), …] :param k: number of basis vectors to keep. Defaults to min(N, D).

Returns:

(D,) base weight vector (average) P: (D, k) projection matrix with top-k basis directions

Return type:

mean_w

on_train_epoch_start(trainer, model)

Track epoch.

on_train_epoch_end(trainer, model)

Main TWA logic: build subspace and update in β space.

# Math: # TWA weight update: # w_twa = mean_w + P * beta # mean_w = (1/n) * sum_i w_i (SWA baseline) # beta <- (1 - eta * lambda) * beta - eta * P^T * g # g = gradient of training loss w.r.t. full model weights # eta = learning rate, lambda = ridge regularization # P = orthonormal basis spanning sampled checkpoints {w_i}

on_fit_end(trainer, model)

Replace with TWA model at the end.