Expand description
§Quaternion (k=4) rotational state — the non-abelian generalisation of RoPE
Mamba-3’s data-dependent RoPE realises a complex-valued SSM: the state
transition factors as a per-head scalar decay times a block-diagonal of
2×2 rotations (paper Prop. Complex-to-Real SSM Equivalence), and because
SO(2) ≅ U(1) is abelian the cumulative rotation collapses to a
cumsum of angles and is absorbed into B/C (the “RoPE trick”, Prop.
Complex SSM, Data-Dependent RoPE Equivalence). See
crate::mamba3::rotation::rope::apply_rope.
This module implements the next rung of the ladder: a quaternion
(k = 4) rotational state, i.e. the transition’s rotation lives in the
left-isoclinic subgroup SU(2) ⊂ SO(4) instead of SO(2). Unit
quaternions under multiplication are SU(2), which is non-abelian and
contains non-solvable finite subgroups (the binary icosahedral group
2I = SL(2,5), a double cover of A₅). By Barrington’s theorem this lifts
the layer’s reachable state-tracking from the solvable/TC⁰ regime (parity,
mod-k) toward NC¹, which abelian rotations provably cannot reach.
§What survives, what changes
The key fact (derivable purely from telescoping + orthogonality, without
commutativity — see the crate discussion) is that the RoPE factoring
survives intact: with the ordered cumulative rotation
Pₜ = Rₜ Rₜ₋₁ ⋯ R₁,
Cₜᵀ (Rₜ⋯Rᵢ₊₁) Bᵢ = (Pₜᵀ Cₜ)ᵀ (Pᵢᵀ Bᵢ) = C̄ₜᵀ B̄ᵢ ,so the scalar-decay SSD core (L ⊙ C̄B̄ᵀ) is unchanged — only the
projections B̄ᵢ = Pᵢᵀ Bᵢ, C̄ₜ = Pₜᵀ Cₜ are rotated. What is lost is the
closed-form cumsum: the cumulative rotation must be built by an
associative scan over the per-step quaternions (quat_cumprod) rather
than a sum of angles. Because a product of unit quaternions is again a unit
quaternion, the scan stays exactly orthogonal (no drift, no wrap_angle
needed), and the cross-chunk carry is a single quaternion per block/head —
the exact analogue of cum_angle in the existing caches.
SO(2) (today’s apply_rope) is the abelian collapse: restricting each
quaternion to a single fixed axis makes them commute and reduces
quat_cumprod to a cumsum of half-angles (asserted in the tests).
§Pipeline (the k = 4 instantiation of the rotation block)
per-step unit quaternion qₜ (materialise from the in-projection; caller)
│ quat_cumprod (assoc. scan, + cross-chunk carry)
▼
cumulative rotation Qₜ
│ rotate_state_rank_blocks(B, conj(Qₜ)) , rotate_state_rank_blocks(C, conj(Qₜ))
▼
B̄, C̄ ──► standard scalar-decay SSD (unchanged)For RotationKind::Rotor4D the same pipeline runs with the two factors
stacked along the block axis, and the last step becomes the two-sided
rotate_state_rank_blocks_two_sided(B, conj(Qₜ), Tₜ).
§SO(4): the whole rotation group of a block (RotationKind::Rotor4D)
SU(2) is only half of what a 4-block can turn by. The general element of
SO(4) ≅ (SU(2)×SU(2))/±1 is the two-sided product
Rₜ(v) = qₜ ⊗ v ⊗ p̄ₜ (a rotor; left factor q, right factor p)and everything above survives it, because the factoring never used more than “the per-step maps compose, and each is orthogonal”. Composing,
Pₜ(v) = Qₜ ⊗ v ⊗ T̄ₜ , Qₜ = qₜ⊗⋯⊗q₁ , Tₜ = pₜ⊗⋯⊗p₁
Pₜ⁻¹(v) = Qₜ* ⊗ v ⊗ Tₜ ⇒ B̄ᵢ = Qᵢ* ⊗ Bᵢ ⊗ Tᵢ , C̄ₜ likewise— note the conjugation reverses the right-hand order twice, so T
accumulates by the same left fold as Q: one quat_cumprod over a
doubled block axis, not a second, reversed scan. The cost over
Quaternion4D is twice the generator
channels, twice the scan’s block axis, and one extra quat_mul per
B/C application; the SSD core is still untouched.
Why bother, when left and right factors commute with each other and so add
no “more non-abelianness”: left multiplication is isoclinic — L_q turns
both invariant planes of the block by the same angle — so SU(2) cannot
produce two independent plane angles, and in particular does not contain the
abelian SO(2)² rotation it was introduced to generalise. The right factor
is exactly what opens the maximal torus (plane angles a−b, a+b for the
half-angles of q, p), making the ladder
Complex2D ⊂ Rotor4D ⊃ Quaternion4D a real one. It also contains the
adjoint action v ↦ q ⊗ v ⊗ q̄, i.e. a faithful SO(3) on the block’s
imaginary part — so a group like A₅ can be tracked as itself rather than
through its double cover 2I, where ±g denote one element but two
different states. And as a representation of SU(2)×SU(2), ℍ is the
irreducible tensor product (½,½), not (½,0) ⊕ (0,½): no arrangement of
left-only blocks reproduces it.
SO(4) is the ceiling for k = 4 (the largest norm-preserving transition
group of a block), and k = 4 is the last rung with a cheap closed form —
at k = 8 the octonions are non-associative and the scan itself breaks.
§The bottom rung: RotationKind::Real1D
k = 1, the trivial group: no rotation, a purely real transition. It is the
ablation the ladder is measured against, and it is structural rather than a
zeroed knob — the in-projection spends no channels on rotation, the cache
carries no accumulator (RotationState::Real), and B/C reach the SSD
core untouched. Switching the rotation off is therefore a choice of kind;
rope_fraction only ever narrows a rotation that exists.
Quaternion layout: the last axis has size 4 and holds (w, x, y, z) with
w the real part. A state_rank of r = 4·J is treated as J independent
quaternion blocks; the rotation acts within each block, exactly as RoPE acts
within each 2-pair. Mamba3 selects it
with RotationKind::Quaternion4D and drives it through one
RotationSpec (the SSD kernels themselves need no edits).
Modules§
- rope
- Rotary (RoPE) application: the mechanical pairwise rotation the abelian pathway factors into B/C. Rotary (RoPE) application helpers for the abelian rotation pathway.
Structs§
- NoRotation
- The stateless payload of
RotationState::Real— a [Module] holding no tensors, so a real transition’s cache slot allocates nothing and converts between backends byClone. - Rotation
Spec - Everything the rotation needs from the block: which algebra, how much of
state_rankit turns, and how far one step may turn it.
Enums§
- Rotation
Kind - Which rotational-state algebra the block uses for the data-dependent
transition rotation absorbed into
B/C. - Rotation
State - The cumulative-rotation accumulator carried between calls in a Mamba-3 cache
— the variant matching the block’s
RotationKind.
Functions§
- angle_
increment - The abelian per-step angle increment
θ̂ₜ = Δₜ · range·π·tanh(ϑₜ), one per head and rotation pair. - bound_
rotation_ vector - Bound a rotation vector’s magnitude, leaving its direction alone:
returns
max_angle · tanh(‖r‖) · r̂. - generator_
increment - The quaternion per-step rotation vector
gₜ = Δₜ · range·π·tanh(‖r‖)·r̂, one per head and quaternion block (feed it toquat_from_scaled_axis). - quat_
conj - Quaternion conjugate
q* = (w, −x, −y, −z)(shape[..., 4]). - quat_
cumprod - Cumulative (ordered, left-accumulating) quaternion product along the sequence axis, with a cross-chunk carry.
- quat_
from_ scaled_ axis - Materialise a unit quaternion from a scaled rotation vector
g ∈ ℝ³(axis · angle) via the exponential map — the data-dependent “materialiseRₜ” step, analogous to RoPE’sΔₜ · π · tanh(θₜ)angle. - quat_
mul - Hamilton product
a ⊗ bof two quaternion tensors. - quat_
normalize - Normalise quaternions to unit norm along the last axis (shape
[..., 4]). - quat_
to_ rot4 - Materialise the
4×4orthogonal matrix of left-multiplication byq. - rotate_
bc_ forward - Rotate
B/Cfor a full sequence by the data-dependent transition rotation, returning the rotated projections and the new cumulativeRotationStateto store in the cache. - rotate_
bc_ step - Single-token counterpart of
rotate_bc_forwardfor the recurrentstep. - rotate_
blocks_ partial - Apply a per-block quaternion rotation to the first
rope_widthentries of thestate_rankaxis (a multiple of 4); the remainder passes through. The quaternion analogue ofapply_rope_partial. - rotate_
blocks_ two_ sided_ partial - Two-sided counterpart of
rotate_blocks_partial: rotates the firstrope_widthentries of thestate_rankaxis byv ↦ ql ⊗ v ⊗ qr, passing the remainder through. - rotate_
state_ rank_ blocks - Apply a per-block quaternion rotation to the
state_rankaxis ofv. - rotate_
state_ rank_ blocks_ two_ sided - Apply a per-block two-sided quaternion rotation to the
state_rankaxis ofv:v ↦ ql ⊗ v ⊗ qrper 4-block. - safe_
norm 🔒 - Euclidean norm over the last axis, formed scale-free so it cannot overflow: the components are divided by their (detached) largest magnitude before squaring, and the result is scaled back.
- split_
rotor - Split a stacked
RotationState::Rotoraccumulator[..., 2·J, 4]into its(left, right)factors, each[..., J, 4].