GRIN: GMM-Adapted Reversible Normalization

• The paper introduces GRIN, a GMM-adapted reversible instance normalization framework that enhances deep neural forecasting through exact inversion.
• It leverages global GMM estimates to perform per-instance normalization independent of batch statistics, effectively handling non-stationary and multi-modal data.
• Empirical results in the TimeGMM framework demonstrate significant improvements in CRPS and NMAE over traditional normalization methods.

GMM-Adapted Reversible Instance Normalization (GRIN) is a normalization methodology tailored for deep neural architectures, designed to model and adapt to multi-modal and non-stationary data distributions. The GRIN formulation enables exact reversibility and independence from batch statistics by integrating global Gaussian Mixture Model (GMM) estimates into the normalization and denormalization process. This mechanism, as instantiated in the TimeGMM forecasting framework, provides robust adaptation to temporal-probabilistic shifts and improves probabilistic prediction accuracy in applications suffering from non-stationary feature distributions and distributional mismatch (Liu et al., 18 Jan 2026, Kalayeh et al., 2018).

1. Foundational Principles

GRIN extends prior normalization strategies by leveraging Gaussian Mixture Models in the normalization transform. Traditional Batch Normalization (BN) assumes all samples in a batch are drawn from a unimodal Gaussian. BN’s transform $x \mapsto \hat{x} = (x - \mu)/\sigma$ can be interpreted as a Fisher vector corresponding to the Gaussian likelihood gradient. However, most deep learning activations, especially under non-linearities, exhibit heavy-tailed and asymmetric characteristics incompatible with a single-mode Gaussian (Kalayeh et al., 2018). Mixture Normalization (MN) replaces the Gaussian assumption by estimating mixtures from batch data but remains batch-dependent. GRIN generalizes MN by using a global (typically per-channel) GMM fit from running averages and performs normalization at the instance level, removing batch coupling and guaranteeing reversibility.

2. Mathematical Formulation and Workflow

Forward Transform

Given input $x_t^{(i)}$ for channel $i=1,\ldots,V$, time index $t=1,\ldots,L_h$, GRIN computes window-level statistics

$$\mu_t^{(i)} = \frac{1}{L_h} \sum_{s=1}^{L_h} x_s^{(i)} ,\quad \sigma_t^{(i)\,2} = \frac{1}{L_h} \sum_{s=1}^{L_h}\left(x_s^{(i)} - \mu_t^{(i)}\right)^2$$

with numerical stabilizer $\epsilon > 0$. The normalized value is then

$$\tilde x_t^{(i)} = a^{(i)} \frac{x_t^{(i)} - \mu_t^{(i)}}{\sqrt{\sigma_t^{(i)\,2} + \epsilon}} + b^{(i)}$$

where $a^{(i)}$, $b^{(i)}$ are trainable affine parameters per channel. For the full GMM-adapted formulation, normalization scale and shift are computed by integrating mixture responsibilities $\nu_k(x)$ arising from the global GMM parameters $\lambda_k$, $\mu_k$, $\sigma_k$ via

$$s = \sum_{k=1}^K \frac{\nu_k(x)}{\sqrt{\lambda_k}} \cdot \frac{1}{\sigma_k + \epsilon} ,\quad b = -\sum_{k=1}^K \frac{\nu_k(x)}{\sqrt{\lambda_k}} \cdot \frac{\mu_k}{\sigma_k + \epsilon}$$

Yielding the normalized output

$\tilde x = s\cdot x + b$

Inverse Transform

GRIN is exactly reversible, so denormalization applies the inverse affine

$x = \frac{\tilde x - b}{s}$

If trainable parameters $\gamma$, $\beta$ (as in standard normalization layers) are included, the forward is $y = \gamma \tilde x + \beta$ and the inverse uses saved buffers $A = s\cdot \gamma$, $B = \beta + \gamma b$: $x = \frac{y - B}{A}$

Decoder-Level Usage

In forecasting architectures (e.g., TimeGMM), the decoder emits normalized GMM parameters $(\tilde\mu, \tilde\sigma, w)$. Final predicted means and scales are denormalized using the same window-level statistics, guaranteeing that outputs are mapped back to the original data domain: $$\mu_{u,k}^{(i)} = \sqrt{\sigma_t^{(i)\,2} + \epsilon} \frac{\tilde\mu_{u,k}^{(i)} - b^{(i)}}{a^{(i)}} + \mu_t^{(i)}$$

$$\sigma_{u,k}^{(i)} = \sqrt{\sigma_t^{(i)\,2} + \epsilon} \frac{\tilde\sigma_{u,k}^{(i)}}{a^{(i)}}$$

Mixture weights are recovered via softmax normalization.

3. Mechanisms for Temporal-Probabilistic Adaptation

GRIN addresses distribution shift by recomputing window-level statistics for each incoming historical input, instantly tracking changes in series level and volatility. For each forecast, the normalization and denormalization are parameterized by statistics of the most recent input window, ensuring adaptation to transient trends and seasonality changes. This approach obviates the need for the network to relearn scale or location effects and aligns mixture prediction with the true data scale at inference. The same set of statistics used for normalization is re-employed at the output, meaning no additional parameters or tracking are required between normalization and denormalization steps (Liu et al., 18 Jan 2026).

A plausible implication is that by conditioning the entire prediction cascade on up-to-date window statistics, GRIN enables the model to maintain calibration and sharpness under realistic non-stationarity, without explicit retraining or model recalibration.

4. Implementation Details and Pseudocode

The GRIN layer is efficiently implemented with minimal overhead. A PyTorch-style algorithm comprises:

eps = 1e-5 mu_t = X.mean(dim=2, keepdim=True) # [B,V,1] var_t = X.var(dim=2, unbiased=False, keepdim=True) # [B,V,1] std_t = torch.sqrt(var_t + eps) # [B,V,1] a, b = model.grin_a, model.grin_b # each of shape [V,1] X_norm = a * ((X - mu_t) / std_t) + b

For the inverse:

mu = std_t * ((tilde_mu - b) / a) + mu_t # broadcasts over L_f, K sig = std_t * (tilde_sigma / a) w = torch.softmax(w_logits, dim=-1)

In the general GRIN algorithm (using global GMM parameters):

• For each instance and location, compute GMM responsibilities $\nu_k(x)$.
• Accumulate scale and shift using mixture-adapted terms.
• Apply per-channel affine transformation and store scale/shift buffers for exact inversion during back-propagation. GRIN backward pass multiplies the upstream gradient by the stored scale buffer, matching BN’s computational cost. Pre-fitting a GMM (typically $K=3-5$ components) per channel over training statistics is required, but the transform itself is instance-based and reversible (Kalayeh et al., 2018).

5. Empirical Performance and Impact

Ablation studies within the TimeGMM framework demonstrate substantial quantitative gains attributable to GRIN’s normalization. In representative datasets (ETTm1, ETTm2, ETTh2, Weather), removing GRIN degrades the Continuous Ranked Probability Score (CRPS) by 15–20% relative, while the removal of the mixture itself causes smaller performance drops. Over six benchmark datasets and four forecast horizons, the full TimeGMM system (with GRIN) outperforms the next-best probabilistic forecaster by up to 22.5% in CRPS and 21.2% in NMAE (Liu et al., 18 Jan 2026). This suggests GRIN’s normalization and denormalization to be a critical enabling factor for robust probabilistic forecasting.

Dataset	Full TimeGMM (CRPS)	–GMM (single Gaussian)	–GRIN (no normalization)
ETTm1	0.2378	0.2407	0.2743
ETTm2	0.1409	0.1440	0.1588
ETTh2	0.1677	0.1681	0.2036
Weather	0.0585	0.0613	0.0611

Removing GRIN leads to the most significant degradation, confirming its centrality under regime shift.

6. Comparative Analysis and Theoretical Context

GRIN generalizes several normalization frameworks:

• BatchNorm employs unimodal, batch statistics per feature.
• Mixture Normalization (MN) fits a GMM on each batch for improved representation of batch-wise multi-modality.
• GRIN utilizes a global, precomputed GMM and is fully per-instance and reversible.

Unlike BN and MN, no batch-dependent reestimation or coupling is needed; GRIN can operate on batch-size 1 and supports exact inversion for downstream tasks requiring reversibility (e.g., probabilistic modeling, forecasting, generative networks). The underlying Fisher kernel analysis shows that normalizing with respect to mixture responsibilities aligns network inputs with latent subpopulation geometries, enhancing representation learning irrespective of batch composition (Kalayeh et al., 2018).

A plausible implication is that GRIN’s independence from batch statistics makes it particularly suitable for non-i.i.d. settings, fine-tuning, and low-data regimes.

7. Applications, Limitations, and Extensions

GRIN is deployed within the TimeGMM architecture for single-pass probabilistic time series forecasting, particularly in domains with pronounced distribution shifts such as energy and finance. Its guarantees of invertibility and multi-modal adaptation extend to other domains where feature non-stationarity and batch independence are paramount—such as image and speech models affected by covariate drift. Since GMM parameters are fit globally and not adapted online, GRIN does not back-propagate through mixture parameters during training (these are constant). This restricts dynamic adaptation of the mixture itself but simplifies computational practice and accelerates training. Potential future directions include adaptive online GMM fitting and hierarchical mixture normalization for higher-dimensional latent features.

In sum, GMM-Adapted Reversible Instance Normalization provides a principled, invertible, and multi-modal normalization framework for deep networks, enabling robust forward and inverse transformation under non-stationary and multi-modal data regimes (Liu et al., 18 Jan 2026, Kalayeh et al., 2018).