Skip to main content

Module step_constant

Module step_constant 

Source
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) b

P 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 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 by div_eps.
complex_mul 🔒
Component-wise complex product (ar + i·ai)(br + i·bi).
cos_sin 🔒
(cos, sin) of an angle tensor, reduced mod first (value-exact).
mul_complex_partial 🔒
Multiply the rotation-active entries of x by the complex scalar (re, im) per pair — same pairing conventions as apply_rope_partial — and the pass-through entries by the real scalar tail.
mul_quat_partial 🔒
Left-multiply the first rope_width state-rank entries of x by the (not necessarily unit) quaternion f per block, and scale the pass-through entries by the real scalar tail — the quaternion analogue of mul_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 by div_eps. (Used on 1 − α q, whose norm is bounded below by 1 − α > 0.)
quat_one_minus 🔒
1 − sq for a (scalar-scaled) quaternion sq: negate, add 1 to the real part.
quat_pow 🔒
exp(k·g/2) — the per-step rotation quaternion raised to the real power k, built like quat_from_scaled_axis but with the half-angle reduced mod so large k stays fp-accurate (the reduction is value-exact and the wrap offset is detached, so d/dg matches the unrolled product).
quat_scalar_affine 🔒
a + b·q for per-head scalars a, b ([batch, nheads, 1, 1]) and a quaternion tensor q.