Learning Discriminative Geometry for Drifting Models

Doudou Zhang; Wenwen Hou; Yilin Chen; Qi Chen

cs.LGOct 3, 2026

Learning Discriminative Geometry for Drifting Models

Authors: Doudou Zhang, Wenwen Hou, Yilin Chen, Qi Chen

Organizations: ELLIS Institute Finland · Aalto University

Abstract

Recently proposed Drifting Models shift iterative distribution refinement from inference to training, enabling effective one-step generation. However, their performance on complex image datasets depends strongly on the representation used to construct the drifting field: pixel-space drifting performs poorly, whereas pretrained feature spaces substantially improve sample quality for reasons that remain unclear. We trace this gap to the discriminative geometry of the representation, which determines sample weighting in kernel density estimation (KDE) and, consequently drift. We introduce persistent representation learning, which continuously learns a more discriminative representation geometry as the generator evolves across batches. We further establish a current-step gradient equivalence between the KDE ratio loss and drift regression loss under matched conditions, connecting density-ratio-based generator optimization to empirical drifting and motivating direct control of the drifting velocity. Across multiple datasets, our method learns effective discriminative representations directly from pixels and reduces FID by approximately $82-95\%$ over the original pixel-space Drifting Models, without pretrained encoders. Adapting pretrained representations and applying velocity clipping provide further gains.

Figures & tables

Figure 1: Persistent Representation Learning Improves Drift Alignment. (a) Drift field constructed in the raw space is disorganized by the nuisance coordinates, while the persistent representation field aligns with the ring shaped target structure. (b) Persistent representation space field reaches substantially higher alignment with the reference field than raw space as training progresses.

Figure 2: Persistent Representations Improve Discrimination and Generation. Pixel-space KDE discrimination remains near chance, while a frozen MoCo-v2 representation improves both discrimination and sample quality, and persistent adaptation improves them further. Dashed lines indicate the loss for indistinguishable distributions and chance-level accuracy, respectively.

Figure 3: Qualitative Comparisons on CIFAR-10. Our method produces substantially clearer and more recognizable samples than standard drifting, with more coherent object structure than SNGAN.

Figure 4: FID over Training Steps. Our method converges to substantially lower FID than standard drifting on SVHN, CIFAR-10, and CIFAR-100.

	SVHN		CIFAR-10		CIFAR-100
Method	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$	IS $\uparrow$	FID $\downarrow$
Standard drifting	2.91	61.56	3.18	124.94	3.46	113.14
Ours	3.24	2.88	8.69	12.81	8.52	20.61
SNGAN	3.04	13.16	7.22	28.59	7.96	30.33

Table 1: Quantitative Comparisons. Our method substantially improves both FID and IS over standard drifting on SVHN, CIFAR-10, and CIFAR-100, and also outperforms SNGAN.

Figure 5: Effect of Persistence in Different Components. Among the four configurations, placing persistence only in the representation achieves substantially lower FID and higher IS than the other configurations, whereas introducing a persistent scalar critic degrades performance.

Table 7

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Generator
	Ablation default	Standard drifting	Ours
Architecture	U-Net
Base channels, mult.	$128$ , $(1,2,2,2)$
Res. blocks / attn. res.	$2$ per stage; $16{\times}16$
Generator Optimizer
Optimizer	AdamW

Appendix

Table 4: Configurations for the Main Benchmark and Ablation Studies.

Panel of Figure 6 — Figure 6: Persistent Representation Learning with a Residual Parameterization.

Figure 7: Representation Similarity across Training. Linear CKA between encoder checkpoints shows a gradual departure from the identity initialization that stabilizes among late-training checkpoints.

Figure 8: Velocity Budgets Relative to Drift Magnitudes. Distribution of

\|\widehat{V}_{\varphi}(u_{i})\|_{2}

at steps 500, 1000, and 2000 of the

V_{\max}=0.005

run. Dashed lines mark the five

V_{\max}

budgets compared in Table 2 .

Figure 9: Effect of Kernel Bandwidth on Pixel-space Drifting. FID over training steps on CIFAR-10 for four KDE temperatures

\tau\in\{0.002,0.02,0.2,2\}

in pixel space. None approaches the

27.41

FID achieved by persistent representation learning within the same 10k-step budget.

Figure 10: FID under the Same Training Time on CIFAR-10. FID against training time within 5 hours.

Figure 11: Additional Qualitative Results on SVHN. Samples generated by our method at the checkpoint attaining the best FID of $2.88$ .

Drifting Models

Drifting models learn transport fields that move generated samples toward a target distribution, often enabling one-step generation.

See all 53 papers

DriftXpress: Faster Drifting Models via Projected RKHS Fields
May 12, 2026Ali Falahati, Elliot Creager, Gautam Kamath +1Image Generation One-Step Generative Modeling
Cheriton School of Computer Science, University of Waterloo and Vector Institute. · Department of Electrical and Computer Engineering, University of Waterloo and Vector Institute. · A*STAR Centre for Frontier AI Research, Singapore.
Lookahead Drifting Model
Apr 10, 2026Guoqiang Zhang, Kenta Niwa, W. Bastiaan KleijnDrifting Models Image Generation
University of Exeter · NTT Communication Science Laboratories · Victoria University of Wellington
Sinkhorn-Drifting Generative Models
Mar 12, 2026Ping He, Om Khangaonkar, Hamed Pirsiavash +2Sinkhorn Divergence Drifting Models
Department of Computer Science, Vanderbilt University · Department of Computer Science, University of California, Davis · Department of Computer Science, Purdue University

May 12, 2026J/K move · Enter open · S save

DriftXpress: Faster Drifting Models via Projected RKHS Fields

Ali Falahati, Elliot Creager, Gautam Kamath, Shubhankar Mohapatra

Cheriton School of Computer Science, University of Waterloo and Vector Institute. · Department of Electrical and Computer Engineering, University of Waterloo and Vector Institute. · A*STAR Centre for Frontier AI Research, Singapore.

Image Generation One-Step Generative Modeling Kernel Methods Drifting Models

Open paper AskPDF arXiv

Drifting Models have emerged as a new paradigm for one-step generative modeling, achieving strong image quality without iterative inference. The premise is to replace the iterative denoising process in diffusion models with a single evaluation of a generator. However, this creates a different trade-off: drifting reduces inference cost by moving much of the computation into training. We introduce DriftXpress, an accelerated formulation of drifting models based on projected RKHS fields. DriftXpress approximates the drifting kernel in a low-rank feature space. This preserves the attraction-repulsion structure of the original drifting field while reducing the cost of field evaluation. Across image-generation benchmarks, DriftXpress achieves comparable FID to standard drifting while reducing wall-clock training cost. These results show that the training-inference trade-off of drifting models can be pushed further without giving up their one-step inference advantage.

Abstract

Drifting Models

Explore similar work

DriftXpress: Faster Drifting Models via Projected RKHS Fields

Lookahead Drifting Model

Sinkhorn-Drifting Generative Models