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 is a first-order (Jacobian) linearization – the same “delta method” used for LayerNorm’s covariance push-forward.

  • 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 contract (streaming / layer-local / constant memory)

propagate_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 rather than materializing everything at once. Concretely, at any point during the walk only (a) the CURRENT running law (mu, Sigma) – an O(d_model^2) object – and (b) the ONE block currently being processed are resident; once a block’s output law and closure-error receipt are recorded, its intermediate quantities (the attention MGF terms, the GELU Jacobian, the local Monte Carlo samples used for the receipt, etc.) are dropped. This is the moment-propagation analogue of a real forward pass that would otherwise have to materialize per-layer activations for the whole depth of the network simultaneously (or lean on gradient checkpointing to avoid it, as CausalLM.gradient_checkpointing does for the real module). Peak memory therefore does not scale with network depth – only with d_model – which is verified directly in mixle/tests/moment_propagation_test.py.

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

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

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)[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 so the marginal moments stay exact even though the correlation structure is linearized.

Parameters:

law (MultivariateGaussianDistribution)

Return type:

tuple[MultivariateGaussianDistribution, ndarray]

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]

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

Streaming, layer-local, constant-memory Gaussian-law propagation through a real mixle.models.transformer.CausalLM.

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

Execution contract: this walks model.blocks one block at a time, then the final model.ln and model.head. At each step only the CURRENT running law and the layer being processed are resident; see the module docstring for the full memory-contract discussion and its relation to mixle.inference.precision_plan.recommend_compute_precision()’s inspect-then-decide walking pattern. Verified empirically (peak memory vs. depth) in mixle/tests/moment_propagation_test.py.

Returns a list of LayerMoments, one per block plus one for the final LayerNorm and one for the (weight-tied) output head, in execution order.

Parameters:

  • model (Any)

  • input_law (MultivariateGaussianDistribution)

  • n_mc (int)

  • seq_len (int)

  • seed (int)

Return type:

list[LayerMoments]