Mh/diffusion full pipeline latent rmse curve

[moritzhauschulz]

Description

This PR enables plotting of latent RMSE during inference of the diffusion model. Simply enable via test_config.latent_rollout_rmse=True.

NB: This involved adding a new batch to the sampler (see changes). I dropped the idea of reusing the num_steps_input as this would mess with the indexing of the entire pipeline, and there are many ways that it could fail silently. This implementation is a bit more intrusive, but all changes are gated with a flag, so that the standard path is minimally disrupted. If this is to become a permanent feature on develop, I think a broader discussion should be had about what is the best way.

Issue Number

Is this PR a draft? Mark it as draft.

Checklist before asking for review

• I have performed a self-review of my code
• My changes comply with basic sanity checks:
  • I have fixed formatting issues with ./scripts/actions.sh lint
  • I have run unit tests with ./scripts/actions.sh unit-test
  • I have documented my code and I have updated the docstrings.
  • I have added unit tests, if relevant
• I have tried my changes with data and code:
  • I have run the integration tests with ./scripts/actions.sh integration-test
  • (bigger changes) I have run a full training and I have written in the comment the run_id(s): launch-slurm.py --time 60
  • (bigger changes and experiments) I have shared a hegdedoc in the github issue with all the configurations and runs for this experiments
• I have informed and aligned with people impacted by my change:
  • for config changes: the MatterMost channels and/or a design doc
  • for changes of dependencies: the MatterMost software development channel

[clessig]

Why do we need this? How do you distinguish between latent (R)MSE and physical (R)MSE in the training config?

[moritzhauschulz]

Why do we need this? How do you distinguish between latent (R)MSE and physical (R)MSE in the training config?

@clessig We noticed that the problem with analyzing the diffusion training using physical RMSE is heavily dependent on the decoder. Looking at the latent RMSE will give us direct insights on when roll-out diverges in latent space. Afaik there is currently no way of saving the latents for eval, so it needed to be implemented. Note that this is at inference time, not at training time. Training is already via latent loss (as before) unless you put a weight on the physical space – so this is purely a diagnostics tool, and it is triggered via a separate arg in the config (test_config.latent_rollout_rmse=True).