azula.guidance.mmps

Moment Matching Posterior Sampling (MMPS) internals.

References

Learning Diffusion Priors from Observations by Expectation Maximization (Rozet et al., 2024)
https://arxiv.org/abs/2405.13712

Classes

MMPSDenoiser

Creates a MMPS denoiser module.

Descriptions

class azula.guidance.mmps.MMPSDenoiser(denoiser, y, A, cov_y, solver='gmres', iterations=1)[source]

Creates a MMPS denoiser module.

Parameters:

  • denoiser (Denoiser) – A denoiser \(q_\phi(X \mid X_t)\).

  • y (Tensor) – An observation \(y \sim \mathcal{N}(A(x), \Sigma_y)\), with shape \((*, D)\).

  • A (Callable[[Tensor], Tensor]) – The forward operator \(x \mapsto A(x)\).

  • cov_y (Covariance) – The noise covariance \(\Sigma_y\).

  • solver (str) – The linear solver name ("cg" or "gmres").

  • iterations (int) – The number of solver iterations.