Difference between JAX and PyTorch Training

[ehonig]

Hi,

I have had a problem for some time now, similar to #6325.

Using seemingly identical setups, including exactly copied model initialization weights and identical stochasticity, Jax/Flax has stable training for this toy Variational Diffusion Model (VDM) example, while PyTorch training consistently fails (tested on 50+ PyTorch seeds).

Below are the two implementations:
Jax/Flax
PyTorch

In particular, it can be seen from the training loop cell outputs that the KL-divergence loss term (bpd_latent) goes to 0 in PyTorch training, despite its computation being identical to that of the Jax/Flax version.

I believe I have accounted for the differences between optax and torch.optim regarding AdamW, and have been considering this problem for a while.

Any help or comments are greatly appreciated; thanks!

[ehonig]

Solved: gradients weren't being propagated through the diffusion loss term properly.

To replace jax.jvp in PyTorch: one should use functorch.jvp to properly propagate gradients, or use torch.autograd.functional.jvp with create_graph=True.