Repository: pyg-team/pytorch_geometric
Pull request #8347 (2023-11-08 20:14:08)

Add noise scheduler utility for diffusion based graph generative models
We will use this utility to reimplement Permutation Invariant Graph Generation via Score-Based Generative Modeling (https://arxiv.org/abs/2003.00638) with PyG

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8347?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8347](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8347?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (e35007e) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/dadfeec2e38c8a220e478455cbc6d552a8f40ab5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (dadfeec) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

> :exclamation: Current head e35007e differs from pull request most recent head 49a19ed. Consider uploading reports for the commit 49a19ed to get more accurate results

```diff
@@           Coverage Diff           @@
##           master    #8347   +/-   ##
=======================================
  Coverage   87.17%   87.17%           
=======================================
  Files         473      474    +1     
  Lines       28756    28773   +17     
=======================================
+ Hits        25067    25084   +17     
  Misses       3689     3689           
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8347?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/noise\_scheduler.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8347?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL25vaXNlX3NjaGVkdWxlci5weQ==) | `100.00% <100.00%> (ø)` | |

:mega: Codecov offers a browser extension for seamless coverage viewing on GitHub. Try it in [Chrome](https://chrome.google.com/webstore/detail/codecov/gedikamndpbemklijjkncpnolildpbgo) or [Firefox](https://addons.mozilla.org/en-US/firefox/addon/codecov/) today!


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index ab342b3e5..a8e6599fc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added noise scheduler utility for diffusion based graph generative models ([#8347](https://github.com/pyg-team/pytorch_geometric/pull/8347))
 - Added the equivariant `ViSNet` model ([#8287](https://github.com/pyg-team/pytorch_geometric/pull/8287))
 - Added temporal-related capabilities to `Data` ([#8454](https://github.com/pyg-team/pytorch_geometric/pull/8454))
 - Added support for returning multi graphs in `to_networkx` ([#8575](https://github.com/pyg-team/pytorch_geometric/pull/8575))
diff --git a/torch_geometric/utils/noise_scheduler.py b/torch_geometric/utils/noise_scheduler.py
new file mode 100644
index 000000000..9908d9115
--- /dev/null
+++ b/torch_geometric/utils/noise_scheduler.py
@@ -0,0 +1,109 @@
+import math
+from typing import Literal, Optional
+
+import torch
+from torch import Tensor
+
+
+def get_smld_sigma_schedule(
+    sigma_min: float,
+    sigma_max: float,
+    num_scales: int,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    r"""Generates a set of noise values on a logarithmic scale for "Score
+    Matching with Langevin Dynamics" from the `"Generative Modeling by
+    Estimating Gradients of the Data Distribution"
+    <https://arxiv.org/abs/1907.05600>`_ paper.
+
+    This function returns a vector of sigma values that define the schedule of
+    noise levels used during Score Matching with Langevin Dynamics.
+    The sigma values are determined on a logarithmic scale from
+    :obj:`sigma_max` to :obj:`sigma_min`, inclusive.
+
+    Args:
+        sigma_min (float): The minimum value of sigma, corresponding to the
+            lowest noise level.
+        sigma_max (float): The maximum value of sigma, corresponding to the
+            highest noise level.
+        num_scales (int): The number of sigma values to generate, defining the
+            granularity of the noise schedule.
+        dtype (torch.dtype, optional): The output data type.
+            (default: :obj:`None`)
+        device (torch.device, optional): The output device.
+            (default: :obj:`None`)
+    """
+    return torch.linspace(
+        math.log(sigma_max),
+        math.log(sigma_min),
+        num_scales,
+        dtype=dtype,
+        device=device,
+    ).exp()
+
+
+def get_diffusion_beta_schedule(
+    schedule_type: Literal['linear', 'quadratic', 'constant', 'sigmoid'],
+    beta_start: float,
+    beta_end: float,
+    num_diffusion_timesteps: int,
+    dtype: Optional[torch.dtype] = None,
+    device: Optional[torch.device] = None,
+) -> Tensor:
+    r"""Generates a schedule of beta values according to the specified strategy
+    for the diffusion process from the `"Denoising Diffusion Probabilistic
+    Models" <https://arxiv.org/abs/2006.11239>`_ paper.
+
+    Beta values are used to scale the noise added during the diffusion process
+    in generative models. This function creates an array of beta values
+    according to a pre-defined schedule, which can be either :obj:`"linear"`,
+    :obj:`"quadratic"`, :obj:`"constant"`, or :obj:`"sigmoid"`.
+
+    Args:
+        schedule_type (str): The type of schedule to use for beta values.
+        beta_start (float): The starting value of beta.
+        beta_end (float): The ending value of beta.
+        num_diffusion_timesteps (int): The number of timesteps for the
+            diffusion process.
+        dtype (torch.dtype, optional): The output data type.
+            (default: :obj:`None`)
+        device (torch.device, optional): The output device.
+            (default: :obj:`None`)
+    """
+    if schedule_type == 'linear':
+        return torch.linspace(
+            beta_start,
+            beta_end,
+            num_diffusion_timesteps,
+            dtype=dtype,
+            device=device,
+        )
+
+    if schedule_type == 'quadratic':
+        return torch.linspace(
+            beta_start**0.5,
+            beta_end**0.5,
+            num_diffusion_timesteps,
+            dtype=dtype,
+            device=device,
+        )**2
+
+    if schedule_type == 'constant':
+        return torch.full(
+            (num_diffusion_timesteps, ),
+            fill_value=beta_end,
+            dtype=dtype,
+            device=device,
+        )
+
+    if schedule_type == 'sigmoid':
+        return torch.linspace(
+            -6,
+            6,
+            num_diffusion_timesteps,
+            dtype=dtype,
+            device=device,
+        ).sigmoid() * (beta_end - beta_start) + beta_start
+
+    raise ValueError(f"Found invalid 'schedule_type' (got '{schedule_type}')")


Test patch:
diff --git a/test/utils/test_noise_scheduler.py b/test/utils/test_noise_scheduler.py
new file mode 100644
index 000000000..c5fba67dd
--- /dev/null
+++ b/test/utils/test_noise_scheduler.py
@@ -0,0 +1,34 @@
+import pytest
+import torch
+
+from torch_geometric.utils.noise_scheduler import (
+    get_diffusion_beta_schedule,
+    get_smld_sigma_schedule,
+)
+
+
+def test_get_smld_sigma_schedule():
+    expected = torch.tensor([
+        1., 0.59948425, 0.35938137, 0.21544347, 0.12915497, 0.07742637,
+        0.04641589, 0.02782559, 0.01668101, 0.01
+    ])
+    out = get_smld_sigma_schedule(
+        sigma_min=0.01,
+        sigma_max=1.0,
+        num_scales=10,
+    )
+    assert torch.allclose(out, expected)
+
+
+@pytest.mark.parametrize(
+    'schedule_type',
+    ['linear', 'quadratic', 'constant', 'sigmoid'],
+)
+def test_get_diffusion_beta_schedule(schedule_type):
+    out = get_diffusion_beta_schedule(
+        schedule_type,
+        beta_start=0.1,
+        beta_end=0.2,
+        num_diffusion_timesteps=10,
+    )
+    assert out.size() == (10, )