model/graphcast/dpm_solver_plus_plus_2s.py

# Copyright 2024 DeepMind Technologies Limited.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS-IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""DPM-Solver++ 2S sampler from https://arxiv.org/abs/2211.01095."""

from typing import Optional

from . import casting
from . import denoisers_base
from . import samplers_base as base
from . import samplers_utils as utils
from . import xarray_jax
import haiku as hk
import jax.numpy as jnp
import xarray


class Sampler(base.Sampler):
  """Sampling using DPM-Solver++ 2S from [1].

  This is combined with optional stochastic churn as described in [2].

  The '2S' terminology from [1] means that this is a second-order (2),
  single-step (S) solver. Here 'single-step' here distinguishes it from
  'multi-step' methods where the results of function evaluations from previous
  steps are reused in computing updates for subsequent steps. The solver still
  uses multiple steps though.

  [1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic
  Models, https://arxiv.org/abs/2211.01095
  [2] Elucidating the Design Space of Diffusion-Based Generative Models,
  https://arxiv.org/abs/2206.00364
  """

  def __init__(self,
               denoiser: denoisers_base.Denoiser,
               max_noise_level: float,
               min_noise_level: float,
               num_noise_levels: int,
               rho: float,
               stochastic_churn_rate: float,
               churn_min_noise_level: float,
               churn_max_noise_level: float,
               noise_level_inflation_factor: float
               ):
    """Initializes the sampler.

    Args:
      denoiser: A Denoiser which predicts noise-free targets.
      max_noise_level: The highest noise level used at the start of the
        sequence of reverse diffusion steps.
      min_noise_level: The lowest noise level used at the end of the sequence of
        reverse diffusion steps.
      num_noise_levels: Determines the number of noise levels used and hence the
        number of reverse diffusion steps performed.
      rho: Parameter affecting the spacing of noise steps. Higher values will
        concentrate noise steps more around zero.
      stochastic_churn_rate: S_churn from the paper. This controls the rate
        at which noise is re-injected/'churned' during the sampling algorithm.
        If this is set to zero then we are performing deterministic sampling
        as described in Algorithm 1.
      churn_min_noise_level: Minimum noise level at which stochastic churn
        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
      churn_max_noise_level: Maximum noise level at which stochastic churn
        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
      noise_level_inflation_factor: This can be used to set the actual amount of
        noise injected higher than what the denoiser is told has been added.
        The motivation is to compensate for a tendency of L2-trained denoisers
        to remove slightly too much noise / blur too much. S_noise from the
        paper. Only used if stochastic_churn_rate > 0.
    """
    super().__init__(denoiser)
    self._noise_levels = utils.noise_schedule(
        max_noise_level, min_noise_level, num_noise_levels, rho)
    self._stochastic_churn = stochastic_churn_rate > 0
    self._per_step_churn_rates = utils.stochastic_churn_rate_schedule(
        self._noise_levels, stochastic_churn_rate, churn_min_noise_level,
        churn_max_noise_level)
    self._noise_level_inflation_factor = noise_level_inflation_factor

  def __call__(
      self,
      inputs: xarray.Dataset,
      targets_template: xarray.Dataset,
      forcings: Optional[xarray.Dataset] = None,
      **kwargs) -> xarray.Dataset:

    dtype = casting.infer_floating_dtype(targets_template)  # pytype: disable=wrong-arg-types
    noise_levels = jnp.array(self._noise_levels).astype(dtype)
    per_step_churn_rates = jnp.array(self._per_step_churn_rates).astype(dtype)

    def denoiser(noise_level: jnp.ndarray, x: xarray.Dataset) -> xarray.Dataset:
      """Computes D(x, sigma, y)."""
      bcast_noise_level = xarray_jax.DataArray(
          jnp.tile(noise_level, x.sizes['batch']), dims=('batch',))
      # Estimate the expectation of the fully-denoised target x0, conditional on
      # inputs/forcings, noisy targets and their noise level:
      return self._denoiser(
          inputs=inputs,
          noisy_targets=x,
          noise_levels=bcast_noise_level,
          forcings=forcings)

    def body_fn(i: jnp.ndarray, x: xarray.Dataset) -> xarray.Dataset:
      """One iteration of the sampling algorithm.

      Args:
        i: Sampling iteration.
        x: Noisy targets at iteration i, these will have noise level
          self._noise_levels[i].

      Returns:
        Noisy targets at the next lowest noise level self._noise_levels[i+1].
      """
      def init_noise(template):
        return noise_levels[0] * utils.spherical_white_noise_like(template)

      # Initialise the inputs if i == 0.
      # This is done here to ensure both noise sampler calls can use the same
      # spherical harmonic basis functions. While there may be a small compute
      # cost the memory savings can be significant.
      # TODO(dominicmasters): Figure out if we can merge the two noise sampler
      # calls into one to avoid this hack.
      maybe_init_noise = (i == 0).astype(noise_levels[0].dtype)
      x = x + init_noise(x) * maybe_init_noise

      noise_level = noise_levels[i]

      if self._stochastic_churn:
        # We increase the noise level of x a bit before taking it down again:
        x, noise_level = utils.apply_stochastic_churn(
            x, noise_level,
            stochastic_churn_rate=per_step_churn_rates[i],
            noise_level_inflation_factor=self._noise_level_inflation_factor)

      # Apply one step of the ODE solver to take x down to the next lowest
      # noise level.

      # Note that the Elucidating paper's choice of sigma(t)=t and s(t)=1
      # (corresponding to alpha(t)=1 in the DPM paper) as well as the standard
      # choice of r=1/2 (corresponding to a geometric mean for the s_i
      # midpoints) greatly simplifies the update from the DPM-Solver++ paper.
      # You need to do a bit of algebraic fiddling to arrive at the below after
      # substituting these choices into DPMSolver++'s Algorithm 1. The simpler
      # update we arrive at helps with intuition too.

      next_noise_level = noise_levels[i + 1]
      # This is s_{i+1} from the paper. They don't explain how the s_i are
      # chosen, but the default choice seems to be a geometric mean, which is
      # equivalent to setting all the r_i = 1/2.
      mid_noise_level = jnp.sqrt(noise_level * next_noise_level)

      mid_over_current = mid_noise_level / noise_level
      x_denoised = denoiser(noise_level, x)
      # This turns out to be a convex combination of current and denoised x,
      # which isn't entirely apparent from the paper formulae:
      x_mid = mid_over_current * x + (1 - mid_over_current) * x_denoised

      next_over_current = next_noise_level / noise_level
      x_mid_denoised = denoiser(mid_noise_level, x_mid)  # pytype: disable=wrong-arg-types
      x_next = next_over_current * x + (1 - next_over_current) * x_mid_denoised

      # For the final step to noise level 0, we do an Euler update which
      # corresponds to just returning the denoiser's prediction directly.
      #
      # In fact the behaviour above when next_noise_level == 0 is almost
      # equivalent, except that it runs the denoiser a second time to denoise
      # from noise level 0. The denoiser should just be the identity function in
      # this case, but it hasn't necessarily been trained at noise level 0 so
      # we avoid relying on this.
      return utils.tree_where(next_noise_level == 0, x_denoised, x_next)

    # Init with zeros but apply additional noise at step 0 to initialise the
    # state.
    noise_init = xarray.zeros_like(targets_template)
    return hk.fori_loop(
        0, len(noise_levels) - 1, body_fun=body_fn, init_val=noise_init)