This project explores advanced sampling techniques for diffusion models that go beyond the standard reverse diffusion process. By implementing search-based methods like Best-of-N sampling and Beam Search, we aim to improve image generation quality by allowing the model to explore more of the generation space and select optimal paths.
Traditional diffusion models follow a deterministic reverse sampling process, generating images step-by-step from noise to final output. This project investigates whether search-based sampling methods can improve generation quality by:
- Exploring multiple candidate paths at each denoising step
- Evaluating path quality using forward likelihood estimation
- Selecting optimal trajectories through the generation space
We developed a novel Forward Negative Log-Likelihood (NLL) metric to evaluate the quality of diffusion model predictions at each step:
NLL = 0.5 * log(2π * variance_t) + 0.5 * (||ε_pred||² / d)
Where:
ε_predis the predicted noisevariance_tis the noise variance at timestep tdis the dimensionality of the latent space- Lower values indicate better predictions
This metric assesses how plausible the predicted noise is under the forward diffusion process, providing a reliable proxy for prediction confidence and model performance.
- Baseline implementation using the default Stable Diffusion 1.5 pipeline
- Single deterministic path through the reverse process
- Minimal computational overhead
- Manual implementation of the diffusion forward process
- Foundation for implementing custom sampling strategies
- Direct control over each denoising step
- Implementation of forward NLL computation
- Visualization of prediction quality throughout the sampling process
- Analysis of score function estimates and their statistics
- Generate N candidate samples at each step
- Select the candidate with the lowest forward NLL
- Excellent results with minimal computational overhead
- Maintain multiple beams (candidate paths) throughout generation
- Use lookahead evaluation for better path selection
- More sophisticated but computationally intensive approach
| Method | Generation Time | Quality Improvement | Computational Overhead |
|---|---|---|---|
| Standard Pipeline | ~20 seconds | Baseline | 1x |
| Best-of-N (N=10) | ~2 min | ✅ Significant | ~6x |
| Beam Search | ~1 hour | ❌ Minimal | ~100x |
- ✅ Excellent image quality improvement
- ✅ Minimal computational overhead (~50% increase)
- ✅ Easy to implement and tune
- ✅ Consistent results across different prompts
- ❌ Drastically increased computation time (~1 hour per image on laptop GPU RTX 4050)
- ❌ No visible quality improvement over Best-of-N
- ❌ Complex implementation with lookahead evaluation
- ❓ May benefit from better hyperparameter tuning
pip install diffusers transformers torch torchvision pillow tqdm numpy matplotlib-
Standard Generation:
python 1_standard_pipeline.py
-
Best-of-N Sampling (Recommended):
python 4_best_of_N_sampling.py
-
Beam Search (Experimental):
python 5_beam_search_diffusion_sampling.py
-
Likelihood Analysis:
python 3_step_linkelihood_estimation.py
All scripts use the same prompt by default:
prompt = "A fantasy landscape with mountains and rivers"Key parameters can be modified in each script:
num_inference_steps: Number of denoising steps (default: 50)guidance_scale: Classifier-free guidance strength (default: 7.5)N: Number of candidates for Best-of-N (default: 10)beam_width: Number of beams for beam search (default: 3)
The project generates several image files:
generated_image.png- Standard pipeline outputgenerated_image_for.png- Forward process implementationgenerated_image_lh.png- Likelihood estimation outputgenerated_image_best_of_N.png- Best-of-N sampling resultgenerated_image_beam_search.png- Beam search resultintermediate_x0_predictions.png- Visualization of intermediate predictions
-
Dataset Generation & Metrics
- Generate a comprehensive dataset using different sampling methods
- Implement CLIP-based automatic evaluation metrics
- Conduct human evaluation studies for subjective quality assessment
-
Advanced Search Methods
- Monte Carlo Tree Search (MCTS) for diffusion sampling
- Adaptive lookahead strategies
- Multi-objective optimization (quality vs. speed)
-
Performance Optimization
- GPU memory optimization for batch processing
- Parallel candidate evaluation
- Early stopping criteria based on convergence
-
Theoretical Analysis
- Mathematical foundation for forward likelihood estimation
- Convergence guarantees for search-based methods
- Optimal exploration-exploitation trade-offs
-
Model Improvements
- Integration with newer diffusion architectures (SDXL, SD3)
- Learned evaluation functions instead of hand-crafted NLL
- Adaptive sampling strategies based on content complexity
-
Applications
- Extension to other domains (audio, video, 3D)
- Integration with fine-tuned models
- Real-time generation with search-based quality enhancement
This project demonstrates that Best-of-N sampling provides an excellent trade-off between generation quality and computational cost. While Beam Search shows the potential for more sophisticated approaches, it requires significant optimization to be practical.
The forward likelihood estimation method proves to be a valuable tool for evaluating and guiding the diffusion sampling process, opening new avenues for improving generative model performance.
- Base Model: Stable Diffusion 1.5 (
runwayml/stable-diffusion-v1-5) - Scheduler: DDPM for stochastic sampling
- Precision: Float32 for accurate NLL computation
- Resolution: 512x512 pixels
- Minimum: CUDA-capable GPU with 8GB VRAM
- Recommended: RTX 3080/4070 or better for reasonable beam search performance
- CPU Fallback: Available but significantly slower
This project is part of ongoing research into advanced sampling methods for diffusion models. Contributions and feedback are welcome!