Discriminative geometry improves one-step image generation from pixels
Learning Discriminative Geometry for Drifting Models
Machine LearningComputer Vision and Pattern Recognition
Summary
Generating images quickly by drifting models is hard when using raw pixel data because the way pixels represent images is not good for guiding the changes. The authors show that this problem comes from the way the model weighs examples based on their geometry in the representation space. They introduce a method that learns better representations continuously during training, making the generation step more reliable and improving image quality significantly without needing pretrained encoders. This technique helps the model learn useful shapes and features directly from pixel data, leading to much clearer generated images.
What this means in practice
- •For image generation engineers: Generate higher-quality images in one step without relying on pretrained feature extractors by learning representations directly from raw pixels.
- •For computer vision developers: Improve kernel density estimations in visual models for applications needing fast training and inference, like real-time image synthesis.
Authors
Doudou Zhang, Wenwen Hou, Yilin Chen, Qi Chen
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.