issues training TiTok & TA-TiTok with the single stage loss

[SnakeOnex]

Hello,
first of all thanks for amazing research and for open sourcing the code & weights.

I had issues with training the TiTok & TA-TiTok with Single Stage loss, I tried:

1. training TiTok with the Single Stage loss
2. training TA-TiTok with Single Stage loss, but with placeholder text guidance prompt on ImageNet

In both cases I got very bad results: loss goes down, but the reconstructed image is just pale noise with not any real resemblence of the right colors or shapes. I trained with bs=32 on a single A100 GPU for 10 hours. After which I would expect to start seeing signs of convergence. I noticed the grad norms are all in the 1e-7 to 1e-9 range.

Below are linked a minimal repro (~5 lines changed in config) and the resulting wandb training run with results.

minimal diff repro
wandb training run report

[TACJu]

Hi,

Thank you for bringing this issue to our attention. We’ve identified that the problem was caused by the Perceptual Loss not being updated to align with the latest configuration. This has now been fixed in the latest update. We’ve verified that with the fix, the model begins reconstructing reasonable images around 25k steps with a total batch size of 256 on 8 A100 GPUs. Please give it a try and let us know if you encounter any further issues.