Flow Along the K-Amplitude for Generative Modeling

TL;DR

K-Flow introduces a flow-matching framework that flows along the $K$-amplitude to model data across multiple scales. By defining the $K$-amplitude decomposition (Fourier, Wavelet, PCA) and treating the scaling parameter $k$ as a time-like dimension, it builds localized, frequency-aware vector fields and trainable interpolants to progressively reconstruct data from frequency bands. The approach achieves competitive unconditional and class-conditioned image generation, enables explicit scale steerability for controlled outputs, and attains state-of-the-art performance in molecular assembly tasks, highlighting the practical value of explicit multi-scale, frequency-domain modeling. Overall, K-Flow offers a principled, steerable, and scalable framework for multi-scale generative modeling with broad applicability to images and scientific data.

Abstract

In this work, we propose a novel generative learning paradigm, K-Flow, an algorithm that flows along the $K$-amplitude. Here, $k$ is a scaling parameter that organizes frequency bands (or projected coefficients), and amplitude describes the norm of such projected coefficients. By incorporating the $K$-amplitude decomposition, K-Flow enables flow matching across the scaling parameter as time. We discuss three venues and six properties of K-Flow, from theoretical foundations, energy and temporal dynamics, and practical applications, respectively. Specifically, from the practical usage perspective, K-Flow allows steerable generation by controlling the information at different scales. To demonstrate the effectiveness of K-Flow, we conduct experiments on unconditional image generation, class-conditional image generation, and molecule assembly generation. Additionally, we conduct three ablation studies to demonstrate how K-Flow steers scaling parameter to effectively control the resolution of image generation.

Figures (12)

• Figure 1: Unconditional generation using K-Flow using three types of $K$-amplitude decomposition: Fourier, Wavelet, and PCA.
• Figure 2: Pipeline of K-Flow. In this figure, we have a bat figure as the input and three inverted images after three transformations at different granularities.
• Figure 3: On the low-scaling hypothesis. The graph illustrates the relative norm distribution for each scaling component as defined by the wavelet decomposition in the latent space. It can be observed that the low-scaling component exhibits a significantly higher norm (energy), nearly twice that of the high-scaling component.
• Figure 4: Projection Error Comparison with Different Models. The graph illustrates the PCA projection errors of two models throughout the entire flow process, with distinct segments marked by dashed lines. The red and blue lines represent the original latent flow matching (LFM) and the K-Flow with two amplitude components, respectively. The projection error is quantified by the reconstruction error for each generation step from the PCA compression, using the first two principal components. Owing to the scaling-aware nature of our flow, the low-amplitude portion (the initial part of the curve) resides in a relatively high-dimensional space, resulting in higher projection errors for the two-dimensional PCA projection.
• Figure 5: Comparison of multi-scale modeling: pixel data space and K-Amplitude space.