patchfm 3.2.0__tar.gz → 3.3.0__tar.gz

{patchfm-3.2.0/src/patchfm.egg-info → patchfm-3.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 3.2.0
+Version: 3.3.0
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License

{patchfm-3.2.0 → patchfm-3.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "patchfm"
-version = "3.2.0"
+version = "3.3.0"
 authors = [
   { name="Samy-Melwan Vilhes", email="samy-melwan.vilhes@insa-rouen.fr" },
 ]

{patchfm-3.2.0 → patchfm-3.3.0}/src/patchfm/configs/model_config.py

@@ -5,13 +5,13 @@ from dataclasses import asdict, dataclass, field
 class PatchFMConfig:
     max_seq_len: int = 1024
     patch_len: int = 32
-    d_model: int = 2048
-    n_heads: int = 32
-    n_layers_encoder: int = 6
+    d_model: int = 1024
+    n_heads: int = 16
+    n_layers_encoder: int = 20
     quantiles: list[float] = field(default_factory=lambda: [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9])
     
-    compile: bool = True
-    full_leakage: bool = True
+    compile: bool = False
+    full_leakage: bool = False
 
     def __getitem__(self, key):
         return getattr(self, key)

{patchfm-3.2.0 → patchfm-3.3.0}/src/patchfm/inference/forecaster.py

@@ -4,7 +4,8 @@ import torch.nn as nn
 from einops import rearrange
 
 from patchfm.inference.modules import (CausalRevIN, PatchFM, ResidualBlock,
-                                       SeqTypeConverter, TransformerEncoder)
+                                       SeqTypeConverter, TransformerEncoder, 
+                                       quantile_widths, monotone_quantiles)
 
 
 # --- Forecaster Model ---
@@ -22,6 +23,9 @@ class Forecaster(nn.Module):
         self.n_quantiles = len(self.quantiles)
         self.max_patches = self.max_seq_len // self.patch_len
 
+        widths, self.median_idx = quantile_widths(self.quantiles)
+        self.register_buffer("quantile_widths", widths, persistent=False)
+
         if config["full_leakage"]:
             print("Loading full leakage model from HuggingFace Hub...")
             model_name = "vilhess/PatchFM-Leakage"
@@ -48,7 +52,7 @@ class Forecaster(nn.Module):
         self.revin = CausalRevIN()
         self.proj_embedding = ResidualBlock(
             in_dim=self.patch_len, 
-            hid_dim=4 * self.d_model, 
+            hid_dim=self.d_model, 
             out_dim=self.d_model
         )
         self.transformer_encoder = TransformerEncoder(
@@ -103,6 +107,19 @@ class Forecaster(nn.Module):
         x = x[:, -1:, :]  # Keep only the last patch for autoregressive forecasting
 
         forecasting = self.proj_output(x)
+        forecasting = rearrange(
+            forecasting,
+            "b 1 (pl q) -> b 1 pl q",
+            pl=self.patch_len,
+            q=self.n_quantiles,
+        )
+        forecasting = monotone_quantiles(
+            forecasting, self.quantile_widths, self.median_idx, width_factor=2.0
+        )
+        forecasting = rearrange(
+            forecasting,
+            "b 1 pl q -> b 1 (pl q)",
+        )
         forecasting = self.revin(forecasting, mode="denorm")
 
         # Reshape to (bs, patch_num, patch_len, n_quantiles)
@@ -122,6 +139,19 @@ class Forecaster(nn.Module):
             x = self.transformer_encoder(x)
             x = x[:, -1:, :]  # Keep only the last patch for autoregressive forecasting
             forecasting = self.proj_output(x)
+            forecasting = rearrange(
+                forecasting,
+                "b 1 (pl q) -> b 1 pl q",
+                pl=self.patch_len,
+                q=self.n_quantiles,
+            )
+            forecasting = monotone_quantiles(
+                forecasting, self.quantile_widths, self.median_idx, width_factor=2.0
+            )
+            forecasting = rearrange(
+                forecasting,
+                "b 1 pl q -> b 1 (pl q)",
+            )
             forecasting = self.revin(forecasting, mode="denorm")
 
             # Reshape to (bs*n_quantiles, patch_num, patch_len, n_quantiles)

{patchfm-3.2.0 → patchfm-3.3.0}/src/patchfm/inference/modules.py

@@ -87,9 +87,8 @@ def nanstd(o, dim, keepdim=False):
     return std
 
 class CausalRevIN(nn.Module):
-    def __init__(self, eps=1e-5):
+    def __init__(self):
         super().__init__()
-        self.eps = eps
         self.cached_mean = None
         self.cached_std = None
 
@@ -159,7 +158,7 @@ class CausalRevIN(nn.Module):
         self.cached_cumsum_x2 = cumsum_x2[:, -1:, :]
 
         var = (cumsum_x2 - 2 * mean * cumsum_x + counts * mean**2) / counts
-        std = torch.sqrt(var + 1e-5)
+        std = torch.sqrt(var.clamp_min(0)).clamp_min(1e-3)
 
         return mean, std
     
@@ -205,9 +204,10 @@ class MultiHeadAttention(nn.Module):
         super().__init__()
         assert d_model%n_heads==0, f"d_model ({d_model}) must be divisible by n_heads ({n_heads})"
 
-        self.WQ = nn.Linear(d_model, d_model)
-        self.WK = nn.Linear(d_model, d_model)
-        self.WV = nn.Linear(d_model, d_model)
+        self.WQ = nn.Linear(d_model, d_model, bias=False)
+        self.WK = nn.Linear(d_model, d_model, bias=False)
+        self.WV = nn.Linear(d_model, d_model, bias=False)
+        self.gate = nn.Linear(d_model, d_model, bias=False)
 
         self.out_proj = nn.Linear(d_model, d_model)
 
@@ -238,6 +238,8 @@ class MultiHeadAttention(nn.Module):
             is_causal = False
             offset += (context - 1)
 
+        gated = self.gate(q)
+
         q = self.WQ(q).reshape(bs, -1, self.n_heads, self.head_dim).transpose(1, 2)
         k = self.WK(k).reshape(bs, -1, self.n_heads, self.head_dim).transpose(1, 2)
         v = self.WV(v).reshape(bs, -1, self.n_heads, self.head_dim).transpose(1, 2)
@@ -261,9 +263,10 @@ class MultiHeadAttention(nn.Module):
         k = self.k_norm(k)
         q = self.per_dim_scale(q)
 
-        values = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal)
+        values = nn.functional.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=1.0)
 
         values = values.transpose(1, 2).reshape(bs, -1, dim)
+        values = gated.sigmoid() * values
         values = self.out_proj(values)
         return values
     
@@ -293,9 +296,9 @@ class FeedForward(nn.Module):
 class TransformerEncoderLayer(nn.Module):
     def __init__(self, d_model, n_heads, last=False):
         super().__init__()
-        self.ln1 = nn.LayerNorm(d_model)
+        self.ln1 = nn.RMSNorm(d_model)
         self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, last=last)
-        self.ln2 = nn.LayerNorm(d_model)
+        self.ln2 = nn.RMSNorm(d_model)
         self.ff = FeedForward(d_model=d_model)
     
     def forward(self, x):
@@ -314,12 +317,34 @@ class TransformerEncoder(nn.Module):
             ]
         )
         self.layers.append(TransformerEncoderLayer(d_model=d_model, n_heads=n_heads, last=True))
-        self.norm = nn.LayerNorm(d_model)
+        self.norm = nn.RMSNorm(d_model)
 
     def forward(self, x):
         for layer in self.layers:
             x = layer(x)
         return self.norm(x)
+
+def quantile_widths(quantiles):
+    m = quantiles.index(0.5)
+    widths = [0.0] * len(quantiles)
+    for i in range(m + 1, len(quantiles)):
+        widths[i] = quantiles[i] - quantiles[i - 1]
+    for i in range(m):
+        widths[i] = quantiles[i + 1] - quantiles[i]
+    return torch.tensor(widths), m
+
+
+def monotone_quantiles(raw, widths, m, width_factor=2.0):
+    inc = width_factor * widths.to(raw.dtype) * torch.nn.functional.softplus(raw)
+    median = raw[..., m : m + 1]
+    up = torch.cumsum(inc[..., m + 1 :], dim=-1)
+    down = torch.flip(
+        torch.cumsum(torch.flip(inc[..., :m], dims=[-1]), dim=-1), dims=[-1]
+    )
+    big = torch.finfo(raw.dtype).max
+    return torch.cat([median - down, median, median + up], dim=-1).clamp(
+        min=-big, max=big
+    )
     
 class PatchFM(nn.Module, PyTorchModelHubMixin): 
     def __init__(self, config):
@@ -337,7 +362,7 @@ class PatchFM(nn.Module, PyTorchModelHubMixin):
         self.revin = CausalRevIN()
         self.proj_embedding = ResidualBlock(
             in_dim=self.patch_len, 
-            hid_dim=4 * self.d_model, 
+            hid_dim=self.d_model, 
             out_dim=self.d_model
         )
         self.transformer_encoder = TransformerEncoder(

{patchfm-3.2.0 → patchfm-3.3.0/src/patchfm.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: patchfm
-Version: 3.2.0
+Version: 3.3.0
 Summary: a Foundation Model for Univariate Time Series Forecasting
 Author-email: Samy-Melwan Vilhes <samy-melwan.vilhes@insa-rouen.fr>
 License: MIT License