PNDMScheduler combines raw v_prediction outputs in PRK and PLMS instead of noise predictions #14816

Describe the bug

PNDMScheduler with prediction_type="v_prediction" applies the PRK and PLMS combinations to the raw v outputs and only converts the combined result to a noise prediction afterwards, inside _get_prev_sample, using the alphas and the sample of the current step:

The PLMS formula (55 e_t - 59 e_{t-1} + 37 e_{t-2} - 9 e_{t-3}) / 24 and the four PRK stages are derived for noise predictions. v depends on the sample and on alpha_prod_t of the step it was computed at, so combining v values from different timesteps and then converting the sum with the current alpha_prod_t and sample does not give the combined noise prediction. The epsilon path is not affected.

Two consequences:

1. The same denoiser written as a v_prediction model and as an epsilon model gives different samples with PNDM. With DDIMScheduler, DPMSolverMultistepScheduler, UniPCMultistepScheduler and DEISMultistepScheduler the two forms agree to float32 precision (script below).
2. For a model whose noise prediction is exact, epsilon PNDM recovers the clean sample to about 1e-6 at any step count, while v_prediction PNDM leaves an error that shrinks roughly like 1 / steps**2: max abs error 0.28 to 0.30 at 5 steps, 0.043 at 10, 0.011 at 20 and 0.0016 at 50, for both skip_prk_steps=True and False.

Converting each model output to a noise prediction with its own timestep and sample when it arrives in step_prk and step_plms, and having _get_prev_sample take noise predictions only, makes the v_prediction and epsilon runs agree to float32 precision and makes the exact model error vanish for v_prediction as well. I have that change and a regression test ready and would like to open the PR once this is acknowledged. The existing test_full_loop_with_v_prediction golden values change with it, since they were recorded with the current behaviour.

Reproduction

import torch

from diffusers import DDIMScheduler, DPMSolverMultistepScheduler, PNDMScheduler, UNet2DModel, UniPCMultistepScheduler

torch.manual_seed(0)
unet = UNet2DModel(
    sample_size=8,
    in_channels=3,
    out_channels=3,
    layers_per_block=1,
    block_out_channels=(8, 16),
    down_block_types=("DownBlock2D", "AttnDownBlock2D"),
    up_block_types=("AttnUpBlock2D", "UpBlock2D"),
    norm_num_groups=4,
).eval()


def sample_with(scheduler_cls, prediction_type, num_steps=20):
    scheduler = scheduler_cls(prediction_type=prediction_type)
    scheduler.set_timesteps(num_steps)
    alphas_cumprod = scheduler.alphas_cumprod
    sample = torch.randn(1, 3, 8, 8, generator=torch.Generator().manual_seed(1))
    for t in scheduler.timesteps:
        with torch.no_grad():
            eps = unet(sample, t).sample
        if prediction_type == "v_prediction":
            # the exact v output of the same network: v = sqrt(a) * eps - sqrt(1 - a) * x0
            a = alphas_cumprod[t]
            x0 = (sample - (1 - a) ** 0.5 * eps) / a**0.5
            model_output = a**0.5 * eps - (1 - a) ** 0.5 * x0
        else:
            model_output = eps
        sample = scheduler.step(model_output, t, sample).prev_sample
    return sample


for cls in (DDIMScheduler, DPMSolverMultistepScheduler, UniPCMultistepScheduler, PNDMScheduler):
    eps_out = sample_with(cls, "epsilon")
    v_out = sample_with(cls, "v_prediction")
    print(f"{cls.__name__}: max |sample| = {eps_out.abs().max():.3g}, max |epsilon run - v run| = {(eps_out - v_out).abs().max():.3e}")

Logs

DDIMScheduler: max |sample| = 1, max |epsilon run - v run| = 1.583e-04
DPMSolverMultistepScheduler: max |sample| = 565, max |epsilon run - v run| = 2.441e-04
UniPCMultistepScheduler: max |sample| = 565, max |epsilon run - v run| = 2.136e-04
PNDMScheduler: max |sample| = 349, max |epsilon run - v run| = 2.984e+01

With the fix described above the last line becomes max |epsilon run - v run| = 7.629e-05.

System Info

• diffusers 0.41.0.dev0, main at a3e0b8e
• torch 2.14.0+cpu, numpy 2.5.2
• Python 3.13, Linux

AI disclosure: I used an AI coding agent to help find this, to write the reproduction and this report. I have read and checked the report and the scripts myself and I will answer questions personally.

Who can help?

@yiyixuxu @dg845