Breaking the Likelihood-Quality Trade-off in Diffusion Models by Merging Pretrained Experts

AI summary

Overview

• Research area: Generative modeling with diffusion models — specifically, the tension between data likelihood (density estimation) and perceptual sample quality (FID) in image generation.
• Technical level: Intermediate. The paper assumes familiarity with diffusion model notation (noise schedules, SNR, score matching, PF ODE, BPD/FID evaluation), but the core idea is simple enough to follow without it.
• Scope: A training-free, plug-and-play sampling method that switches between two pretrained diffusion models along the denoising trajectory, evaluated on CIFAR-10 and ImageNet32.

What This Paper Is About

Diffusion models for images face a well-documented trade-off: objectives that emphasize high-noise denoising steps produce realistic-looking images but poor likelihoods, while likelihood-oriented training overweights low-noise steps and degrades visual fidelity. Prior work (notably Kim et al., 2021) reports an inverse correlation between likelihood and FID, and the authors cite evidence that perceptual quality is governed by high noise levels while likelihood is highly sensitive to low-level pixel statistics. The goal of this paper is to get both at once, without retraining anything, by merging two pretrained experts that were each optimized for one side of the trade-off.

Key Contributions

1. A general adaptation procedure for reusing a pretrained score model in a different stochastic process. By matching negative log-SNR between the expert's training schedule and a target schedule, a pretrained model can be used at any noise level it was trained on. For variance-preserving (VP) processes, where α² + σ² = 1, this reduces to a simple time remapping: 𝐬_θ(𝐳_t, t) = 𝐬̃_θ(𝐳_t, γ̃⁻¹(γ_t)).
2. A hard-switch expert merging scheme. Rather than training a new network, the method assigns a weight of 1 to exactly one expert at any given time, switching from an image-quality expert at high noise to a likelihood expert at low noise at a chosen intermediate step η.
3. Empirical demonstration on CIFAR-10 and ImageNet32 that the merged model matches or outperforms both base components (pretrained EDM and VDM) on likelihood, FID, or both, with no retraining or fine-tuning.
4. A comparison against other approaches targeting both metrics (Soft Truncation and Consistency Trajectory Models), showing the merged model outperforms Soft Truncation while relying only on standard denoising objectives.

Main Findings

• The trade-off is broken on CIFAR-10 at η = 0.3. This is reported as the best overall operating point: NLL of 2.63 BPD and FID of 2.01, compared with EDM's 3.21 BPD / 2.02 FID and VDM's 2.64 BPD / 9.37 FID. NFE at this setting is 169.
• Pushing η to 0.4 improves likelihood beyond the VDM baseline. NLL reaches 2.62 BPD (versus VDM's 2.64 under the authors' evaluation) with only a slight FID degradation, from 2.02 to 2.14, at 173 NFE.
• On ImageNet32, η = 0.5 matches VDM in likelihood while surpassing EDM in FID. NLL is 3.72 BPD (tying VDM's 3.72) and FID is 6.58 (versus EDM's 7.38 and VDM's 9.85), at 180 NFE.
• The switching threshold η smoothly mediates the trade-off. Across both datasets, likelihood improves monotonically as η increases and FID degrades once η grows large; the reported feasible range is η_min ≈ 0.0394 to η_max ≈ 0.8294.
• The two base models were trained over different noise ranges. VDM over γ^VDM ∈ [−13.3, 5] and EDM over γ^EDM ∈ [−12.43, 8.764], giving the merged model a combined operating range of γ ∈ [−13.3, 8.764], with a linear schedule γ_t = γ_0^VDM + t(γ_1^EDM − γ_0^VDM).
• Comparison with prior work on CIFAR-10. The merged model (η = 0.3: 2.63 NLL, 2.01 FID) outperforms Soft Truncation (3.01 NLL, 3.96 FID). CTM reports better numbers (2.43 NLL, 1.87 FID at NFE 2) but combines multiple loss functions including GAN-based objectives and data augmentation, whereas this method uses only standard denoising objectives. The authors explicitly state they do not claim state-of-the-art performance.
• Qualitative behavior. Using the ODE sampler on CIFAR-10 with a shared initial noise sample, all trajectories are identical up to the switching point. Early switching (larger η) improves likelihood while perceptual quality stays largely unchanged, until excessive reliance on VDM begins to degrade fidelity. The authors also note the two base models frequently produce remarkably similar samples from the same initial noise, despite differing architectures, training objectives, and ELBO weighting.
• NFE increases modestly relative to EDM. On CIFAR-10, NFE rises from 125 for pure EDM to 169 at η = 0.3; on ImageNet32, from 120 to 180 at η = 0.5. These are small compared to what the paper reports for the VDM baselines in its comparison table (206 on CIFAR-10, 158 on ImageNet32 under the authors' evaluation).

Methodology in Plain English

The researchers took two already-trained diffusion models: EDM, which produces visually strong images, and VDM, which is optimized for likelihood. Both were trained with their own noise schedules and their own notion of "how noisy" a sample is at a given point.

The trick is a common currency for noise: the negative log signal-to-noise ratio, γ. Every noise level in every model maps to a value of γ. So when the merged sampler is at time t in its own target schedule, it looks up which time u in EDM's or VDM's schedule has the same γ, and feeds the sample to that model at that corresponding time. For VP processes, the rescaling factors cancel out entirely, so this is literally just plugging a different time value into the network.

Sampling then proceeds straightforwardly from noise: run EDM until reaching a chosen switching threshold η, then hand the partially denoised sample to VDM for the remaining, low-noise steps. Only one number must be chosen — η — and it is selected empirically. Nothing is retrained, no extra supervision is used.

For evaluation, the authors measure FID on 50k generated samples following the EDM evaluation protocol, and report test log-likelihood in bits per dimension two ways: via the standard variational lower bound, and via the probability flow ODE integrated with the log-density tracked through the instantaneous change-of-variables formula, combined with truncated normal dequantization (yielding a tighter bound). The divergence in the ODE likelihood is estimated with the Skilling–Hutchinson trace estimator using a Rademacher distribution, and integration uses the RK45 solver with atol=1e-5 and rtol=1e-5. VDM was trained for 10 million steps on 8×A100 (40GB) GPUs with no data augmentation and batch size 128 for CIFAR-10; for ImageNet32 with 256 channels and a total batch size of 512 for 2 million steps. Since no pretrained EDM checkpoint exists for ImageNet32, the authors trained one using the official EDM code (--cond 0 --arch ddpmpp --duration 1000) for 1000M images at total batch size 1024, with no hyperparameter tuning.

Why This Matters

The paper shows that two models which each lose on one metric can be combined into a system that beats both, using only existing checkpoints and a single scalar hyperparameter. It reframes the likelihood–quality trade-off not as an inherent property of diffusion training, but as something addressable at sampling time.

Real-world applications implied by the paper's framing:

• Text-to-image and text-to-video generation, which the paper lists among the domains where diffusion models have achieved state-of-the-art results, and where both realism and statistical fidelity matter.
• Density estimation and likelihood-based tasks, also explicitly named by the authors as a diffusion application area, where quality-optimized models currently perform poorly.
• Reusing published checkpoints. The method works with publicly available models (the paper points to the public EDM CIFAR-10 checkpoint and a VDM PyTorch re-implementation), so practitioners can upgrade existing models without retraining.
• Compute-constrained deployment, since the approach adds no training cost and only a modest increase in NFE relative to the quality expert.

Industry relevance: training competitive diffusion models is expensive, and the paper's approach converts an expensive retraining problem into a cheap inference-time decision. Modular composition of specialized checkpoints — rather than training one model to do everything — is a practical path to improving deployed generative systems.

Future Directions

• Automated or learned switching mechanisms. The optimal threshold η must currently be determined empirically, and the authors list this as a limitation and a target for future work.
• Integration with advanced samplers. The paper identifies combining the merging scheme with more sophisticated sampling methods as an open direction.
• Extensions to latent and consistency-based diffusion models. All experiments are restricted to pixel-space diffusion models, leaving adaptation to other architectures and training regimes unexplored.
• Generalizing beyond a two-expert hard switch. The framework as described allows score estimates from multiple adapted experts to be combined in an ensemble fashion, for example by convex combination; only the hard two-model switch is instantiated here. The authors also note that performance depends on the characteristics of the merged models, and the method's behavior when the experts disagree more strongly is not characterized.

Target Audience

Researchers and practitioners working on diffusion and flow-based generative models who care about both likelihood and sample quality, and who want to improve existing pretrained systems without retraining. It is also relevant to anyone studying model merging, mixture-of-experts, or efficient reuse of pretrained checkpoints. Readers without background in diffusion notation may find the preliminaries section heavy, but the core method and results are accessible at an intermediate level.

Authors’ abstract

Diffusion models for image generation often exhibit a trade-off between perceptual sample quality and data likelihood: training objectives emphasizing high-noise denoising steps yield realistic images but poor likelihoods, whereas likelihood-oriented training overweights low-noise steps and harms visual fidelity. We introduce a simple plug-and-play sampling method that combines two pretrained diffusion experts by switching between them along the denoising trajectory. Specifically, we apply an image-quality expert at high noise levels to shape global structure, then switch to a likelihood expert at low noise levels to refine pixel statistics. The approach requires no retraining or fine-tuning -- only the choice of an intermediate switching step. On CIFAR-10 and ImageNet32, the merged model consistently matches or outperforms its base components, improving or preserving both likelihood and sample quality relative to each expert alone. These results demonstrate that expert switching across noise levels is an effective way to break the likelihood-quality trade-off in image diffusion models.