Lighten TimesNet exposure backbone

backbones.py

@@ -151,82 +151,27 @@ class AgeSinusoidalEncoding(nn.Module):
         return output
 
 
-class GRN2d(nn.Module):
-    """Global Response Normalization from ConvNeXt V2 for NCHW tensors.
+class LiteTimesBackbone2d(nn.Module):
+    """Cheap local feature extractor for a TimesNet period image."""
 
-    Reference:
-        Woo et al., "ConvNeXt V2: Co-designing and Scaling ConvNets with
-        Masked Autoencoders", CVPR 2023. https://arxiv.org/abs/2301.00808
-    """
-
-    def __init__(self, dim: int, eps: float = 1e-6):
+    def __init__(self, dim: int, kernel_size: int = 5,
+                 expansion: float = 2.0, dropout: float = 0.0):
         super().__init__()
-        self.gamma = nn.Parameter(torch.zeros(1, dim, 1, 1))
-        self.beta = nn.Parameter(torch.zeros(1, dim, 1, 1))
-        self.eps = eps
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        gx = torch.norm(x, p=2, dim=(2, 3), keepdim=True)
-        nx = gx / (gx.mean(dim=1, keepdim=True) + self.eps)
-        return x + self.gamma * (x * nx) + self.beta
-
-
-class LayerNorm2d(nn.Module):
-    """Channel-wise LayerNorm for NCHW tensors."""
-
-    def __init__(self, dim: int, eps: float = 1e-6):
-        super().__init__()
-        self.weight = nn.Parameter(torch.ones(dim))
-        self.bias = nn.Parameter(torch.zeros(dim))
-        self.eps = eps
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = x.permute(0, 2, 3, 1)
-        x = F.layer_norm(x, (self.weight.numel(),), self.weight, self.bias, self.eps)
-        return x.permute(0, 3, 1, 2)
-
-
-class ConvNeXtV2Block2d(nn.Module):
-    """Lightweight ConvNeXt V2-style 2D block for TimesNet period images.
-
-    This is intentionally a block, not the full ConvNeXt V2 image backbone:
-    TimesNet's 2D tensors are reshaped time-series period maps rather than
-    natural images, so aggressive visual downsampling would destroy axis
-    semantics.
-    """
-
-    def __init__(
-        self,
-        dim: int,
-        kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
-        dropout: float = 0.0,
-    ):
-        super().__init__()
-        padding = kernel_size // 2
-        hidden_dim = int(dim * mlp_ratio)
+        if kernel_size <= 0 or kernel_size % 2 == 0:
+            raise ValueError("kernel_size must be a positive odd integer")
+        if expansion <= 0:
+            raise ValueError("expansion must be > 0")
+        hidden_dim = max(dim, int(round(dim * expansion)))
         self.dwconv = nn.Conv2d(
-            dim,
-            dim,
-            kernel_size=kernel_size,
-            padding=padding,
-            groups=dim,
+            dim, dim, kernel_size=kernel_size,
+            padding=kernel_size // 2, groups=dim,
         )
-        self.norm = LayerNorm2d(dim)
+        self.norm = nn.GroupNorm(1, dim)
         self.pwconv1 = nn.Conv2d(dim, hidden_dim, kernel_size=1)
         self.act = nn.GELU()
-        self.grn = GRN2d(hidden_dim)
         self.pwconv2 = nn.Conv2d(hidden_dim, dim, kernel_size=1)
         self.drop = nn.Dropout(dropout)
-        self.reset_parameters()
-
-    def reset_parameters(self) -> None:
-        nn.init.normal_(self.dwconv.weight, mean=0.0, std=0.02)
-        nn.init.zeros_(self.dwconv.bias)
-        nn.init.normal_(self.pwconv1.weight, mean=0.0, std=0.02)
-        nn.init.zeros_(self.pwconv1.bias)
-        nn.init.normal_(self.pwconv2.weight, mean=0.0, std=0.02)
-        nn.init.zeros_(self.pwconv2.bias)
+        self.layer_scale = nn.Parameter(torch.full((1, dim, 1, 1), 1e-2))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         residual = x
@@ -234,14 +179,12 @@ class ConvNeXtV2Block2d(nn.Module):
         x = self.norm(x)
         x = self.pwconv1(x)
         x = self.act(x)
-        x = self.grn(x)
         x = self.pwconv2(x)
-        x = self.drop(x)
-        return residual + x
+        return residual + self.layer_scale * self.drop(x)
 
 
 class TimesNetBlock(nn.Module):
-    """TimesNet block with ConvNeXt V2-style 2D extraction.
+    """TimesNet block with lightweight depthwise-separable 2D extraction.
 
     The block follows TimesNet's idea: discover dominant periods with FFT,
     reshape a 1D sequence into period-wise 2D maps, run a 2D convolutional
@@ -255,25 +198,27 @@ class TimesNetBlock(nn.Module):
     def __init__(
         self,
         d_model: int,
-        top_k: int = 3,
-        n_convnext_blocks: int = 2,
-        conv_kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
+        top_k: int = 2,
+        n_backbone_blocks: int = 1,
+        backbone_kernel_size: int = 5,
+        backbone_expansion: float = 2.0,
         dropout: float = 0.0,
     ):
         super().__init__()
         if top_k <= 0:
             raise ValueError(f"top_k must be > 0, got {top_k}")
+        if n_backbone_blocks <= 0:
+            raise ValueError("n_backbone_blocks must be > 0")
         self.top_k = top_k
         self.norm = nn.LayerNorm(d_model)
         self.extractor = nn.Sequential(*[
-            ConvNeXtV2Block2d(
+            LiteTimesBackbone2d(
                 dim=d_model,
-                kernel_size=conv_kernel_size,
-                mlp_ratio=mlp_ratio,
+                kernel_size=backbone_kernel_size,
+                expansion=backbone_expansion,
                 dropout=dropout,
             )
-            for _ in range(n_convnext_blocks)
+            for _ in range(n_backbone_blocks)
         ])
 
     def _select_periods(self, x: torch.Tensor) -> tuple[list[int], torch.Tensor]:
@@ -326,10 +271,10 @@ class TimesNetEncoder(nn.Module):
         input_dim: int,
         d_model: int,
         n_layers: int = 2,
-        top_k: int = 3,
-        n_convnext_blocks: int = 2,
-        conv_kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
+        top_k: int = 2,
+        n_backbone_blocks: int = 1,
+        backbone_kernel_size: int = 5,
+        backbone_expansion: float = 2.0,
         dropout: float = 0.0,
         append_observed_mask: bool = True,
     ):
@@ -342,9 +287,9 @@ class TimesNetEncoder(nn.Module):
             TimesNetBlock(
                 d_model=d_model,
                 top_k=top_k,
-                n_convnext_blocks=n_convnext_blocks,
-                conv_kernel_size=conv_kernel_size,
-                mlp_ratio=mlp_ratio,
+                n_backbone_blocks=n_backbone_blocks,
+                backbone_kernel_size=backbone_kernel_size,
+                backbone_expansion=backbone_expansion,
                 dropout=dropout,
             )
             for _ in range(n_layers)
@@ -415,10 +360,10 @@ class TimesNetExposureEncoder(nn.Module):
         monthly_input_dim: int = 2,
         d_model: int | None = None,
         n_layers: int = 2,
-        top_k: int = 3,
-        n_convnext_blocks: int = 2,
-        conv_kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
+        top_k: int = 2,
+        n_backbone_blocks: int = 1,
+        backbone_kernel_size: int = 5,
+        backbone_expansion: float = 2.0,
         dropout: float = 0.0,
         use_gate: bool = True,
     ):
@@ -429,9 +374,9 @@ class TimesNetExposureEncoder(nn.Module):
             d_model=d_model,
             n_layers=n_layers,
             top_k=top_k,
-            n_convnext_blocks=n_convnext_blocks,
-            conv_kernel_size=conv_kernel_size,
-            mlp_ratio=mlp_ratio,
+            n_backbone_blocks=n_backbone_blocks,
+            backbone_kernel_size=backbone_kernel_size,
+            backbone_expansion=backbone_expansion,
             dropout=dropout,
             append_observed_mask=True,
         )
@@ -440,9 +385,9 @@ class TimesNetExposureEncoder(nn.Module):
             d_model=d_model,
             n_layers=n_layers,
             top_k=top_k,
-            n_convnext_blocks=n_convnext_blocks,
-            conv_kernel_size=conv_kernel_size,
-            mlp_ratio=mlp_ratio,
+            n_backbone_blocks=n_backbone_blocks,
+            backbone_kernel_size=backbone_kernel_size,
+            backbone_expansion=backbone_expansion,
             dropout=dropout,
             append_observed_mask=True,
         )
@@ -485,10 +430,10 @@ class TimesNetSequenceDecoder(nn.Module):
         latent_dim: int,
         d_model: int,
         n_layers: int = 2,
-        top_k: int = 3,
-        n_convnext_blocks: int = 2,
-        conv_kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
+        top_k: int = 2,
+        n_backbone_blocks: int = 1,
+        backbone_kernel_size: int = 5,
+        backbone_expansion: float = 2.0,
         dropout: float = 0.0,
     ):
         super().__init__()
@@ -498,9 +443,9 @@ class TimesNetSequenceDecoder(nn.Module):
             TimesNetBlock(
                 d_model=d_model,
                 top_k=top_k,
-                n_convnext_blocks=n_convnext_blocks,
-                conv_kernel_size=conv_kernel_size,
-                mlp_ratio=mlp_ratio,
+                n_backbone_blocks=n_backbone_blocks,
+                backbone_kernel_size=backbone_kernel_size,
+                backbone_expansion=backbone_expansion,
                 dropout=dropout,
             )
             for _ in range(n_layers)
@@ -538,10 +483,10 @@ class TimesNetExposureAutoencoder(nn.Module):
         monthly_input_dim: int = 2,
         d_model: int | None = None,
         n_layers: int = 2,
-        top_k: int = 3,
-        n_convnext_blocks: int = 2,
-        conv_kernel_size: int = 7,
-        mlp_ratio: float = 4.0,
+        top_k: int = 2,
+        n_backbone_blocks: int = 1,
+        backbone_kernel_size: int = 5,
+        backbone_expansion: float = 2.0,
         dropout: float = 0.0,
     ):
         super().__init__()
@@ -553,9 +498,9 @@ class TimesNetExposureAutoencoder(nn.Module):
             d_model=d_model,
             n_layers=n_layers,
             top_k=top_k,
-            n_convnext_blocks=n_convnext_blocks,
-            conv_kernel_size=conv_kernel_size,
-            mlp_ratio=mlp_ratio,
+            n_backbone_blocks=n_backbone_blocks,
+            backbone_kernel_size=backbone_kernel_size,
+            backbone_expansion=backbone_expansion,
             dropout=dropout,
             use_gate=True,
         )
@@ -564,9 +509,9 @@ class TimesNetExposureAutoencoder(nn.Module):
             d_model=d_model,
             n_layers=n_layers,
             top_k=top_k,
-            n_convnext_blocks=n_convnext_blocks,
-            conv_kernel_size=conv_kernel_size,
-            mlp_ratio=mlp_ratio,
+            n_backbone_blocks=n_backbone_blocks,
+            backbone_kernel_size=backbone_kernel_size,
+            backbone_expansion=backbone_expansion,
             dropout=dropout,
         )
         self.encoder = TimesNetExposureEncoder(**encoder_kwargs)