Improve lead time support for diffusion models #980
Conversation
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
|
/blossom-ci |
There was a problem hiding this comment.
All changes proposed look good to me. Just a few details would require some improvements:
-
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...)
-
The new flexibility to independently use lead-time and positional embeddings should be clearly explained in the docstrings.
-
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_labelcan be done whileself.lt_embdis notNone, which leads to an error- Conversely,
lead_time_labelcould be a user-provided tensor, whileself.lt_embdisNone, which leads tolead_time_labelbeing silently ignored.
I strongly support better parameters validation to eliminate this failure modes, either in the forward method or the __Init__ when possible.
|
Hi @CharlelieLrt,
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.
As they were already implemented independently, I don't think it's a new flexibility, but I can improve the docstrings in that regard.
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). |
|
/blossom-ci |
| if (lead_time_channels is None) or (lead_time_channels <= 0): | ||
| raise ValueError( | ||
| "`lead_time_channels` must be >= 1 if `lead_time_mode` is enabled." | ||
| ) |
There was a problem hiding this comment.
Agree with this validation. However, there is clearly redundancy between lead_time_mode and lead_time_channels. In addition, for positional embeddings, we can disable them by setting N_grid_channels = 0 (there is no boolean parameter positional_embedding_mode), so we expect the same mechanism to work for lead-time embeddings (i.e. disable by setting lead_time_channels = 0).
If your concern is about backward compatibility if removing lead_time_mode:
- This parameters was introduced only very recently (<2 months), so it could be okay to remove it
- If still a concern, can we keep
lead_time_modebut add a deprecation warning in this if statement:
if self.lead_time_mode:
warnings.warn(
f"The parameter `lead_time_mode` will be deprecated in a future version. "
f"The recommended way to enable (disable) lead-time embeddings is to set `lead_time_channels > 0` (`lead_time_channels = 0`)")
[...]| embeddings = self.pos_embd | ||
| elif lead_time_label is None: # lead time embedding only | ||
| embeddings = self.lt_embd[lead_time_label[0].int()] |
There was a problem hiding this comment.
Why using only lead_time_label[0] here? I know it follows the original implementation, but isn't it an error?
There was a problem hiding this comment.
Unrelated, I think there is a second bug here. The elif clauses should read elif self.pos_embd is not None and elif self.lead_time_label is not None, right?
There was a problem hiding this comment.
I confirmed with @tge25 that the [0] is some prototyping code, so we should get rid of it.
This elif self.lead_time_label is not None is not possible, because lead_time_label is not an attribute. But I agree that the elif clause should be different. I think it should be elif self.lt_embd is not None, right?
| @@ -245,6 +253,10 @@ def __call__(self, net, images, condition=None, labels=None, augment_pipe=None): | |||
| An optional data augmentation function that takes images as input and | |||
| returns augmented images. If not provided, no data augmentation is applied. | |||
|
|
|||
| lead_time_label: torch.Tensor, optional | |||
| Lead-time labels to pass to the model, shape (batch_size, 1). | |||
There was a problem hiding this comment.
Are you sure about the shape (batch_size, 1) of lead_time_label? I've seen other places (not in your PR), where it's (batch_size,), some other places where it's (batch_size, 1, 1, 1)
PhysicsNeMo Pull Request
Description
Adds/fixes support for lead-time labels in various places where it was missing or not working:
SongUNetPosEmbdnow 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 toNone.deterministic_samplernow accepts lead-time labels and passes them through to the model, if given.EDMLossalso now supports lead-time labels.Checklist
Dependencies
No new dependencies needed.