From 7ca64e33596031311d66dc21cc6e44718c8aa4d9 Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Thu, 15 Oct 2026 08:42:33 -0700 Subject: [PATCH 1/5] Add Kumo-Anomaly DDPM reverse step Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/sampling.py | 69 +++++++++++ .../kumo/timeseries/anomaly/test_sampling.py | 112 ++++++++++++++++++ 2 files changed, 181 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 0000000000..555f7db665 --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/sampling.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import torch +from torch import Tensor + +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule + + +class DDPMStep(torch.nn.Module): + """Apply one Kumo-Anomaly DDPM reverse-diffusion update. + + This component consumes predicted noise, not a denoiser. Callers control + the reverse-step loop, conditioning, and random draws. No normalization, + observation restoration, or anomaly scoring is applied. + + Args: + schedule: Diffusion coefficients used to derive the mean and noise + scales. Coefficients retain the schedule's device and dtype and + are excluded from the state dict. Construct a new step module + when changing the schedule. + """ + + mean_scale: Tensor + noise_scale: Tensor + std: Tensor + + def __init__(self, schedule: DiffusionSchedule) -> None: + super().__init__() + alpha = schedule.alpha + previous_alpha = torch.cat([alpha.new_ones(1), alpha[:-1]]) + for name, value in ( + ("mean_scale", schedule.alpha_hat.rsqrt()), + ("noise_scale", (1 - schedule.alpha_hat) / (1 - alpha).sqrt()), + ( + "std", + ((1 - previous_alpha) / (1 - alpha) * schedule.beta).sqrt(), + ), + ): + self.register_buffer(name=name, tensor=value, persistent=False) + + def forward( + self, + x: Tensor, + predicted_noise: Tensor, + step: Tensor, + noise: Tensor, + ) -> Tensor: + """Take a reverse step using explicitly supplied Gaussian noise. + + Args: + x: Current normalized sample with shape ``[B, K, L]``. + predicted_noise: Denoiser prediction with the shape and dtype of + ``x``. + step: Integer schedule indices with shape ``[B]`` or ``[1]``; + a scalar also broadcasts across the batch. Zero is the final + reverse step. All inputs must be on the coefficient device. + noise: Standard-normal noise with the same shape and dtype as + ``x``. The noise contribution is zero at step zero. Random + draws and their generator are controlled by the caller. + + Returns: + Updated samples with the shape and dtype of ``x``. Coefficients + are cast to the sample dtype before the update. + """ + mean_scale = self.mean_scale[step].to(x.dtype)[..., None, None] + noise_scale = self.noise_scale[step].to(x.dtype)[..., None, None] + std = self.std[step].to(x.dtype)[..., None, None] + return mean_scale * (x - noise_scale * predicted_noise) + std * noise 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 0000000000..20c4b56b8a --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_sampling.py @@ -0,0 +1,112 @@ +# 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.sampling import DDPMStep +from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule +from sdm.testing import withCUDA + + +@withCUDA +@pytest.mark.parametrize( + ("schedule", "num_steps", "indices", "dtype"), + [ + ("quad", 7, [6, 0, 3, 3], torch.float32), + ("linear", 7, [3], torch.float64), + ("quad", 1, [0], torch.float32), + ("linear", 12, [11, 2, 2, 0], torch.float16), + ("quad", 6, [5], torch.bfloat16), + ("quad", 100, [99, 0, 50, 99], torch.float32), + ], +) +def test_ddpm_step( + device: torch.device, + schedule: Literal["quad", "linear"], + num_steps: int, + indices: list[int], + dtype: torch.dtype, +) -> None: + """Match Kumo-TS reverse coefficients for shared and per-window steps.""" + coefficients = DiffusionSchedule( + num_steps=num_steps, + beta_start=1e-4, + beta_end=0.5, + schedule=schedule, + device=device, + ) + module = DDPMStep(coefficients) + x = torch.randn(4, 2, 5, device=device, dtype=dtype) + predicted = torch.randn_like(x) + noise = torch.randn_like(x) + step = torch.tensor(indices, device=device) + out = module(x=x, predicted_noise=predicted, step=step, noise=noise) + if schedule == "quad": + beta = np.linspace(1e-4**0.5, 0.5**0.5, num_steps) ** 2 + else: + beta = np.linspace(1e-4, 0.5, num_steps) + alpha_hat = 1 - beta + alpha = np.clip(np.cumprod(alpha_hat), 1e-6, 1 - 1e-6) + expected = torch.empty_like(x) + for batch in range(x.size(0)): + t = indices[batch] if len(indices) > 1 else indices[0] + coeff1 = x.new_tensor(1 / alpha_hat[t] ** 0.5) + coeff2 = x.new_tensor((1 - alpha_hat[t]) / (1 - alpha[t]) ** 0.5) + expected[batch] = coeff1 * (x[batch] - coeff2 * predicted[batch]) + if t > 0: + sigma = ((1 - alpha[t - 1]) / (1 - alpha[t]) * beta[t]) ** 0.5 + expected[batch] += x.new_tensor(sigma) * noise[batch] + torch.testing.assert_close(out, expected) + assert out.dtype == dtype + assert out.device == device + + +@withCUDA +def test_ddpm_step_final_and_gradients(device: torch.device) -> None: + """Ignore noise at the final step while retaining sample gradients.""" + schedule = DiffusionSchedule( + num_steps=3, beta_start=0.1, beta_end=0.3, device=device + ) + module = DDPMStep(schedule) + x = torch.randn( + size=(2, 1, 4), device=device, dtype=torch.float64, requires_grad=True + ) + predicted = torch.randn_like(x, requires_grad=True) + noise = torch.randn_like(x, requires_grad=True) + step = torch.tensor(0, device=device) + out = module(x=x, predicted_noise=predicted, step=step, noise=noise) + other = module( + x=x, predicted_noise=predicted, step=step, noise=noise * 100 + ) + torch.testing.assert_close(out, other) + out.sum().backward() + coeff1 = schedule.alpha_hat[0].rsqrt() + coeff2 = (1 - schedule.alpha_hat[0]) / (1 - schedule.alpha[0]).sqrt() + torch.testing.assert_close(x.grad, coeff1.expand_as(x)) + torch.testing.assert_close(predicted.grad, (-coeff1 * coeff2).expand_as(x)) + torch.testing.assert_close(noise.grad, torch.zeros_like(noise)) + + +@withCUDA +def test_ddpm_step_conversion(device: torch.device) -> None: + """Move derived coefficients with the module without checkpoint entries.""" + schedule = DiffusionSchedule(num_steps=3, beta_start=0.1, beta_end=0.3) + module = DDPMStep(schedule) + x = torch.randn(2, 1, 4) + predicted = torch.randn_like(x) + noise = torch.randn_like(x) + step = torch.tensor([2, 1]) + expected = module(x=x, predicted_noise=predicted, step=step, noise=noise) + module.to(device=device, dtype=torch.float32) + actual = module( + x=x.to(device), + predicted_noise=predicted.to(device), + step=step.to(device), + noise=noise.to(device), + ) + torch.testing.assert_close(actual, expected.to(device)) + assert not module.state_dict() From d58787af67e036f0637c8115e1236d703dfcdca1 Mon Sep 17 00:00:00 2001 From: Aditi Gautam Date: Thu, 15 Oct 2026 09:17:22 -0700 Subject: [PATCH 2/5] Replace DDPM step with DPM-Solver noise schedule Signed-off-by: Aditi Gautam --- .../kumo/timeseries/anomaly/noise_schedule.py | 100 ++++++++++++ .../kumo/timeseries/anomaly/sampling.py | 69 -------- .../timeseries/anomaly/test_noise_schedule.py | 153 ++++++++++++++++++ .../kumo/timeseries/anomaly/test_sampling.py | 112 ------------- 4 files changed, 253 insertions(+), 181 deletions(-) create mode 100644 sdm/models/kumo/timeseries/anomaly/noise_schedule.py delete mode 100644 sdm/models/kumo/timeseries/anomaly/sampling.py create mode 100644 test/models/kumo/timeseries/anomaly/test_noise_schedule.py delete mode 100644 test/models/kumo/timeseries/anomaly/test_sampling.py 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..0a96b4baac --- /dev/null +++ b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py @@ -0,0 +1,100 @@ +# 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 -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. + """ + + _log_alpha: Tensor + _time: Tensor + + def __init__(self, betas: Tensor) -> 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. + clipped = int( + torch.searchsorted((log_alpha - log_sigma).flip(0), -5.1) + ) + self.total_N = log_alpha.numel() - 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/sdm/models/kumo/timeseries/anomaly/sampling.py b/sdm/models/kumo/timeseries/anomaly/sampling.py deleted file mode 100644 index 555f7db665..0000000000 --- a/sdm/models/kumo/timeseries/anomaly/sampling.py +++ /dev/null @@ -1,69 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -import torch -from torch import Tensor - -from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule - - -class DDPMStep(torch.nn.Module): - """Apply one Kumo-Anomaly DDPM reverse-diffusion update. - - This component consumes predicted noise, not a denoiser. Callers control - the reverse-step loop, conditioning, and random draws. No normalization, - observation restoration, or anomaly scoring is applied. - - Args: - schedule: Diffusion coefficients used to derive the mean and noise - scales. Coefficients retain the schedule's device and dtype and - are excluded from the state dict. Construct a new step module - when changing the schedule. - """ - - mean_scale: Tensor - noise_scale: Tensor - std: Tensor - - def __init__(self, schedule: DiffusionSchedule) -> None: - super().__init__() - alpha = schedule.alpha - previous_alpha = torch.cat([alpha.new_ones(1), alpha[:-1]]) - for name, value in ( - ("mean_scale", schedule.alpha_hat.rsqrt()), - ("noise_scale", (1 - schedule.alpha_hat) / (1 - alpha).sqrt()), - ( - "std", - ((1 - previous_alpha) / (1 - alpha) * schedule.beta).sqrt(), - ), - ): - self.register_buffer(name=name, tensor=value, persistent=False) - - def forward( - self, - x: Tensor, - predicted_noise: Tensor, - step: Tensor, - noise: Tensor, - ) -> Tensor: - """Take a reverse step using explicitly supplied Gaussian noise. - - Args: - x: Current normalized sample with shape ``[B, K, L]``. - predicted_noise: Denoiser prediction with the shape and dtype of - ``x``. - step: Integer schedule indices with shape ``[B]`` or ``[1]``; - a scalar also broadcasts across the batch. Zero is the final - reverse step. All inputs must be on the coefficient device. - noise: Standard-normal noise with the same shape and dtype as - ``x``. The noise contribution is zero at step zero. Random - draws and their generator are controlled by the caller. - - Returns: - Updated samples with the shape and dtype of ``x``. Coefficients - are cast to the sample dtype before the update. - """ - mean_scale = self.mean_scale[step].to(x.dtype)[..., None, None] - noise_scale = self.noise_scale[step].to(x.dtype)[..., None, None] - std = self.std[step].to(x.dtype)[..., None, None] - return mean_scale * (x - noise_scale * predicted_noise) + std * noise 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..c17dada8a4 --- /dev/null +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -0,0 +1,153 @@ +# 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", ["quad", "linear"]) +@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +def test_noise_schedule( + device: torch.device, + schedule: Literal["quad", "linear"], + dtype: torch.dtype, +) -> 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) + 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 >= -5.1).sum()) + assert module.total_N == retained + assert module.T == 1.0 + assert module.schedule == "discrete" + assert (retained < 1000) == (schedule == "quad") + + 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: + """Extrapolate boundary segments and flatten scalar/batched inputs.""" + 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 + ) + 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(), + ) + torch.testing.assert_close( + actual=module.marginal_alpha(t[0, 0]), expected=expected[:1].exp() + ) + empty = torch.empty(0, device=device, dtype=betas.dtype) + assert module.marginal_alpha(empty).shape == (0,) + assert module.inverse_lambda(empty).shape == (0,) + + +@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() + + +@withCUDA +def test_noise_schedule_gradients(device: torch.device) -> None: + """Differentiate time and half-logSNR away from interpolation knots.""" + betas = torch.tensor( + data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64 + ) + module = NoiseScheduleVP(betas) + t = torch.tensor( + data=[0.3, 0.6, 0.9], + device=device, + dtype=betas.dtype, + requires_grad=True, + ) + for method in ( + module.marginal_log_mean_coeff, + module.marginal_alpha, + module.marginal_std, + module.marginal_lambda, + ): + assert torch.autograd.gradcheck(method, (t,)) + lamb = module.marginal_lambda(t).detach().requires_grad_() + assert torch.autograd.gradcheck(module.inverse_lambda, (lamb,)) + + +@pytest.mark.parametrize("betas", [[], [0.1], [0.99999, 0.99999]]) +def test_noise_schedule_too_short(betas: list[float]) -> None: + """Reject schedules without two usable interpolation knots.""" + with pytest.raises(ValueError, match="At least two diffusion steps"): + NoiseScheduleVP(torch.tensor(betas)) diff --git a/test/models/kumo/timeseries/anomaly/test_sampling.py b/test/models/kumo/timeseries/anomaly/test_sampling.py deleted file mode 100644 index 20c4b56b8a..0000000000 --- a/test/models/kumo/timeseries/anomaly/test_sampling.py +++ /dev/null @@ -1,112 +0,0 @@ -# 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.sampling import DDPMStep -from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule -from sdm.testing import withCUDA - - -@withCUDA -@pytest.mark.parametrize( - ("schedule", "num_steps", "indices", "dtype"), - [ - ("quad", 7, [6, 0, 3, 3], torch.float32), - ("linear", 7, [3], torch.float64), - ("quad", 1, [0], torch.float32), - ("linear", 12, [11, 2, 2, 0], torch.float16), - ("quad", 6, [5], torch.bfloat16), - ("quad", 100, [99, 0, 50, 99], torch.float32), - ], -) -def test_ddpm_step( - device: torch.device, - schedule: Literal["quad", "linear"], - num_steps: int, - indices: list[int], - dtype: torch.dtype, -) -> None: - """Match Kumo-TS reverse coefficients for shared and per-window steps.""" - coefficients = DiffusionSchedule( - num_steps=num_steps, - beta_start=1e-4, - beta_end=0.5, - schedule=schedule, - device=device, - ) - module = DDPMStep(coefficients) - x = torch.randn(4, 2, 5, device=device, dtype=dtype) - predicted = torch.randn_like(x) - noise = torch.randn_like(x) - step = torch.tensor(indices, device=device) - out = module(x=x, predicted_noise=predicted, step=step, noise=noise) - if schedule == "quad": - beta = np.linspace(1e-4**0.5, 0.5**0.5, num_steps) ** 2 - else: - beta = np.linspace(1e-4, 0.5, num_steps) - alpha_hat = 1 - beta - alpha = np.clip(np.cumprod(alpha_hat), 1e-6, 1 - 1e-6) - expected = torch.empty_like(x) - for batch in range(x.size(0)): - t = indices[batch] if len(indices) > 1 else indices[0] - coeff1 = x.new_tensor(1 / alpha_hat[t] ** 0.5) - coeff2 = x.new_tensor((1 - alpha_hat[t]) / (1 - alpha[t]) ** 0.5) - expected[batch] = coeff1 * (x[batch] - coeff2 * predicted[batch]) - if t > 0: - sigma = ((1 - alpha[t - 1]) / (1 - alpha[t]) * beta[t]) ** 0.5 - expected[batch] += x.new_tensor(sigma) * noise[batch] - torch.testing.assert_close(out, expected) - assert out.dtype == dtype - assert out.device == device - - -@withCUDA -def test_ddpm_step_final_and_gradients(device: torch.device) -> None: - """Ignore noise at the final step while retaining sample gradients.""" - schedule = DiffusionSchedule( - num_steps=3, beta_start=0.1, beta_end=0.3, device=device - ) - module = DDPMStep(schedule) - x = torch.randn( - size=(2, 1, 4), device=device, dtype=torch.float64, requires_grad=True - ) - predicted = torch.randn_like(x, requires_grad=True) - noise = torch.randn_like(x, requires_grad=True) - step = torch.tensor(0, device=device) - out = module(x=x, predicted_noise=predicted, step=step, noise=noise) - other = module( - x=x, predicted_noise=predicted, step=step, noise=noise * 100 - ) - torch.testing.assert_close(out, other) - out.sum().backward() - coeff1 = schedule.alpha_hat[0].rsqrt() - coeff2 = (1 - schedule.alpha_hat[0]) / (1 - schedule.alpha[0]).sqrt() - torch.testing.assert_close(x.grad, coeff1.expand_as(x)) - torch.testing.assert_close(predicted.grad, (-coeff1 * coeff2).expand_as(x)) - torch.testing.assert_close(noise.grad, torch.zeros_like(noise)) - - -@withCUDA -def test_ddpm_step_conversion(device: torch.device) -> None: - """Move derived coefficients with the module without checkpoint entries.""" - schedule = DiffusionSchedule(num_steps=3, beta_start=0.1, beta_end=0.3) - module = DDPMStep(schedule) - x = torch.randn(2, 1, 4) - predicted = torch.randn_like(x) - noise = torch.randn_like(x) - step = torch.tensor([2, 1]) - expected = module(x=x, predicted_noise=predicted, step=step, noise=noise) - module.to(device=device, dtype=torch.float32) - actual = module( - x=x.to(device), - predicted_noise=predicted.to(device), - step=step.to(device), - noise=noise.to(device), - ) - torch.testing.assert_close(actual, expected.to(device)) - assert not module.state_dict() From 1cf0c6165f6122b41fad1b2a79515e50b77f778e Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:22:01 +0900 Subject: [PATCH 3/5] Expose noise schedule cutoff and focus inference tests --- .../kumo/timeseries/anomaly/noise_schedule.py | 13 +++--- .../timeseries/anomaly/test_noise_schedule.py | 43 ++++++------------- 2 files changed, 22 insertions(+), 34 deletions(-) diff --git a/sdm/models/kumo/timeseries/anomaly/noise_schedule.py b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py index 0a96b4baac..e47f3880bb 100644 --- a/sdm/models/kumo/timeseries/anomaly/noise_schedule.py +++ b/sdm/models/kumo/timeseries/anomaly/noise_schedule.py @@ -14,7 +14,8 @@ class NoiseScheduleVP(torch.nn.Module): 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 -5.1 and assign the retained steps times ``(i + 1) / total_N``. + 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. @@ -29,22 +30,24 @@ class NoiseScheduleVP(torch.nn.Module): 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) -> None: + 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. - clipped = int( - torch.searchsorted((log_alpha - log_sigma).flip(0), -5.1) + num_clipped = int( + torch.searchsorted((log_alpha - log_sigma).flip(0), clipped_lambda) ) - self.total_N = log_alpha.numel() - clipped + self.total_N = log_alpha.numel() - num_clipped if self.total_N < 2: raise ValueError( "At least two diffusion steps must remain after truncation" diff --git a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py index c17dada8a4..31018a42d0 100644 --- a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -13,12 +13,21 @@ @withCUDA -@pytest.mark.parametrize("schedule", ["quad", "linear"]) -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) +@pytest.mark.parametrize( + ("schedule", "dtype", "clipped_lambda"), + [ + ("quad", torch.float32, -5.1), + ("quad", torch.float64, -5.1), + ("linear", torch.float32, -5.1), + ("quad", torch.float32, -2.0), + ("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( @@ -29,14 +38,14 @@ def test_noise_schedule( dtype=dtype, device=device, ) - module = NoiseScheduleVP(discrete.beta) + 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 >= -5.1).sum()) + retained = int((half_log_snr >= clipped_lambda).sum()) assert module.total_N == retained assert module.T == 1.0 assert module.schedule == "discrete" - assert (retained < 1000) == (schedule == "quad") + assert (retained < 1000) == (schedule == "quad" or clipped_lambda == -2.0) grid = torch.linspace(0, 1, retained + 1, dtype=torch.float32)[1:].to( discrete.beta @@ -122,30 +131,6 @@ def test_noise_schedule_conversion(device: torch.device) -> None: assert not module.state_dict() -@withCUDA -def test_noise_schedule_gradients(device: torch.device) -> None: - """Differentiate time and half-logSNR away from interpolation knots.""" - betas = torch.tensor( - data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64 - ) - module = NoiseScheduleVP(betas) - t = torch.tensor( - data=[0.3, 0.6, 0.9], - device=device, - dtype=betas.dtype, - requires_grad=True, - ) - for method in ( - module.marginal_log_mean_coeff, - module.marginal_alpha, - module.marginal_std, - module.marginal_lambda, - ): - assert torch.autograd.gradcheck(method, (t,)) - lamb = module.marginal_lambda(t).detach().requires_grad_() - assert torch.autograd.gradcheck(module.inverse_lambda, (lamb,)) - - @pytest.mark.parametrize("betas", [[], [0.1], [0.99999, 0.99999]]) def test_noise_schedule_too_short(betas: list[float]) -> None: """Reject schedules without two usable interpolation knots.""" From ccf8e81217d02620114cf0f901650b91c61d557e Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:50:29 +0900 Subject: [PATCH 4/5] Check query gradients in noise schedule interpolation test --- .../kumo/timeseries/anomaly/test_noise_schedule.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py index 31018a42d0..85a8b02986 100644 --- a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -92,7 +92,10 @@ def test_noise_schedule_extrapolation(device: torch.device) -> None: 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 + data=[[0.2, 0.25], [0.375, 1.1]], + device=device, + dtype=betas.dtype, + requires_grad=True, ) expected = torch.stack( tensors=[ @@ -113,6 +116,11 @@ def test_noise_schedule_extrapolation(device: torch.device) -> None: empty = torch.empty(0, device=device, dtype=betas.dtype) assert module.marginal_alpha(empty).shape == (0,) assert module.inverse_lambda(empty).shape == (0,) + gradient = torch.autograd.grad(module.marginal_log_mean_coeff(t).sum(), t)[ + 0 + ] + assert torch.isfinite(gradient).all() + assert (gradient != 0).any() @withCUDA From ad16003349869ca19f31c27df83fb6ba972db2cf Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:55:21 +0900 Subject: [PATCH 5/5] Trim redundant noise schedule test cases --- .../timeseries/anomaly/test_noise_schedule.py | 20 +++++-------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py index 85a8b02986..5cfb867ccb 100644 --- a/test/models/kumo/timeseries/anomaly/test_noise_schedule.py +++ b/test/models/kumo/timeseries/anomaly/test_noise_schedule.py @@ -19,7 +19,6 @@ ("quad", torch.float32, -5.1), ("quad", torch.float64, -5.1), ("linear", torch.float32, -5.1), - ("quad", torch.float32, -2.0), ("linear", torch.float32, -2.0), ], ) @@ -43,8 +42,6 @@ def test_noise_schedule( 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 module.T == 1.0 - assert module.schedule == "discrete" assert (retained < 1000) == (schedule == "quad" or clipped_lambda == -2.0) grid = torch.linspace(0, 1, retained + 1, dtype=torch.float32)[1:].to( @@ -85,7 +82,7 @@ def test_noise_schedule( @withCUDA def test_noise_schedule_extrapolation(device: torch.device) -> None: - """Extrapolate boundary segments and flatten scalar/batched inputs.""" + """Preserve boundary extrapolation and query-time gradients.""" betas = torch.tensor( data=[0.1, 0.2, 0.3, 0.4], device=device, dtype=torch.float64 ) @@ -110,15 +107,9 @@ def test_noise_schedule_extrapolation(device: torch.device) -> None: actual=module.inverse_lambda(module.marginal_lambda(t)), expected=t.flatten(), ) - torch.testing.assert_close( - actual=module.marginal_alpha(t[0, 0]), expected=expected[:1].exp() + (gradient,) = torch.autograd.grad( + module.marginal_log_mean_coeff(t).sum(), t ) - empty = torch.empty(0, device=device, dtype=betas.dtype) - assert module.marginal_alpha(empty).shape == (0,) - assert module.inverse_lambda(empty).shape == (0,) - gradient = torch.autograd.grad(module.marginal_log_mean_coeff(t).sum(), t)[ - 0 - ] assert torch.isfinite(gradient).all() assert (gradient != 0).any() @@ -139,8 +130,7 @@ def test_noise_schedule_conversion(device: torch.device) -> None: assert not module.state_dict() -@pytest.mark.parametrize("betas", [[], [0.1], [0.99999, 0.99999]]) -def test_noise_schedule_too_short(betas: list[float]) -> None: +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(betas)) + NoiseScheduleVP(torch.tensor([0.99999, 0.99999]))