Add EDM-style noise schedules for VEScoreEstimator#1754
Open
Conversation
❌ 1 Tests Failed:
View the full list of 1 ❄️ flaky test(s)
To view more test analytics, go to the Test Analytics Dashboard |
Implements configurable noise schedules from Karras et al. (2022) "Elucidating the Design Space of Diffusion-Based Generative Models": - train_schedule: "uniform" (default) | "lognormal" - solve_schedule: "uniform" (default) | "power_law" - Parameters: lognormal_mean, lognormal_std, power_law_exponent Includes input validation and numerical stability: - Validates sigma bounds (sigma_min > 0, sigma_max > sigma_min) - Log-space clamping in lognormal schedule to prevent NaN - Warning when >5% of samples require clamping - Explicit boundary assignment in power-law schedule References: https://arxiv.org/abs/2206.00364
- Extract sigma_min, sigma_max, and schedule params for VE estimators - Update docstrings for solve_schedule usage - Add comment clarifying VE simulation budget in tests
- test_ve_edm_schedules: all schedule combinations - test_ve_lognormal_no_nan_with_extreme_params: numerical stability
janfb
force-pushed
the
add-edm-noise-schedules
branch
from
d1e2f68 to
f625a41
Compare
| std_0: float = 1.0, | ||
| t_min: float = 1e-3, | ||
| t_max: float = 1.0, | ||
| train_schedule: Literal["uniform", "lognormal"] = "uniform", |
Contributor
There was a problem hiding this comment.
Depending on how the benchmark goes, I think we should even make lognormal and power_law the default.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Extends
VEScoreEstimatorwith configurable noise schedules from Karras et al. (2022) "Elucidating the Design Space of Diffusion-Based Generative Models". This follows up on #1736, which refactored schedule methods into the base class but deferred EDM-specific schedules pending benchmarking.This applies only to VE score estimators. The EDM papers suggests to use log-normal / geometric time schedules for the training and power-law time schedules for the ODE / SDE solving schedules. So, instead of a uniform schedule, we use:
exp(P_mean + P_std * z), converts to time via VE's geometric relationshipσ_i = (σ_max^(1/ρ) + i/(N-1) * (σ_min^(1/ρ) - σ_max^(1/ρ)))^ρ, concentrates steps at low noiseBecause these schedules can lead to numerical instabilities in edge cases, this PR also adds:
log(negative)The hyperparameters are exposed to the user via the build function and
posterior_nn:build_vector_field_estimatornow passessigma_min,sigma_max, and all schedule parameters to VE estimatorsposterior_score_nnaccepts these parameters whensde_type="ve"All new parameters have defaults matching previous uniform behavior. Existing code is unaffected.
Example
Benchmarking
I am currently running benchmarks to compare against uniform schedules, will post them here once done.
References