Sequence Modeling with Spectral Mean Flows

TL;DR

The paper addresses learning nonlinear and probabilistic state dynamics in sequences by embedding the entire sequence distribution into a tensor-product RKHS and formulating learning as a maximum mean discrepancy (MMD) gradient flow over sequences. It introduces spectral mean flows, combining spectral decompositions of hidden Markov model (HMM) operators with a tensor-network embedding to yield a scalable and trainable framework. The authors extend MMD flows to time-dependent RKHS and connect them to flow matching to achieve finite-time convergence, enabling simulation-free training and faster sampling. Empirically, the method achieves competitive or superior performance on synthetic and diverse time-series benchmarks, including irregular and physics-informed settings, while maintaining scalable computation and end-to-end trainability.

Abstract

A key question in sequence modeling with neural networks is how to represent and learn highly nonlinear and probabilistic state dynamics. Operator theory views such dynamics as linear maps on Hilbert spaces containing mean embedding vectors of distributions, offering an appealing but currently overlooked perspective. We propose a new approach to sequence modeling based on an operator-theoretic view of a hidden Markov model (HMM). Instead of materializing stochastic recurrence, we embed the full sequence distribution as a tensor in the product Hilbert space. A generative process is then defined as maximum mean discrepancy (MMD) gradient flow in the space of sequences. To overcome challenges with large tensors and slow sampling convergence, we introduce spectral mean flows, a novel tractable algorithm integrating two core concepts. First, we propose a new neural architecture by leveraging spectral decomposition of linear operators to derive a scalable tensor network decomposition of sequence mean embeddings. Second, we extend MMD gradient flows to time-dependent Hilbert spaces and connect them to flow matching via the continuity equation, enabling simulation-free learning and faster sampling. We demonstrate competitive results on a range of time-series modeling datasets. Code is available at https://github.com/jw9730/spectral-mean-flow.

Figures (9)

• Figure 1: A tensor diagram bridgeman2017handahle2024tensorcookbook of the inner product $\langle \phi({\bf x}^1)\otimes\cdots\otimes\phi({\bf x}^N), \mu_\rho\rangle$\ref{['eq:sequence_inner_product']} for $N=5$, showing how its evaluation involving a tensor $\mu_\rho$ decomposes into matrix and vector multiplications.
• Figure 2: Peak GPU memory of inner product $\langle {\bf x}^1\otimes\cdots\otimes{\bf x}^N, \mu\rangle$ depending on the use of tensor network decomposition \ref{['eq:time_dependent_sequence_inner_product']}.
• Figure 3: 2D checkerboard experiment. Left: Intermediate distributions over sampling timesteps (zoom in for a better view). Right: MMD between the intermediate and target distributions over sampling timesteps, measured with an RBF kernel of bandwidth $1$.
• Figure 4: Examples of real and generated data for the Sines dataset.
• Figure 5: Examples of real and generated data for the Stocks dataset.

Theorems & Definitions (21)

• Remark A.1
• Lemma B.1
• proof
• Proposition B.2
• proof
• Lemma B.3
• proof
• Corollary B.4
• Lemma B.5
• proof