From 9bd692bd4bc8cf01f954de8ee1bfdf6de1e32629 Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Mon, 19 Oct 2026 11:06:07 -0700 Subject: [PATCH] Add Kumo-Anomaly DPM-Solver++ sampler Co-authored-by: Aditi Gautam --- .../kumo/timeseries/anomaly/sampling.py | 149 ++++++++++++++++++ .../kumo/timeseries/anomaly/test_sampling.py | 66 ++++++++ 2 files changed, 215 insertions(+) create mode 100644 sdm/models/kumo/timeseries/anomaly/sampling.py create mode 100644 test/models/kumo/timeseries/anomaly/test_sampling.py diff --git a/sdm/models/kumo/timeseries/anomaly/sampling.py b/sdm/models/kumo/timeseries/anomaly/sampling.py new file mode 100644 index 000000000..7b35362d0 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/sampling.py @@ -0,0 +1,149 @@ +# Portions derived from DPM-Solver: +# Copyright (c) 2022 Cheng Lu +# Licensed under the MIT License; see third_party/dpm-solver/LICENSE. + +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# ruff: noqa: D205 + +from collections.abc import Callable + +import torch +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.noise_schedule import NoiseScheduleVP + + +class DPMSolver(torch.nn.Module): + r"""The second-order, single-step sampler from the `"DPM-Solver++: Fast + Solver for Guided Sampling of Diffusion Probabilistic Models" + `__ paper. + + Time runs uniformly from ``T = 1`` to ``1 / total_N`` of + :class:`NoiseScheduleVP`, and the solver integrates clean-data + predictions without final denoising or clipping. + + The denoiser receives training step indices: each solver time ``t`` + maps to ``round(t * total_N - 1)``, clamped to ``[0, total_N - 1]``. + + Args: + betas: Discrete training variances on CPU, in training order. + num_steps: Number of denoiser evaluations, at least one. Odd counts + end with a first-order update. + """ + + alpha: Tensor + sigma: Tensor + delta: Tensor + fraction: Tensor + steps: Tensor + + def __init__(self, betas: Tensor, num_steps: int = 20) -> None: + super().__init__() + schedule = NoiseScheduleVP(betas) + self.orders = (2,) * (num_steps // 2) + (1,) * (num_steps % 2) + time = torch.linspace( + start=schedule.T, + end=1 / schedule.total_N, + steps=num_steps + 1, + device="cpu", + dtype=torch.float32, + ) + intervals = [] + deltas = [] + fractions = [] + offset = 0 + for order in self.orders: + start, end = time[offset], time[offset + order] + # Place the intermediate time at the midpoint of the interval, + # expressed as a fraction of its half-logSNR step: + inner = torch.linspace( + start=start, + end=end, + steps=3, + device="cpu", + dtype=torch.float32, + ) + lamb = schedule.marginal_lambda(inner) + delta = lamb[2] - lamb[0] + fraction = (lamb[1] - lamb[0]) / delta + mid = schedule.inverse_lambda(lamb[0] + fraction * delta)[0] + intervals.append(torch.stack([start, mid, end])) + deltas.append(delta) + fractions.append(fraction) + offset += order + times = torch.stack(intervals) # [num_intervals, 3] + steps = (times * schedule.total_N - 1).round().long() + steps = steps.clamp(0, schedule.total_N - 1) + for name, value in ( + ("alpha", schedule.marginal_alpha(times).reshape(-1, 3)), + ("sigma", schedule.marginal_std(times).reshape(-1, 3)), + ("delta", torch.stack(deltas)), + ("fraction", torch.stack(fractions)), + ("steps", steps), + ): + self.register_buffer(name=name, tensor=value, persistent=False) + + def forward( + self, + x: Tensor, + predict_noise: Callable[[Tensor, Tensor], Tensor], + ) -> Tensor: + """Integrate initial noise using a noise-predicting denoiser. + + Args: + x: Initial noise, for example ``[batch, features, time]``, on the + module's device. + predict_noise: Callable accepting the current sample and a + one-element training step index tensor, returning predicted + noise with the sample's shape. + + Returns: + Reconstructed data with ``x``'s shape. + """ + for i, order in enumerate(self.orders): + alpha, sigma = self.alpha[i], self.sigma[i] + delta, fraction = self.delta[i], self.fraction[i] + eps = predict_noise(x, self.steps[i, :1]) + data = (x - sigma[0] * eps) / alpha[0] + data_mid = None + if order == 2: + mid = _update( + x=x, + data=data, + alpha_end=alpha[1], + sigma_start=sigma[0], + sigma_end=sigma[1], + delta=fraction * delta, + ) + eps_mid = predict_noise(mid, self.steps[i, 1:2]) + data_mid = (mid - sigma[1] * eps_mid) / alpha[1] + x = _update( + x=x, + data=data, + alpha_end=alpha[2], + sigma_start=sigma[0], + sigma_end=sigma[2], + delta=delta, + data_mid=data_mid, + fraction=fraction, + ) + return x + + +def _update( + x: Tensor, + data: Tensor, + alpha_end: Tensor, + sigma_start: Tensor, + sigma_end: Tensor, + delta: Tensor, + data_mid: Tensor | None = None, + fraction: Tensor | float = 0.5, +) -> Tensor: + factor = alpha_end * (-delta).expm1() + out = (sigma_end / sigma_start) * x - factor * data + if data_mid is not None: + out = out - (0.5 / fraction) * factor * (data_mid - data) + return out diff --git a/test/models/kumo/timeseries/anomaly/test_sampling.py b/test/models/kumo/timeseries/anomaly/test_sampling.py new file mode 100644 index 000000000..9c7276f02 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_sampling.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.noise_schedule import NoiseScheduleVP +from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule +from sdm.testing import withCUDA + +BETAS = DiffusionSchedule(500, 0.0001, 0.01, "linear").beta.float() + + +def test_dpm_solver_steps() -> None: + steps = [] + + def predict_noise(sample: Tensor, step: Tensor) -> Tensor: + steps.append(step) + return torch.zeros_like(sample) + + DPMSolver(BETAS, num_steps=3)(torch.zeros(2, 3), predict_noise) + torch.testing.assert_close(torch.cat(steps), torch.tensor([499, 333, 166])) + + +@withCUDA +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_dpm_solver_constant_data( + device: torch.device, + dtype: torch.dtype, +) -> None: + # A constant clean signal and noise define the exact ODE trajectory + # alpha_t * clean + sigma_t * noise, which the solver follows exactly. + schedule = NoiseScheduleVP(BETAS) + time = torch.tensor([1.0, 1 / schedule.total_N]) + alpha_start, alpha_end = schedule.marginal_alpha(time).tolist() + sigma_start, sigma_end = schedule.marginal_std(time).tolist() + clean = torch.full((2, 3, 4), 0.25, device=device, dtype=dtype) + noise = torch.randn_like(clean) + x = alpha_start * clean + sigma_start * noise + # Three steps take one second-order and one first-order update. + solver = DPMSolver(BETAS, num_steps=3).to(device) + + def predict_noise(sample: Tensor, step: Tensor) -> Tensor: + assert step.device == device + return noise + + out = solver(x, predict_noise) + expected = alpha_end * clean + sigma_end * noise + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +def test_dpm_solver_linear_noise() -> None: + # Predicting the sample itself as noise has the exact solution + # x_end / x_start = (alpha_end / alpha_start) * exp(u_end - u_start), + # where u = asinh(sigma / alpha). The solver stays within 1%, while + # first-order updates alone miss it by 5%. + schedule = NoiseScheduleVP(BETAS) + time = torch.tensor([1.0, 1 / schedule.total_N]) + alpha = schedule.marginal_alpha(time) + u = (schedule.marginal_std(time) / alpha).asinh() + expected = (alpha[1] / alpha[0]) * (u[1] - u[0]).exp() + + out = DPMSolver(BETAS, num_steps=21)(torch.ones(1), lambda x, step: x) + torch.testing.assert_close(out, expected.view(1), rtol=1e-2, atol=0.0)