Annealed Langevin Dynamics Scheduler

AnnealedLangevinDynamicsScheduler is a scheduler that uses Langevin dynamics to sample from the posterior distribution of the model parameters. The scheduler anneals the temperature of the Langevin dynamics over time, starting from a high temperature and gradually decreasing it to a low temperature. The scheduler is based on the paper Generative Modeling by Estimating Gradients of the Data Distribution by Yang Song and Stefano Ermon. Stanford AI Lab.

The abstract of the paper is the following:

We introduce a new generative model where samples are produced via Langevin dynamics using gradients of the data distribution estimated with score matching. Because gradients can be ill-defined and hard to estimate when the data resides on low-dimensional manifolds, we perturb the data with different levels of Gaussian noise, and jointly estimate the corresponding scores, i.e., the vector fields of gradients of the perturbed data distribution for all noise levels. For sampling, we propose an annealed Langevin dynamics where we use gradients corresponding to gradually decreasing noise levels as the sampling process gets closer to the data manifold. Our framework allows flexible model architectures, requires no sampling during training or the use of adversarial methods, and provides a learning objective that can be used for principled model comparisons. Our models produce samples comparable to GANs on MNIST, CelebA and CIFAR-10 datasets, achieving a new state-of-the-art inception score of 8.87 on CIFAR-10. Additionally, we demonstrate that our models learn effective representations via image inpainting experiments.

AnnealedLangevinDynamicsScheduler

class ncsn.scheduler.AnnealedLangevinDynamicsScheduler(num_train_timesteps, num_annealed_steps, sigma_min, sigma_max, sampling_eps)[source]

Annealed Langevin Dynamics scheduler for Noise Conditional Score Networks (NCSN).

This scheduler inherits from SchedulerMixin. Check the superclass documentation for it’s generic methods implemented for all schedulers (such as downloading or saving).

Parameters:

  • num_train_timesteps (int) – Number of training timesteps.

  • num_annealed_steps (int) – Number of annealed steps.

  • sigma_min (float) – Minimum standard deviation for the isotropic Gaussian noise.

  • sigma_max (float) – Maximum standard deviation for the isotropic Gaussian noise.

  • sampling_eps (float) – Sampling epsilon for the Langevin dynamics.

__init__(num_train_timesteps, num_annealed_steps, sigma_min, sigma_max, sampling_eps)[source]
add_noise(original_samples, noise, timesteps)[source]

Add noise to the original samples.

Parameters:

  • original_samples (torch.Tensor) – The original samples.

  • noise (torch.Tensor) – The noise to be added.

  • timesteps (torch.Tensor) – The timesteps.

Returns:

The noisy samples.

Return type:

torch.Tensor

set_sigmas(num_inference_steps, sigma_min=None, sigma_max=None, sampling_eps=None)[source]

Sets the sigmas and step sizes for the scheduler (to be run before inference).

Parameters:

  • num_inference_steps (int) – The number of diffusion steps used when generating samples with a pre-trained model. If used, sigmas and step_size must be None.

  • sigma_min (float, optional) – The minimum standard deviation for the isotropic Gaussian noise. If None, the default value is used.

  • sigma_max (float, optional) – The maximum standard deviation for the isotropic Gaussian noise. If None, the default value is used.

  • sampling_eps (float, optional) – The sampling epsilon for the Langevin dynamics. If None, the default value is used.

Return type:

None

set_timesteps(num_inference_steps, sampling_eps=None, device=None)[source]

Sets the timesteps for the scheduler (to be run before inference).

Parameters:

  • num_inference_steps (int) – The number of diffusion steps used when generating samples with a pre-trained model. If used, timesteps must be None.

  • device (str or torch.device, optional) – The device to which the timesteps should be moved to. If None, the timesteps are not moved. Defined to maintain compatibility with other pipelines, but this argument is not actually used.

  • sampling_eps (float, optional) – The sampling epsilon for the Langevin dynamics. If None, the default value is used.

  • timesteps (List[int], optional) – Custom timesteps used to support arbitrary spacing between timesteps. If None, then the default timestep spacing strategy of equal spacing between timesteps is used. If timesteps is passed, num_inference_steps must be None.

Return type:

None

step(model_output, timestep, sample, return_dict=True, **kwargs)[source]

Perform one step following Langevin dynamics. Annealing must be done separately.

Parameters:

  • model_output (torch.Tensor) – The score output from learned neural network-based score function.

  • timestep (int) – The current timestep.

  • sample (torch.Tensor) – The current sample.

  • return_dict (bool, optional) – Whether or not to return AnnealedLangevinDynamicsOutput or tuple.

Returns:

if return_dict is True, AnnealedLangevinDynamicsOutput is returned, otherwise a tuple is returned where the first element is the updated sample.

Return type:

AnnealedLangevinDynamicsOutput or tuple

AnnealedLangevinDynamicsOutput

class ncsn.scheduler.AnnealedLangevinDynamicsOutput(prev_sample)[source]

Annealed Langevin Dynamics output class.