Sequence Modeling with Spectral Mean Flows
Jinwoo Kim, Max Beier, Petar Bevanda, Nayun Kim, Seunghoon Hong
A fundamental 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 a useful yet 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 whole sequence distribution as a tensor in the product Hilbert space. A generative process is then defined as a 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 that leverages the spectral decomposition of linear operators to derive a scalable tensor network decomposition of the sequence mean embedding. Second, we extend MMD gradient flows to time-dependent Hilbert spaces, which might be of independent interest. The resulting technique is connected to flow matching via the continuity equation, enabling end-to-end, simulation-free learning via double backpropagation. We demonstrate state-of-the-art results on time-series modeling.