Skip to main content

Module rotation

Module rotation 

Source
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 isoclinicL_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 by Clone.
RotationSpec
Everything the rotation needs from the block: which algebra, how much of state_rank it turns, and how far one step may turn it.

Enums§

RotationKind
Which rotational-state algebra the block uses for the data-dependent transition rotation absorbed into B/C.
RotationState
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 to quat_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 “materialise Rₜ” step, analogous to RoPE’s Δₜ · π · tanh(θₜ) angle.
quat_mul
Hamilton product a ⊗ b of two quaternion tensors.
quat_normalize
Normalise quaternions to unit norm along the last axis (shape [..., 4]).
quat_to_rot4
Materialise the 4×4 orthogonal matrix of left-multiplication by q.
rotate_bc_forward
Rotate B/C for a full sequence by the data-dependent transition rotation, returning the rotated projections and the new cumulative RotationState to store in the cache.
rotate_bc_step
Single-token counterpart of rotate_bc_forward for the recurrent step.
rotate_blocks_partial
Apply a per-block quaternion rotation to the first rope_width entries of the state_rank axis (a multiple of 4); the remainder passes through. The quaternion analogue of apply_rope_partial.
rotate_blocks_two_sided_partial
Two-sided counterpart of rotate_blocks_partial: rotates the first rope_width entries of the state_rank axis by v ↦ ql ⊗ v ⊗ qr, passing the remainder through.
rotate_state_rank_blocks
Apply a per-block quaternion rotation to the state_rank axis of v.
rotate_state_rank_blocks_two_sided
Apply a per-block two-sided quaternion rotation to the state_rank axis of v: v ↦ ql ⊗ v ⊗ qr per 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::Rotor accumulator [..., 2·J, 4] into its (left, right) factors, each [..., J, 4].