mixle.models.self_distillation module

Self-distillation during training (roadmap J3): EMA-teacher consistency + stochastic-depth targets, wired into normal training as loss-hooks – not a separate post-hoc distillation pass, and not a new trainer either.

The idea, concretely

Two self-consistency pressures applied during ordinary next-token training:

  1. EMA-teacher consistency (EMATeacher): maintain an exponential-moving-average copy of the model’s own weights, updated every step (teacher = decay * teacher + (1 - decay) * student, the standard mean-teacher/BYOL/DINO pattern). The actively-trained (student) model is pushed to agree with this temporally-smoothed version of itself via consistency_loss() – an implicit regularizer with no extra labels or extra model.

  2. Stochastic-depth consistency (stochastic_depth_forward()): each step, run the SAME input through the model twice – once at full depth, once with a random subset of blocks skipped entirely (the standard stochastic-depth / drop-path regularizer) – and add a consistency term pulling the partial-depth output toward the full-depth output. This directly trains the model to tolerate missing blocks, which is exactly the redundancy G3’s mixle.models.coarsening depth-merge exploits.

Why this belongs at the loss-hook level, not a new trainer

mixle.models.grad_leaf already establishes the “compose via wrapping” pattern for this codebase’s M-step: a training loop is generic, and custom OBJECTIVES are a loss(module, x, w) -> scalar hook, not a subclass tree (see GradLeaf/GradEstimator). CausalLM doesn’t fit GradLeaf directly (it has no log_density; its own dense-teacher-forcing loop lives in mixle.models.language_model and mixle.models.streaming_transformer_leaf), so train_with_self_distillation() mirrors THOSE loops’ own conventions (F.cross_entropy over (context, next_token) micro-batches from mixle.data.stream_token_source.stream_token_source(), a plain torch.optim.Adam M-step) and adds the two consistency terms as extra, addable loss components on top of the same per-step cross-entropy – the loss-hook composition pattern, applied at the place this model family’s training loop actually lives.

class EMATeacher(model, decay=0.999)[source]

Bases: object

An exponential-moving-average copy of a model’s own weights: the standard mean-teacher/BYOL/DINO self-distillation teacher.

update(student_model) applies teacher = decay * teacher + (1 - decay) * student to every tensor in the teacher’s state_dict (parameters AND buffers, so e.g. non-trainable statistics stay consistent too) – called once per training step, AFTER the optimizer step, so the teacher always tracks a temporally-smoothed trailing average of the student. The teacher is a real, independent, forward-passable module (forward/predict), held in eval mode with gradients disabled: it is a read-only distillation TARGET, never itself directly optimized.

Parameters:

  • model (Any)

  • decay (float)

update(student_model)[source]

One EMA step: pull every teacher tensor toward the student’s current value.

CausalLM ties head.weight to tok.weight (weight tying), so several state_dict() keys alias the SAME underlying storage – updating each key naively would apply the EMA formula to that storage more than once per step (a double update). seen dedupes by storage identity (data_ptr()) so every real tensor is updated exactly once, however many names alias it.

Parameters:

student_model (Any)

Return type:

None

forward(x)[source]

Run x through the EMA-teacher weights (eval mode, no grad).

Parameters:

x (Any)

Return type:

Any

predict(x)

Run x through the EMA-teacher weights (eval mode, no grad).

Parameters:

x (Any)

Return type:

Any

class TrainStats(ce_loss=<factory>, stochastic_depth_loss=<factory>, ema_consistency_loss=<factory>, total_loss=<factory>)[source]

Bases: object

Per-step telemetry from train_with_self_distillation() – cross-entropy, stochastic-depth consistency, and EMA-teacher consistency losses, kept separately so a caller can see which pressure is doing what (and the combined total actually optimized).

Parameters:

  • ce_loss (list)

  • stochastic_depth_loss (list)

  • ema_consistency_loss (list)

  • total_loss (list)

consistency_loss(student_output, teacher_output, mode='mse')[source]

The self-distillation consistency term between a student prediction and a teacher/target prediction on the SAME input – mode="mse" (default, plain squared-error between logits, the mean-teacher convention) or mode="kl" (KL(teacher_softmax || student_log_softmax), the classic soft-target distillation loss).

Parameters:

  • student_output (Any)

  • teacher_output (Any)

  • mode (str)

Return type:

Any

stochastic_depth_forward(model, x, drop_prob, generator=None)[source]

Run model on x twice: once at full depth, once with each block independently dropped with probability drop_prob (at least one block is always kept, so the partial pass never degenerates to the bare embedding/head). Returns (full_output, partial_output) – the pair train_with_self_distillation() feeds to consistency_loss().

At drop_prob == 0 both passes keep every block, so the two outputs are IDENTICAL (no dropout elsewhere in Block) – the degenerate-case sanity check pinned in mixle/tests/self_distillation_test.py.

Parameters:
Return type:

tuple

train_with_self_distillation(model, data, steps, *, ema_decay=0.999, drop_prob=0.1, consistency_weight=1.0, ema_weight=None, stochastic_depth_weight=None, consistency_mode='mse', lr=3e-3, device='cpu', optimizer=None, seed=0, log=None)[source]

Train model (a CausalLM, trained in place and also returned) for steps next-token cross-entropy steps, with EMA-teacher consistency and stochastic-depth consistency added as extra loss terms on top of the SAME per-step batch – both self-distillation pressures happen DURING training, not as a separate post-hoc pass.

data yields (context, next_token) micro-batches shaped exactly like mixle.data.stream_token_source.stream_token_source() (context: (batch, block) float ids, next_token: (batch,) int ids). data may be:

  • a zero-arg CALLABLE returning a fresh iterator each time (e.g. lambda: stream_token_source(ids, block=64, batch_size=32)) – restarted automatically whenever it runs dry before steps is reached, so training can outlast one epoch; or

  • a plain iterable/iterator (e.g. a list of batches, or a single generator object) – consumed once, sized to yield at least steps batches (a bare generator can’t be rewound).

Per step: loss = cross_entropy(full_depth_logits, target) + stochastic_depth_weight * consistency(partial_depth_logits, full_depth_logits.detach()) + ema_weight * consistency(full_depth_logits, ema_teacher(context)), then one optimizer step, then one EMA-teacher update. ema_weight/stochastic_depth_weight each default to consistency_weight when unset.

Parameters:
Return type:

Any