Skip to content

Fix operator precedence in x0 LVLB weights and NaN assertion#404

Open
Mr-Neutr0n wants to merge 1 commit intoCompVis:mainfrom
Mr-Neutr0n:fix/lvlb-weights-nan-check
Open

Fix operator precedence in x0 LVLB weights and NaN assertion#404
Mr-Neutr0n wants to merge 1 commit intoCompVis:mainfrom
Mr-Neutr0n:fix/lvlb-weights-nan-check

Conversation

@Mr-Neutr0n
Copy link

Summary

Fixes two related bugs in DDPM.register_schedule() (ldm/models/diffusion/ddpm.py):

  • Operator precedence bug in x0 LVLB weights (line 163): The expression 2. * 1 - torch.Tensor(alphas_cumprod) is parsed as (2.0) - alphas_cumprod instead of 2.0 * (1 - alphas_cumprod). This silently produces incorrect variational lower bound weights when using x0 parameterization, degrading training quality. Added parentheses to enforce the correct grouping.

  • NaN assertion uses .all() instead of .any() (line 169): assert not torch.isnan(...).all() only triggers if every element is NaN. A single NaN value — which is sufficient to corrupt the loss and propagate through the model — passes undetected. Changed to .any() so the assertion catches any NaN.

Test plan

  • Verified the diff is minimal and limited to the two fixes
  • Run model initialization with parameterization="x0" and confirm LVLB weights match the expected formula
  • Run model initialization and confirm the NaN assertion correctly fires when any weight is NaN

Two bugs in DDPM.register_schedule():

1. The x0 parameterization formula `2. * 1 - torch.Tensor(alphas_cumprod)`
   is evaluated as `(2. * 1) - torch.Tensor(alphas_cumprod)` due to
   operator precedence, yielding `2.0 - alphas_cumprod` instead of the
   intended `2.0 * (1 - alphas_cumprod)`. This produces incorrect LVLB
   weights that silently degrade training when using x0 parameterization.

2. The NaN guard `assert not isnan(...).all()` only fires when *every*
   element is NaN. A single NaN — which is enough to corrupt the loss —
   passes undetected. Changed to `.any()` so any NaN triggers the
   assertion.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant