Improve lead time support for diffusion models

[jleinonen]

PhysicsNeMo Pull Request

Description

Adds/fixes support for lead-time labels in various places where it was missing or not working:

• SongUNetPosEmbd now works properly using either, both or neither of positional embedding and lead-time embedding. In the previous version some pieces of code could try to access properties of these even when set to None.
• deterministic_sampler now accepts lead-time labels and passes them through to the model, if given.
• EDMLoss also now supports lead-time labels.
• Added tests for the above features.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

No new dependencies needed.

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[CharlelieLrt]

All changes proposed look good to me. Just a few details would require some improvements:

1. As far as I understand, this PR decouples lead-time embeddings from positional embeddings, in order to allow more flexibility in using them independently from each other. This new functionality does not seem to be used in any of the training recipes/examples. It could be useful to detail in the PR description the broader context (e.g. which applications is it going to be applied to? Will there be a follow-up PR? etc...)

2. The new flexibility to independently use lead-time and positional embeddings should be clearly explained in the docstrings.

3. IMO the current implementation of the lead-time embeddings has too many failure modes to be safely exposed to broader applications. For example, in positional_embedding_indexing:

• lead_time_label can be done while self.lt_embd is not None, which leads to an error
• Conversely, lead_time_label could be a user-provided tensor, while self.lt_embd is None, which leads to lead_time_label being silently ignored.

I strongly support better parameters validation to eliminate this failure modes, either in the forward method or the __Init__ when possible.

[jleinonen]

Hi @CharlelieLrt,

1. As far as I understand, this PR decouples lead-time embeddings from positional embeddings, in order to allow more flexibility in using them independently from each other. This new functionality does not seem to be used in any of the training recipes/examples. It could be useful to detail in the PR description the broader context (e.g. which applications is it going to be applied to? Will there be a follow-up PR? etc...)

I would say lead-time embeddings were already decoupled from positional embeddings before the PR. This PR just includes some fixes to make sure that they can be enabled when positional embeddings are disabled, or vice versa.

2. The new flexibility to independently use lead-time and positional embeddings should be clearly explained in the docstrings.

As they were already implemented independently, I don't think it's a new flexibility, but I can improve the docstrings in that regard.

3. IMO the current implementation of the lead-time embeddings has too many failure modes to be safely exposed to broader applications. For example, in positional_embedding_indexing:

• lead_time_label can be done while self.lt_embd is not None, which leads to an error
• Conversely, lead_time_label could be a user-provided tensor, while self.lt_embd is None, which leads to lead_time_label being silently ignored.

I strongly support better parameters validation to eliminate this failure modes, either in the forward method or the __Init__ when possible.

I'll add some checks to make sure the inputs conform with the model configuration (but note that as far as I understand, these failure modes already existed before the PR).

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[CharlelieLrt]

LGTM

[CharlelieLrt]

/blossom-ci

[CharlelieLrt]

/blossom-ci