[Bug] Asymmetric scale guard in revin causes roundtrip invariant violation and silent forecast annihilation for low-variance / quiescent series

Summary

In both PyTorch (src/timesfm/torch/util.py) and Flax (src/timesfm/flax/util.py), the Reversible Instance Normalization (revin) implementation contains an asymmetry between forward normalization and reverse denormalization when handling near-zero variance inputs ($\sigma < \text{_TOLERANCE}$).

While forward normalization guards against division-by-zero by clamping small $\sigma$ to $1.0$, reverse denormalization uses raw $\sigma$ directly:

• Forward: (x - mu) / safe_sigma where safe_sigma = where(sigma < 1e-6, 1.0, sigma)
• Reverse: x * sigma + mu (raw sigma without guard)

This breaks the roundtrip identity: $$\text{revin}(\text{revin}(x, \mu, \sigma, \text{reverse=False}), \mu, \sigma, \text{reverse=True}) \neq x$$

For any input series with zero or near-zero variance ($\sigma < 10^{-6}$, common in telemetry, quiescent sensors, piecewise-constant metrics, or idle systems), the transformer backbone's predicted forecast deltas are multiplied by $\approx 0.0$ in denormalization, silently annihilating all model predictions and collapsing the forecast into a flat line ($\mu$).

---

Mathematical Invariant Violation

As stated in tests/test_torch_utils.py (lines 190–196):

"This is the fundamental invariant of reversible normalization... Any deviation here directly corrupts the final forecast values."

Kim et al. (ICLR 2022) defines RevIN such that normalization and denormalization must be exact mathematical inverses. When variance is below threshold, forward pass maps $x \mapsto x - \mu$ (treating effective $\sigma = 1.0$). Therefore, the inverse mapping must also treat effective $\sigma = 1.0$ ($x \mapsto x \cdot 1.0 + \mu$).

Currently, multiplying by the raw unguarded $\sigma$ causes:

1. When $\sigma = 0$: The forecast delta $\Delta = \hat{y} - \mu$ is zeroed out entirely ($\Delta \times 0 = 0$).
2. When $0 < \sigma < 10^{-6}$: Forecast deltas are attenuated by up to $1,000,000\times$.

---

Minimal Reproducible Example

Run the following standalone script (reproducing src/timesfm/torch/util.py:revin):

import torch

_TOLERANCE = 1e-6

def revin_current(x, mu, sigma, reverse=False):
    if not reverse:
        safe_sigma = torch.where(sigma < _TOLERANCE, 1.0, sigma)
        return (x - mu) / safe_sigma
    return x * sigma + mu

def revin_symmetric(x, mu, sigma, reverse=False):
    safe_sigma = torch.where(sigma < _TOLERANCE, 1.0, sigma)
    if not reverse:
        return (x - mu) / safe_sigma
    return x * safe_sigma + mu

# Quiescent series (e.g. constant sensor reading of 5.0, sigma = 0.0)
x = torch.tensor([5.0, 5.0, 5.0])
mu = torch.tensor([5.0])
sigma = torch.tensor([0.0])

# 1. Forward normalization
norm = revin_current(x, mu, sigma, reverse=False)

# 2. Model outputs forecast with expected variations: [-0.5, 0.0, +0.5]
forecast_latent = norm + torch.tensor([-0.5, 0.0, 0.5])

# 3. Reverse denormalization
output_current = revin_current(forecast_latent, mu, sigma, reverse=True)
output_symmetric = revin_symmetric(forecast_latent, mu, sigma, reverse=True)

print("Target Forecast Expected :", (mu + torch.tensor([-0.5, 0.0, 0.5])).tolist())
print("Output with Current Code :", output_current.tolist())
print("Output with Symmetric Fix:", output_symmetric.tolist())

Output:

Target Forecast Expected : [4.5, 5.0, 5.5]
Output with Current Code : [5.0, 5.0, 5.0]    <-- Forecast annihilated to flat line!
Output with Symmetric Fix: [4.5, 5.0, 5.5]    <-- Exact reconstruction preserved

---

Proposed Fix

In both src/timesfm/torch/util.py and src/timesfm/flax/util.py, apply the tolerance guard symmetrically to reverse=True:

PyTorch (src/timesfm/torch/util.py):

def revin(x: torch.Tensor, mu: torch.Tensor, sigma: torch.Tensor, reverse: bool = False) -> torch.Tensor:
    safe_sigma = torch.where(sigma < _TOLERANCE, 1.0, sigma)
    if not reverse:
        return (x - mu) / safe_sigma
    return x * safe_sigma + mu

Flax (src/timesfm/flax/util.py):

def revin(x: jnp.ndarray, mu: jnp.ndarray, sigma: jnp.ndarray, reverse: bool = False) -> jnp.ndarray:
    safe_sigma = jnp.where(sigma < _TOLERANCE, 1.0, sigma)
    if not reverse:
        return (x - mu) / safe_sigma
    return x * safe_sigma + mu

I have verified this fix locally on the full test suite (all 18 test modules pass). I am happy to submit a PR with this fix and regression tests if the maintainers agree.