Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix langevin corrector alpha to general num timesteps #81

Merged
merged 1 commit into from
Feb 24, 2025

Conversation

danielzuegner
Copy link
Contributor

@danielzuegner danielzuegner commented Feb 24, 2025

Addresses issue raised in #79 where the $\alpha$ coefficient for the Langevin corrector was not generalizing for different numbers of integration time steps.

The way it's generalized now works as follows:

In the score SDE paper, the scaling coefficient in the noise distribution $p(x_t | x_0)$ is defined as $\sqrt{\bar{\alpha}} = \exp\left(-0.5\int_0^t \beta(s) ds\right)$, so $\bar{\alpha} = \exp\left(-\int_0^t \beta(s) ds\right)$ (Eq. 29). In the discrete case $\bar{\alpha} = \prod_{t=1}^T \alpha_t$, i.e., the cumulative product of the per-step $\alpha_t$ coefficients. Thus, for the continuous case we obtain the equivalent by dividing the cumulative products for $t$ and $t + dt$, i.e., $\alpha_t = \frac{\bar{\alpha_t}}{\bar{\alpha}_{t + dt}}$. This happens in these lines in the code:

alpha_bar = sde._marginal_mean_coeff(t) ** 2
alpha_bar_before = sde._marginal_mean_coeff(t + dt) ** 2
alpha = alpha_bar / alpha_bar_before

where sde._marginal_mean_coeff(t) corresponds to $\sqrt{\bar{\alpha}_t}$.

Also minor fixes to tests and MP20 config.

@danielzuegner danielzuegner merged commit 87ccbce into main Feb 24, 2025
4 checks passed
@danielzuegner danielzuegner deleted the fix-langevin-alphas branch February 24, 2025 16:12
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.

2 participants