mixle.experimental.structure_edit_schedule module

H3: structure-edit schedule during training – the neural half of ConditionalJIT (roadmap H).

D5 (mixle.inference.conditional_jit_controller) built a generic learned ActionType registry and explicitly left STRUCTURE_EDIT as a documented EXTENSION POINT, not implemented (see its module docstring and ACTION_TYPE_REGISTRY). This module is that wiring: a real action space of architecture edits –

• grow (H1, mixle.experimental.growth_operators) – net2net_widen/widen_block (width) and insert_block (depth);

• prune / depth-merge (G3, mixle.models.coarsening) – depth_merge;

• rank change (G2, mixle.models.sigma_weighted_projection) – sigma_weighted_low_rank;

• 2:4 sparsity (I4 – no standalone I4 PR had landed when this module was built; the underlying 2:4 projection primitive already exists as part of G2’s own module, sigma_weighted_block_sparse() with pattern="2:4", so a snapshot (non-ramped) 2:4 projection IS wired here – see STRUCTURE_EDIT_REGISTRY’s note on "sparsity_2_4" for exactly what is and is not covered);

• MoE expert add/merge (H2 – not landed when this module was built) – SCAFFOLDED ONLY: the action-type name is registered and raises a clear NotImplementedError from apply_structure_edit(), per the roadmap item’s explicit “optional, document don’t block” instruction.

under one uniform interface (apply_structure_edit()), gated by an F4-style training-health check plus a real function-preservation/output-parity check (should_apply_edit()), driven by a StructureEditController that extends D5’s LearnedController/ActionType machinery with a REAL STRUCTURE_EDIT arm (reusing D5’s own mixle.task.bandit wiring pattern, per that module’s “reusable brain” note), and exercised end-to-end by train_with_adaptive_structure() – a real training loop that starts small and grows/edits structure as training proceeds, per the round’s controller decision.

Note on mixle.inference.conditional_jit_controller.ACTION_TYPE_REGISTRY: that dict’s own STRUCTURE_EDIT entry is left untouched here (D5’s own test pins its “EXTENSION POINT” text) – this module’s STRUCTURE_EDIT_REGISTRY is a SEPARATE, more detailed registry of the actual edit-type strings apply_structure_edit() accepts ("grow_insert", "grow_widen_block", "prune_depth_merge", "rank_reduce", "sparsity_2_4", "moe_expert_add"), not a replacement for D5’s coarser action-type-level registry.

class AdaptiveTrainingResult(model, total_flops, steps, final_loss, reached_target, edits_applied=<factory>, edits_rejected=<factory>, health_report=<factory>)

Bases: object

Output of train_with_adaptive_structure(): the final (possibly grown/edited) model, the REAL measured total compute (sum of F4’s own theoretical_flops_per_iter over every step, at whatever the model’s shape was AT that step – so growth rounds correctly cost more from the round they take effect, not before), and the edit/health bookkeeping.

Parameters:

• model (Any)

• total_flops (float)

• steps (int)

• final_loss (float)

• reached_target (bool)

• edits_applied (list[tuple[int, str]])

• edits_rejected (list[tuple[int, str, str]])

• health_report (dict[str, Any])

class StructureEditController(*, edit_moves=_DEFAULT_EDIT_MOVES, max_layer=None, ucb_c=1.0, seed=None)

Bases: LearnedController[StructureEditState, ControllerAction]

The real ActionType.STRUCTURE_EDIT arm D5 left as an extension point (see this module’s docstring): an online bandit – reusing mixle.task.bandit exactly as D5’s own BanditController does, per that module’s “reusable brain” note – over a small discrete set of edit “moves” (default: {no_edit, grow_insert}; any apply_structure_edit()-shaped {"edit_type": ..., **params} dict may be added). select_action returns a ControllerAction tagged ActionType.STRUCTURE_EDIT whose payload carries the chosen move (budget_fraction is unused by this action type, set to 1.0 for interface symmetry with D5’s other actions).

At capacity (state.n_layer >= max_layer, when max_layer is set) only "none" is legal, so growth arms are skipped without consulting/perturbing the bandit – a forced move never counts as an exploration pull.

Parameters:

• edit_moves (tuple[dict[str, Any], ...])

• max_layer (int | None)

• ucb_c (float)

• seed (int | None)

select_action(state)

Return this round’s action given state.

Parameters:

state (StructureEditState)

Return type:

ControllerAction

update(state, action, realized_gain, realized_cost)

Feed back the REALIZED gain/cost of action taken in state – the online-learning signal every concrete controller trains from.

Parameters:

• state (StructureEditState)

• action (ControllerAction)

• realized_gain (float)

• realized_cost (float)

Return type:

None

class StructureEditReceipt(edit_type, parity, detail=None)

Bases: object

One structure edit’s receipt: which edit, the real forward-pass ParityReceipt used by the function-preservation gate (see should_apply_edit()), and the edit-specific detail object (a GrowthReceipt, ScaleReceipt, or ProjectionReceipt-shaped record, whichever the underlying H1/G3/G2 op returns) for anyone wanting the edit’s own native receipt too.

Parameters:

• edit_type (str)

• parity (ParityReceipt | None)

• detail (Any)

class StructureEditState(round_index, loss_ema, loss_slope, n_layer, healthy)

Bases: object

One round’s controller-visible state for the structure-edit decision: the running loss EMA and its recent slope (the plateau signal), current depth, and whether F4 currently reports healthy – small and specific to “should I consider editing structure right now”, mirroring D5’s own ControllerState role but for this different decision.

Parameters:

• round_index (int)

• loss_ema (float)

• loss_slope (float)

• n_layer (int)

• healthy (bool)

apply_structure_edit(model, edit_type, params=None)

Apply one structure edit to model and return (new_model, receipt) – the uniform interface every STRUCTURE_EDIT action funnels through, wrapping H1/G3/G2’s real ops (see STRUCTURE_EDIT_REGISTRY for exactly what each edit_type does and does not cover).

Every edit type except "grow_widen_block" (block-scoped, see the registry note) returns a full, forward-passable model and a real forward-pass ParityReceipt (computed via verify_output_parity() on the SAME random batch, or params["parity_batch"] if supplied) – the function-preservation half of should_apply_edit()’s gate.

Parameters:

• model (Any)

• edit_type (str)

• params (dict[str, Any] | None)

Return type:

tuple[Any, StructureEditReceipt]

health_report_from_monitor(monitor, lookback=5)

Build the health_report should_apply_edit() expects from a real F4 TrainingHealthMonitor: healthy iff no anomaly was raised in the last lookback observed steps (an anomaly from steps ago should not permanently block future edits; a RECENT one – loss spiking, NaN/Inf grads, a restart discontinuity – should).

Parameters:

• monitor (TrainingHealthMonitor)

• lookback (int)

Return type:

dict[str, Any]

should_apply_edit(health_report, parity_check)

The H3 gate: commit to a structure edit only if BOTH hold –

1. health_report (see health_report_from_monitor()) reports no recent F4 anomaly (don’t structurally edit a model mid-anomaly: a loss spike, NaN/Inf grad, or restart discontinuity means the current state is not trustworthy to branch a structural decision from);

2. parity_check (a real ParityReceipt from apply_structure_edit(), per H1/D6’s established output-divergence pattern) reports the edit is within its stated tolerance.

Otherwise the edit is skipped for this round – the caller keeps training the UNedited model and may try again (a different edit, or the same one) at a later round.

Parameters:

• health_report (dict[str, Any])

• parity_check (ParityReceipt | None)

Return type:

bool

train_with_adaptive_structure(initial_model, make_batch, target_loss, *, max_steps=2000, max_layer=3, batch_size=64, lr=5e-3, min_steps_before_edit=80, plateau_window=40, plateau_eps=0.01, parity_tolerance=1e-4, health_lookback=5, seed=0, controller=None)

Train initial_model (expected small) toward target_loss, letting a StructureEditController decide when/how to grow structure as training proceeds.

Each step: one AdamW step on a batch from make_batch(batch_size, rng) (real cross-entropy loss, real backward pass), fed into a real F4 TrainingHealthMonitor (loss, grad-norm). A loss-EMA PLATEAU DETECTOR (no improvement over the last plateau_window steps, past min_steps_before_edit steps since the last edit, and below max_layer) is what decides WHEN to even consider a structure edit – the controller is consulted only at plateau moments, mirroring how a real scheduler would not burn an edit decision every single step. When consulted, the controller’s chosen move is applied via apply_structure_edit() and gated by should_apply_edit() (a real F4 health check plus the edit’s own real output-parity receipt) before being committed – a rejected edit is simply skipped, the run keeps training the unedited model, and the plateau window resets so a fresh signal is required before trying again.

Stops as soon as the loss EMA drops below target_loss (after a short warmup), or at max_steps. Returns an AdaptiveTrainingResult with the REAL measured total compute.

Parameters:

• initial_model (Any)

• make_batch (Callable[[int, Generator], tuple[Any, Any]])

• target_loss (float)

• max_steps (int)

• max_layer (int)

• batch_size (int)

• lr (float)

• min_steps_before_edit (int)

• plateau_window (int)

• plateau_eps (float)

• parity_tolerance (float)

• health_lookback (int)

• seed (int)

• controller (StructureEditController | None)

Return type:

AdaptiveTrainingResult