Expand description
Pooled SSM-state moments (StateMoments) for state participation ratios.
§Pooled SSM-state moments (for state participation ratios)
StateMoments carries the exact first and second moments of every
per-token SSM state hₜ ∈ ℝ^{per_head_dim × state_rank} of a forward
pass, treating each state row (one (token, per_head_dim) pair) as a
sample in ℝ^{state_rank} — the same sample convention a token-by-token
step loop reading the cache would produce.
The moments are all a participation ratio (PR, a differentiable effective
rank) needs: with Σ the sample covariance, PR = (tr Σ)² / tr(Σ²) and
both traces derive from Σ hhᵀ, Σ h, and the sample count. Storing raw
sums (not averages) makes moments composable: StateMoments::merge
pools across forward calls (streaming chunks, eval batches) and the PR of
the merged moments is the exact pooled PR.
Structs§
- State
Moments - Raw (un-normalised) first/second moments of the per-token SSM states of
one block’s
forwardpass, pooled over tokens andper_head_dimrows.
Enums§
- State
Pairing - How the
state_rankaxis of a (realified) SSM state groups into complex / quaternionic coordinates — the realification layout the Mamba-3 rotation applies to B/C (and hence to the state). Consumed byStateMoments::pr_complex.