MolecularDiffusion.modules.models.nextmol.scheduler

Continuous-time VP noise schedule (NoiseScheduleVPV2, DPM-Solver).

Ported from others/NExT-Mol/data_provider/diffusion_scheduler.py:24.

Only the cosine and linear branches are here. The discrete / discrete_poly branches need a 1000-step table and an interpolation helper, and every released DMT checkpoint was trained with noise_scheduler: cosine and discrete_schedule: False – so porting them would be dead code today. Note the cosine schedule’s T = 0.9946, not 1.0: t=1 has numerical issues.

Classes

NoiseScheduleVPV2

marginal_prob(t) -> (alpha_t, sigma_t) for a VP forward SDE.

Module Contents

class MolecularDiffusion.modules.models.nextmol.scheduler.NoiseScheduleVPV2(schedule: str = 'cosine', continuous_beta_0: float = 0.1, continuous_beta_1: float = 20.0, discrete_mode: bool = False)

marginal_prob(t) -> (alpha_t, sigma_t) for a VP forward SDE.

marginal_alpha(t)
marginal_log_mean_coeff(t)
marginal_prob(t)
marginal_std(t)
T = 0.9946
beta_0 = 0.1
beta_1 = 20.0
cosine_log_alpha_0
cosine_s = 0.008
discrete_mode = False
schedule = 'cosine'
total_N = 1000