Contrastive noise alignment improves generative flow models quality and speed
Beyond Random Couplings: Contrastive Noise Alignment in Generative Flows
Computer Vision and Pattern RecognitionMachine Learning
Summary
Generating images using flow-based models usually involves adding random noise in a way that can make learning difficult and slow. The authors introduce Contrastive Noise Alignment (CNA), a method that actively adjusts the noise used during training to better match the data. This makes the model's job easier, leading to improved image quality and faster generation with fewer steps. Their method maintains the math properties needed for smooth operation while significantly reducing errors compared to earlier approaches.
What this means in practice
- •For machine learning engineers: Improve generative image models for faster and higher-quality outputs with fewer sampling steps.
- •For graphics programmers: Develop more efficient image synthesis pipelines by reducing complexity in noise-data alignment during model training.
Authors
Lennart Wittke, Vinicius Azevedo
Abstract
Diffusion and flow-matching models are typically trained by corrupting data through independently sampled Gaussian noise. While simple and scalable, this forward process induces arbitrary data-noise couplings, forcing the network to learn high-curvature transports between unrelated endpoints. Existing optimal-transport methods reduce this burden by reassigning fixed noise samples to data, but the source noise distribution itself remains passive. To address this, we introduce Contrastive Noise Alignment (CNA), a training-time method that creates dynamic, contrastive couplings by optimizing the noise representations directly. By modeling the noise batch as an interacting particle system, CNA employs a cross-modal InfoNCE objective to align noise particles with their paired data targets. To prevent spatial collapse, this alignment is regularized using an angular entropy term and a radial norm penalty. We show theoretically that this equilibrium asymptotically preserves Gaussian structures, maintaining tractability during inference. Empirically, CNA improves the alignment between noise and data, reduces flow curvature, and provides better generation quality with fewer required sampling steps. For few-step, pixel-space generation (2-4 NFEs), CNA reduces FID by over 50\% compared to standard rectified flow, and by at least 24\% against Optimal Transport baselines.