[consistency distillation] LoRA scripts omits EMA update

[jon-chuang]

Describe the bug

Consistency distillation scripts omit EMA update of teacher unet

I think that this will limit the quality in the training limit, as the teacher will not be updated to produce better one-step denoising.

I think the idea of EMA is that it can mimic something like progressive distillation, which in itself is shown to improve few-step diffusion over baseline.

Reproduction

Check the scripts

Logs

No response

System Info

main

Who can help?

cc @shuminghu regarding whether EMA was necessary to achieve high quality results for LCM.

[dg845]

The LCM distillation scripts should do an EMA update of target_unet here (e.g. for SD 1.X/2.X):

diffusers/examples/consistency_distillation/train_lcm_distill_sd_wds.py

Lines 1344 to 1346 in a551ddf

	if accelerator.sync_gradients:
	# 12. Make EMA update to target student model parameters (`target_unet`)
	update_ema(target_unet.parameters(), unet.parameters(), args.ema_decay)

Note that teacher_unet never gets modified during distillation. The script maintains online and target student unet models during distillation, following the original consistency models (CM) paper and the LCM paper; I believe the reasoning is that if we used only one student unet the consistency distillation target value would depend on the same parameters we want to optimize (since it will also be calculated using unet), causing training instability (see Section 4 of the CM paper).

Note that the LCM-LoRA distillation scripts do not maintain a separate target_unet since we found this to be unnecessary during training.

cc @luosiallen, @patil-suraj

[shuminghu]

@jon-chuang You might have misunderstood teacher_unet. As @dg845 mentioned, teacher_unet is never modified. It's used to get x_prev with DDIM solver to feed into target_unet(which is an EMA update of unet) in the distillation script linked above.

[jon-chuang]

Sorry guys, my observation was limited to the LoRA scripts only.

I figured out a way to apply EMA as follows:

@torch.no_grad()
def update_ema(target_params, source_params, rate=0.99):
    """
    Update target parameters to be closer to those of source parameters using
    an exponential moving average.

    :param target_params: the target parameter sequence.
    :param source_params: the source parameter sequence.
    :param rate: the EMA rate (closer to 1 means slower).
    """
    for targ in target_params:
        src = None
        while (src is None or src.shape != targ.shape):
            src = next(source_params, None)
            assert src is not None
        targ.detach().mul_(rate).add_(src, alpha=1 - rate)
    assert next(source_params, None) is None

I also do the following:

unet.fuse_lora()
update_ema(teacher_unet.parameters(), unet.parameters(), rate=0.99)
unet.unfuse_lora()

Note that teacher_unet never gets modified during distillation.

I see, so I guess this above usage is wrong, I need to maintain an additional target_unet.

Note that the LCM-LoRA distillation scripts do not maintain a separate target_unet since we found this to be unnecessary during training.

I see. Did you find conversely that EMA was unnecessary for non-LoRA as well? Maybe it wasn't that important after all.

[jon-chuang]

causing training instability

Btw, you know what guys, I've seen a lot of training instability for my LCM-LoRA runs (leading to results that resemble overly high CFG), so EMA may be important after all, just saying.

[shuminghu]

Btw, you know what guys, I've seen a lot of training instability for my LCM-LoRA runs (leading to results that resemble overly high CFG), so EMA may be important after all, just saying.

Interesting. LoRA has much fewer number of params need to be updated. It is likely more stable than full FT?

By the way, Improved training on CM paper removed EMA for CT and presents "a theoretical argument indicating that the EMA decay rate for the teacher network should always be zero for CT, although it can be nonzero for CD". Thus, theorectially, we can still use EMA since we are doing CD, though I never delved into the maths to undersand the theory :P

Upon re-inspecting the paper, I also realized where the confusion in naming is from. "teacher network" from the paper correponds to target_unet not teacher_unet in current code space.

[dg845]

I have opened a PR to allow a target model to be optionally used in the LCM-LoRA scripts at #6537.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.