Expand description
§Constant-input stepping shortcuts — step_n_approx / step_infinite
When the same token is fed to Mamba3::step over and over, every
data-dependent quantity is the same at each step: the trapezoid
coefficients α, β, γ, the QK-normed b/c, and the per-step
rotation increment (angle θ̂ for Complex2D, unit quaternion q for
Quaternion4D). Only the cumulative rotation moves — by that same
increment each step — so with P the per-step rotation operator and R₁
the cumulative rotation at the first constant step,
Bₜ = R₁ Pᵗ⁻¹ b , Cₜ = R₁ Pᵗ⁻¹ c ,
hₜ = α hₜ₋₁ + x ⊗ (β Bₜ₋₁ + γ Bₜ) (t ≥ 2)and K = n − 1 unrolled steps collapse to a matrix geometric series in
α P⁻¹ (spectral radius α < 1, since a_floor > 0 and Δ > 0):
h_n = α^K h₁ + x ⊗ R₁ P^{K−1} (I − (αP⁻¹)^K)(I − αP⁻¹)⁻¹ (β + γP) bP is block-diagonal: per RoPE pair it is the complex scalar e^{iθ̂}, per
quaternion block left-multiplication by w = q*. In both cases all the
factors live in one abelian subalgebra (powers of a single rotation
commute — for quaternions, span{1, û} ≅ ℂ), so the series is a handful of
scalar complex / quaternion ops per (head, pair/block): n steps cost O(1).
In the readout y_n = Cₙᵀ h_n + D x the unbounded phase cancels
(⟨R c, R v⟩ = ⟨c, v⟩), leaving only the relative rotation P⁻¹. As
n → ∞ the h₁ term and the partial-sum correction decay like αⁿ, so the
output converges even though the state h orbits forever:
y_∞ = xᵀ · cᵀ (γ + β P⁻¹)(I − α P⁻¹)⁻¹ b + D x— the block’s stationary fixed point, independent of any starting cache.
Mamba3::step_infinite evaluates exactly this (and therefore takes and
returns no cache); Mamba3::step_n_approx evaluates the finite-n form,
returning the exact step-n output and cache.
Numerics: α^K is exp(K·ΔA); every accumulated rotation power has its
(half-)angle reduced mod 2π with the value-exact wrap_angle; the
geometric denominators satisfy |1 − α e^{−iθ̂}|² ≥ (1 − α)² and are floored
by div_eps. When α → 1 and θ̂ → 0 the series
value (β+γ)(1−α^K)/(1−α) → K·(β+γ)-ish stays finite, but its fp32
evaluation loses precision once 1 − α nears the epsilon floor — the same
regime where the unrolled recurrence itself accumulates K near-undamped
terms.
Functions§
- complex_
div 🔒 - Component-wise complex quotient;
|den|²floored bydiv_eps. - complex_
mul 🔒 - Component-wise complex product
(ar + i·ai)(br + i·bi). - cos_sin 🔒
(cos, sin)of an angle tensor, reduced mod2πfirst (value-exact).- mul_
complex_ 🔒partial - Multiply the rotation-active entries of
xby the complex scalar(re, im)per pair — same pairing conventions asapply_rope_partial— and the pass-through entries by the real scalartail. - mul_
quat_ 🔒partial - Left-multiply the first
rope_widthstate-rank entries ofxby the (not necessarily unit) quaternionfper block, and scale the pass-through entries by the real scalartail— the quaternion analogue ofmul_complex_partial. - per_
step_ 🔒angle θ̂ = Δ · π·tanh(rot)— the constant per-step RoPE angle increment.- per_
step_ 🔒generator g = Δ · π·tanh(rot)per quaternion block — the constant per-step rotation generator (q = exp(g/2)).- quat_
inv 🔒 - Quaternion inverse
q⁻¹ = q* / ‖q‖², with‖q‖²floored bydiv_eps. (Used on1 − α q, whose norm is bounded below by1 − α > 0.) - quat_
one_ 🔒minus 1 − sqfor a (scalar-scaled) quaternionsq: negate, add 1 to the real part.- quat_
pow 🔒 exp(k·g/2)— the per-step rotation quaternion raised to the real powerk, built likequat_from_scaled_axisbut with the half-angle reduced mod2πso largekstays fp-accurate (the reduction is value-exact and the wrap offset is detached, sod/dgmatches the unrolled product).- quat_
scalar_ 🔒affine a + b·qfor per-head scalarsa,b([batch, nheads, 1, 1]) and a quaternion tensorq.