diff --git a/sdm/models/kumo/timeseries/anomaly/noise_schedule.py b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py new file mode 100644 index 0000000000..e47f3880bb --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py @@ -0,0 +1,103 @@ +# 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 + +import torch +from torch import Tensor + + +class NoiseScheduleVP(torch.nn.Module): + """Continuous-time VP coefficients for Kumo-Anomaly's DPM-Solver++. + + Interpolate ``log(alpha_t) = 0.5 * cumsum(log(1 - beta))`` between + discrete steps. As in Kumo-TS, discard the tail where half-logSNR is + below ``clipped_lambda`` (default -5.1) and assign the retained steps + times ``(i + 1) / total_N``. + ``total_N`` is the retained step count and ``T = 1`` is the final time. + Unlike :class:`DiffusionSchedule`, cumulative alphas are not clamped. + + Methods accept scalar or batched tensors on the module's device and + return flattened, one-dimensional tensors. Outside the retained time + interval ``[1 / total_N, T]``, interpolation uses the boundary segment + for linear extrapolation. The derived tables are nonpersistent buffers. + + Args: + betas: One-dimensional diffusion variances strictly between zero + and one, with at least two steps remaining after truncation. + Coefficients retain the input device and dtype (float32 or + float64). Pass ``DiffusionSchedule.beta.float()`` to match + Kumo-TS inference. + clipped_lambda: Minimum retained half-logSNR. Defaults to -5.1, + matching Kumo-TS's numerical-stability cutoff. + """ + + _log_alpha: Tensor + _time: Tensor + + def __init__(self, betas: Tensor, *, clipped_lambda: float = -5.1) -> None: + super().__init__() + self.schedule = "discrete" + self.T = 1.0 + log_alpha = 0.5 * (1 - betas.detach()).log().cumsum(dim=0) + log_sigma = 0.5 * (1 - (2 * log_alpha).exp()).log() + # A single host synchronization at construction determines table size. + num_clipped = int( + torch.searchsorted((log_alpha - log_sigma).flip(0), clipped_lambda) + ) + self.total_N = log_alpha.numel() - num_clipped + if self.total_N < 2: + raise ValueError( + "At least two diffusion steps must remain after truncation" + ) + self.register_buffer( + name="_log_alpha", + tensor=log_alpha[: self.total_N], + persistent=False, + ) + # Match upstream's float32 CPU grid before dtype/device conversion. + time = torch.linspace( + start=0, + end=self.T, + steps=self.total_N + 1, + device="cpu", + dtype=torch.float32, + )[1:].to(betas) + self.register_buffer("_time", time, persistent=False) + + def marginal_log_mean_coeff(self, t: Tensor) -> Tensor: + """Return interpolated log signal coefficients at times ``t``.""" + return _interpolate(t.reshape(-1), self._time, self._log_alpha) + + def marginal_alpha(self, t: Tensor) -> Tensor: + """Return signal coefficients ``alpha_t`` at continuous times ``t``.""" + return self.marginal_log_mean_coeff(t).exp() + + def marginal_std(self, t: Tensor) -> Tensor: + """Return noise standard deviations ``sqrt(1 - alpha_t**2)``.""" + return (1 - (2 * self.marginal_log_mean_coeff(t)).exp()).sqrt() + + def marginal_lambda(self, t: Tensor) -> Tensor: + """Return half-logSNR ``log(alpha_t) - log(sigma_t)``.""" + log_alpha = self.marginal_log_mean_coeff(t) + return log_alpha - 0.5 * (1 - (2 * log_alpha).exp()).log() + + def inverse_lambda(self, lamb: Tensor) -> Tensor: + """Return continuous times for the supplied half-logSNR values.""" + log_alpha = -0.5 * torch.logaddexp(torch.zeros_like(lamb), -2 * lamb) + return _interpolate( + x=log_alpha.reshape(-1), + xp=self._log_alpha.flip(0), + yp=self._time.flip(0), + ) + + +def _interpolate(x: Tensor, xp: Tensor, yp: Tensor) -> Tensor: + """Interpolate sorted 1D knots, extrapolating at either boundary.""" + right = torch.searchsorted(xp, x.contiguous()).clamp(1, xp.numel() - 1) + left = right - 1 + return yp[left] + (x - xp[left]) * (yp[right] - yp[left]) / ( + xp[right] - xp[left] + ) diff --git a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py new file mode 100644 index 0000000000..5cfb867ccb --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -0,0 +1,136 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import Literal + +import numpy as np +import pytest +import torch + +from sdm.models.kumo.timeseries.anomaly.noise_schedule import NoiseScheduleVP +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize( + ("schedule", "dtype", "clipped_lambda"), + [ + ("quad", torch.float32, -5.1), + ("quad", torch.float64, -5.1), + ("linear", torch.float32, -5.1), + ("linear", torch.float32, -2.0), + ], +) +def test_noise_schedule( + device: torch.device, + schedule: Literal["quad", "linear"], + dtype: torch.dtype, + clipped_lambda: float, +) -> None: + """Match VP coefficients, tail truncation, and interpolation.""" + discrete = DiffusionSchedule( + num_steps=1000, + beta_start=1e-4, + beta_end=0.5 if schedule == "quad" else 0.02, + schedule=schedule, + dtype=dtype, + device=device, + ) + module = NoiseScheduleVP(discrete.beta, clipped_lambda=clipped_lambda) + log_alpha = 0.5 * (1 - discrete.beta).log().cumsum(0) + half_log_snr = log_alpha - 0.5 * (1 - (2 * log_alpha).exp()).log() + retained = int((half_log_snr >= clipped_lambda).sum()) + assert module.total_N == retained + assert (retained < 1000) == (schedule == "quad" or clipped_lambda == -2.0) + + grid = torch.linspace(0, 1, retained + 1, dtype=torch.float32)[1:].to( + discrete.beta + ) + torch.testing.assert_close( + actual=module.marginal_log_mean_coeff(grid), + expected=log_alpha[:retained], + ) + # Unsorted queries, repeated values, and both interval endpoints. + t = torch.tensor( + data=[1.0, 0.435, 1 / retained, 0.731, 0.435], + device=device, + dtype=dtype, + ) + expected = torch.tensor( + data=np.interp( + x=t.cpu().numpy(), + xp=grid.cpu().numpy(), + fp=log_alpha[:retained].cpu().numpy(), + ), + device=device, + dtype=dtype, + ) + alpha = expected.exp() + sigma = (1 - (2 * expected).exp()).sqrt() + for actual, target in ( + (module.marginal_log_mean_coeff(t), expected), + (module.marginal_alpha(t), alpha), + (module.marginal_std(t), sigma), + (module.marginal_lambda(t), expected - sigma.log()), + (module.inverse_lambda(module.marginal_lambda(t)), t), + ): + assert actual.device == device + assert actual.dtype == dtype + torch.testing.assert_close(actual, target) + + +@withCUDA +def test_noise_schedule_extrapolation(device: torch.device) -> None: + """Preserve boundary extrapolation and query-time gradients.""" + betas = torch.tensor( + data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64 + ) + module = NoiseScheduleVP(betas) + logs = 0.5 * (1 - betas).log().cumsum(0) + t = torch.tensor( + data=[[0.2, 0.25], [0.375, 1.1]], + device=device, + dtype=betas.dtype, + requires_grad=True, + ) + expected = torch.stack( + tensors=[ + logs[0] - 0.2 * (logs[1] - logs[0]), + logs[0], + (logs[0] + logs[1]) / 2, + logs[-1] + 0.4 * (logs[-1] - logs[-2]), + ], + ) + torch.testing.assert_close(module.marginal_log_mean_coeff(t), expected) + torch.testing.assert_close( + actual=module.inverse_lambda(module.marginal_lambda(t)), + expected=t.flatten(), + ) + (gradient,) = torch.autograd.grad( + module.marginal_log_mean_coeff(t).sum(), t + ) + assert torch.isfinite(gradient).all() + assert (gradient != 0).any() + + +@withCUDA +def test_noise_schedule_conversion(device: torch.device) -> None: + """Move coefficients while keeping outputs and checkpoint contents.""" + betas = torch.tensor( + data=[0.1, 0.2, 0.3, 0.4], dtype=torch.float64, requires_grad=True + ) + module = NoiseScheduleVP(betas) + t = torch.tensor([0.3, 0.6, 0.9], dtype=betas.dtype) + expected = module.marginal_lambda(t).detach() + module = module.to(device=device, dtype=torch.float32) + actual = module.marginal_lambda(t.to(device).float()) + torch.testing.assert_close(actual, expected.to(device).float()) + assert not actual.requires_grad + assert not module.state_dict() + + +def test_noise_schedule_too_short() -> None: + """Reject schedules without two usable interpolation knots.""" + with pytest.raises(ValueError, match="At least two diffusion steps"): + NoiseScheduleVP(torch.tensor([0.99999, 0.99999]))