mixle.experimental.tying_discovery module

Tying discovery via arrangement similarity – the neural half of ConditionalJIT (H4).

The R1 copula note (see the roadmap doc, R1 -> G4, F6, I2, H4) observes that any weight tensor’s flattened value vector v decomposes as v = P . s: a sorted profile s (the empirical quantile function – the values in sorted order) composed with a permutation P (the arrangement – which sorted rank each original position maps to). That is exactly Sklar’s theorem one level down: a joint distribution factors into marginals (here: the profile / value distribution) and a dependence structure (here: the arrangement). This module uses only the marginal half of that decomposition as a discovery signal: two tensors whose PROFILES are close are, value-distribution-wise, “the same numbers, differently arranged” – exactly the situation a weight tie (two module attributes sharing one nn.Parameter, the tensor analogue of this codebase’s keys= mechanism for tying distribution parameters across combinator children, see mixle.stats.combinator) can exploit for a real, measurable parameter reduction.

This module deliberately stays on the “discovery” side of the R1 note. Training a permutation (differentiable OT / Sinkhorn / torchsort) is explicitly out of scope here – profile comparison via sorted-vector L2 distance is closed-form and needs no optimization loop. See profile_distance() for why sorted-vector L2 already is the right 1-D optimal-transport quantity for this purpose.

Workflow:

1. tensor_profile() – extract a fixed-length quantile-function profile from a weight tensor.

2. profile_distance() – an L2 distance between two profiles (= 1-D Wasserstein-2 distance between the corresponding empirical distributions).

3. propose_ties() – rank all pairs of named tensors by profile distance and return the most promising tying candidates.

4. apply_tie() – actually replace two tensors with one shared nn.Parameter and report a measured output-parity receipt (the model is NOT guaranteed function-preserving by a tie; the receipt is the honest, measured delta, not an assumed zero).

Everything here lives in mixle/experimental/ per F7: it graduates out once the mechanism has field mileage and (once merged) should be registered against E0’s graduation ledger – a trivial follow-up, not done here since mixle/experimental/graduation.py does not yet exist on this branch.

tensor_profile(tensor, n_quantiles=256)

Return the fixed-length quantile-function profile of a tensor’s values.

The profile is the empirical quantile function of the tensor’s flattened values: sort the values, then resample that sorted curve onto n_quantiles evenly spaced quantile positions in [0, 1] by linear interpolation. Fixing n_quantiles makes profiles from differently shaped/sized tensors directly comparable (a (64, 64) weight and a (32, 128) weight both reduce to a length-n_quantiles curve).

This throws away the ARRANGEMENT entirely (sorting is exactly “forget the permutation P in v = P . s”) and keeps only the value distribution – by design, since arrangement similarity is not what tying discovery is asking about here: two tensors that hold “the same numbers” but scattered into different spatial slots are still excellent tying candidates once one of them is permuted (or once the model is simply insensitive to which of the two arrangements it uses, which the parity receipt in apply_tie() checks directly).

Parameters:

• tensor (Any) – A torch tensor of any shape (or anything torch.as_tensor accepts).

• n_quantiles (int) – Number of fixed quantile positions to resample the sorted values onto. Must be >= 2. Default 256 is small enough to be cheap and large enough that two profiles from differently-shaped tensors compare fairly.

Returns:

A 1-D torch tensor of length n_quantiles, dtype float32, containing the interpolated quantile function.

Return type:

Any

profile_distance(profile_a, profile_b)

L2 distance between two equal-length quantile-function profiles.

Sorted-vector-vs-sorted-vector L2 distance between two empirical quantile functions IS the (discretized) 1-D Wasserstein-2 distance between the two underlying empirical distributions – a standard, well-known fact worth stating explicitly rather than leaning on silently: for 1-D distributions, the optimal transport coupling between two empirical measures is exactly the sorted-to-sorted (rank-to-rank) matching, so W2(mu, nu) = ||sort(x) - sort(y)||_2 up to the 1/sqrt(n_quantiles) normalization used here. That is also the v = P . s decomposition’s marginal half compared directly: two tensors with identical profiles (profile_distance == 0) hold the same multiset of values up to arrangement, i.e. one is some permutation of the other – the exact condition a weight tie wants.

Parameters:

• profile_a (Any) – A 1-D tensor, as returned by tensor_profile().

• profile_b (Any) – A 1-D tensor of the same length as profile_a.

Returns:

The (root-mean-square-normalized) L2 distance between the two profiles, >= 0.

Return type:

float

class TyingCandidate(name_a, name_b, distance, shape_a, shape_b)

Bases: object

A single proposed weight tie between two named tensors.

Parameters:

• name_a (str)

• name_b (str)

• distance (float)

• shape_a (tuple)

• shape_b (tuple)

name_a

Name of the first tensor (as passed into propose_ties()).

Type:

str

name_b

Name of the second tensor.

Type:

str

distance

Profile distance between the two tensors (lower = more similar = better candidate).

Type:

float

shape_a

Shape of the first tensor.

Type:

tuple[int, …]

shape_b

Shape of the second tensor.

Type:

tuple[int, …]

propose_ties(named_tensors, n_quantiles=256, max_distance=None, top_k=None)

Propose weight-tying candidates by pairwise profile similarity across a set of named tensors.

Computes a tensor_profile() for every tensor, then ranks all pairs by profile_distance() (ascending – most similar first). This is the discovery step of H4: the analogue of scanning distribution combinators for parameters that could share a keys= tag, applied to torch tensors instead.

Parameters:

• named_tensors (dict[str, Any]) – Mapping from a tensor name (e.g. "layer0.head2.q") to the tensor itself. Tensors may have different shapes – profiles fix that via resampling.

• n_quantiles (int) – Passed through to tensor_profile().

• max_distance (float | None) – If given, only pairs with distance <= max_distance are returned. Left unset (None) by default so callers can inspect the full ranked list and pick a threshold.

• top_k (int | None) – If given, only the top_k closest pairs are returned (after max_distance filtering, if any).

Returns:

Candidates sorted by ascending distance (most similar / best tying candidate

first).

Return type:

list[TyingCandidate]

class ParityReceipt(max_abs_diff, relative_l2, params_before, params_after)

Bases: object

Measured output-parity receipt for a function-preserving-edit attempt.

Parameters:

• max_abs_diff (float)

• relative_l2 (float)

• params_before (int)

• params_after (int)

max_abs_diff

Max absolute elementwise difference between pre- and post-edit outputs.

Type:

float

relative_l2

||after - before||_2 / ||before||_2 (0 if before is identically zero).

Type:

float

params_before

Total parameter count before the edit.

Type:

int

params_after

Total parameter count after the edit.

Type:

int

property params_reduced: int

Absolute parameter-count reduction (>= 0 for a real tie; can be 0 if no sharing occurred).

property params_reduced_fraction: float

Fractional parameter-count reduction, in [0, 1].

apply_tie(module, name_a, name_b, inputs, strategy='average')

Apply a proposed weight tie to two nn.Parameter attributes on module and measure the parity cost.

Replaces the two named parameters with ONE shared nn.Parameter (both attributes point at the same tensor object afterward, so a gradient step on either updates both – an actual, literal tie, not merely initializing them equal). strategy picks the shared value:

• "average" (default): elementwise mean of the two original tensors. Requires identical shapes. This is the natural choice when the two tensors are similar-valued-but-not-identical (the realistic case tying discovery is meant to catch, per the H4 acceptance note) – it minimizes the summed squared perturbation to both tensors simultaneously.

• "keep_a": keep name_a’s tensor verbatim and point name_b at it. Useful when one tensor is trusted more (e.g. it trained longer, or name_a is canonical by convention).

Weight tying is NOT assumed to be output-preserving – it is exactly a bet that it will be nearly so, which is what the profile-similarity threshold in propose_ties() is for. This function does not enforce any tolerance itself; it measures and returns the actual delta via a forward pass on inputs before and after the edit, so the caller can decide whether the receipt is acceptable.

Parameters:

• module (Any) – A torch.nn.Module whose forward pass is deterministic given inputs (caller should .eval() it first if it contains dropout/batchnorm-type layers).

• name_a (str) – Dotted attribute path to the first nn.Parameter (e.g. "blocks.0.attn.qkv.weight").

• name_b (str) – Dotted attribute path to the second nn.Parameter. Must have the same shape as name_a.

• inputs (Any) – Whatever module(inputs) accepts; used for the before/after forward pass.

• strategy (str) – "average" or "keep_a"; see above.

Returns:

measured output delta and parameter-count change.

Return type:

ParityReceipt