mixle.models.sigma_weighted_projection module

Sigma-weighted structured projections (roadmap G2): thin solvers over borrowed primitives.

Given a weight matrix W (out_dim x in_dim) and the covariance Sigma (in_dim x in_dim, PSD) of the activations it will actually be multiplied against – e.g. the propagated-law covariance coming out of mixle.models.moment_propagation (roadmap G1) – the objective that matters for preserving downstream behavior is NOT plain Frobenius compression (||W - What||_F^2, which treats every input direction as equally important) but the SIGMA-WEIGHTED version

min_What tr((W - What) @ Sigma @ (W - What)^T)

which penalizes reconstruction error in input directions the real data varies along more, and tolerates more error in directions the data barely explores (the “optimal brain damage” / Fisher-weighted pruning idea, generalized from a diagonal Hessian approximation to a full covariance weighting).

Per the roadmap’s build-vs-borrow note, this module BORROWS the heavy primitives rather than reimplementing them:

  • the low-rank case has a real closed-form solution via a whiten/SVD/un-whiten reduction to plain Eckart-Young truncated SVD (sigma_weighted_low_rank()) – no iterative solver needed;

  • the block-sparse / 2:4 case has no closed form, so sigma_weighted_block_sparse() uses a textbook projected-gradient (“alternating projection”) scheme: a gradient step on the (convex, quadratic-in-What) weighted objective, alternated with a hard projection onto the structural constraint set (a fixed support mask, or a dynamically re-selected 2:4 pattern);

  • the permutation case is solved with Sinkhorn’s algorithm – the standard entropic-OT relaxation of a linear assignment problem. torchsort/POT/geomloss were checked (see the PR description) and are not installed in this environment; Sinkhorn itself is a well-known ~10-line fixed-point iteration (alternating row/column normalization of a Gibbs kernel in log-domain for numerical stability), so it is implemented directly here rather than pulling in a heavy optional dependency for a few lines of numpy. This is also the differentiable “profile . arrangement” building block roadmap item G4 (later, not this item) reuses for permutation x profile quantization.

sigma_weighted_error(w, w_hat, sigma)[source]

tr((W - What) @ Sigma @ (What - W)^T) – the Sigma-weighted reconstruction objective itself.

Used both as the convergence check inside the iterative solvers below and as the metric the tests compare solvers against each other with. Sigma is assumed PSD (a covariance matrix); this function does not itself validate that – callers pass a real covariance (e.g. from mixle.models.moment_propagation.propagate_moments()) or a synthetic A @ A.T construction.

Parameters:
Return type:

float

sigma_weighted_low_rank(w, sigma, rank)[source]

Exact closed-form solver for min_{rank(What)<=rank} tr((W-What) Sigma (W-What)^T).

Derivation (generalized SVD via whitening): for symmetric PSD Sigma with symmetric square root Sigma^(1/2) (Sigma = Sigma^(1/2) Sigma^(1/2), itself symmetric so (Sigma^(1/2))^T = Sigma^(1/2)),

tr((W-What) Sigma (W-What)^T) = tr((W-What) Sigma^(1/2) Sigma^(1/2) (W-What)^T)

= || (W-What) @ Sigma^(1/2) ||_F^2 .

Substituting B = W @ Sigma^(1/2) and Bhat = What @ Sigma^(1/2), the constraint rank(What) <= rank becomes (for full-rank Sigma^(1/2)) exactly rank(Bhat) <= rank, and the objective becomes the PLAIN (unweighted) Frobenius low-rank problem min ||B - Bhat||_F^2, whose exact global optimum is the truncated SVD of B (Eckart-Young). Un-whitening What = Bhat @ Sigma^(1/2)^+ (pseudo-inverse, needed if Sigma is rank-deficient) recovers the optimal What in the ORIGINAL objective. When Sigma has a null space, those input directions contribute nothing to the objective regardless of What’s value there, so the pseudo-inverse un-whitening (which zeroes What on that null space) is one particular optimum among many – still provably attaining the true minimum objective value, which is all the stated objective can see.

This is the SAME closed form used for Fisher/Hessian-weighted low-rank compression (a diagonal special case of this is “optimal brain damage”-style weighted SVD); here it is implemented for a full (non-diagonal) Sigma.

Parameters:
Return type:

ndarray

sigma_weighted_block_sparse(w, sigma, block_pattern_or_2_4, max_iter=200, tol=1e-10)[source]

Alternating-projection solver for min_{What in S} tr((W-What) Sigma (W-What)^T) where S is a structural constraint set with no closed form: a fixed block-sparse/arbitrary support mask, or the 2:4 semi-structured pattern.

block_pattern_or_2_4:

  • the literal string "2:4" – 2:4 semi-structured sparsity (pattern re-selected every step, see _project_2_4());

  • a boolean array shaped like W – an explicit (e.g. block-sparse) fixed support pattern.

Algorithm: projected gradient descent on the (convex, quadratic-in-What) objective – grad_What = -2 (W - What) Sigma – with step size 1 / (2 * lambda_max(Sigma)) (the standard Lipschitz-safe step for a quadratic with Hessian 2*Sigma acting on the right), alternated with a HARD projection onto the structural constraint set after every gradient step. Convergence contract: for a FIXED mask the constraint set is a linear subspace, so this is plain convex projected gradient descent and converges to the GLOBAL optimum of that subspace-constrained problem (matches the closed-form per-row constrained-least-squares solution – see the optimality test). For 2:4 the constraint set is a finite, non-convex union of subspaces (the pattern itself is re-chosen every step), so only convergence to a LOCAL optimum of the alternating scheme is guaranteed – NOT global optimality over all possible 2:4 masks (that combinatorial problem is not attempted here).

Parameters:
Return type:

ndarray

sigma_weighted_permutation(w, sigma, target_profile, temperature=0.1, max_iter=100)[source]

Sinkhorn-based soft-permutation solver for What = P @ target_profile – the “profile o arrangement” pattern (roadmap H4/R1, feeding G4’s permutation x profile quantization): find the ROW-permutation P of a fixed canonical target_profile that best matches W under the Sigma-weighted objective min_P tr((W - P @ target_profile) Sigma (W - P @ target_profile)^T).

This is a linear assignment problem in disguise: it decomposes over ROW-PAIRS (row i of W matched to row j of target_profile) with pairwise cost cost[i,j] = (W_i - profile_j) @ Sigma @ (W_i - profile_j)^T, so the discrete problem min_{P permutation} sum_i cost[i, perm(i)] is EXACTLY a linear assignment problem. We solve it with the differentiable Sinkhorn relaxation the roadmap asks for (a Gibbs kernel K = exp(-cost/temperature), alternately row/column normalized in log-domain – this is the torchsort/POT-style soft-permutation building block G4 later reuses for a jointly-differentiable profile+arrangement objective), THEN round the converged soft doubly-stochastic coupling to a hard permutation via linear-sum-assignment (Hungarian) on the SAME cost matrix – a standard, exact final-rounding step (the Sinkhorn relaxation supplies the differentiable pattern; committing to a hard answer is a separate, exact combinatorial step, not claimed to itself be “the Sinkhorn solution”). Convergence contract: the returned answer is the exact optimum of the linear assignment problem (Hungarian rounding is exact for assignment problems); the SINKHORN PLAN itself only converges to the true permutation in the low-temperature limit – what “converged Sinkhorn solution” means here is that repeated Sinkhorn normalization has converged to a fixed doubly-stochastic coupling for the given temperature, not that temperature itself has been annealed to zero.

Parameters:
Return type:

ndarray

sigma_weighted_butterfly(w, sigma, n_stages=None, n_sweeps=4)[source]

Sigma-weighted BUTTERFLY structured projection: constrain What to be (the top-left d_out x d_in block of) an N x N “butterfly matrix” – a product of L sparse factors, each with exactly 2 nonzeros per row connecting index j to j XOR stride (stride doubling stage to stage: 1, 2, 4, …) – the SAME block-diagonal-then-permute connectivity pattern FFT’s radix-2 decimation uses. This gives O(N log N) free parameters (2 * N per stage, L = log2(N) stages) instead of O(N^2) for a dense matrix, where N is the next power of two >= max(d_out, d_in).

Solved by ALTERNATING LEAST SQUARES over the L stage factors, per the roadmap card’s Steps: reusing the SAME whiten-by-Sigma^(1/2) reduction sigma_weighted_low_rank() uses (via _symmetric_sqrt_and_pinv_sqrt()) to turn each per-stage subproblem into a plain (unweighted) linear least-squares problem in that stage’s 2*N free parameters (closed-form, via numpy.linalg.lstsq()) given every OTHER stage held fixed – a genuine block-coordinate solve, monotonically non-increasing in the Sigma-weighted objective per stage update (see sigma_weighted_error(), the same convergence metric the other three solvers already use).

Two SIMPLIFICATIONS versus a textbook FFT butterfly, both bounded and stated here rather than hidden:

  1. Each stage’s two taps per row are FREE real parameters fit to the data, not fixed unitary FFT twiddle factors – this follows the “butterfly matrices for structured compression” line of work (generalizing FFT’s O(n log n) connectivity to a learnable factorization), not a literal (inverse) Fourier transform.

  2. Rectangular W is handled by zero-padding W/Sigma up to the square N x N problem and reading off the top-left d_out x d_in block at the end. Sigma’s padded rows/columns are zero (those input directions cost nothing, same convention as _symmetric_sqrt_and_pinv_sqrt()’s null-space handling), but padded OUTPUT rows (beyond d_out, when d_out is not already a power of two) are fit toward zero using the SAME shared stage parameters as the real rows – a mild, honest dilution of fitting capacity for non-power-of-two d_out, not a hidden bug.

  3. n_sweeps bounds the number of ALS passes over all L stages rather than iterating to convergence – the “fixed number of butterfly stages/sweeps” bounded-fix simplification the roadmap card allows for. It does not make the family a no-op or fold it into another family: each stage solve is a real, distinct least-squares fit and the returned What has the genuine sparse butterfly parameter count, not a dense low-rank or block-sparse structure.

Parameters:
Return type:

ndarray

class ProjectionReport(structure, sigma_weighted_error, stats=<factory>)[source]

Bases: object

Uniform report shape project() returns alongside What, regardless of structure.

sigma_weighted_error is always the SAME objective (sigma_weighted_error()) every solver in this module already minimizes/reports, computed once here so callers get one consistent number to compare across families. stats carries whatever structure-specific numbers that solver’s own return value already lets a caller compute (rank, sparsity fraction, stage/parameter counts, …) – nothing new is invented here, this just wraps numbers each solver already makes derivable.

Parameters:
project(w, sigma, structure, **kw)[source]

Unified front door for roadmap G2’s four structure families: structure in {"low_rank", "block_sparse", "butterfly", "perm_profile"}.

Dispatches to this module’s existing standalone solvers (sigma_weighted_low_rank(), sigma_weighted_block_sparse(), sigma_weighted_butterfly(), sigma_weighted_permutation()) – this function does not reimplement any solver, it only picks one by name and wraps its result (plus the shared sigma_weighted_error() metric and a few structure-specific stats already derivable from that result) into a single ProjectionReport shape, so callers that want to pick a structure by string (e.g. a search/schedule over structures) do not need a per-family if/elif of their own.

**kw per structure (forwarded to the underlying solver; see each solver’s docstring for details):

  • "low_rank": rank (int, required).

  • "block_sparse": pattern ("2:4" or a boolean mask shaped like W; required), optional max_iter, tol.

  • "butterfly": optional n_stages, n_sweeps.

  • "perm_profile": target_profile (required), optional temperature, max_iter.

Parameters:
Return type:

tuple[ndarray, ProjectionReport]