Rethinking Refinement: Correcting Generative Bias without Noise Injection

TL;DR

Rethinking Refinement identifies systematic bias in iterative generative models and presents Bi-stage Flow Refinement (BFR) as a post-hoc solution that preserves deterministic sampling. It introduces Data-space Flow Refinement (DFR) and Latent-space Flow Refinement (LFR), both formulated as ODE-based refinements, to correct bias without retraining the base model. Empirical results show consistent gains in fidelity and coverage across MNIST, CIFAR-10, and FFHQ, with a remarkable MNIST FID of 1.46 achieved using only 1 function evaluation (1-NFE). The framework also demonstrates robustness to base-model shifts and extends to molecular generation, highlighting BFR's generality and practical impact for improving diffusion/flow-based generative systems.

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.

Figures (7)

• Figure 1: Overview of bi-stage flow refinement (BFR). (A) Baseline generation directly maps Gaussian noise $z$ to data samples $x$ via the base generator $G$. (B) Data-space refinement: samples $x_1 = G(z)$ are further refined by a data-space flow $F$ to correct residual bias. (C) Latent-space refinement: noise $z$ is first transformed into a bias-corrected latent $z'$ by a latent flow $F$, followed by generation through $G$.
• Figure 2: Illustration of Bi-stage Flow Refinement (BFR) for training and inference. Left: Training stage, showing how DFR refines generated samples $\hat{x}_1$ and LFR refines latent variables $z_1$. Right: Inference stage, where the trained refiners $F_\psi$ (data-space) and $F_\phi$ (latent-space) are applied via ODE solvers to produce corrected samples.
• Figure 3: Qualitative comparison of MNIST samples before and after refinement. Each row corresponds to a refinement method. DFR and LFR significantly improve visual quality.
• Figure 4: Qualitative comparison of CIFAR-10 samples before and after refinement. Each row corresponds to a refinement method. DFR and LFR improve visual fidelity over the base model.
• Figure 5: Qualitative comparison of FFHQ 256$\times$256 samples before and after refinement. Each row corresponds to a refinement method. DFR and LFR improve visual fidelity over the base model.