Technical Summary: Generative Modeling via Drifting

Problem Statement

Generative modeling is fundamentally the task of learning a mapping $f$ such that the pushforward of a prior distribution $p_{prior}$ (e.g., Gaussian noise) matches the data distribution $p_{data}$, denoted as $f\#p_{prior} \approx p_{data}$. While prevailing paradigms like Diffusion Models and Flow Matching achieve this by iteratively refining samples at inference time (decomposing the complex pushforward map into a chain of feasible transformations), this approach necessitates multiple neural network evaluations (NFE) per sample, leading to high computational costs.

Existing attempts to create one-step generators often rely on distilling pre-trained multi-step models or approximating the trajectories of stochastic differential equations (SDEs) or ordinary differential equations (ODEs) during training. The authors propose a new paradigm that fundamentally shifts the iterative process from inference to training time, enabling high-quality generation with a single neural function evaluation (1-NFE) without relying on SDE/ODE formulations.

Methodology: Drifting Models

The core contribution is the Drifting Model, which formulates generative modeling as the evolution of the pushforward distribution during the training process itself.

1. Training-Time Evolution

Instead of viewing training as merely updating weights to minimize a static loss, the authors view the sequence of models $\{f_i\}$ produced by iterative optimization (e.g., SGD) as a sequence of pushforward distributions $\{q_i\}$. At each iteration, a sample $x_i = f_i(\epsilon)$ is implicitly "drifted" to $x_{i+1} = f_{i+1}(\epsilon)$ due to parameter updates. The goal is to define a mechanism where this drift naturally guides the distribution $q$ toward $p_{data}$.

2. The Drifting Field

To govern this evolution, the authors introduce a drifting field $V_{p,q}(x)$, a function that computes a displacement vector $\Delta x$ for a sample $x$ based on the current generated distribution $q$ and the target data distribution $p$.

• Equilibrium Condition: The field is designed such that $V_{p,q}(x) = 0$ if and only if $p = q$. When the distributions match, samples stop drifting, reaching an equilibrium.
• Anti-Symmetry: The field satisfies $V_{p,q}(x) = -V_{q,p}(x)$. This property ensures that if the distributions are swapped, the direction of drift reverses, guaranteeing that $V=0$ when $p=q$.

3. Instantiation via Attraction and Repulsion

The specific drifting field used in the paper is inspired by the mean-shift algorithm and is constructed using positive samples from $p$ and negative samples from $q$:
$$V_{p,q}(x) = V^+_p(x) - V^-_q(x)$$
Where:

• $V^+_p(x)$ is an attractive force pulling $x$ toward the data distribution $p$.
• $V^-_q(x)$ is a repulsive force pushing $x$ away from the generated distribution $q$.
• These forces are computed using a kernel function $k(x, y)$ (specifically an exponential kernel based on $\ell_2$ distance) and normalized via softmax operations to ensure stability.

4. Training Objective

The training objective is derived from the fixed-point relation at equilibrium: $f(\epsilon) = f(\epsilon) + V_{p,q}(f(\epsilon))$. The loss function minimizes the distance between the current prediction and a "frozen target" (the prediction plus the drift):
$$L = \mathbb{E}_{\epsilon} \left[ \| f_\theta(\epsilon) - \text{stopgrad}(f_\theta(\epsilon) + V_{p,q}(f_\theta(\epsilon))) \|^2 \right]$$
This formulation allows the optimizer to evolve the distribution by moving samples toward their drifted versions. Crucially, the gradient does not backpropagate through the distribution $q$ itself, avoiding the complexities of differentiating through distributions.

5. Feature Space and Guidance

• Feature Space: To handle high-dimensional data, the drifting loss is computed in a feature space using a pre-trained self-supervised encoder (e.g., ResNet-style MAE). This ensures that semantically similar samples remain close, providing richer gradient signals.
• Classifier-Free Guidance (CFG): The method naturally supports CFG by mixing unconditional data samples into the negative sample distribution during training, allowing for controllable generation at inference without retraining.

Key Contributions

1. New Paradigm: Introduces "Drifting Models," a framework where the pushforward distribution evolves during training time, eliminating the need for iterative inference.
2. Drifting Field: Proposes a novel loss mechanism based on a drifting field that achieves equilibrium when the generated and data distributions match, governed by attraction to data and repulsion from the current generation.
3. State-of-the-Art One-Step Generation: Demonstrates that a single-pass network can achieve performance competitive with, or superior to, multi-step diffusion and flow-based models.
4. Robustness to Mode Collapse: Toy experiments suggest the method is robust to mode collapse; if the generated distribution collapses to one mode, the attraction from other modes in the data distribution continues to drive the drift, pushing the distribution to evolve.

Experimental Results

The authors evaluate Drifting Models on ImageNet 256×256 in both latent and pixel spaces.

• Latent Space: Using a standard SD-VAE tokenizer, the model achieves a 1-NFE FID of 1.54 (with a larger L/2 model). This outperforms all previous one-step methods (e.g., iMeanFlow, AdvFlow) and is competitive with multi-step diffusion models (e.g., SiT-XL/2 with 250 steps achieves 2.06 FID).
• Pixel Space: Without a latent tokenizer, the model achieves a 1-NFE FID of 1.61, substantially outperforming previous pixel-space one-step methods (e.g., StyleGAN-XL at 2.30 FID) and competing with multi-step pixel-space models.
• Robotics Control: In the Diffusion Policy setting, the one-step Drifting Model matches or exceeds the performance of multi-step diffusion policies (100 NFE) on various robotic manipulation tasks.

Significance and Claims

The paper claims to open new opportunities for high-quality, efficient generative modeling. By reframing iterative neural network training as a mechanism for distribution evolution, the work challenges the necessity of SDE/ODE formulations for generative tasks.

The authors emphasize that their method:

• Naturally admits one-step inference (1-NFE), significantly reducing computational cost at deployment.
• Does not rely on distillation from pre-trained multi-step models, training from scratch.
• Offers a conceptually distinct alternative to diffusion and flow matching, focusing on the update rule ($x_{i+1} = x_i + \Delta x_i$) during training rather than inference.

The authors remain modest regarding theoretical guarantees, noting that while $p=q \implies V=0$ is proven, the converse ($V \to 0 \implies p \to q$) holds empirically and under specific non-degeneracy assumptions but is not guaranteed for arbitrary vector fields. They view their work as a starting point for exploring other realizations of distribution evolution mechanisms.