Flash STU: Fast Spectral Transform Units

Flash STU is a sequence model built around spectral filtering, a method for learning the behavior of dynamical systems without ever estimating their parameters. The work pairs a careful theoretical treatment of the Spectral Transform Unit with an efficient implementation, interleaves it with a small amount of local attention, and shows across three very different domains that the resulting architecture matches or surpasses Transformers and the leading state space models at a fixed parameter budget while remaining close to linear in sequence length.

Learning in linear dynamical systems

Many neural architectures for sequence modeling can be read through the lens of linear dynamical systems, a foundational object in control theory and time series analysis. Such a system carries a hidden state that evolves under a fixed set of matrices and emits an observation at every step,

x_{t+1} = A x_t + B u_t + w_t, \qquad y_t = C x_t + D u_t + \zeta_t,

where the input enters through the matrix B, the state propagates through A, and the noise terms absorb everything left unmodeled. Unrolling the recursion in the noiseless case writes each output purely in terms of past inputs,

y_t = \sum_{i=1}^{t} C A^i B \, u_{t-i}.

The largest absolute eigenvalue of the transition matrix, its spectral radius, controls how slowly the influence of the past decays. When that radius approaches one the system has very long memory, and this is exactly the regime where the classical recipes falter. System identification recovers the matrices directly but is non-convex and unreliable near the boundary, autoregressive fitting needs a parameter count that grows with the length of the memory, and Kalman filtering is fragile under adversarial inputs.

Spectral filtering takes a different route. Rather than estimating the system matrices at all, it represents the dynamics through a small set of fixed convolutional filters drawn from the dominant spectral components of the problem. The construction has three properties that matter in practice. It avoids the non-convex recovery of parameters entirely, its filter count does not grow with the hidden dimension of the system, and it stays effective precisely when the spectral radius is close to one, which is the long-context regime. Because the operation is a convolution, it can be carried out with the fast Fourier transform, which gives subquadratic cost in the sequence length.

The Spectral Transform Unit

The Spectral Transform Unit, or STU, turns spectral filtering into a neural network layer. An input sequence is first lifted into a working dimension by a learned map, then convolved against a bank of fixed filters, and finally projected back out by a learned transformation before an optional nonlinearity. The whole layer can be written compactly as

\hat{y}_t = \sigma\left( \sum_{i=1}^{k} M_i \cdot \langle \Phi_i, u_{t:t-L} \rangle \right),

where the filters from the first through the kth are fixed and known in advance, and the matrices that carry the index i are the only learned projections.

The Spectral Transform Unit. An input is lifted into a working dimension, convolved against a fixed bank of spectral filters, projected by the learned matrices, and passed through an optional nonlinearity.

The filters are not chosen by hand. They are the leading eigenvectors of a single Hankel matrix that encodes every fixed-length impulse response,

Z = \int_{0}^{1} \mu_\alpha \mu_\alpha^\top \, d\alpha, \qquad \mu_\alpha = (1, \alpha, \ldots, \alpha^L),

whose entries also admit a simple closed form,

Z(i,j) = \frac{2}{(i+j)^3 - (i+j)}.

The reason this works is a striking fact about such matrices. The spectrum of any real Hankel matrix decays exponentially, so a handful of eigenvectors already captures almost all of the relevant structure and the error left behind by keeping only a few filters is exponentially small. In practice the loss stops improving once the number of filters passes roughly fifteen, which is why the entire bank can be precomputed and stored before a single token is seen.

The first twenty spectral filters used in the model. Each one is an eigenvector of the fixed Hankel matrix, computed once and reused unchanged across training and inference.

A further optimization, the tensordot approximation, factors each projection into two thinner matrices,

M_i \approx M_i^1 \times M_i^2, \qquad M_i^1 \in \mathbb{R}^{d_{\text{model}} \times k}, \quad M_i^2 \in \mathbb{R}^{d_{\text{model}} \times d_{\text{out}}},

which lets the projection be applied before the convolution and cuts the cost by about a factor of the filter count. The variant that uses it is called STU-T. It trades a little expressivity for a large saving in memory and time, and it is the form that makes scaling to billions of parameters feasible.

The Flash STU architecture

Flash STU is the hybrid that gives the paper its name. A pure stack of spectral layers is efficient but, like other state space models, it can struggle on tasks that demand exact recall of a specific earlier token. The remedy is to alternate STU-T layers with a small amount of local attention, each followed by a feedforward block, so that the spectral layers carry the long-range structure cheaply while the attention layers supply precise lookups. To keep the attention itself inexpensive, the global mechanism is replaced by sliding window attention whose window spans one eighth of the sequence, and the convolutions are computed with a specialized fast Fourier transform kernel. The normalization, the gated feedforward layers, and the loss all run through the same optimized kernels used for the Transformer baseline. Without the tensordot factorization the spectral models ran out of memory beyond a billion parameters, so it is that approximation which ultimately unlocks scale.

Optimization and synthetic tasks

The first set of experiments isolates the STU layer on data drawn from a known linear system whose spectral radius is tuned to give an effective memory of about one hundred steps. Trained alongside attention, a diagonal state space layer, and a selective state space layer, the spectral layer reaches a low error quickly and repeatably while the others vary from seed to seed. The width of the layer does not need to match the hidden dimension of the underlying system in order to recover its dynamics, which is exactly what the theory predicts.

Mean squared prediction error on a linear dynamical system with long memory. Both the spectral layer and its factored variant converge quickly and reliably, while the attention and state space baselines fluctuate across random seeds.

Part of the reason for this reliability is geometric. Under the squared error the STU layer is a convex function of its parameters, whereas any model that learns the system matrices directly faces a non-convex objective whose difficulty grows with the sequence length. Visualizing the loss surface after a few steps of training makes the difference plain. The spectral layer sits in a smooth basin with clear directions of progress, while the diagonal state space layer in particular shows a rough and broken landscape.

Local loss surfaces for the spectral layer, a diagonal state space layer, a selective state space layer, and attention, taken after a few steps of training. The spectral surface is a smooth bowl, whereas the diagonal state space surface is jagged.

The picture on the more nonlinear synthetic tasks is honest about the trade-offs. On selection style problems such as induction heads and selective copy, the selective state space model has a built-in advantage from its input-dependent gating, and attention is the natural fit for associative recall, where many key and value pairs must be retrieved in parallel. Even so, the spectral layers consistently beat the diagonal state space baseline across these tasks and find their solutions more reliably, which marks the STU out as a simple and easily optimized primitive to build on.

Robotics control

The second domain is next-state prediction for simulated robots in the MuJoCo physics engine. Here the task is to predict the next configuration of a moving body from its current state and the applied torques,

x_{t+1} = f(x_t, u_t),

a map that is nonlinear and non-smooth, unlike the clean linear system of the synthetic study. Each model is held to the same half-million-parameter budget and trained to minimize the squared distance between the predicted and the true next state,

L(\theta) = \frac{1}{2}\left\| \hat{x}_{t+1}(\theta) - x_{t+1} \right\|_2^2.

Across three controllers the spectral models excel at this kind of dynamical prediction. The factored STU-T variant gives the best accuracy of any model tested, including the selective state space model, while running faster per step, and it descends more smoothly early in training. A consistent pattern is that the spectral models gain more from added depth, whereas the Transformer leans on width and saturates quickly beyond a couple of layers.

Next-state prediction error on the Ant controller in MuJoCo, with every model held to the same parameter budget. The factored spectral model reaches the lowest validation loss.

Language modeling

The final and most demanding domain is language. Every model is pretrained on roughly ten billion tokens of curated web text under a shared recipe in the style of the LLaMA family, and the spectral models are compared against a strong Transformer, a selective state space model, and a hybrid of the two. At the compute-optimal size of about five hundred million parameters, Flash STU reaches the lowest validation loss of the group, 3.40 against the Transformer at 3.92, and it leads on the average of a broad suite of downstream benchmarks. Scaled up to two billion parameters it again comes out ahead, with a validation loss of 3.09 and a perplexity of 22.08 against 3.18 and 24.11 for the Transformer, at roughly equal time per step.

Validation loss for Flash STU and a matched Transformer during language model pretraining on curated web text. The spectral hybrid settles to a lower loss.

Two qualitative observations accompany these numbers. The Transformer proved more sensitive to the choice of learning rate and more prone to sudden spikes in its loss, while Flash STU trained smoothly throughout. The authors are also careful to note that the Transformer benefits from years of hardware-level tuning that the spectral implementation does not yet match, so the fact that the two already run at comparable wall-clock speed, with the spectral model holding an asymptotic advantage, is itself the encouraging result.

Outcome

Taken together the three studies make a consistent case. A model built on fixed spectral filters, given just enough local attention to handle exact recall, is a genuine competitor to the Transformer and to the leading state space models, and at a fixed parameter budget it tends to come out ahead. Its standout practical quality is how cleanly it optimizes, a property that follows from the convex foundation of spectral filtering and that grows more valuable as models scale. The full implementation and a distributed pretraining pipeline are released openly, and the work was published in the proceedings of the IEEE Conference on Decision and Control. The code is available on GitHub.