PaperlayerView exact revision on arXiv

Source-linked AI summary

Generative Modeling via Drifting

Mingyang Deng, He Li, Tianhong Li, Yilun Du, Kaiming He

arXiv:2602.04770v2cs.LGcs.CV

TL;DR

Generative modeling seeks mappings whose pushforward distributions match data, but common diffusion and flow approaches rely on iterative inference. Drifting Models instead evolve the pushforward during training with a drifting field and achieve one-step ImageNet generation with FID 1.54 in latent space and 1.61 in pixel space.

  • Problem

    Generative modeling requires learning a mapping whose pushforward distribution matches the data, while prevailing diffusion and flow approaches realize this mapping through iterative inference.

  • Method

    Drifting Models use a single-pass network, a distribution-dependent drifting field, and a drift-minimizing objective to evolve the pushforward distribution during iterative training.

  • Results

    1-NFE FID 1.54 on ImageNet 256×256 in latent space and 1.61 in pixel space demonstrate strong one-step generation results.

  • Takeaways & Limitations

    The approach reframes iterative training as distribution evolution and naturally supports one-step inference for high-quality, efficient generation.

  • Takeaways & Limitations

    The method was unable to work on ImageNet without a feature encoder, and zero drift does not generally imply matching distributions for arbitrary vector fields.

Abstract

from arXiv · show

Generative modeling can be formulated as learning a mapping f such that its pushforward distribution matches the data distribution. The pushforward behavior can be carried out iteratively at inference time, for example in diffusion and flow-based models. In this paper, we propose a new paradigm called Drifting Models, which evolve the pushforward distribution during training and naturally admit one-step inference. We introduce a drifting field that governs the sample movement and achieves equilibrium when the distributions match. This leads to a training objective that allows the neural network optimizer to evolve the distribution. In experiments, our one-step generator achieves state-of-the-art results on ImageNet at 256 x 256 resolution, with an FID of 1.54 in latent space and 1.61 in pixel space. We hope that our work opens up new opportunities for high-quality one-step generation.

1. Introduction

Drifting Models shift the iterative evolution of generative distributions from inference into training, enabling a single-pass network and one-step generation. A drifting field guides samples toward equilibrium when generated and data distributions match, supporting strong ImageNet results.

  • Core idea: Diffusion and Flow Matching progressively transform noisier samples toward the data distribution through inference-time chains of feasible transformations.Their iterative computation decomposes a complex pushforward map into multiple transformations.
  • Core idea: Drifting Models learn a pushforward map that evolves during training, replacing iterative inference with a single-pass, non-iterative network.Training updates the map and thereby evolve the pushforward distribution.
  • Training mechanism: A drifting field depends on generated and data distributions and becomes zero when they match, defining an equilibrium with no further sample drift.The field governs sample movement during training.
  • Training mechanism: The training objective minimizes sample drift, inducing movements that evolve the pushforward distribution through iterative optimization such as SGD.The paper also specifies the drifting field, neural network, and training algorithm.
  • Results: 1-NFE FID 1.54 on ImageNet 256×256 establishes state-of-the-art performance among single-step methods under the standard latent-space protocol.The result remains competitive with multi-step diffusion- and flow-based models.

2. Related Work

Related approaches include iterative diffusion and flow models, one-step distillation and training-from-scratch methods, and classical one-step or distribution-matching generators. Drifting Models differ by explicitly governing sample movement during training without adversarial optimization.

  • Diffusion-/Flow-based Models: Diffusion and flow-based models formulate noise-to-data mappings through SDEs or ODEs and generate samples with multiple network evaluations.Their inference-time updates can use numerical solvers such as Euler.
  • One-step generation: Distillation methods compress pretrained multi-step models into one-step generators, while other methods train one-step diffusion or flow models from scratch.Training-from-scratch approaches approximate trajectories induced by SDE or ODE dynamics.
  • Classical generators: Unlike GANs, Drifting Models use a single-pass generator whose goodness is evaluated by a loss function without adversarial optimization.Both approaches map noise to data with a network f.
  • Classical generators: Classical VAEs can generate in one step from a Gaussian prior, although prevailing applications often use priors learned from diffusion or autoregressive models.In those applications, VAEs effectively act as tokenizers.
  • Classical generators: Normalizing Flows conceptually generate in one step through an inverse network, but require invertible architectures and computable Jacobians.They learn mappings from data to noise and optimize sample log-likelihood.
  • Distribution matching: Moment-matching methods minimize MMD between generated and data distributions, whereas Drifting Models explicitly govern sample drifts at training time.Both approaches use kernel-related concepts, including positive and negative samples.
  • Distribution matching: Drifting fields use positive data samples and negative generated samples, creating a conceptual connection to contrastive representation learning.This connection has also appeared in generative models such as GANs and Flow Matching.

3. Drifting Models for Generation

Drifting Models evolve the generator’s pushforward distribution during training, replacing iterative inference with a single-pass network. A drifting field moves samples toward equilibrium, where matching generated and data distributions produce zero drift.

  • Core paradigm: Drifting Models evolve the pushforward distribution during training, enabling one-step generation with a single-pass, non-iterative network.The training process iteratively updates the network and thereby evolves its pushforward distribution.
  • Pushforward at training time: The generator output distribution q is the pushforward of the noise distribution pϵ under network f.The network maps noise ϵ to samples x, whose distribution is denoted q.
  • Drifting field: A drifting field V computes sample changes from the current generated distribution q and data distribution p, governing training-time sample movement.The field is defined as a function that computes ∆x from x and depends on both distributions.
  • Equilibrium: When q = p, an anti-symmetric drifting field becomes zero for every sample, establishing an equilibrium.The paper notes that the converse does not generally hold for arbitrary drifting fields.
  • Training objective: The training objective minimizes drift indirectly by moving network predictions toward frozen, drifted targets rather than back-propagating through the distribution-dependent field.The loss equals the expected squared norm of the drifting field, while stop-gradient supplies the frozen target.
  • Drifting-field design: The proposed kernelized field combines attraction toward data samples and repulsion from generated samples through weighted mean-shift vectors.Its instantiated form uses the vector difference y+ − y− and jointly normalized kernel weights while preserving anti-symmetry.
  • Feature-space drifting: The formulation extends to feature space, where the drifting loss aims to match the feature-space pushforward distributions ϕ#q and ϕ#p without requiring paired samples.A feature extractor can operate on real or generated samples before applying the loss.

4. Implementation for Image Generation

The image-generation implementation primarily operates in a 32×32×4 latent space using a DiT-like generator and feature-space drifting loss. It also supports pixel-space generation and configurable conditioning.

  • Training configuration: Training on ImageNet targets 256×256 image generation and primarily uses a feature-space drifting loss.The feature extractor is an image encoder, with ResNet-style models drawing features from multiple stages and combining losses across scales.
  • Latent-space generation: The default ImageNet setup generates 32×32×4 latents using the standard SD-VAE tokenizer.Generation is performed in latent space before decoding to images.
  • Architecture: The generator uses a DiT-like architecture with 32×32×4 Gaussian noise input, same-dimensional latent output, and patch size 2.adaLN-zero handles class conditioning or other additional conditioning.
  • Conditioning: CFG conditioning samples a scale α during training and allows α to be freely varied at inference without retraining.Negative samples are prepared according to the sampled conditioning scale.
  • Batching: Each training epoch is defined by generated-sample count, corresponding to Ndata/B iterations for a dataset of size Ndata.The effective batch size with class labels is B = Nc×N.
  • Toy evolution experiment: A 2D toy experiment shows q evolving toward a bimodal target p from three initializations without mode collapse.The initializations begin between the modes, far from both modes, or collapsed onto one mode.
  • Pixel-space generation: Pixel-space generation is supported with 256×256×3 inputs and outputs, patch size 16, and a pixel-space feature extractor.The primary experiments nevertheless focus on latent-space generation.

5. Experiments

Experiments show that Drifting Models evolve generated distributions toward targets without mode collapse, and perform strongly in one-step ImageNet and robotic-control generation.

  • 5.1. Toy Experiments: The toy experiment approximates a bimodal target without mode collapse, even when initialized in a collapsed single-mode state.Other target modes attract samples, allowing the generated distribution to continue evolving.
  • 5.1. Toy Experiments: The toy training loss decreases as the generated distribution converges to the target, supporting drift reduction as an equilibrium signal.The loss equals ∥V∥2 in the reported 2D cases.
  • 5.1. Toy Experiments: Breaking anti-symmetry causes catastrophic failure because the attraction from data and repulsion from generated samples no longer cancel at equilibrium.The anti-symmetric configuration is the default ablation case.
  • 5.2. ImageNet Experiments: Under a fixed compute budget, increasing positive and negative sample counts improves generation quality.The study fixes 100 epochs and batch size B = Nc×Npos = 4096.
  • 5.2. ImageNet Experiments: The customized latent-MAE encoder performs best among compared feature encoders, with wider encoders, longer pre-training, and classifier fine-tuning improving quality to 3.36 FID.Standard SimCLR and MoCo encoders also achieve competitive results.
  • 5.2. ImageNet Experiments: On ImageNet, the method could not be made to work without a feature encoder, even when using a latent VAE.The paper leaves further study of this limitation for future work.
  • 5.2. ImageNet Experiments: The method achieves 1.54 FID with native 1-NFE latent-space generation and 1.61 FID in pixel space.The latent result outperforms previous 1-NFE methods, while the pixel-space result uses 87G FLOPs versus StyleGAN-XL’s 1574G FLOPs at 2.30 FID.
  • 5.3. Experiments on Robotic Control: In robotic control, the 1-NFE Drifting Policy matches or exceeds Diffusion Policy using 100 NFE across single-stage and multi-stage tasks.The evaluation includes four single-stage tasks and two multi-stage tasks.

6. Discussion and Conclusion

Drifting Models shift iterative distribution evolution from inference into training, enabling one-step generation while leaving theoretical and practical questions open.

  • Discussion and Conclusion: The method models pushforward-distribution evolution during iterative training rather than iterative inference, enabling one-step generation.The update rule xi+1 = xi + ∆xi is applied during training optimization.
  • Discussion and Conclusion: The theoretical converse from zero drifting field to matching distributions does not generally hold, although sufficient conditions are provided for the kernelized formulation.The paper also reports that decreasing drift norm correlates empirically with improved generation quality.
  • Discussion and Conclusion: Several practical design choices may remain sub-optimal, including the drifting-field kernels, feature encoder, and generator architecture.These components remain open for future exploration.
  • Discussion and Conclusion: The broader perspective reframes iterative neural-network training as a mechanism for distribution evolution, contrasting it with diffusion- and flow-based differential equations.The authors suggest exploring other realizations of this mechanism.

A.2. Generator Architecture

The generator is a single-pass DiT-style Transformer that maps noise and conditioning to latent- or pixel-space outputs while supporting continuous and discrete random variables.

  • Generator Architecture: The generator receives random noise with conditioning and produces outputs in either SD-VAE latent space or directly in pixel space.Latent outputs have shape 32×32×4, while pixel outputs have shape 256×256×3.
  • Generator Architecture: The architecture uses a DiT-style Transformer with patchified Gaussian noise, adaLN and in-context conditioning, and unpatchification to the target shape.The implementation also uses SwiGLU, RoPE, RM-SNorm, and QK-Norm.
  • Generator Architecture: The drifting-field computation is organized as an algorithm that takes generated samples, positive samples, negative samples, and a temperature.Pairwise distances are computed between x and the positive or negative samples, excluding self-comparisons when needed.
  • Generator Architecture: The implementation forms logits from negative distances, normalizes them along both sample dimensions, and splits the resulting weights into positive and negative components.The positive and negative drifts are weighted matrix products with the corresponding sample sets.
  • Generator Architecture: Random style embeddings augment the conditioning vector with discrete codebook indices, and the framework permits noise distributions beyond Gaussians.The generator models the pushforward of the joint continuous and discrete random variables.

A.3. Implementation of ResNet-style MAE

The drifting loss uses multi-scale ResNet-style MAE features, normalized feature distances, and per-feature drift terms that are summed across representations.

  • Implementation of ResNet-style MAE: A customized convolutional ResNet-style MAE serves as the feature encoder for the drifting loss, producing multi-scale features through a U-Net-style encoder-decoder.Only the ResNet-style encoder is used for feature extraction during drifting-loss computation.
  • Implementation of ResNet-style MAE: The encoder follows a four-stage ResNet design with GroupNorm and outputs feature maps at four spatial scales and channel widths.The architecture uses residual blocks with 3×3 convolutions and downsampling in later stages.
  • Implementation of ResNet-style MAE: Latent inputs use shape 32×32×4, while pixel inputs are patchified from 256×256×3 into 32×32×192 before ResNet processing.The decoder returns outputs to the original input shape using deconvolutions and skip connections.
  • Implementation of ResNet-style MAE: The MAE zeroes independently masked 2×2 patches with 50% probability and minimizes ℓ2 reconstruction loss on masked regions.Training uses AdamW, batch size 8192, EMA decay 0.9995, and random resized-crop augmentation.
  • Implementation of ResNet-style MAE: Feature extraction evaluates multiple encoders and computes separate drifting losses across scales and spatial locations before summing them.The study includes MoCo, SimCLR, ConvNeXt-V2, and the customized MAE encoder.
  • Implementation of ResNet-style MAE: Richer feature sets improve performance while adding little computational overhead after the feature encoder has run.Feature and drift normalization reduce sensitivity to feature magnitude and support different encoders and feature collections.

A.9. Training Loop

Each training iteration samples labels, conditioning, noise, and comparison samples, computes feature-based drifting loss, and updates the generator parameters.

  • Training Loop: The training loop samples class labels, a CFG scale, and noise, then feeds the conditioned noise to the generator.The resulting generated samples are compared with same-class positive samples and unconditional samples for CFG.
  • Training Loop: Generated, positive, and unconditional samples are passed through the feature encoder, and their features are used to compute the drifting loss.The loss is evaluated after feature extraction for all sample groups.
  • Training Loop: Backpropagation and parameter updates evolve the model during training, while the ablation and final settings are documented for pixel-space generation.The cited tables distinguish ablation settings from final pixel-space configurations.

B.1. Ablations on Pixel-Space Generation

Pixel-space performance depends strongly on feature-encoder strength, while longer training and larger models improve results to an FID of 1.61.

  • Feature encoder: FID 32.11 with a weaker MAE encoder improves to FID 9.35 with a stronger MAE encoder.The encoder choice has a more significant effect on pixel-space generation quality than in the latent setting.
  • Training duration and model size: FID 1.61 is achieved for pixel-space generation after training pixel-space models for 640 epochs.This is fewer epochs than the 1280 used for the latent counterpart, and longer training is expected to improve performance further.
  • Training duration and model size: The pixel-space FID of 1.61 is the result reported in the main paper’s Table 6.

B.2. Ablation on Kernel Normalization

Kernel normalization affects performance, with two-axis softmax performing best, while the method remains reasonably effective without normalization.

  • Normalization variants: 8.46 FID is obtained with softmax normalization over both the x and y axes, improving over y-only softmax at 8.92 FID.The ablation uses a B/2 model trained for 100 epochs.
  • Normalization variants: Without normalization, performance remains decent, demonstrating robustness to the normalization design.The no-normalization variant corresponds to setting Z = 1.
  • Equilibrium condition: All three normalization variants satisfy the equilibrium condition Vp,q(x) = 0 when p = q.This equilibrium property explains why even the destructive no-normalization setting avoids catastrophic failure.

B.3. Ablation on CFG

CFG increases image-quality scores but can worsen distributional coverage beyond the FID sweet spot; the best L/2 FID occurs at α=1.0.

  • CFG behavior: Increasing CFG scale raises IS, while increases beyond the FID sweet spot worsen FID.The observed trade-off is between distributional coverage measured by FID and per-image quality measured by IS.
  • CFG behavior: The L/2 model achieves its optimal FID at α=1.0, commonly treated as without CFG in diffusion and flow models.Unlike standard CFG, inference with this method does not require an unconditional model, although training uses unconditional real samples as negatives.
  • Sample analysis: Generated samples are visually distinct from their top-10 nearest ImageNet training images retrieved using CLIP features.The nearest-neighbor visualization suggests novelty rather than direct memorization of training samples.
  • Sample analysis: The paper compares uncurated samples with improved MeanFlow, identified as the current state-of-the-art one-step method.Figures 7–10 show the model’s samples, while Figures 11–15 provide side-by-side comparisons.
  • Zero-drift identifiability: Under the stated non-degeneracy setup, zero drift implies p = q for the analyzed finite-dimensional distribution class.The argument uses linearly independent interaction vectors and probability-density normalization.

C.2. The Drifting Field of MMD

The paper relates drifting fields to MMD gradients, showing that MMD is a special case while the drifting-field framework supports broader design choices.

  • General connection: A discrepancy loss minimized at p = q can induce a drifting field proportional to the negative gradient with respect to generated samples.The paper’s formulation only requires the induced field to vanish when the distributions match.
  • MMD-derived field: The gradient of the MMD loss yields the underlying drifting field corresponding to the MMD loss L_MMD^2.For a radial kernel, the field uses the kernel derivative with respect to sample distance.
  • Normalization caveat: Without normalization, the Gaussian-kernel drift no longer satisfies the assumptions underlying Algorithm 2, and its mean-shift interpretation breaks down.
  • Differences from MMD: Normalized kernels distinguish the general drifting formulation from standard MMD-based methods.The normalized kernel divides k(x,y) by an expectation-based factor Z(x), and the paper states this choice cannot be naturally derived from the MMD perspective.
  • Differences from MMD: The V-centric formulation supports flexible step sizes and naturally extends to classifier-free guidance.The paper presents these capabilities as practical possibilities enabled by the drifting-field formulation.
  • Empirical comparison: Although a special case reduces to MMD, the authors report that they could not obtain reasonable experimental results using the MMD framework.
Loading 2602.04770v2…