MolecularDiffusion.modules.models.jodo.noise_schedule¶
JODO’s VP noise schedule.
Ported from others/JODO/diffusion/noise_schedule.py (itself from DPM-Solver). Only the two continuous schedules JODO’s own configs use are kept – cosine (every released JODO config) and linear. The discrete/discrete_poly branches and the DPM-Solver-only helpers (inverse_lambda, numerical_clip_alpha, interpolate_fn) are dropped with the fast sampler they exist for; ancestral sampling needs only marginal_prob.
Classes¶
Continuous-time VP forward SDE, alpha_t / sigma_t of a label t. |
Module Contents¶
- class MolecularDiffusion.modules.models.jodo.noise_schedule.NoiseScheduleVP(schedule: str = 'cosine', continuous_beta_0: float = 0.1, continuous_beta_1: float = 20.0)¶
Continuous-time VP forward SDE, alpha_t / sigma_t of a label t.
- marginal_alpha(t: torch.Tensor) torch.Tensor¶
- marginal_log_mean_coeff(t: torch.Tensor) torch.Tensor¶
log(alpha_t).
- marginal_prob(t: torch.Tensor)¶
(alpha_t, sigma_t).
- marginal_std(t: torch.Tensor) torch.Tensor¶
- T = 0.9946¶
- beta_0 = 0.1¶
- beta_1 = 20.0¶
- cosine_log_alpha_0¶
- cosine_s = 0.008¶
- schedule = 'cosine'¶
- total_N = 1000¶