GenMD4 uses twice as many zt samples then MD4 per gradient step

Hi, congrats on the great paper!

One thing that caught my attention: diffusion loss in GenMD4 seems to be computed with twice as many `zt` samples than MD4 per gradient step (due to averaging):

https://github.com/google-deepmind/md4/blob/b9fbf29216f818cccff46cb5727199f86f35593e/md4/models/diffusion/genmd4.py#L281

Am I correct here?

This has the potential to lower the variance in diffusion loss, which improves learning.

But this also means that the improvement of GenMD4 over MD4 can't be solely attributed to the state-dependent noise.

Any thoughts on this?