Rethinking Refinement: Correcting Generative Bias without Noise Injection

TL;DR

This paper introduces a post-hoc bias correction method called Bi-stage Flow Refinement (BFR) for generative models, improving sample quality without noise injection or multi-step resampling, demonstrated on multiple datasets.

Contribution

The paper proposes BFR, a novel flow-matching-based refinement framework that corrects generative bias deterministically at different stages without perturbing sampling dynamics.

Findings

BFR improves sample fidelity and coverage across datasets.

Latent space refinement achieves state-of-the-art FID on MNIST.

Method maintains sample diversity while enhancing quality.

Abstract

Generative models, including diffusion and flow-based models, often exhibit systematic biases that degrade sample quality, particularly in high-dimensional settings. We revisit refinement methods and show that effective bias correction can be achieved as a post-hoc procedure, without noise injection or multi-step resampling of the sampling process. We propose a flow-matching-based \textbf{Bi-stage Flow Refinement (BFR)} framework with two refinement strategies operating at different stages: latent space alignment for approximately invertible generators and data space refinement trained with lightweight augmentations. Unlike previous refiners that perturb sampling dynamics, BFR preserves the original ODE trajectory and applies deterministic corrections to generated samples. Experiments on MNIST, CIFAR-10, and FFHQ at 256x256 resolution demonstrate consistent improvements in fidelity and coverage; notably, starting from base samples with FID 3.95, latent space refinement achieves a \textbf{state-of-the-art} FID of \textbf{1.46} on MNIST using only a single additional function evaluation (1-NFE), while maintaining sample diversity.