mixle.models.continual module

Continual / multi-stage fine-tuning helpers: parameter snapshot + diagonal Fisher + EWC for neural leaves.

Continued pretraining (CPT) without catastrophic forgetting = continue the same module on new data plus an EWC penalty lambda * sum_i F_i (theta_i - theta*_i)^2 anchoring to the pretrained params theta* weighted by the diagonal Fisher F (how much each parameter mattered for the old task). The Fisher is the same curvature mixle uses for posterior approximation; here it is per-parameter importance for anti-forgetting. Use it as a declarative stage in the pipeline:

pre  = Categorical(logits=Net(out=K)).fit(yA, given={"x": XA})
F    = fisher_diagonal(pre.dist, XA, yA)
cpt  = Categorical(logits=Net(out=K)).fit(
    yB, given={"x": XB}, init=pre, ewc=ewc(snapshot(pre.dist), F, lam=100, floor=1e-3))

floor is not decoration. A stage that fit well leaves per-example gradients near zero, so the diagonal empirical Fisher is near zero almost everywhere and the penalty it defines is rank-deficient: minimizing it walks the module through the Fisher’s null space instead of holding it near theta*, and the old task is lost with the penalty itself satisfied to four decimal places. Raising lam does not fix that and eventually diverges. See ewc() for the measurements behind both statements.

class ParameterBundle(values, names)[source]

Bases: list

Tensor list carrying the ordered parameter identity manifest.

Parameters:
snapshot(leaf_or_module)[source]

Detached clones of the module’s parameters – the anchor theta* for an EWC penalty.

Parameters:

leaf_or_module (Any)

Return type:

ParameterBundle

fisher_diagonal(leaf, x, y, *, samples=512, device='cpu', seed=0)[source]

Diagonal empirical Fisher computed on an isolated evaluation-mode copy.

Parameters:
Return type:

ParameterBundle

ewc(anchor, fisher, lam=1.0, floor=0.0)[source]

Bundle (anchor, fisher, lambda) for .fit(..., ewc=...) (the EWC anti-forgetting penalty).

floor adds a constant to every Fisher entry, so the penalty weights are F_i + floor and the penalty is lambda * sum_i (F_i + floor) (theta_i - theta*_i)^2. It defaults to 0.0, which is EWC exactly as defined – and which, on its own, usually does not retain the old task, for a reason that is structural rather than a matter of tuning lam.

The diagonal empirical Fisher is estimated from per-example gradients at the end of the previous stage. A stage that fit well leaves those gradients near zero, so almost every entry of F is near zero and the quadratic it defines is badly rank-deficient. Minimizing that penalty then does not hold the module near theta*; it moves the module through the Fisher’s null space, where the penalty is cheap and the old task’s function is not preserved. Measured on a 4-feature 3-class pair of unrelated tasks (see neural_ppl_test.test_cpt_ewc_retains_the_old_task), raising lam from 10 to 8000 drives the Fisher-weighted drift down by four orders of magnitude, 3.6e-1 to 3.7e-5, while the raw distance from the anchor nearly doubles, 81 to 147, and retained accuracy on the old task stays flat around 0.21-0.24 against 0.19 for plain continuation. The penalty is satisfied; the capability it exists to protect is not. Past lam of about 1e4 the quadratic is stiffer than the routed optimizer’s step size and the objective diverges instead.

A small positive floor closes that null space – it is EWC plus L2-to-anchor, a Tikhonov regularization of the importance estimate rather than of the parameters. On the same task pair floor=1e-3, lam=1e2 retains 0.59-0.74 on the old task against 0.19-0.34 for plain continuation, at 0.62-0.66 on the new one: a real stability/plasticity trade rather than a satisfied constraint. Prefer a floor over a large lam; the floor buys retention where lam alone only buys stiffness.

Parameters:

  • anchor (ParameterBundle)

  • fisher (ParameterBundle)

  • lam (float)

  • floor (float)

Return type:

tuple