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() – build an isolated candidate with one shared nn.Parameter, measure its output delta, and return it only when the caller’s explicit parity budget accepts it. The source model is never mutated.

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. A fixed-length resampled profile is intentionally lossy: zero distance is a strong discovery signal, but is not an exact permutation certificate. propose_ties() records that stricter fact separately by comparing the full sorted tensors.

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, exact_value_multiset)

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)

• exact_value_multiset (bool)

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, …]

exact_value_multiset

Whether the full tensors contain exactly the same values with the same multiplicities. Unlike a resampled profile comparison, this is an exact permutation certificate.

Type:

bool

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. Only shape/dtype/device-compatible pairs are proposed because a direct parameter tie cannot represent incompatible tensors.

• 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, max_relative_l2, max_abs_tolerance, accepted)

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_relative_l2 (float)

• max_abs_tolerance (float | None)

• accepted (bool)

max_abs_diff

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

Type:

float

relative_l2

||after - before||_2 / ||before||_2. This is infinity when the baseline is identically zero but the candidate is not, and zero only when both are zero.

Type:

float

params_before

Total parameter count before the edit.

Type:

int

params_after

Total parameter count after the edit.

Type:

int

max_relative_l2

Caller-supplied relative-L2 acceptance limit.

Type:

float

max_abs_tolerance

Optional caller-supplied maximum-absolute-error limit.

Type:

float | None

accepted

Whether the candidate satisfied every supplied limit and was returned.

Type:

bool

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', *, max_relative_l2, max_abs_tolerance=None)

Transactionally evaluate a weight tie and return (selected_model, receipt).

The edit is made only on a deep-copied candidate. Both candidate attributes point to one nn.Parameter (an actual tie, not merely equal initialization). The source model and all parameter identities referenced by its optimizer remain untouched. The candidate is returned only when its measured error is within the explicit budget; otherwise the original module object is returned. A caller that adopts an accepted candidate must construct or transactionally migrate an optimizer for that new model. 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. Profile similarity is only a discovery signal; this function gates adoption on a real forward comparison under the caller’s error budget.

Parameters:

• module (Any) – A torch.nn.Module. Parity is measured in evaluation mode, then every module’s prior local training/evaluation flag is restored.

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

• max_relative_l2 (float) – Required finite, non-negative relative-L2 acceptance limit.

• max_abs_tolerance (float | None) – Optional finite, non-negative maximum-absolute-error limit.

Returns:

The isolated candidate when accepted, otherwise the original module, plus the

measured decision receipt.

Return type:

tuple[Any, ParityReceipt]