agedi.diffusion.samplers.heun

Stochastic Heun sampler — second-order reverse-SDE integration.

Classes

HeunSampler

Second-order stochastic sampler (Karras et al. 2022).

Module Contents

class agedi.diffusion.samplers.heun.HeunSampler(score_fn: Callable[[agedi.data.AtomsGraph], agedi.data.AtomsGraph], noisers: List[agedi.diffusion.noisers.Noiser])

Bases: agedi.diffusion.samplers.base.Sampler

Second-order stochastic sampler (Karras et al. 2022).

Uses two score evaluations per step to achieve second-order accuracy for the stochastic reverse SDE:

  1. First score at (x_t, t): s_1 = s_θ(x_t, t)

  2. Deterministic EM predictor (no noise) for SDE noisers: x̃ = x_t + dt·(g(t)²·s_1 - f(x_t, t))

  3. Second score at (x̃, t - dt): s_2 = s_θ(x̃, t - dt)

  4. Average scores: s_avg = 0.5·(s_1 + s_2)

  5. Stochastic EM step using averaged score: x_{t-dt} = x_t + dt·(g(t)²·s_avg - f(x_t, t)) + √dt·g(t)·z (no noise when last=True)

The score averaging provides second-order accuracy in the SDE sense, analogous to the stochastic Heun method from Karras et al. (2022) “Elucidating the Design Space of Diffusion-Based Generative Models”.

For noisers with a .sde attribute (continuous SDE-based noisers such as PositionsNoiser), the two-step procedure is applied. For noisers without .sde (e.g. discrete types noiser), a single EM step is applied using the first score evaluation, since score averaging is less principled for discrete diffusion.

Parameters:

  • score_fn (callable) – Score-model forward function.

  • noisers (list of Noiser) – Noisers in forward order.

Notes

noiser.denoise() reads positions and scores directly from tensor attributes and does not require an up-to-date neighbour list. This allows the original positions x_t to be restored (step 5) without rebuilding the graph before the final denoising call in step 5 — the graph is only rebuilt once at the very end of the step.

step(batch: agedi.data.AtomsGraph, dt: torch.Tensor, last: bool) → agedi.data.AtomsGraph

Second-order stochastic Heun step.