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=200))
- snapshot(leaf_or_module)[source]
Detached clones of the module’s parameters – the anchor
theta*for an EWC penalty.
- fisher_diagonal(leaf, x, y, *, samples=512, device='cpu', seed=0)[source]
Diagonal empirical Fisher of a classification leaf’s module on
(x, y): mean of(d log p(y|x)/dtheta)^2.