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 an exact linear-assignment solve over profile rows. It intentionally does not claim to expose an entropic-transport relaxation or differentiable Sinkhorn object.

sigma_weighted_error(w, w_hat, sigma)[source]

tr((W - What) @ Sigma @ (W - What)^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.

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)[source]

Exact row-permutation solver for What = P @ target_profile.

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. The returned hard permutation is therefore computed directly with the Hungarian algorithm. This function exposes no temperature or iteration controls because it does not compute or return an entropic soft plan.

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:

ButterflyProjection

class ButterflyProjection(a, b, strides, d_out, d_in, n)[source]

Bases: object

Compact executable product of sparse butterfly factors.

apply(x) accepts vectors or batches whose last dimension is d_in and returns the equivalent of x @ projection.to_dense().T without materializing an N x N matrix. projection @ x provides ordinary matrix-multiplication orientation for arrays whose first dimension is d_in. NumPy consumers that require a dense matrix remain compatible through np.asarray(projection).

Parameters:
property shape: tuple[int, int]

Logical dense matrix shape.

property parameter_count: int

Number of stored floating-point factor parameters.

property parameter_nbytes: int

Physical bytes occupied by stored floating-point factors.

property serialized_nbytes: int

Actual protocol-5 pickle size of the executable representation.

apply(x)[source]

Apply the projection to vectors/batches stored along their final axis.

Parameters:

x (Any)

Return type:

ndarray

to_dense()[source]

Materialize the logical dense matrix only when a caller explicitly requests it.

Return type:

ndarray

exception CovarianceAdjustmentWarning[source]

Bases: RuntimeWarning

A covariance had only roundoff-scale negative eigenvalues and was projected to the PSD cone.

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).

Parameters:
Return type:

tuple[Any, ProjectionReport]