the results of edm2-imagenet-64-s is all-zero

[PeiqinSun]

Hi, thank you for your nice work!

I'm currently trying to port the logic of this code into the EDM2 repository to replicate the ImageNet-64 experiment results. However, I'm encountering issues with poor performance, as the initial training outputs are just black images. Do you happen to have any experience with training a consistency model on ImageNet-64? Specifically, did you observe black images in the early stages of training (less than 1,000 iterations with a batch size of 2,048 per iteration), and did they gradually improve over time? Or did you see meaningful content from the beginning?

Thank you for your help!

[xandergos]

Hi,

EDM2 is not inherently compatible with consistency distillation, and the paper gives a detailed explanation of why and how to fix it. The main changes are using positional embedding instead of Fourier embedding, and performing pixel-norm on the AdaGN layer (see paper for details).

Once these changes are made, there shouldn't be any problems training with EDM2. In fact, I've used a very similar script to train with EDM2 here, which may be useful as a reference.

[PeiqinSun]

I really appreciate your reply.
I did indeed replace the Fourier embedding and perform pixel-norm on the AdaGN layer as mentioned. After comparing with your EDM2 code, most of it seems to match. I'll take a closer look at the code you provided. By the way, have you tried training a consistency model from scratch based on the EDM2 architecture?
It seems like your code is focused on distillation.
I've been tracking the value of out_gain, which is a parameter in the out_conv of EDM2. I found that it's very small, which is causing the images to appear black. My learning rate follows the paper's recommendation of 1e-4. Is this too low? Should I increase it to 0.01?

[xandergos]

My code is focused on distillation. I have no idea what training looks like for EDM2, although I expect it would take much longer to converge.

[PeiqinSun]

Thanks for your reply. I have another question regarding your EDM2 distillation code: "Using positional embedding instead of Fourier embedding, and performing pixel normalization on the AdaGN layer"—does this also apply to model_m_pretrained, which has no gradients? In this case, can the teacher still provide meaningful information?

[xandergos]

In my case I apply the changes to all models. This is I think the most reasonable way to do it, as usually the pretrained model is the same one used for distillation. In theory though it shouldn't matter how the pretrained model is defined since its gradients are never actually used.

[PeiqinSun]

OK. I followed your approach, but encountered F_theta_grad is nan. Have you come across this issue? This isn't particularly easy to debug and seems somewhat similar to the overflow discussed in section 5.1 of the paper. Did you use FP16 or FP32 for your training?

[xandergos]

Sorry for the long wait (lots of stuff going on in life), but in case you're still working on this, the problem might be that you are using weight = 1 / sigma. I found this to cause issues even though its what the paper suggests. Is that relevant to you?