Progressive Distillation for Accelerated Sampling

Module: Generative AI | Difficulty: Advanced

Algorithm

1. Train teacher with steps
2. For : a. Initialize student from b. Train to take 2 steps per step c.

Student Loss

where is the 2-step estimate from .

def progressive_distill_step(teacher, student, x0, t, scheduler):
    t1 = scheduler.time_shift(t, mu=0.0)
    t2 = scheduler.time_shift(t-1, mu=0.0)
    x_t = scheduler.add_noise(x0, t1)
    with torch.no_grad():
        x_t1 = scheduler.step(teacher, x_t, t1, t2)
    s1 = student(x_t, t1)
    s2 = student(x_t1, t2)
    return ((s1-x0)**2).mean() + ((s2-x0)**2).mean()

Research Insight: Progressive distillation achieves near-lossless quality down to 4-8 steps. Below 4 steps, the quality degrades because the ODE trajectory becomes too non-linear for 2-step approximation.