mixle.models.coarsening module

Coarsening operator R with per-scale receipts (roadmap G3): depth grouping plus structure projection 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 an experimental width-representation helper); 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 implemented model moves and representation helper

  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. Experimental width representation (experimental_width_merge_representation()): constructs a residual-stream map 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) 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.

    This helper does not rewire a model and is not invoked by coarsen().

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

Depth grouping and structure projection are applied by coarsen() under a divergence BUDGET and a local TRUST REGION. The width representation is deliberately separate until normalization, attention, embedding, and output-head rewiring can produce a genuine narrower executable model.

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, divergence_basis='full_rank_gaussian', regularized_kl_divergence=None, assumptions=(), sensitivity=<factory>, scope='local_candidate', artifact_digest=None)[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)

  • divergence_basis (str)

  • regularized_kl_divergence (float | None)

  • assumptions (tuple[str, ...])

  • sensitivity (dict[str, float])

  • scope (str)

  • artifact_digest (str | None)

class ProjectionReceipt(name, mode, sigma_weighted_error, accepted=True, final_kl_divergence=None, artifact_digest=None)[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:

  • name (str)

  • mode (str)

  • sigma_weighted_error (float)

  • accepted (bool)

  • final_kl_divergence (float | None)

  • artifact_digest (str | None)

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 experimental_width_merge_representation()). 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 CoarseningMetrics(source_parameter_count, result_parameter_count, source_parameter_bytes, result_parameter_bytes, source_manifest_entries, result_manifest_entries, source_leaf_blocks, result_leaf_blocks, source_estimated_block_evaluations, result_estimated_block_evaluations, latency_seconds=None, latency_basis='unmeasured: no calibration input was supplied')[source]

Bases: object

Structural measurements for the source and returned executable models.

estimated_block_evaluations counts ordinary block-branch evaluations in one forward: a plain block costs one and the current numerical Taylor merge costs four. FLOPs and latency are deliberately not fabricated without shapes or calibration inputs.

Parameters:

  • source_parameter_count (int)

  • result_parameter_count (int)

  • source_parameter_bytes (int)

  • result_parameter_bytes (int)

  • source_manifest_entries (int)

  • result_manifest_entries (int)

  • source_leaf_blocks (int)

  • result_leaf_blocks (int)

  • source_estimated_block_evaluations (int)

  • result_estimated_block_evaluations (int)

  • latency_seconds (float | None)

  • latency_basis (str)

class CoarsenResult(model, receipt_map=<factory>, accepted_pairs=<factory>, rejected_pairs=<factory>, depth_only_kl=0.0, total_kl=0.0, budget=inf, trust_region=inf, within_budget=True, structure_receipts=<factory>, artifact_digest=None, metrics=None)[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 third move’s own receipt list – see _narrow_block_mlp2() – 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. The final_artifact scale receipt is authoritative for the returned model: its end-to-end law, total_kl, and digest are computed after every accepted structure projection.

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]

experimental_width_merge_representation(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.

This is explicitly a representation helper, not a model transformation: it does not rewire transformer weights, normalization, attention heads, or output heads and therefore does not return or claim a smaller executable model.

receipt.kl_divergence reports the true support result. It is infinite for a genuine rank reduction because the full-rank teacher assigns mass outside the rank-deficient round-trip support. For callers that explicitly assume isotropic downstream noise, regularized_kl_divergence and the receipt’s sensitivity values report that separate surrogate.

Parameters:

  • target_width (int)

  • input_law (MultivariateGaussianDistribution)

  • temperature (float)

  • n_iter (int)

Return type:

tuple[WidthMergeRepresentation, ScaleReceipt]

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

Compatibility wrapper for the experimental representation helper.

model is intentionally unused: this function does not transform a model. New code should call experimental_width_merge_representation(), whose name makes that boundary explicit.

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