Repository: huggingface/transformers
Pull request: Fix TimesFM patch normalization instability

This PR fixes patch normalization issue in TimesFM model:
- Numerical instability for low-variance data
- Incorrect clamping of calculated sigma to one instead of configured eps

Issue only affects data with low relative variance and results in incorrect prediction.
Variance is now calculated after mean removal to avoid catastrophic cancellation.

@kashif Please, take a look as you contributed this model.

## Bug demo

<img width="889" height="590" alt="image" src="https://github.com/user-attachments/assets/110bc961-333f-44fc-ae95-366172720271" />

<details>
<summary> Code to reproduce </summary>

```py
import random

import numpy as np
import matplotlib.pyplot as plt
import torch

from transformers import TimesFmModelForPrediction


seed = 42
np.random.seed(seed)
torch.manual_seed(seed)
random.seed(seed)


model = TimesFmModelForPrediction.from_pretrained(
    "google/timesfm-2.0-500m-pytorch",
    attn_implementation="sdpa",
    device_map="auto",
)

model.eval()


X_inp = np.arange(256)
X_pred = np.arange(256, 256 + 128)

fig, axs = plt.subplots(2, 2, figsize=(12, 12))

for ax, sigma in zip(axs.flatten(), [1e-0, 1e-1, 1e-2, 1e-3]):
    rng = np.random.default_rng(seed=seed)
    Y_inp = 1e3 + rng.normal(loc=0, scale=sigma, size=(256,)).cumsum()

    inputs = [
        torch.from_numpy(Y_inp).to(model.device, dtype=model.dtype)
    ]

    with torch.no_grad():
        pred = model(past_values=inputs)
        pred = pred.mean_predictions.cpu().numpy().squeeze()

    ax.plot(X_inp, Y_inp, label='Inputs')
    ax.plot(X_pred, pred, label='Mean Prediction')
    ax.set_title(f'Sigma={sigma}')
    ax.legend()

```
</details>


Changed files:
filename: src/transformers/models/timesfm/modeling_timesfm.py
patch: @@ -341,11 +341,7 @@ def _forward_transform(
     ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         """Input is of shape [B, N, P]."""
         mu, sigma = self._timesfm_masked_mean_std(inputs, patched_pads)
-        sigma = torch.where(
-            sigma < self.config.tolerance,
-            torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),
-            sigma,
-        )
+        sigma = torch.clamp(sigma, min=self.config.tolerance)
 
         # Normalize each patch
         outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
@@ -524,24 +520,16 @@ def _get_patch_index(arr: torch.Tensor):
 
         # Calculate the number of valid elements
         num_valid_elements = torch.sum(mask, dim=1)
-        num_valid_elements = torch.where(
-            num_valid_elements == 0,
-            torch.tensor(1, dtype=num_valid_elements.dtype, device=num_valid_elements.device),
-            num_valid_elements,
-        )
+        num_valid_elements = torch.clamp(num_valid_elements, min=1.0)
 
-        # Calculate the masked sum and squared sum
+        # Calculate the masked sum and mean
         masked_sum = torch.sum(arr * mask, dim=1)
-        masked_squared_sum = torch.sum((arr * mask) ** 2, dim=1)
-
-        # Calculate the masked mean and standard deviation
-        masked_mean = masked_sum / num_valid_elements
-        masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
-        masked_var = torch.where(
-            masked_var < 0.0,
-            torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),
-            masked_var,
-        )
+        masked_mean = masked_sum / num_valid_elements  # [b]
+
+        # Calculate the masked variance using centered values
+        masked_centered_arr = (arr - masked_mean.unsqueeze(-1)) * mask
+        masked_var = torch.sum(masked_centered_arr**2, dim=1) / num_valid_elements
+        masked_var = torch.clamp(masked_var, min=0.0)
         masked_std = torch.sqrt(masked_var)
 
         return masked_mean, masked_std

filename: src/transformers/models/timesfm/modular_timesfm.py
patch: @@ -297,11 +297,7 @@ def _forward_transform(
     ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
         """Input is of shape [B, N, P]."""
         mu, sigma = self._timesfm_masked_mean_std(inputs, patched_pads)
-        sigma = torch.where(
-            sigma < self.config.tolerance,
-            torch.tensor(1.0, dtype=sigma.dtype, device=sigma.device),
-            sigma,
-        )
+        sigma = torch.clamp(sigma, min=self.config.tolerance)
 
         # Normalize each patch
         outputs = (inputs - mu[:, None, None]) / sigma[:, None, None]
@@ -480,24 +476,16 @@ def _get_patch_index(arr: torch.Tensor):
 
         # Calculate the number of valid elements
         num_valid_elements = torch.sum(mask, dim=1)
-        num_valid_elements = torch.where(
-            num_valid_elements == 0,
-            torch.tensor(1, dtype=num_valid_elements.dtype, device=num_valid_elements.device),
-            num_valid_elements,
-        )
+        num_valid_elements = torch.clamp(num_valid_elements, min=1.0)
 
-        # Calculate the masked sum and squared sum
+        # Calculate the masked sum and mean
         masked_sum = torch.sum(arr * mask, dim=1)
-        masked_squared_sum = torch.sum((arr * mask) ** 2, dim=1)
-
-        # Calculate the masked mean and standard deviation
-        masked_mean = masked_sum / num_valid_elements
-        masked_var = masked_squared_sum / num_valid_elements - masked_mean**2
-        masked_var = torch.where(
-            masked_var < 0.0,
-            torch.tensor(0.0, dtype=masked_var.dtype, device=masked_var.device),
-            masked_var,
-        )
+        masked_mean = masked_sum / num_valid_elements  # [b]
+
+        # Calculate the masked variance using centered values
+        masked_centered_arr = (arr - masked_mean.unsqueeze(-1)) * mask
+        masked_var = torch.sum(masked_centered_arr**2, dim=1) / num_valid_elements
+        masked_var = torch.clamp(masked_var, min=0.0)
         masked_std = torch.sqrt(masked_var)
 
         return masked_mean, masked_std