Spectral Flow Matching: Stabilizing Stochastic GFlowNets via Frequency-Domain Regularization
Abstract
Generative Flow Networks (GFNs) offer a powerful paradigm for diverse sampling, yet they often exhibit instability and poor convergence when applied to stochastic or sparse-reward environments. To mitigate the high variance inherent in these settings, we propose a fundamental re-framing of the GFlowNet training objective within the frequency domain. We present \textbf{Spectral Time-Dependent GFlowNets (ST-GFNs)}, a framework that leverages Fourier analysis to enforce smoothness and stability in learned policies. Our theoretical analysis proves that our proposed spectral loss is mathematically equivalent to regularized value iteration, acting as a principled low-pass filter that separates signal from noise. Furthermore, we tackle the challenge of exploration in sparse landscapes by introducing a novel autocorrelated intrinsic reward derived from the Wiener-Khinchin theorem. Through extensive experiments ranging from adversarial games and noisy sequence generation to high-dimensional single-cell perturbation modelling, we demonstrate that ST-GFNs significantly outperform existing baselines in terms of robustness, sample efficiency, and mode discovery.
Lay Summary
Artificial Intelligence (AI) models known as Generative Flow Networks are highly effective at exploring vast search spaces to discover diverse, high-quality solutions—such as identifying promising new molecular candidates for life-saving drugs or designing complex biological sequences. However, when these AI agents are deployed in "noisy" or highly unpredictable environments where rewards are sparse and chaotic, they suffer from extreme training instability. To solve this problem, this work introduces a novel training framework that views AI learning through the lens of audio and signal processing. Just as noise-canceling headphones filter out chaotic, high-frequency background static to preserve a smooth, clear audio track, our framework introduces a mathematical "spectral filter" into the AI’s learning objective. By shifting the problem into the frequency domain, we can explicitly identify and dampen the chaotic, high-frequency gradient noise that destabilizes training, enforcing a fundamental layer of smoothness and stability. We demonstrate the power of this method across multiple challenging, non-smooth landscapes, including biological single-cell simulation tasks. Our approach allows the AI to stay perfectly stable and efficiently discover hidden global peaks without getting distracted or misdirected by environmental noise. Ultimately, this bridging of signal processing and generative modeling provides a robust new blueprint for training reliable, AI agents across highly unpredictable real-world domains.