mixle.experimental.growth_operators module

Growth operators (roadmap H1): net2net widening + progressive depth stacking – G3’s mixle.models.coarsening run backwards, over the real transformer in mixle.models.transformer.

Where G3’s coarsen() FOLDS capacity down (depth-merge, width-merge, structure-projection) under a divergence budget, H1 SPLITS capacity up, function-preservingly, so that a rung-(k) checkpoint can be grown into a rung-(k+1) initialization and continued-trained rather than retrained from scratch. Concretely, the two moves here are genuine inverses of G3’s shrink operators:

  1. Net2Net widening (net2net_widen(), widen_block() – Chen, Goodfellow & Shlens, “Net2Net: Accelerating Learning via Knowledge Transfer”, 2016): the real duplication rule. Given two consecutive nn.Linear layers L_in: R^d -> R^h and L_out: R^h -> R^o composed as L_out(act(L_in(x))), widen the hidden dimension h -> h' by choosing a mapping g: {0,...,h'-1} -> {0,...,h-1} that is the identity on the first h indices and, for each new index, COPIES a (systematically or randomly chosen) existing column j. The new L_in copies row j of the old L_in (and its bias) into the new row – the widened hidden unit computes literally the same pre-activation as its source, so any pointwise activation between the two layers is also identical on the new unit. The new L_out copies column j of the old L_out into the new column, and then DIVIDES every column of the new L_out by its post-widening replication count (1 for never-duplicated original columns, k+1 for a column duplicated k extra times) – this is the step that keeps the SUM L_out(h) invariant despite now having k+1 identical copies of that unit’s activation feeding into it. This is the textbook Net2WiderNet rule, not an approximation: the pre-widening and post-widening compositions compute the exact same function on any input, up to floating-point round-off (verified below by an actual forward-pass receipt, per this track’s stated invariant).

  2. Progressive depth stacking (insert_block()): the direct inverse of G3’s depth_merge – instead of folding two adjacent blocks x -> x + f(x), x -> x + g(x) into one merged block via a second-order Taylor approximation of f and g’s composition, SPLIT capacity by inserting a brand new block at position whose residual branches are zero-initialized. A Block has TWO separate residual adds (x = x + attn(ln1(x)); x = x + mlp(ln2(x))), so BOTH final linears that get summed back onto the residual stream – attention’s proj and the MLP’s second Linear – have their weight and bias set to exactly zero, so the new block computes x -> x + 0 + 0 = x, an exact identity, immediately after insertion. The residual connection is what makes this trivial: unlike depth_merge’s Taylor approximation (needed because composing two already-nonlinear branches has no closed form), inserting an literal no-op block needs no approximation at all – the new block’s output is bit-for-bit identical to its input by construction, so the WHOLE model’s output is unchanged by insertion, exactly (not to second order).

Function-preservation receipt. Per the track’s stated invariant (“all function-preserving edits carry an output-parity receipt at the moment of the edit”), every growth move here is checked with a REAL forward-pass comparison (verify_output_parity()) of the model before vs. after growth on the same input batch – not a closed-form law-level argument the way G3’s KL receipts are. Net2Net widening and zero-init depth stacking are ALGEBRAICALLY exact (not second-order truncations like depth_merge), so the measured max-abs/max-rel differences here are expected to sit at plain float round-off, not at some larger truncation-controlled scale.

Symmetric naming with G3, deliberately: net2net_widen()/widen_block() undo width_merge, insert_block() undoes depth_merge, and verify_output_parity() plays the same role here that gaussian_kl() plays there – a real, verifiable, reported number, not a hand-wave.

class ParityReceipt(max_abs_diff, max_rel_diff, tolerance, within_tolerance, batch_shape)[source]

Bases: object

A real, measured before/after forward-pass comparison – the function-preservation receipt this track requires “at the moment of the edit”. max_abs_diff/max_rel_diff are computed over every output element on the SAME input batch; within_tolerance is max_abs_diff <= tolerance.

Parameters:
class GrowthReceipt(name, parity=None)[source]

Bases: object

Receipt for one growth operation: its own name, the width/depth change it made, and the ParityReceipt from a real forward-pass comparison – filled in by the caller (widen_block, insert_block) once the grown module/model exists, via verify_output_parity().

Parameters:

  • name (str)

  • parity (ParityReceipt | None)

verify_output_parity(model_before, model_after, batch, tolerance=1e-5)[source]

The actual before/after forward-pass comparison used by every growth operator below: run model_before and model_after on the SAME batch (in eval() mode, no-grad, so dropout/BN – not present in Block today, but this keeps the receipt honest if that ever changes – doesn’t inject spurious stochastic differences) and report the real measured max absolute and max relative difference between their outputs.

This is the module’s one hard requirement per the track’s stated invariant: “every growth operation must be verified via a real forward-pass comparison… showing the output is unchanged to a stated bitwise/numerical tolerance.” Both net2net_widen (an exact algebraic identity) and insert_block (an exact zero-residual identity) are expected to pass at a tight tolerance close to float precision.

Parameters:
Return type:

ParityReceipt

net2net_widen(linear_in, linear_out, new_width, seed=0, systematic=True)[source]

Widen the hidden dimension between two consecutive nn.Linear layers linear_in: R^d -> R^h, linear_out: R^h -> R^o (composed as linear_out(act(linear_in(x))) for any pointwise activation act, or no activation at all) to new_width = h' via the real Net2Net duplication rule (see module docstring): a mapping g picks, for each new hidden unit, an existing unit to copy; linear_in’s new row is copied verbatim from its source row (so the new unit’s pre-activation, and hence any pointwise post-activation, is IDENTICAL to its source); linear_out’s new column is copied from its source column and then every column (including the original, now-duplicated ones) is divided by its total replication count, so the SUM linear_out(h) is invariant.

Returns (new_linear_in, new_linear_out, receipt) – receipt.parity is left None here (no model to run a forward pass on at this granularity); callers needing the parity number should use widen_block() or verify_output_parity() directly on an assembled module.

Parameters:
Return type:

tuple[Any, Any, GrowthReceipt]

widen_block(block, new_d_model, seed=0, systematic=True)[source]

Widen an entire transformer Block’s d_model from its current width old_d to new_d_model, applying net2net-style widening COHERENTLY across every d_model-shaped surface (both LayerNorms, attention qkv/proj, and the MLP’s two Linears) so the block’s overall function is unchanged, not just one isolated Linear pair.

new_d_model must be an exact integer multiple of old_d (r = new_d_model // old_d), and every source coordinate is duplicated EXACTLY r times (see _residual_widen_mapping()). This uniform-ratio restriction is what makes LayerNorm exact under duplication: LayerNorm’s mean/variance are population statistics over the FULL width, and duplicating a value v a NON-uniform number of times changes the weighted mean/variance relative to the original – with a UNIFORM r-fold duplication, every source value contributes the same relative weight it always did (r copies out of r * old_d total, same as 1 copy out of old_d), so the widened mean and (population) variance are algebraically identical to the pre-widening ones, and ln1/ln2’s elementwise weight/bias need only a plain coordinate-duplicate (no division) to match.

Every OTHER layer that reads the widened residual stream (qkv, proj, the MLP’s first Linear) divides its incoming weights by r to correct for the fact that a plain matmul SUMS over its (now r-times-duplicated) input axis – exactly analogous to net2net’s own outgoing-weight correction, just needed on the input side here because d_model is read repeatedly through the block rather than produced once. Attention gets one further correction (see _widen_attention()) for the Q . K dot-product scale. The MLP’s OWN internal hidden width (4 * d_model) widens by ordinary (non-uniform, arbitrary) net2net duplication, exactly as in net2net_widen(), since it is not part of the residual stream and has no LayerNorm-style population-statistic constraint.

Returns (new_block, receipt) where receipt.parity is filled in by an actual forward-pass comparison (verify_output_parity()) of block vs. new_block on the SAME old_d-wide random batch (both blocks accept the same input width at the moment of growth).

Parameters:
Return type:

tuple[Any, GrowthReceipt]

insert_block(model, position, seed=0)[source]

Progressive depth stacking: insert a brand-new, near-identity Block into model.blocks at index position (0 <= position <= len(model.blocks)), function-preservingly, by ZERO-INITIALIZING BOTH of the new block’s residual branches – attention’s proj (the last thing summed onto the residual stream in x + self.attn(self.ln1(x))) AND the MLP’s second Linear (the last thing summed on in x + self.mlp(self.ln2(x))), weight and bias both – so the new block computes x -> (x + 0) + 0 = x, an EXACT identity (not a second-order approximation the way G3’s depth_merge needs, since there is no composition of two nonlinear branches to linearize here – a literal no-op block needs no approximation at all). Zeroing only one of the two branches is NOT sufficient: a Block has two separate residual adds, and either one left non-zero would perturb x before the block’s output is reached.

This is the direct inverse move of depth_merge: where that operator folds two adjacent blocks into one via a Taylor-approximated composition, this operator splits capacity by inserting one new block that starts as a no-op and is left for the optimizer to grow into useful capacity during continued training – exactly the roadmap’s “rung-(k) checkpoints initialize rung-(k+1) function-preservingly” role.

Returns (new_model, receipt) – a fresh CausalLM-shaped module (built by shallow-copying the input model’s embeddings/head and inserting the new block into a fresh blocks list, mirroring CoarsenedLM’s own approach on the shrink side) whose forward is identical to model’s at the moment of insertion, verified by an actual forward pass on real token ids via verify_output_parity().

Parameters:
Return type:

tuple[Any, GrowthReceipt]