mixle.models.coarsening module

Coarsening operator R with per-scale receipts (roadmap G3): depth-merge + width-merge + structure- projection, iterated under a divergence budget and a trust region, over the real transformer in mixle.models.transformer.

Build vs. borrow: this module builds only what the landscape check found unoccupied for G3 itself (the depth-merge Taylor-composition machinery and the width-merge OT-based near-duplicate pairing); it BORROWS everything else –

  • the Gaussian LAW representation and per-layer propagation primitives (mixle.models.moment_propagation, roadmap G1) – linear_law, layernorm_law, gelu_law, attention_law, and G1’s own per-block closure-error receipt (_closure_error_block);

  • the “structure-projection” move itself, which is exactly roadmap G2 (mixle.models.sigma_weighted_projection) called directly, not reimplemented.

The three moves

  1. Depth-merge (depth_merge()): folds two adjacent Blocks x -> x + f(x) and x -> x + g(x) into one merged block via a SECOND-ORDER Taylor approximation of their residual-flow composition x -> x + f(x) + g(x + f(x)):

    g(x + f(x)) ~= g(x) + Dg(x)[f(x)] + O(||f(x)||^2)

    so the merged branch is h(x) = f(x) + g(x) + Dg(x)[f(x)], accurate to second order in the (typically small) per-block residual magnitude – f and g themselves are NOT linearized (both keep their full real attention/LayerNorm/GELU nonlinearity); only the CROSS-TERM introduced by composing them is approximated, which is exactly the “residual is a small perturbation” regime a pre-norm residual stack is designed to live in. At the LAW level this is computed analytically by chaining G1’s own per-branch Jacobians (see _block_branch()), which is also literally how the closed-form per-scale receipt below is obtained – both the teacher (exact sequential G1 propagation through both blocks) and the student (the merged, second-order approximation) end up as Gaussian laws, so their divergence is a KNOWN CLOSED FORM (gaussian_kl()), not an estimate. At the REAL forward-pass level (for actual token sequences, not laws), MergedBlock evaluates the identical algebraic expression using a genuine, per-input Jacobian-vector product (not a single frozen linearization anchor).

  2. Width-merge (width_merge()): reduces the residual-stream width d_model -> target_width by finding near-duplicate directions of the (Sigma-weighted) residual-stream covariance – the same “functionally near-duplicate, once permutation-aligned, can be merged/averaged” idea as neuron-permutation (“git re-basin”) symmetries – via an entropic-OT (Sinkhorn) plan, then projecting down. G2’s own sigma_weighted_permutation() was checked first (see its docstring discussion below in width_merge()) but solves a different-shaped problem (aligning two SAME-shape weight matrices via a square permutation against a fixed target_profile), not the many-to-few d -> target_width reduction needed here, so a small companion RECTANGULAR Sinkhorn is implemented locally, reusing the identical log-domain fixed-point structure G2 uses for its square case.

  3. Structure-projection (structure_project()): a thin wrapper directly around G2’s sigma_weighted_low_rank() / sigma_weighted_block_sparse() – no reimplementation.

These are iterated by coarsen() under a divergence BUDGET (stop once the accumulated closed-form KL between teacher and student exceeds it) and a local TRUST REGION (any individual merge whose own local KL exceeds the trust region is rejected and the original blocks are kept instead).

H1 is this operator inverted

Roadmap H1 (growth operators, not built here) is the natural INVERSE of coarsen(): instead of folding two blocks into one under a divergence budget, it would SPLIT one block into two (or widen d_model) under a capacity/EIG budget, re-using the exact same closed-form Gaussian-law receipt machinery in reverse – gaussian_kl() doesn’t care which direction the model size changes, and ScaleReceipt already records both a teacher and a student law symmetrically enough that swapping which one is called “teacher” is the whole difference between coarsening and growing. Concretely, a hypothetical depth_split(block, budget) would invert the linearization here: given a merged block’s branch Jacobian J_h, find an (f, g) pair whose second-order composition reconstructs h to within budget – the same receipt formula, run backwards. Nothing in this module’s interfaces (plain (law) -> (representation, receipt) functions, laws as ordinary MultivariateGaussianDistribution objects) assumes the direction of size change, which is deliberate.

class ScaleReceipt(name, teacher_law, student_law, kl_divergence, surrogate_closure_error, accepted=True)[source]

Bases: object

One per-scale receipt: the CLOSED-FORM teacher/student divergence at this coarsening step, plus (separately) G1’s own closure-error signal for how much the Gaussian-surrogate assumption itself is trusted at this point in the network (nan where no G1 block closure applies, e.g. width-merge, which never runs a real Block forward and so has nothing for G1’s Monte-Carlo closure check to compare against).

Parameters:

  • name (str)

  • teacher_law (MultivariateGaussianDistribution)

  • student_law (MultivariateGaussianDistribution)

  • kl_divergence (float)

  • surrogate_closure_error (float)

  • accepted (bool)

class ProjectionReceipt(name, mode, sigma_weighted_error)[source]

Bases: object

Receipt for a structure_project() call – reports G2’s own Sigma-weighted reconstruction error directly (there is no Gaussian law on either side of a weight-space projection, so this is not a KL divergence; it is the SAME sigma_weighted_error objective G2’s solvers themselves minimize).

Parameters:
class WidthMergeRepresentation(merge, unmerge, target_width, d_model)[source]

Bases: object

Data-free width-reduction representation: a (target_width, d_model) merge operator and its (d_model, target_width) (pseudo-inverse) reconstruction, built from an entropic-OT near-duplicate pairing of residual-stream coordinates (see width_merge()). Kept as an explicit linear map rather than folded into new per-layer weight matrices – conjugating every qkv/proj/mlp weight in the real model by this map is a real but separable engineering step this representation is designed to make straightforward (W_new = merge @ W @ unmerge for a weight whose BOTH axes are d_model, W_new = W @ unmerge / merge @ W for one-sided cases), left to callers that need an actually smaller CausalLM.

Parameters:
class CoarsenResult(model, receipt_map=<factory>, accepted_pairs=<factory>, rejected_pairs=<factory>, total_kl=0.0, budget=inf, trust_region=inf, within_budget=True, structure_receipts=<factory>)[source]

Bases: object

Output of coarsen(): the new (shallower) model, the full per-scale receipt map, and the bookkeeping needed to see exactly which merges were accepted vs. rejected and why.

structure_receipts is the (separate, additive) third move’s own receipt list – see _narrow_block_linears() – kept OUT of receipt_map deliberately: receipt_map values are ScaleReceipt (closed-form KL against a Gaussian law, consumed as-is by hybrid’s surrogate_closure_error-keyed stage ranking in mixle.models.compress), while structure- projection’s own receipt is a ProjectionReceipt (a Sigma-weighted reconstruction error, not a KL) – mixing the two dataclasses into one dict would silently break that attribute lookup.

Parameters:
gaussian_kl(p, q)[source]

Closed-form KL(p || q) for two multivariate Gaussians – the standard textbook formula

``KL(p||q) = 0.5 * ( tr(Sigma_q^-1 Sigma_p) + (mu_q - mu_p)^T Sigma_q^-1 (mu_q - mu_p)

  • k + ln(det Sigma_q / det Sigma_p) )``

computed ANALYTICALLY, not via Monte Carlo – both p (the “teacher” law) and q (the “student” law) are already MultivariateGaussianDistribution objects, which cache inv_covar and log_det from a (self-healing) Cholesky factorization at construction time, so this reuses those cached quantities directly rather than re-deriving them. Clipped at 0 to absorb float round-off on (near-)identical laws (KL is exactly 0 there, mathematically).

Parameters:

  • p (MultivariateGaussianDistribution)

  • q (MultivariateGaussianDistribution)

Return type:

float

depth_merge(block_a, block_b, input_law, n_mc=64, seed=0)[source]

Fold two adjacent Blocks into one via the second-order Taylor composition documented at module level.

Returns (merged_block, receipt) where merged_block is a real, forward-passable MergedBlock and receipt is a ScaleReceipt whose teacher_law/student_law are the EXACT-per-G1 sequential composition (block_a then block_b, propagated exactly as mixle.models.moment_propagation.propagate_moments() would) vs. the second-order MERGED composition, both Gaussian, so kl_divergence is the closed-form gaussian_kl() between them – the receipt for this individual (local) merge step, i.e. what a caller’s TRUST REGION check compares against.

Parameters:

  • block_a (Any)

  • block_b (Any)

  • input_law (MultivariateGaussianDistribution)

  • n_mc (int)

  • seed (int)

Return type:

tuple[Any, ScaleReceipt]

width_merge(model, target_width, input_law, temperature=0.1, n_iter=200)[source]

Reduce the residual-stream width from d_model to target_width by pairing near-duplicate coordinates of the (Sigma-weighted) residual-stream covariance and merging/averaging them.

sigma_weighted_permutation (G2) was checked first per the roadmap note (see the module docstring): it solves min_P tr((W - P @ target_profile) Sigma (W - P @ target_profile)^T) for a SQUARE permutation P matching two SAME-shape objects (W against a fixed target_profile) – the classic one-to-one “git re-basin” alignment. Width reduction needs a genuinely MANY-TO-FEW map (d_model -> target_width, generally target_width < d_model so there is no permutation at all, square or otherwise), so it is not directly reusable here; _rectangular_sinkhorn() reuses the SAME log-domain Sinkhorn fixed-point idea for the rectangular marginals this problem actually has, rather than pulling in a separate heavy OT solver.

Data-free: the only input is input_law.covar (the propagated residual-stream covariance from G1), used to build a correlation-distance cost cost[i, j] = Sigma[i,i] + Sigma[a_j,a_j] - 2*Sigma[i, a_j] between every source coordinate i and target_width anchor coordinates a_j (the highest-variance coordinates, chosen as informative anchors) – cost[i, j] is exactly Var(x_i - x_{a_j}), so a near-zero cost means coordinate i is functionally redundant with anchor a_j and should be merged into it. The resulting Sinkhorn plan, column-normalized into convex combinations, is the merge operator; its pseudo-inverse is the reconstruction (“unmerge”) map.

Returns (representation, receipt) where receipt.teacher_law is input_law itself and receipt.student_law is input_law round-tripped through merge-then-unmerge, both d_model- dimensional so gaussian_kl() applies directly as the (closed-form) width-merge receipt.

Parameters:

  • model (Any)

  • target_width (int)

  • input_law (MultivariateGaussianDistribution)

  • temperature (float)

  • n_iter (int)

Return type:

tuple[WidthMergeRepresentation, ScaleReceipt]

structure_project(weight, sigma, mode='low_rank', rank=None, pattern='2:4')[source]

Thin wrapper calling G2’s mixle.models.sigma_weighted_projection solvers directly – NOT a reimplementation, per the roadmap’s build-vs-borrow note. mode="low_rank" calls sigma_weighted_low_rank() (requires rank); mode="block_sparse" calls sigma_weighted_block_sparse() (uses pattern, either the literal "2:4" or an explicit boolean mask, exactly as G2 documents).

Parameters:
Return type:

tuple[ndarray, ProjectionReceipt]

coarsen(model, budget, trust_region, input_law, n_mc=64, seed=0)[source]

The iterated coarsening operator R: walk model.blocks pairwise, attempting a depth_merge() at each adjacent pair. A merge is ACCEPTED only if BOTH hold:

  • TRUST REGION – its own LOCAL closed-form KL (receipt.kl_divergence) is at most trust_region;

  • BUDGET – accepting it would not push the ACCUMULATED KL (summed over all accepted merges so far) past budget.

A rejected (or budget-exhausted) pair is left UNMERGED – both original blocks are kept, and the running law is propagated through them individually (via G1’s own per-layer laws, reusing _block_branch() plus the outer residual add) so later merge attempts still see the correct running law regardless of whether earlier pairs were merged. This makes the whole pass data-free: the running “receipt map” is built entirely from propagated LAWS, never real data.

Returns a CoarsenResult wrapping a new, real, forward-passable CoarsenedLM (so a caller can still measure REAL per-layer error against the original model by literally running both models on sampled token sequences – see mixle/tests/coarsening_test.py).

Parameters:

  • model (Any)

  • budget (float)

  • trust_region (float)

  • input_law (MultivariateGaussianDistribution)

  • n_mc (int)

  • seed (int)

Return type:

CoarsenResult