--- title: 'TOSE: Tiled Operator-Space Evolution in SSMs' url: https://www.emergentmind.com/topics/tiled-operator-space-evolution-tose type: topic --- # TOSE: Tiled Operator-Space Evolution in SSMs Tiled Operator-Space Evolution (TOSE) is a core algorithmic primitive for O(1)-memory exact analytical differentiation through linear state-space models (SSMs), enabling efficient sensitivity analysis for large-scale sequence modeling tasks. TOSE reframes the gradient computation as the forward evolution of an augmented dynamical system, exploiting the associative structure of operator evolution and leveraging a tiling and scan-based protocol to eliminate the memory overhead inherent in conventional backpropagation. TOSE supports applications such as chromosome-scale modeling in genomics, where conventional Autograd approaches encounter prohibitive memory costs. ## 1. Theoretical Basis and Role in Phase Gradient Flow TOSE is underpinning the Phase Gradient Flow (PGF) framework, which addresses the memory bottleneck in gradient-based sensitivity analysis for large SSMs. PGF interprets backpropagation through a linear recurrence as the forward evolution of a “tangent” dynamical system. For a discrete-time SSM defined by \[ x_{t+1} = A x_t + B u_t,\qquad y_t = C x_t + D u_t, \] the Fréchet derivative product $\nabla y_t = D\mathcal{F}[u] \cdot \nabla u_t$ evolves under a linear recurrence isomorphic to the original dynamics. TOSE leverages this dynamical isomorphism by tiling the sequence into fixed-size blocks. Within each tile, both the primal and tangent systems are evolved via a parallel scan over associative operator products. The computational graph is detached at tile boundaries, resulting in memory usage that does not scale with sequence length $L$. This enables strictly O(1) differentiation graph memory in $L$ [2512.23068]. ## 2. Mathematical Foundation ### 2.1 Primal and Tangent Recursions Given $x_t \in \mathbb{R}^n$, $u_t, y_t \in \mathbb{R}^d$ and allowing for diagonal or block-diagonal, possibly time-varying system matrices (as in Mamba/S4), the tangent dynamics for sensitivity analysis (the variation $\nabla x_t := \partial x_t / \partial u \cdot \nabla u$) satisfy \[ \nabla x_{t+1} = A \nabla x_t + (\partial A / \partial u_t \cdot \nabla u_t)x_t + (\partial B / \partial u_t \cdot \nabla u_t)u_t + B \nabla u_t, \] \[ \nabla y_t = C \nabla x_t + D \nabla u_t, \] where, in the diagonal SSM setting, all derivatives and _per-step_ operators remain element-wise linear. ### 2.2 Augmented Operator-Space The triplet $(x_t, \nabla x_t, 1)$ is encoded as a single augmented vector, and an augmented block-matrix recurrence is defined: \[ \begin{bmatrix} x_t \\ \nabla x_t \\ 1 \end{bmatrix} = M_t \begin{bmatrix} x_{t-1} \\ \nabla x_{t-1} \\ 1 \end{bmatrix}, \] where each $M_t$ is \[ M_t = \begin{bmatrix} A & 0 & B u_t \\ K_t & A & j_t \\ 0 & 0 & 1 \end{bmatrix}, \] with $K_t = \partial A / \partial u_t \cdot \nabla u_t$ and $j_t = (\partial B / \partial u_t \cdot \nabla u_t)u_t + B \nabla u_t$. The associativity $M_{t+1} \cdot M_t = M_{t:t+1}$ allows parallel prefix scan computation of these operator products. The diagonal structure reduces the computation, for each coordinate, to products of compact ($3\times3$) blocks. ## 3. Tiling and Parallel Scan Protocol Sequence length $L$ is partitioned into tiles of size $B$. For each tile $k$ spanning steps $[(k-1)B+1,\; kB]$, the per-step operator $M_t$ is constructed and a parallel scan (prefix product) is performed over the block. Only the state at the tile boundary is retained for subsequent computation, and intermediate autograd history within the tile is detached, yielding memory complexity independent of $L$. The TOSE protocol is formalized as follows: | Step | Operation | Memory Effect | |------------------------------------|-------------------------------------------------|----------------------------| | a) Load block inputs | $u_{\text{block}},\; \nabla u_{\text{block}}$ | No graph growth | | b) Build per-step operators | Construct $M_t$ for tile | Per-step, local | | c) ParallelScan over $M_t$ | Compute $[x_{\text{out}}, \nabla x_{\text{out}}]$| Intra-tile only | | d) Output sensitivity | $\nabla y_{\text{block}} = C \nabla x_{\text{out}} + D \nabla u_{\text{block}}$ | Store external output | | e) Boundary state detach | $x_{\text{in}}, \nabla x_{\text{in}}$ updated via .detach() | Limits recursive history | TOSE relies on an arbitrary associative parallel prefix-product implementation (e.g., tree-based scan) to combine per-step operators within each tile. ## 4. Memory Complexity and Differentiation Graph Architecture By detaching computational history at tile boundaries, the memory footprint for automatic differentiation is proportional to $B \cdot (\mathrm{dim}\, x + \mathrm{dim}\, \nabla x)$, a strict constant in $L$ as $L \rightarrow \infty$. All input and output tensors (e.g., $u,\; \nabla u,\; y,\; \nabla y$) retain linear scaling in $L$, but these are not part of the recursive differentiation graph. TOSE thus guarantees strictly O(1) memory scaling for the differentiation graph with respect to sequence length. ## 5. Numerical Stability in Stiff ODE Regimes Conventional parallel scan algorithms exhibit catastrophic underflow or overflow for very stiff ODEs (e.g., $A_t$ with eigenvalues $\exp(-1000)$) in finite precision arithmetic. TOSE incorporates a "log-shifting stabilizer," normalizing intermediate products within each tile by subtracting the local maximum log-magnitude prior to exponentiation and restoring the shift post-scan. Analytical error analysis demonstrates that the relative error after $L$ steps is bounded by \[ \|\text{error}\| \leq C \epsilon_{\mathrm{mach}}, \] with $C$ governed by the spectral radius of $A$ ($\leq 1$ in typical scenarios). Empirical regression of error versus $L$ shows no slope ($p \gg 0.05$), confirming no error accumulation with increasing sequence length. This distinguishes TOSE from naive parallel prefix schemes, which fail under extreme stiffness for large $L$ [2512.23068]. ## 6. Empirical Performance and Scalability ### 6.1 Memory Utilization Standard Autograd backpropagation incurs $O(L \cdot n \cdot d)$ activation storage; at $L=100{,}000$, a single layer exceeds $10$ GB VRAM. TOSE maintains graph memory at a constant level irrespective of $L$. Across $L \in [10^3,\, 10^5]$, PGF achieves a $94\%$ reduction in peak VRAM usage on NVIDIA RTX 5090 hardware. ### 6.2 Throughput On an RTX 5060 Laptop GPU, PGF demonstrates $11.9\times$ throughput improvement over reverse-mode Autograd at $L=10{,}000$, and sustains $23\times$ higher throughput for $L \approx 128{,}000$ (“chromosome scale”). This efficiency was evaluated via samples per second (see Figures 1 and 3 in [2512.23068]). ### 6.3 Sensitivity Analysis Benchmarks In a $128,000$-step impulse-response (“Ghost Pulse”) benchmark, TOSE precisely recovers micro-perturbations at late timesteps ($t=100{,}000$) without numerical leakage, a regime where Autograd encounters out-of-memory failure in multi-layered models. ## 7. Applications and Significance TOSE enables feasible genome-scale sensitivity analysis ($L > 10^5$) on single 8–12 GB GPUs. This bridges the practical gap between infinite-context theoretical models and existing hardware constraints. The protocol is applicable to any linear SSM and directly supports models with diagonal or block-diagonal system matrices, including Mamba and S4 variants. A plausible implication is that TOSE could generalize to other operator evolution domains beyond SSMs, provided they admit associative and scan-friendly forms. The minimal error accumulation and strict memory guarantees position TOSE as a foundational tool for large-scale differentiable modeling, particularly in genomics, signal processing, and other sequence analysis contexts where computational memory was previously prohibitive [2512.23068]. Source: https://www.emergentmind.com/topics/tiled-operator-space-evolution-tose