mixle.models.moment_propagation module

Gaussian(-mixture) law propagation through the real causal transformer in mixle.models.transformer.

This is a surrogate: instead of running a forward pass on concrete activations, it pushes a Gaussian LAW x ~ N(mu, Sigma) (representing the distribution of a token’s residual-stream vector under some input distribution) analytically through each layer type the real CausalLM is built from – nn.Linear, nn.LayerNorm, nn.GELU, and mixle.models.transformer.CausalAttention – and returns the propagated law at every layer together with a genuine, locally-computed “closure error” receipt: how far the closed-form law is from a small Monte Carlo sample pushed through the REAL torch layer at that point.

Propagated laws are represented with this codebase’s own mixle.stats.multivariate.multivariate_gaussian.MultivariateGaussianDistribution (mean + full covariance), not a parallel numpy-only representation.

What is exact vs. approximate

  • Linear (y = Wx + b): exact. A Gaussian pushed through an affine map is exactly Gaussian.

  • Attention: the output-given-query map is derived to be exactly affine in the query (see attention_law()), via the MGF identity for jointly-Gaussian (K, V). Composed with the exactly-linear qkv/proj projections, the whole attention branch is an exact affine function of its LayerNorm’d input – conditional on treating the key/value population as a single stationary Gaussian (see caveats in attention_law()).

  • LayerNorm: nonlinear and data-dependent; propagated via a first-order Taylor (“re-anchoring”) expansion of the true LayerNorm map around the input mean (see layernorm_law() for the closed-form Jacobian and documented failure modes).

  • GELU: the first two output MOMENTS (mean and per-dimension variance) are exact closed-form expressions in (mu, sigma) (derived via Stein’s lemma + the bivariate normal CDF, see gelu_law()); the OFF-diagonal output covariance starts from a first-order (Jacobian) linearization – the same “delta method” used for LayerNorm’s covariance push-forward – then its approximate correlations are projected to the PSD cone under fixed exact marginal variances, with an explicit correction receipt.

  • Residual connections (x + branch(x)): the two summands are correlated (both are functions of the same x), which the propagation accounts for through a chained JACOBIAN of the branch mapping (composed from the exact/linearized per-layer Jacobians above), giving Cov(x, branch(x)) ~= Sigma_x @ J_branch^T and hence an (approximately) correct Sigma_out = Sigma_x + Sigma_branch + Cov + Cov^T.

Execution and retained-memory contract

iter_moments() mirrors the walking pattern of mixle.inference.precision_plan.recommend_compute_precision(): it inspects the model’s structure and processes it piece by piece and yields each receipt before advancing. Its propagation working set is the current full law plus one block’s intermediates, while closure probes are processed in bounded batches and reduced online rather than materialized as n_mc × seq_len × d_model tensors. Excluding the model itself and any receipts retained by the caller, iterator memory is independent of depth and bounded by O(d_model² + probe_batch_size × seq_len × d_model) (with analogous output-width terms at the head).

propagate_moments() is the convenience materializer: it consumes that iterator into a list and therefore intentionally retains one full mean/covariance law per layer. Its returned-memory cost is O(depth × d_model²) plus the final head law. Use iter_moments() when depth-independent retained memory is required.

References

  • Stein’s lemma (Gaussian integration by parts): E[(X-mu) g(X)] = sigma^2 E[g'(X)] for X ~ N(mu, sigma^2) – used throughout to get closed forms for E[GELU(X)] and its derivative.

  • The Gaussian-product / MGF identity E[Y e^{t^T X}] = M_X(t) (mu_Y + Sigma_YX t) for jointly Gaussian (X, Y) – the exact identity behind attention_law().

  • Data-free quantization (DFQ) BatchNorm-based calibration is the closest prior art for the LayerNorm “re-anchoring” step: both re-derive a cheap closed-form summary of what a normalization layer does to a law, without touching real data.

class LayerMoments(index, name, law, closure_error, diagnostics=<factory>)[source]

Bases: object

One layer’s propagated law plus its locally-computed closure-error receipt.

Parameters:

  • index (int)

  • name (str)

  • law (MultivariateGaussianDistribution)

  • closure_error (float)

  • diagnostics (dict[str, Any])

class GeluCovarianceReceipt(correction_frobenius, min_eigenvalue_before, min_eigenvalue_after, iterations, marginal_variance_max_error)[source]

Bases: object

Measured correction applied to keep linearized GELU correlations PSD at exact marginal variances.

Parameters:

  • correction_frobenius (float)

  • min_eigenvalue_before (float)

  • min_eigenvalue_after (float)

  • iterations (int)

  • marginal_variance_max_error (float)

linear_law(law, weight, bias=None)[source]

Exact Gaussian push-forward of y = Wx + b.

x ~ N(mu, Sigma) => y ~ N(W mu + b, W Sigma W^T) exactly – no approximation. Returns the new law together with the Jacobian dy/dx = W (used to chain residual cross-covariances).

Note on out.covar when weight maps to MORE output dimensions than law.covar’s rank (e.g. an expanding projection like the qkv weight in attention_law()): W Sigma W^T is then mathematically rank-deficient (PSD, not PD), and the returned law – a MultivariateGaussianDistribution, which self-heals a non-PD covariance with a small trace-scaled diagonal jitter so it stays Cholesky-factorizable for log_det/inv_covar/sampling, see _robust_cho_factor in mixle.stats.multivariate.multivariate_gaussian – may therefore differ from the exact W Sigma W^T by that jitter (diagonal-only, ~1e-10 * trace/n). This is NOT an error in the formula above – verified bit-exact whenever the output is genuinely PD (e.g. square invertible weight); see mixle/tests/moment_propagation_test.py::LinearExactTest for both regimes – it is the unavoidable cost of representing an exactly-singular covariance with a factorization that requires one.

Parameters:

  • law (MultivariateGaussianDistribution)

  • weight (ndarray)

  • bias (ndarray | None)

Return type:

tuple[MultivariateGaussianDistribution, ndarray]

layernorm_law(law, weight, bias, eps=1e-5)[source]

Propagate a Gaussian law through LayerNorm: y = weight * (x - m(x)) / sqrt(v(x) + eps) + bias, where m(x) = mean_d(x) and v(x) = mean_d((x - m(x))^2) are the PER-SAMPLE (per-token) statistics LayerNorm computes over the feature axis – this is the nonlinear, data-dependent step in the block.

Derivation (“re-anchoring”)

LayerNorm has no closed-form pushforward of a full law in general (m(x) and v(x) are themselves random, nonlinear functions of x). We use a first-order Taylor expansion around the CURRENT mean mu, i.e. anchor the (unknown, per-sample) normalization statistics at their EXPECTED values under the current law, not merely their values evaluated at the mean vector. m(x) = mean_d(x) is itself LINEAR in x, so E[m(x)] = mean_d(mu) exactly – no correction needed there. But v(x) = mean_d((x - m(x))^2) is QUADRATIC in x, so evaluating it at mu alone (v(mu) = mean((mu - m)^2)) drops a systematic bias term: writing P = I - (1/d) 11^T for the feature-centering projector (symmetric, idempotent), v(x) = (1/d) x^T P x, and the standard quadratic-form expectation identity gives

E[v(x)] = v(mu) + (1/d) trace(P @ Sigma) = v(mu) + (1/d) (trace(Sigma) - (1/d) sum(Sigma)).

The second term is the “spread of x around its own per-token mean” contribution that v(mu) alone misses entirely; it is NOT a higher-order correction that can be dropped once Sigma is non-negligible relative to d – for small d_model (e.g. an 8-wide toy model) it routinely dominates v(mu), which without this correction makes the anchored std = sqrt(v(mu) + eps) far too small and blows the propagated mean/covariance up by an order of magnitude relative to the true LayerNorm output law. We therefore re-anchor at the BIAS-CORRECTED v = E[v(x)] above (still a first-order/delta-method treatment of the covariance push-forward – only the anchor point for v is corrected, not the linearization itself). This is the LLM analogue of BatchNorm-based data-free-quantization (DFQ) calibration, which likewise re-derives cheap closed-form layer statistics (there, running mean/var; here, the expected LayerNorm mean/var under the CURRENT propagated law) without touching real data.

The mean is propagated by evaluating the true (nonlinear) LayerNorm map at mu but with the bias-corrected v:

mu_out = weight * (mu - m) / sqrt(v + eps) + bias.

The covariance is propagated via the JACOBIAN of LayerNorm evaluated at mu with the same bias-corrected v (a standard, textbook LayerNorm-backward-style derivative applied at the corrected anchor):

d(norm_i)/d(x_j) |_{x=mu} = (1/sqrt(v+eps)) * (delta_ij - 1/d - (mu_i - m)(mu_j - m) / (d*(v+eps))) J_ij = weight_i * d(norm_i)/d(x_j) Sigma_out ~= J Sigma J^T.

Known failure modes (feeds the closure-error receipt)

  • Small ``d_model``: m(x) and v(x) are averages over only d_model features, so their sample-to-sample fluctuation around their (now bias-corrected) expectation – which this delta-method approximation still ignores, since the Jacobian itself is frozen at the anchor – is large relative to their magnitude when d_model is small. The approximation is progressively worse as d_model shrinks, even after the mean-bias correction above.

  • Heavy-tailed pre-norm activations: the first-order Taylor expansion is only locally valid; if the input law has high kurtosis / is far from Gaussian in practice (despite being MODELED as Gaussian here), the true per-sample (m, v) can swing far from their Gaussian-law expectations, and the linearization degrades.

  • A deep stack of blocks compounds both effects: even a small per-layer LayerNorm error can accumulate across n_layer re-anchoring steps.

Parameters:
Return type:

tuple[MultivariateGaussianDistribution, ndarray]

gelu_law(law, *, return_receipt=False)[source]

Propagate a Gaussian law through elementwise GELU using the closed-form per-dimension moments in _gelu_scalar_moments().

The per-dimension MEAN and VARIANCE are exact closed-form expressions (no Monte Carlo, no approximation of the GELU functional form). The OFF-diagonal output covariance – Cov(GELU(x_i), GELU(x_j)) for i != j – has no simple closed form (it needs the joint bivariate distribution of (x_i, x_j) through a nonlinearity) and is instead approximated by the standard delta-method / Price’s-theorem-style linearization Cov(y_i, y_j) ~= J_ii * J_jj * Cov(x_i, x_j) with J_ii = d E[GELU(x_i)]/d mu_i the EXACT mean-derivative from Stein’s lemma. The diagonal is then overwritten with the exact closed-form variance. Because that splice need not be PSD, correlations are projected to the PSD cone under a fixed-unit-diagonal constraint and then rescaled by the exact marginal standard deviations. This preserves the scalar variances while correcting only the approximate correlation structure.

Set return_receipt=True to receive a GeluCovarianceReceipt measuring the correction, pre/post minimum correlation eigenvalues, iteration count, and marginal-variance preservation error.

Parameters:

  • law (MultivariateGaussianDistribution)

  • return_receipt (bool)

Return type:

tuple[MultivariateGaussianDistribution, ndarray] | tuple[MultivariateGaussianDistribution, ndarray, GeluCovarianceReceipt]

attention_law(law, qkv_weight, qkv_bias, proj_weight, proj_bias, n_head)[source]

Propagate a Gaussian law through one mixle.models.transformer.CausalAttention layer.

Modeling assumption: the key/value population attended over (across sequence positions) is treated as a single STATIONARY Gaussian – the same joint law as the query’s – rather than tracking per-position laws. This is the “R2 MGF” population-level approximation the roadmap specifies; it does not model the causal mask explicitly (a real causal mask makes early positions attend to a smaller, non-stationary population). That mismatch is a known limitation, checked directly against Monte Carlo softmax attention in mixle/tests/moment_propagation_test.py (tightest for a roughly-stationary population, looser for strongly non-stationary / short sequences).

Derivation

For jointly Gaussian (K, V) (a key vector and its associated value vector from the SAME token) and a fixed query q, the requested MGF identity is

E[exp(q^T K / sqrt(d)) V] = exp(q^T mu_K/sqrt(d) + 0.5 q^T Sigma_KK q / d) * (mu_V + Sigma_VK q / sqrt(d))

which is the standard “E[Y e^{t^T X}] = M_X(t) (mu_Y + Sigma_YX t)” identity for jointly Gaussian (X, Y) with t = q / sqrt(d), X = K, Y = V. The attention DENOMINATOR (the softmax normalizer) is exactly the same MGF evaluated with V replaced by the constant 1:

E[exp(q^T K / sqrt(d))] = exp(q^T mu_K/sqrt(d) + 0.5 q^T Sigma_KK q / d).

Both share the identical exponential prefactor, so it CANCELS in the ratio:

softmax-attention-output(q) ~= E[exp(q^T K/sqrt(d)) V] / E[exp(q^T K/sqrt(d))] = mu_V + (Sigma_VK / sqrt(d)) q.

This is a remarkably clean result: the MGF-approximated attention output, as a function of the query, is EXACTLY AFFINE in q (no exponential term survives). Since q itself has a propagated Gaussian law (the marginal of the joint (Q, K, V) law after the qkv projection), pushing that law through this affine map is exact (reusing linear_law()’s formula) – there is no additional approximation beyond the population-stationarity assumption above and (for the covariance) the affine relation being evaluated with the Sigma_VK estimated from the SAME joint law used for the mean.

Per head, per token position, y_h(q) = mu_{V,h} + (Sigma_{VK,h} / sqrt(d_head)) q_h. Stacking heads into a block-diagonal map A (each head only reads its own query slice) and reusing the FULL (cross-head) query covariance from the qkv projection gives the cross-head covariance of the output “for free” – cross-head correlation in Q (already present in Sigma_QQ’s off-diagonal blocks) propagates through A Sigma_QQ A^T even though A itself only mixes within a head.

Parameters:
Return type:

tuple[MultivariateGaussianDistribution, ndarray]

iter_moments(model, input_law, n_mc=128, seq_len=16, seed=0, probe_batch_size=32)[source]

Yield layer-local Gaussian laws without retaining earlier receipts.

input_law models the distribution of a token’s residual-stream vector ENTERING the block stack (i.e. after token + position embedding) – an N(mu, Sigma) over R^{d_model}.

The iterator walks model.blocks one at a time, then the final model.ln and model.head. Closure probes execute in the target module’s device/dtype and evaluation mode, restore every submodule’s prior mode, and reduce batches of at most probe_batch_size online.

Parameters:

  • model (Any)

  • input_law (MultivariateGaussianDistribution)

  • n_mc (int)

  • seq_len (int)

  • seed (int)

  • probe_batch_size (int)

Return type:

Iterator[LayerMoments]

propagate_moments(model, input_law, n_mc=128, seq_len=16, seed=0, probe_batch_size=32)[source]

Materialize iter_moments() into a list of per-layer full-covariance receipts.

Unlike the iterator, this convenience API intentionally retains O(depth × d_model²) returned law state. It remains useful when callers need random access to every layer; callers requiring depth-independent retained memory should consume iter_moments().

Parameters:

  • model (Any)

  • input_law (MultivariateGaussianDistribution)

  • n_mc (int)

  • seq_len (int)

  • seed (int)

  • probe_batch_size (int)

Return type:

list[LayerMoments]