Research map / Image generation / Fast diffusion sampling
Nearest Neighbour Score Estimators for Diffusion Generative Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a nearest-neighbour self-normalized importance sampling estimator of the diffusion score that uses FAISS k-NN search over the training set, reducing score/posterior-mean variance and improving consistency training (iCT) and PF-ODE integration.
Problem
Score function estimation for diffusion and consistency models is either biased (neural network) or high-variance (single-sample Monte Carlo on the conditional score), and exact marginal score requires summing over the whole training set.
Approach
Use self-normalized importance sampling with a proposal built from the k nearest training neighbours of the noisy input (equivalent to the top-k posterior elements under Gaussian likelihood), with the remaining mass assigned to the k-th neighbour's likelihood, so the posterior mean is estimated from multiple samples cheaply.
Key results
- On CIFAR-10 the KNN estimator reduces posterior-mean and score MSE to near zero, with peak score MSE roughly 100x lower than STF.
- Consistency training with the KNN estimator achieves FID 3.59 versus 3.67 for the reimplemented iCT and 3.79 for iCT+STF (single-step generation on CIFAR-10).
- Replacing the score network in PF-ODE integration (hybrid with EDM after t_switch) reaches FID 0.5, but this matches forward-process initialization, so the method does not generalize on its own.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace the single-sample score estimator with the KNN score estimator during iCT training | iCT · CIFAR-10 single-step unconditional generation | FID: 3.67 -> 3.59 | Helped |
| Replace the single-sample score estimator with the KNN score estimator during iCT training | iCT · CIFAR-10 single-step unconditional generation | Inception Score: 9.45 -> 9.49 | Helped |
| Use the STF score estimator for iCT training instead of the reimplemented iCT estimator | iCT · CIFAR-10 single-step unconditional generation | FID: 3.67 -> 3.79 | Hurt |
| Use the STF score estimator for iCT training instead of the reimplemented iCT estimator | iCT · CIFAR-10 single-step unconditional generation | Inception Score: 9.45 -> 9.50 | Helped |
| Use the KNN score estimator instead of STF | CIFAR-10 score estimation · CIFAR-10 | score MSE: approximately 100 times better | Helped |
| Use the KNN score estimator instead of STF | CIFAR-10 score estimation · CIFAR-10 | posterior mean MSE | Helped |
| Use the KNN score estimator instead of the EDM network estimator | CIFAR-10 posterior mean estimation · CIFAR-10 | posterior mean MSE | Helped |
| Reduce the KNN search size k (ablation) | KNN consistency training · CIFAR-10 | FID | Helped |
| Reduce the estimator sample size n (ablation) | KNN consistency training · CIFAR-10 | FID | Helped |
| Replace EDM with KNN score estimates for PF-ODE integration | Diffusion sampling · CIFAR-10 PF-ODE hybrid sampling | FID: 0.5 | Mixed |
| Use STF instead of KNN for PF-ODE integration when t_switch < 2 | Diffusion sampling · CIFAR-10 PF-ODE hybrid sampling | FID | Hurt |
Surprising findings
- Hybrid PF-ODE sampling with the KNN estimator reaches FID 0.5 but is matched by forward-process initialization, indicating it reproduces training samples rather than generalizing.
- Reducing n and k improved sample FID for consistency training despite higher estimator variance, suggesting variance may have a regularizing effect.
Limitations
- Standalone sampling does not generalize; it essentially reproduces training data.
- Reliance on l2 distance in pixel space does not capture high-level image similarity.
- Gains over reimplemented iCT are small (FID 3.59 vs 3.67) and original iCT results could not be reproduced.
Experimental setup
Datasets: CIFAR-10
Benchmarks: CIFAR-10
Metrics: MSE bias variance FID Inception Score
Baselines: single-sample Monte Carlo estimator STF EDM iCT consistency distillation consistency training
Builds on: consistency training improved consistency training stable target field EDM FAISS
Closest papers by meaning
- Evaluating the design space of diffusion-based generative models (NeurIPS 2024 · similarity 0.90)
- Training Diffusion-based Generative Models with Limited Data (ICML 2025 · similarity 0.90)
- Accelerating Convergence of Score-Based Diffusion Models, Provably (ICML 2024 · similarity 0.90)
- Neural Network-Based Score Estimation in Diffusion Models: Optimization and Generalization (ICLR 2024 · similarity 0.90)
- O(d/T) Convergence Theory for Diffusion Probabilistic Models under Minimal Assumptions (ICLR 2025 · similarity 0.89)
- Dimension-free Score Matching and Time Bootstrapping for Diffusion Models (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.