[Kumo-Anomaly 3/n] Support post-norm and dropout in TransformerBlock

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

@@ -6,7 +6,53 @@
 
 import torch
 from torch import Tensor
-from torch.nn import Conv1d, LayerNorm, Linear, TransformerEncoderLayer
+from torch.nn import GELU, Conv1d, Dropout, LayerNorm, Linear, Sequential
+
+from sdm.nn import TransformerBlock
+
+
+class KumoAnomalyTransformerBlock(TransformerBlock):
+    """Post-norm transformer block of Kumo-Anomaly.
+
+    Args:
+        channels: Input and output width.
+        num_heads: Attention heads.
+        hidden_channels: Feedforward width.
+        dropout: Attention and feedforward dropout probability.
+        device: The device.
+        dtype: The dtype.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        hidden_channels: int,
+        dropout: float,
+        device: torch.device | str | None = None,
+        dtype: torch.dtype | None = None,
+    ) -> None:
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+        super().__init__(
+            channels=channels,
+            num_query_heads=num_heads,
+            mlp=Sequential(
+                Linear(channels, hidden_channels, **factory_kwargs),
+                GELU(),
+                Dropout(dropout),
+                Linear(hidden_channels, channels, **factory_kwargs),
+                Dropout(dropout),
+            ),
+            attn_residual_norm=LayerNorm(channels, **factory_kwargs),
+            mlp_residual_norm=LayerNorm(channels, **factory_kwargs),
+            dropout=dropout,
+            **factory_kwargs,
+        )
+        # Match the initialization of `torch.nn.MultiheadAttention`.
+        torch.nn.init.xavier_uniform_(self.attn.qkv_lin.weight)
+        torch.nn.init.zeros_(self.attn.qkv_lin.bias)
+        self.attn.out_lin.reset_parameters()
+        torch.nn.init.zeros_(self.attn.out_lin.bias)
 
 
 class ResidualBlock(torch.nn.Module):
@@ -84,25 +130,18 @@ def __init__(
             with torch.no_grad():
                 projection.weight.mul_(0.5)
 
-        # Preserve post-residual normalization and attention dropout, which
-        # SDM's TransformerBlock does not currently provide together.
-        # TODO: Switch to SDM's TransformerBlock once it supports both.
-        self.time_layer = TransformerEncoderLayer(
-            d_model=channels,
-            nhead=num_heads,
-            dim_feedforward=hidden_channels,
+        self.time_layer = KumoAnomalyTransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            hidden_channels=hidden_channels,
             dropout=dropout,
-            activation="gelu",
-            batch_first=True,
             **factory_kwargs,
         )
-        self.feature_layer = TransformerEncoderLayer(
-            d_model=channels,
-            nhead=num_heads,
-            dim_feedforward=hidden_channels,
+        self.feature_layer = KumoAnomalyTransformerBlock(
+            channels=channels,
+            num_heads=num_heads,
+            hidden_channels=hidden_channels,
             dropout=dropout,
-            activation="gelu",
-            batch_first=True,
             **factory_kwargs,
         )
         self.norm_after_time = LayerNorm(channels, **factory_kwargs)

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

@@ -5,16 +5,25 @@
 
 from __future__ import annotations
 
+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._warnings import warn_once
 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
 
 
 class KumoAnomaly(Model):
@@ -96,7 +105,7 @@ def forward(
         Args:
             x: The feature tensor with shape ``[..., R, D]`` with ``R`` rows
                 and ``D`` columns.
-            recipe: The custom recipe for pre- and post-processing.
+            recipe: The custom recipe for pre-processing.
             num_samples: The number of reconstruction samples per complementary
                 masking strategy. Reconstructions are combined using the
                 sample median.
@@ -111,13 +120,187 @@ def forward(
             absolute reconstruction error across preprocessed features at each
             time step.
         """
-        raise NotImplementedError
+        callbacks = () if callbacks is None else callbacks
+        requires_grad = self.training
+        requires_grad |= any(callback.requires_grad for callback in callbacks)
+
+        param = next(self.parameters())
+        if not isinstance(x, TableTensor):
+            x = TableTensor.from_tensor(x)
+        x = cast(TableTensor, x.to(param.device))
+
+        recipe = (
+            self.default_recipe() if recipe is None else copy.deepcopy(recipe)
+        )
+        with (
+            torch.amp.autocast(param.device.type, enabled=False),
+            inference_mode("no_grad" if requires_grad else "inference"),
+        ):
+            x = recipe.features.fit_transform(x, generator=generator)
+
+        for callback in callbacks:
+            x, _ = callback.on_query_preprocessing_end(self, x, None)
+
+        invalid = x.active_stypes - self.supported_feature_stypes - {Stype.id}
+        if len(invalid) > 0:
+            stypes = ", ".join(f"{str(stype)!r}" for stype in invalid)
+            warn_once(
+                key="model-unsupported-feature-stypes",
+                message=(
+                    f"{self.__class__.__name__!r} received unsupported "
+                    f"feature stypes {stypes}. Columns with unsupported "
+                    f"feature stypes will not be consumed by the model."
+                ),
+            )
+        if x.numerical.size(-1) > self.model.num_features:
+            raise ValueError(
+                f"{self.__class__.__name__!r} supports at most "
+                f"{self.model.num_features} features after pre-processing "
+                f"(got {x.numerical.size(-1)})"
+            )
+
+        with inference_mode("grad" if requires_grad else "inference"):
+            out = self.model(
+                x.numerical.to(param.dtype),
+                num_samples=num_samples,
+                num_steps=num_steps,
+                generator=generator,
+            )
+            out = TableTensor(
+                columns={Stype.numerical: ["score"]},
+                numerical=out.unsqueeze(-1),
+            )
+            for callback in callbacks:
+                out = callback.on_model_forward_end(self, out)
+        return cast(TableTensor, out.to(x.dtype))
 
 
 class _KumoAnomaly(torch.nn.Module):
     def __init__(
         self,
+        num_features: 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,
+        window_size: int = 100,
+        num_blocks: int = 4,
+        batch_size: int = 32,
         device: torch.device | str | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
+
+        self.num_features = num_features
+        self.num_steps = num_steps
+        self.window_size = window_size
+        self.num_blocks = num_blocks
+        self.batch_size = batch_size
+
+        self.side_embedding = SideEmbedding(
+            num_features=num_features,
+            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, D]
+        num_samples: int,
+        num_steps: int,
+        generator: torch.Generator | None,
+    ) -> Tensor:  # [..., R]
+
+        *size, length, features = x.size()
+        x = x.reshape(-1, length, features)
+        x = F.pad(x, (0, self.num_features - features))
+
+        # Pad the last window by repeating the last row:
+        padding = -length % self.window_size
+        x = torch.cat([x, x[:, -1:].expand(-1, padding, -1)], dim=1)
+        windows = x.reshape(-1, self.window_size, self.num_features)
+
+        schedule = DiffusionSchedule(
+            num_steps=self.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(batch, sampler, num_samples, generator)
+                for batch in windows.split(self.batch_size)
+            ]
+        )
+        out = out.reshape(x.size())[:, :length, :features]
+        out = (out - x[:, :length, :features]).abs().mean(dim=-1)
+        return out.reshape(*size, length)
+
+    def reconstruct(
+        self,
+        x: Tensor,  # [B, L, K]
+        sampler: DPMSolver,
+        num_samples: int,
+        generator: torch.Generator | None,
+    ) -> Tensor:  # [B, L, K]
+
+        x = x.transpose(1, 2)  # [B, K, L]
+        batch, _, length = x.size()
+        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(length, device=x.device).expand(batch, -1)
+        blocks = time // (length // self.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, K, L]
+            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, K, L]
+        step: Tensor,  # [1]
+        *,
+        cond: Tensor,  # [B, K, L]
+        target_mask: Tensor,  # [B, K, L]
+        side_info: Tensor,  # [B, S, K, L]
+        strategy: Tensor,  # [1]
+    ) -> Tensor:  # [B, K, L]
+        return self.denoiser(
+            x=torch.stack([cond, target_mask * x], dim=1),
+            side_info=side_info,
+            step=step,
+            strategy=strategy,
+        )

sdm/nn/attention.py

@@ -458,6 +458,10 @@ class TransformerBlock(torch.nn.Module):
             attention.
         post_attn_norm: Normalization applied to the attention output before
             its residual addition.
+        attn_residual_norm: Normalization applied after the attention residual
+            addition.
+        mlp_residual_norm: Normalization applied after the feedforward
+            residual addition. Requires ``mlp``.
         query_transform: Transformation applied to projected query heads before
             scaled dot-product attention.
         key_transform: Transformation applied to projected key heads before
@@ -468,7 +472,8 @@ class TransformerBlock(torch.nn.Module):
         scale: Scaling factor passed to
             :func:`torch.nn.functional.scaled_dot_product_attention`.
             ``None`` uses ``1 / sqrt(channels_per_head)``.
-        dropout: Attention dropout probability.
+        dropout: Dropout probability applied to attention weights and to the
+            attention output before its residual addition.
         bias: If set to ``False``, the module will not learn an additive bias.
         device: The device.
         dtype: The dtype.
@@ -483,6 +488,8 @@ def __init__(
         query_norm: torch.nn.Module | None = None,
         key_value_norm: torch.nn.Module | None = None,
         post_attn_norm: torch.nn.Module | None = None,
+        attn_residual_norm: torch.nn.Module | None = None,
+        mlp_residual_norm: torch.nn.Module | None = None,
         query_transform: torch.nn.Module | None = None,
         key_transform: torch.nn.Module | None = None,
         query_scaling: QueryScaling | None = None,
@@ -494,12 +501,18 @@ def __init__(
         dtype: torch.dtype | None = None,
     ) -> None:
         super().__init__()
+        if mlp is None and mlp_residual_norm is not None:
+            raise ValueError("'mlp_residual_norm' requires 'mlp'")
+
         factory_kwargs: dict[str, Any] = {"device": device, "dtype": dtype}
 
         self.mlp = mlp
         self.query_norm = query_norm
         self.key_value_norm = key_value_norm
         self.post_attn_norm = post_attn_norm
+        self.attn_residual_norm = attn_residual_norm
+        self.mlp_residual_norm = mlp_residual_norm
+        self.dropout = dropout
 
         self.attn = Attention(
             channels=channels,
@@ -811,20 +824,33 @@ def _forward(
 
         if self.post_attn_norm is not None:
             attn_out = self.post_attn_norm(attn_out)
+        if self.training and self.dropout > 0.0:
+            attn_out = F.dropout(attn_out, p=self.dropout)
 
         if (
             out is not None
             and torch.compiler.is_compiling()
             and not out.is_contiguous()
         ):
             tmp = (attn_out + query).to(out.dtype)
+            if self.attn_residual_norm is not None:
+                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
-                tmp += self.mlp(tmp)
+                tmp = (tmp + self.mlp(tmp)).to(out.dtype)
+                if self.mlp_residual_norm is not None:
+                    tmp = self.mlp_residual_norm(tmp)
             out.copy_(tmp)
         else:
             tmp = torch.add(attn_out, query, out=out)
+            if self.attn_residual_norm is not None:
+                tmp = self.attn_residual_norm(tmp)
             if self.mlp is not None:
                 tmp = torch.add(tmp, self.mlp(tmp), out=out)
+                if self.mlp_residual_norm is not None:
+                    tmp = self.mlp_residual_norm(tmp)
+            # Normalization allocates, so copy its result into `out`.
+            if out is not None and tmp is not out:
+                tmp = out.copy_(tmp)
             out = tmp
 
         return (out, kv) if return_key_value else out