[Kumo-Anomaly 2/n] Add KumoAnomaly reconstruction and scoring

sdm/models/kumo/timeseries/anomaly/denoiser.py

@@ -29,6 +29,7 @@ def __init__(
     ) -> None:
         super().__init__()
         factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        self.num_steps = num_steps
 
         self.diffusion_embedding = DiffusionEmbedding(
             num_steps=num_steps,

sdm/models/kumo/timeseries/anomaly/embedding.py

@@ -87,7 +87,7 @@ class SideEmbedding(torch.nn.Module):
     before conversion to the parameter dtype.
 
     Args:
-        num_features: Number of features in checkpoint order.
+        in_channels: Number of features in checkpoint order.
         time_channels: Even sinusoidal time embedding width, at least two.
         feature_channels: Learned feature embedding width.
         is_unconditional: Append the conditioning mask as the last channel.
@@ -97,7 +97,7 @@ class SideEmbedding(torch.nn.Module):
 
     def __init__(
         self,
-        num_features: int,
+        in_channels: int,
         time_channels: int = 128,
         feature_channels: int = 16,
         is_unconditional: bool = False,
@@ -109,7 +109,7 @@ def __init__(
         self.time_channels = time_channels
         self.is_unconditional = is_unconditional
         self.embed_layer = torch.nn.Embedding(
-            num_embeddings=num_features,
+            num_embeddings=in_channels,
             embedding_dim=feature_channels,
             **factory_kwargs,
         )
@@ -129,7 +129,7 @@ def forward(self, time: Tensor, mask: Tensor) -> Tensor:
             time: Model time coordinates with shape ``[B, L]``. Coordinates
                 are used as supplied, with no normalization or resampling.
             mask: Conditioning mask with shape ``[B, K, L]`` in checkpoint
-                feature order, where ``K`` equals ``num_features``. One marks
+                feature order, where ``K`` equals ``in_channels``. One marks
                 an observed value used for conditioning; zero marks a target.
                 Both inputs must be on the module's device.
 

sdm/models/kumo/timeseries/anomaly/model.py

@@ -7,16 +7,22 @@
 
 import copy
 from collections.abc import Sequence
-from typing import ClassVar
+from functools import partial
+from typing import Any, ClassVar, cast
 
 import torch
+import torch.nn.functional as F
 from torch import Tensor
 
 from sdm import Recipe, Stype, TableTensor
 from sdm._inference import inference_mode
 from sdm.models import Callback, Model
 from sdm.models._huggingface import download_checkpoint
+from sdm.models.kumo.timeseries.anomaly.denoiser import DiffusionDenoiser
+from sdm.models.kumo.timeseries.anomaly.embedding import SideEmbedding
 from sdm.models.kumo.timeseries.anomaly.recipe import default_recipe
+from sdm.models.kumo.timeseries.anomaly.sampling import DPMSolver
+from sdm.models.kumo.timeseries.anomaly.schedule import DiffusionSchedule
 from sdm.processing.execution import MemberQuery, RecipeExecution
 
 
@@ -154,12 +160,14 @@ def forward(
 
             out = TableTensor(
                 columns={Stype.numerical: ("mae",)},
-                numerical=out.to(query.x.numerical),
+                numerical=out.unsqueeze(-1),
             )
 
             for callback in callbacks:
                 out = callback.on_model_forward_end(self, out)
 
+            out = cast(TableTensor, out.to(query.x.dtype))
+
         with (
             torch.amp.autocast(x.device.type, enabled=False),
             inference_mode("grad" if requires_grad else "inference"),
@@ -170,16 +178,137 @@ def forward(
 class _KumoAnomaly(torch.nn.Module):
     def __init__(
         self,
+        in_channels: int = 40,
+        time_channels: int = 256,
+        feature_channels: int = 64,
+        channels: int = 128,
+        num_layers: int = 6,
+        num_heads: int = 8,
+        embedding_channels: int = 256,
+        num_steps: int = 500,
         device: torch.device | str | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        self.in_channels = in_channels
+
+        self.side_embedding = SideEmbedding(
+            in_channels=in_channels,
+            time_channels=time_channels,
+            feature_channels=feature_channels,
+            **factory_kwargs,
+        )
+        self.denoiser = DiffusionDenoiser(
+            channels=channels,
+            side_channels=time_channels + feature_channels,
+            num_steps=num_steps,
+            num_layers=num_layers,
+            num_heads=num_heads,
+            embedding_channels=embedding_channels,
+            **factory_kwargs,
+        )
 
     def forward(
         self,
-        x: Tensor,  # [..., R, C]
+        x: Tensor,  # [..., R, D]
         num_samples: int,
         num_steps: int,
+        window_size: int = 100,
+        num_blocks: int = 4,
+        batch_size: int = 32,
         generator: torch.Generator | None = None,
-    ) -> Tensor:  # [..., R, 1]
-        return x.mean(dim=-1, keepdim=True)
+    ) -> Tensor:  # [..., R]
+
+        if x.size(-1) > self.in_channels:
+            raise ValueError(
+                f"'KumoAnomaly' supports at most {self.in_channels} columns "
+                f"after pre-processing (got {x.size(-1)})"
+            )
+
+        *B, R, D = x.size()
+        x = F.pad(x.reshape(-1, R, D), (0, self.in_channels - D))
+
+        # Pad last window by repeating last row:
+        x = torch.cat([x, x[:, -1:].expand(-1, -R % window_size, -1)], dim=-2)
+        window = x.reshape(-1, window_size, self.in_channels)
+
+        schedule = DiffusionSchedule(
+            num_steps=self.denoiser.num_steps,
+            beta_start=0.0001,
+            beta_end=0.01,
+            schedule="linear",
+            device="cpu",
+        )
+        sampler = DPMSolver(schedule.beta.float(), num_steps).to(x.device)
+
+        out = torch.cat(
+            [
+                self.reconstruct(
+                    x=_x,
+                    sampler=sampler,
+                    num_samples=num_samples,
+                    num_blocks=num_blocks,
+                    generator=generator,
+                )
+                for _x in window.split(batch_size)
+            ],
+            dim=0,
+        )
+
+        out = out.view(x.size())[:, :R, :D]
+        out = (out - x[:, :R, :D]).abs().mean(dim=-1)
+        return out.view(*B, R)
+
+    def reconstruct(
+        self,
+        x: Tensor,  # [B, R, D]
+        sampler: DPMSolver,
+        num_samples: int,
+        num_blocks: int,
+        generator: torch.Generator | None,
+    ) -> Tensor:  # [B, R, D]
+
+        B, R, _ = x.size()
+        x = x.transpose(1, 2)  # [B, D, R]
+        mean = x.mean(dim=(1, 2), keepdim=True)
+        std = x.std(dim=(1, 2), keepdim=True) + 1e-5
+        x = ((x - mean) / std).clamp(-10, 10)
+
+        time = torch.arange(R, device=x.device).expand(B, -1)
+        blocks = time // (R // num_blocks)
+
+        out = x.new_zeros((num_samples, *x.size()))
+        for strategy in range(2):
+            mask = ((blocks % 2) == strategy).to(x.dtype)
+            mask = mask.unsqueeze(1).expand_as(x)  # [B, D, R]
+            predict_noise = partial(
+                self.predict_noise,
+                cond=mask * x,
+                target_mask=1 - mask,
+                side_info=self.side_embedding(time, mask),
+                strategy=torch.tensor([strategy], device=x.device),
+            )
+            for i in range(num_samples):
+                noise = torch.empty_like(x).normal_(generator=generator)
+                sample = sampler(noise, predict_noise) * std + mean
+                out[i] += (1 - mask) * sample
+
+        return out.median(dim=0).values.transpose(1, 2)
+
+    def predict_noise(
+        self,
+        x: Tensor,  # [B, D, R]
+        step: Tensor,  # [1]
+        *,
+        cond: Tensor,  # [B, D, R]
+        target_mask: Tensor,  # [B, D, R]
+        side_info: Tensor,  # [B, S, D, R]
+        strategy: Tensor,  # [1]
+    ) -> Tensor:  # [B, D, R]
+        return self.denoiser(
+            x=torch.stack([cond, target_mask * x], dim=1),
+            side_info=side_info,
+            step=step,
+            strategy=strategy,
+        )

test/models/kumo/timeseries/anomaly/test_model.py

@@ -3,18 +3,18 @@
 
 import torch
 
-from sdm import TableTensor
 from sdm.models import KumoAnomaly
 from sdm.testing import withCUDA
 
 
 @withCUDA
 def test_forward(device: torch.device) -> None:
     model = KumoAnomaly(pretrained=False, device=device)
-    x = TableTensor.from_tensor(torch.randn(32, 5, device=device))
+    x = torch.randn(2, 130, 3, device=device)
 
-    out = model(x)
-    assert out.size() == (32, 1)
+    out = model(x, num_samples=1, num_steps=1)
+    assert out.size() == (2, 130, 1)
+    assert out.columns["numerical"] == ("mae",)
     assert out.device == device
     assert out.numerical.isfinite().all()
-    assert "mae" in out.numerical_columns
+    assert (out.numerical >= 0).all()

test/models/kumo/timeseries/anomaly/test_side_embedding.py

@@ -35,9 +35,9 @@ def test_side_embedding(
         and not torch.cuda.is_bf16_supported()
     ):
         pytest.skip("CUDA bfloat16 not supported")
-    batch, features, length = shape
+    B, D, R = shape
     module = SideEmbedding(
-        num_features=features,
+        in_channels=D,
         time_channels=time_channels,
         feature_channels=4,
         is_unconditional=is_unconditional,
@@ -46,10 +46,10 @@ def test_side_embedding(
     )
     if dtype == torch.float64:
         module = module.to(dtype=dtype)
-    time = torch.arange(batch * length, device=device).reshape(batch, length)
+    time = torch.arange(B * R, device=device).reshape(B, R)
     time = time.float().square() * 0.25 - 2.5
     time[:, -1] = 100000
-    mask = torch.arange(batch * features * length, device=device)
+    mask = torch.arange(B * D * R, device=device)
     mask = mask.reshape(shape).remainder(2).bool()
     actual = module(time, mask)
 
@@ -60,7 +60,7 @@ def test_side_embedding(
         / time_channels,
     )
     phase = time.unsqueeze(-1) * frequencies
-    time_embed = torch.empty(batch, length, time_channels, device=device)
+    time_embed = torch.empty(B, R, time_channels, device=device)
     time_embed[..., 0::2] = phase.sin()
     time_embed[..., 1::2] = phase.cos()
     time_embed = F.layer_norm(
@@ -71,7 +71,7 @@ def test_side_embedding(
         eps=module.time_embed_norm.eps,
     )
     feature_embed = F.embedding(
-        input=torch.arange(features, device=device),
+        input=torch.arange(D, device=device),
         weight=module.embed_layer.weight,
     )
     feature_embed = F.layer_norm(
@@ -83,19 +83,14 @@ def test_side_embedding(
     )
     expected = torch.cat(
         [
-            time_embed[:, :, None].expand(-1, -1, features, -1),
-            feature_embed[None, None].expand(batch, length, -1, -1),
+            time_embed[:, :, None].expand(-1, -1, D, -1),
+            feature_embed[None, None].expand(B, R, -1, -1),
         ],
         dim=-1,
     ).permute(0, 3, 2, 1)
     if is_unconditional:
         expected = torch.cat([expected, mask[:, None].to(dtype)], dim=1)
-    assert actual.shape == (
-        batch,
-        time_channels + 4 + int(is_unconditional),
-        features,
-        length,
-    )
+    assert actual.shape == (B, time_channels + 4 + int(is_unconditional), D, R)
     assert actual.dtype == dtype
     assert actual.device == device
     torch.testing.assert_close(actual, expected)