Why my prior phase trains to Loss inf

[xsddys]

Hi, very glag to learn your work.
Im now train on another dataset, in the vae train phase, it seens don't have any problem.I have the MMD-CD and MMD-EMD both in 1e-4, there is my latent sapce and reconstruction

and

However, when I tried to trian DDM(the prior phase) in my vae model, there is worst loss nan

I tried to lower down the lr, pvd_mse_loss, time_dim, and up the Batchsize. The loss nan then turn to loss inf, a little better than nan.And then, what ever hyperparam I try, it still stay in loss inf. It seems nothing to happen even when I lower down the lr to 1e-8

I also tried to lower down the timesteps lion used, but it still help nothing

Could you give me some advice?

[ZENGXH]

perhaps you could try
train your VAE for longer (and increase your KL loss weight for VAE training a bit, e.g. 5x to 10x), until you observe your latent become noisier (i.e. closer to normal distribution. s.t. the prior is easier to learn). You may see reconstruction getting worse a bit, this is expected as it's tradeoff