arXiv:2604.15838v1 [cs.LG] 17 Apr 2026
Reversible Residual Normalization Alleviates Spatio-Temporal Distribution Shift Zhaobo HU Vincent Gauthier SAMOVAR, Télécom SudParis SAMOVAR, Télécom SudParis Institut Polytechnique de Paris Institut Polytechnique de Paris [email protected] [email protected] Mehdi Naima CNRS – LIP6 Sorbonne Université [email protected]
Abstract Distribution shift severely degrades the performance of deep forecasting models. While this issue is well-studied for individual time series, it remains a significant challenge in the spatio-temporal domain. Effective solutions like instance normalization and its variants can mitigate temporal shifts by standardizing statistics. However, distribution shift on a graph is far more complex, involving not only the drift of individual node series but also heterogeneity across the spatial network where different nodes exhibit distinct statistical properties. To tackle this problem, we propose Reversible Residual Normalization (RRN), a novel framework that performs spatially-aware invertible transformations to address distribution shift in both spatial and temporal dimensions. Our approach integrates graph convolutional operations within invertible residual blocks, enabling adaptive normalization that respects the underlying graph structure while maintaining reversibility. By combining Center Normalization with spectral-constrained graph neural networks, our method captures and normalizes complex Spatio-Temporal relationships in a data-driven manner. The bidirectional nature of our framework allows models to learn in a normalized latent space and recover original distributional properties through inverse transformation, offering a robust and model-agnostic solution for forecasting on dynamic spatio-temporal systems. Keywords: Spatio-Temporal Forecasting, Distribution Shift, Invertible Neural Networks, Graph Neural Networks.
1
Introduction
Spatio-temporal forecasting is fundamental to managing dynamic systems, ranging from urban traffic [1–3] and environmental monitoring [4–6] to public health surveillance. Effective models must capture both the spatial dependencies among interconnected locations and the temporal dynamics governing system evolution. Despite deep learning advancements, a critical challenge persists: distribution shift. This phenomenon occurs when statistical properties 1
change, violating the assumption that training and testing data follow identical distributions. In spatio-temporal systems, this shift is dual-natured. Temporally, observations drift due to evolving dynamics; spatially, locations exhibit statistical heterogeneity despite sharing underlying correlations. This Spatio-Temporal distribution shift [7] severely degrades generalization when networks trained on historical data confront evolved patterns during inference. Current mitigation strategies primarily focus on time series, using methods like instance normalization [8] to standardize sequences. However, these techniques do not naturally extend to graph-structured data. Applying normalization independently to each node neglects spatial heterogeneity and may disrupt topological relationships. Furthermore, most spatio-temporal architectures focus on correlation modeling while implicitly assuming stationarity, creating a gap between model assumptions and the non-stationary nature of real-world data. To bridge this gap, we propose Reversible Residual Normalization (RRN), a novel framework that extends normalization to the spatio-temporal domain via invertible transformations. RRN integrates three key components: Center Normalization [9], a Lipschitzcontinuous alternative [10–12] that removes shifts while maintaining invertibility; graph convolutional operations within invertible residual blocks [12] for spatially-aware normalization; and spectral normalization [13] to constrain the transformation. This bidirectional framework maps data into a stationary latent space for easier learning and uses the inverse transformation to restore location-specific characteristics. Extensive experiments on real-world benchmarks demonstrate that our model-agnostic approach consistently improves performance, validating the necessity of addressing distribution shift in spatio-temporal systems.
2
Related Works
2.1
Distribution Shift in Time Series and Spatio-Temporal Data
Distribution shift arises from non-stationarity, causing significant discrepancies between training and testing distributions [14–17]. To mitigate this, instance normalization and its variants, such as reversible and adaptive normalization effectively standardize temporal sequences by removing non-stationary components, yet they typically process nodes independently. Other advanced approaches explicitly model distribution shifts by learning intra-space and interspace variations, but these remain confined to the temporal dimension. Consequently, the spatial dimension of distribution shift is largely neglected, applying temporal normalization strategies directly to spatio-temporal data ignores the underlying graph structure, failing to account for spatial heterogeneity where distinct nodes exhibit unique statistical properties.
2.2
Invertible Networks
Invertible neural networks learn bijective transformations that allow for exact computation of both forward and inverse mappings, facilitating tasks such as generative modeling and density estimation [11, 18–21]. While normalizing flows typically rely on coupling layers to construct invertible mappings [18, 21–23], Invertible Residual Networks [10–12] extend this capability to standard residual architectures. Rather than relying on partitioning dimensions, these models utilize residual blocks [24] that guarantee invertibility provided the residual function satisfies specific Lipschitz continuity conditions (requiring a constant strictly less than one). This constraint is often enforced via spectral normalization to enable exact inverse 2
computation through fixed-point iteration. Despite their success in time series, the potential of invertible architectures to simultaneously address spatial and temporal distribution shift while preserving graph topology remains underexplored.
3
Problem Formulations
3.1
Spatio-Temporal Forecasting
Let a spatio-temporal system be represented by a graph G = (V, E), where V is a set of N spatial locations (e.g., sensors, regions) with |V | = N , and E is a set of edges representing the spatial relationships between these locations. At each discrete time step t, the system exhibits a feature matrix X (t) ∈ RN ×D , where D is the number of features recorded at each location. This feature matrix is also referred to as a graph signal. The task of spatio-temporal forecasting aims to predict future system states based on historical observations. Given a lookback window of length L, the historical data can be denoted as a tensor Xt−L+1:t = (X (t−L+1) , . . . , X (t) ) ∈ RL×N ×D . The objective is to forecast the subsequent data over a horizon window of length H, denoted as Yt+1:t+H = (X (t+1) , . . . , X (t+H) ) ∈ RH×N ×D . Formally, the goal is to learn a function fθ parameterized by θ that maps the historical observations and the graph structure to the future sequence: Ŷt+1:t+H = fθ (Xt−L+1:t ; G) where Ŷt+1:t+H represents the predicted future sequence.
3.2
Distribution Shift in Spatio-Temporal Series
A primary challenge in real-world spatio-temporal forecasting is the non-stationarity of the data, which manifests as a distribution shift. Many deep learning models implicitly assume that the data-generating process is consistent between the training and test periods, an assumption that is frequently violated in practice. This discrepancy occurs as the joint distribution of the observed spatio-temporal data changes over time, hindering a model’s ability to generalize from past observations to future predictions. Formally, for any two distinct time steps tu and tv , the underlying conditional probability that governs the system’s evolution is not constant. This distribution shift can be expressed as: P (Ytu +1:tu +H |Xtu −L+1:tu ; G) ̸= P (Ytv +1:tv +H |Xtv −L+1:tv ; G) Therefore, our objective is not only to predict future spatio-temporal states but also to explicitly model and mitigate these distribution shifts to enhance forecasting accuracy and generalization. To address this, we propose extending the principles of normalizing flows, which have shown promise in handling temporal distribution shifts, to the more complex spatio-temporal domain.
3
Input Space
Transformed Space
Time
Forecasting Space
Time
Detransformed Space
Time
Time
Figure 1: Overview of the Reversible Residual Normalization framework for spatio-temporal forecasting. The framework operates through four stages: (1) Input Space contains raw spatio-temporal data with distribution shifts across nodes and time; (2) Transformed Space where the invertible transformation Tϕ normalizes the data by removing Spatio-Temporal distribution shifts while preserving correlational structures; (3) Forecasting Space where any forecasting model fθ makes predictions in the stationary latent space; (4) Detransformed Space where the inverse transformation Tϕ−1 restores the original distributional properties to produce final predictions.
4
Methodology
4.1
Invertible Residual Structure
Residual structures [24] have become a fundamental building block in modern deep learning architectures. A residual block takes the form: H(x) = x + G(x)
(1)
Without loss of generality, we define x ∈ Rd is the input, G : Rd → Rd represents a learnable transformation, and the output H(x) combines the identity mapping with the learned residual G(x). This formulation facilitates gradient flow during training and enables the construction of very deep networks. For our purposes, we are particularly interested in when such residual structures are invertible, meaning there exists a unique inverse mapping H −1 such that H −1 (H(x)) = x for all x in the domain. Invertibility is desirable in our context because it allows us to transform data distributions bidirectionally, which is essential for modeling distribution shifts in spatiotemporal forecasting. A sufficient condition for the invertibility of a residual block is provided by the following lemma, which constrains the behavior of the residual function G: Lemma 1 (Sufficient Condition for Invertible Residual Blocks [11, 12]). Let H(x) = x + G(x) be a residual block, where G : Rd → Rd . If the Lipschitz constant of G satisfies: ∥G(x1 ) − G(x2 )∥ <1 ∥x1 − x2 ∥ x1 ̸=x2
L(G) = sup then H is invertible.
4
(2)
The intuition behind this lemma is that when G is a contraction mapping (i.e., L(G) < 1), the residual block H becomes a perturbation of the identity function that preserves bijectivity. Under this condition, the inverse x = H −1 (z) for a given output z can be computed through fixed-point iteration as shown in Algorithm 1. Algorithm 1: Inverse of Residual Block via Fixed-Point Iteration Input: Output x(ℓ) from residual block, residual function G, number of iterations N Output: input of residual block x(ℓ−1) 1 x ← x(ℓ) ; 2 for m = 1, . . . , N do 3 x ← x − G(x); 4 return x
In practice, enforcing the Lipschitz constraint L(G) < 1 requires careful design of the residual function G. When G is composed of linear transformations (such as convolutions) followed by nonlinear activations, a common approach is to constrain the spectral norm of each linear layer. Specifically, if G consists of multiple layers with weight matrices W1 , W2 , . . . , Wi , we can enforce: ( cWi /σ̃i , if c/σ̃i < 1 W̃i = (3) Wi , else where σ̃i = ∥Wi ∥2 denotes the spectral norm (largest singular value) of Wi , and c < 1 are chosen such that the overall Lipschitz constant of G remains below 1. The spectral norm can be efficiently approximated using power iteration methods [13, 25] during training. This framework of invertible residual blocks provides the foundation for our approach to handling distribution shifts in spatio-temporal forecasting, as we will elaborate in the following sections.
4.2
Center Normalization instead of Instance Normalization
Instance Normalization (IN) [8, 14, 15, 17] has demonstrated significant effectiveness in mitigating distribution shift across various domains. By normalizing each sample independently, IN reduces dependence on specific distributional statistics, thereby improving generalization. However, Instance Normalization is not Lipschitz continuous, which contradicts the invertibility requirement in Lemma 1. For input Xt−T +1:t ∈ RT ×N ×D representing temporal features at a spatial location, we omit the time subscripts for brevity and denote it as X . Instance Normalization is defined as: IN(X ) = γ ⊙ z + β, 1 ⊤ y = I − 11 X. T
y where z = , Std(y)
The Jacobian matrix of IN reveals its discontinuity: ∂z 1 1 ⊤ yy ⊤ Jz (X ) = = I − 11 I− . ∂X Std(y) T ∥y∥22 5
(4)
(5)
When Std(y) approaches 0, the Jacobian entries approach ∞, violating Lipschitz continuity. This causes training instability and precludes invertibility. To address this, we introduce Center Normalization[9]: 1 ⊤ CN(X ) = γ ⊙ α I − 11 X + β, (6) T where α is a controllable scaling parameter, γ and β ∈ RN ×D are learnable parameters. CN removes the mean without variance normalization, ensuring Lipschitz continuity. The ) Jacobian of the Center Normalization is ∂CN(X = α I − T1 11⊤ . When γ = 1 and β = 0, ∂X Center Normalization has Lipschitz constant Lip(CNX ) = α, it’s trivial to verify ∥CN(X1 ) − CN(X2 )∥ ≤ α∥X1 −X2 ∥. By setting α ≤ 1, CN satisfies the invertibility condition in Lemma 1. The centering operation preserves distribution shift mitigation by removing mean shifts, a common form of temporal distribution shift.
4.3
Overall Structure
After applying Center Normalization, we incorporate a graph convolution module within the invertible residual block to capture spatial dependencies. Following the framework in Section 3.1, we introduce a Lipschitz-constrained Graph Convolutional Network (GCN) layer [12, 26] to ensure the overall residual block satisfies the invertibility condition. 4.3.1
Spectral Normalization for GCN
A standard GCN layer with residual connection can be formulated as: H(X (t) ) = X (t) + σ(ÂX (t) W ) = X (t) + g(X (t) ),
(7)
where X (t) ∈ RN ×D is the input node representation, we omit the time superscripts for brevity 1 1 and denote it as X,  = D̃− 2 ÃD̃− 2 is the normalized adjacency matrix with à = A + I (adjacency matrix with self-loops) and D̃ is the diagonal degree matrix, W ∈ Rd×d is the learnable weight matrix, and σ(·) is a Lipschitz continuous activation function (e.g., ReLU, tanh). According to Lemma 1, the residual block in Eq. (7) is invertible if the Lipschitz constant of the residual function g(X) = σ(ÂXW ) satisfies: ∥g(X1 ) − g(X2 )∥2 < 1. ∥X1 − X2 ∥2 X1 ̸=X2
Lip(g) = sup
(8)
For contractive activation functions like ReLU and tanh where Lip(σ) < 1, the condition is satisfied if: ∥ÂXW ∥2 sup < 1. (9) ∥X∥2 X̸=0 This supremum is upper bounded by: ∥ÂXW ∥2 ≤ ∥Â∥2 ∥W ∥F , ∥X∥2 X̸=0 sup
6
(10)
where ∥ · ∥2 denotes the spectral norm (largest singular value), and ∥W ∥F is the Frobenius norm of W , while strict invertibility requires bounding the spectral norm ∥W ∥2 , we strictly constrain the Frobenius norm ∥W ∥F in practice as an upper bound (∥W ∥2 ≤ ∥W ∥F ) for computational efficiency, same operation as [12], bypass the inconvenience of power iteration. For the normalized adjacency matrix  in GCN, we have ∥Â∥2 = 1. Therefore, to ensure Lip(g) < c for some c < 1, we normalize the weight matrix W according to Eq. (3) after each gradient descent step. This spectral normalization ensures the invertibility of the residual block throughout training. 4.3.2
Reversible Residual Normalization
We now combine Center Normalization and the Lipschitzconstrained GCN into an invertible residual block. The complete block is defined as: (ℓ)
(ℓ)
(ℓ)
H(Xt−T +1:t ) = Xt−T +1:t + σ(Â · CN(Xt−T +1:t ) · W ), (11)
Center Norm
where CN(·) is Center Normalization from Eq. (6). The Lipschitz constant of this block by chain rules of Lipschitz satisfies: Lip(g) ≤ Lip(σ) · ∥Â∥2 · Lip(CN) · ∥W ∥F .
4.4
(12)
GNN
Bidirectional Transformation Framework
Our approach transforms spatio-temporal data into a stationary latent space, performs forecasting, and transforms predictions back to the original space. Formally, let Xt−L+1:t denote the historical observations. We define: Forward Transformation (to stationary space): Z = Tϕ (Xt−L+1:t ; G),
GNN
(13)
where H (ℓ) denotes the ℓ-th invertible spatio-temporal Figure 2: Invertible Residual Block block from Eq. (11). architecture.
Prediction in stationary space: Ẑt+1:t+H = fθ (Z; G),
(14)
where fθ is an arbitrary spatio-temporal forecasting model operating in the stationary latent space. Inverse Transformation (back to original space): Ŷt+1:t+H = Tϕ−1 (Ẑt+1:t+H ; G),
(15)
where each (H (ℓ) )−1 is computed via fixed-point iteration due to guaranteed invertibility. The key insight is that by transforming to a stationary latent space Z, we mitigate distribution shift: the conditional distribution P (Ẑt+1:t+H |Z; G) becomes approximately invariant across time, enabling more robust forecasting with any backbone model fθ . 7
5
Experiments
Table 1: Forecasting performance comparison across different datasets and models. All results are the average of 5 random seeds. Bold values indicate better performance in each baseline vs. RRN pair (lower is better).
PemsBay
MetrLA
SDWPF
Dataset
5.1
GWavenet
GWavenet + RRN
DCRNN + RRN
Auto-DSTSGN
Auto-DSTSGN+RRN
MAE
RMSE
MAE
RMSE
MAE
DCRNN RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
GMAN RMSE
GMAN + RRN MAE
RMSE
MAE
STPGNN RMSE
STPGNN + RRN MAE
RMSE
58.72
110.61
54.20
105.07
64.06
119.12
60.85
114.68
56.62
105.91
55.26
104.02
58.26
104.58
56.18
104.63
57.03
108.72
55.18
103.88
Horizon
3 6
83.34
151.75
78.81
144.83
90.13
161.25
85.91
158.15
80.52
145.92
78.91
144.81
82.12
145.68
79.69
145.09
82.20
149.41
79.91
144.55
12
114.95
204.17
112.59
197.38
126.31
216.19
123.28
210.47
113.71
198.42
113.32
195.55
119.63
199.07
114.33
197.30
118.37
204.57
114.85
195.78
3
2.75
5.22
2.73
5.19
2.91
5.58
2.84
5.40
2.77
5.35
2.78
5.38
2.81
5.47
2.76
5.32
2.92
5.67
2.92
5.66
6
3.16
6.27
3.16
6.24
3.40
6.74
3.30
6.51
3.13
6.29
3.15
6.39
3.19
6.45
3.15
6.40
3.36
6.75
3.36
6.77
12
3.68
7.44
3.66
7.39
4.13
8.25
3.93
7.90
3.58
7.35
3.57
7.37
3.65
7.51
3.63
7.61
3.87
7.87
3.88
7.91
3
1.33
2.78
1.28
2.77
1.42
2.99
1.39
2.91
1.36
2.83
1.32
2.80
1.37
2.93
1.34
2.82
1.39
2.93
1.34
2.84
6
1.67
3.73
1.63
3.71
1.87
4.19
1.81
4.03
1.72
3.76
1.67
3.71
1.71
3.82
1.67
3.67
1.77
3.93
1.68
3.79
12
2.01
4.59
2.01
4.52
2.43
5.54
2.33
5.21
2.04
4.48
1.99
4.46
2.01
4.55
1.97
4.48
2.15
4.81
2.01
4.59
Experimental setup
Datasets We evaluate our framework on three diverse spatio-temporal benchmarks. METRLA and PEMS-BAY are classic traffic speed datasets containing 207 and 325 sensors from Los Angeles and the San Francisco Bay Area, respectively, with 5-minute sampling frequency over 4–6 months. SDWPF [27] represents an energy domain, providing 24 months of wind power data from 134 turbines, featuring 19 dynamic variables including SCADA system parameters and ERA5 meteorological data. Regarding graph construction, METR-LA and PEMS-BAY utilize distance-based connectivity via Gaussian kernel thresholding, while SDWPF leverages spatial proximity between turbines. All datasets are chronologically split into 70% training, 10% validation, and 20% testing sets to evaluate forecasting performance under temporal distribution shifts. Baselines We evaluate our model against several state-of-the-art methods addressing distribution shift. RevIN [14] normalizes each instance by its mean and variance, symmetrically denormalizing outputs to restore original scales. Dish-TS [16] models intra-space and inter-space shifts using a Dual-CONET framework to learn separate distribution coefficients. SAN [15] adopts a fine-grained approach, normalizing sub-series "slices" while predicting future statistics for adaptive denormalization. ST-Norm [28] utilizes temporal and spatial modules to separately refine high-frequency and local components by removing low-frequency and global trends from raw data. Backbone models To demonstrate the versatility of our model-agnostic framework, we evaluate it on diverse mainstream architectures. DCRNN [29] integrates diffusion convolution with GRUs to model traffic as a directed diffusion process. Graph WaveNet [30] combines dilated TCNs with an adaptive adjacency matrix to capture hidden spatial dependencies. GMAN [31] utilizes an encoder-decoder structure with spatio-temporal attention
8
Table 2: Performance comparison of DCRNN with various enhancement methods across different datasets. All results represent the average of five random seeds (lower is better). DCRNN
PemsBay
MetrLA
SDWPF
Dataset
DCRNN + RevIN
DCRNN + Dish-TS
DCRNN + SAN
DCRNN + STNORM
DCRNN + RRN
Horizon MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
MAE
RMSE
3
64.06
119.12
62.32
120.76
61.50
117.53
70.77
125.27
58.89
113.52
60.85
114.68
6
90.13
161.25
88.91
164.25
88.11
158.89
96.88
164.92
85.06
155.72
85.91
158.15
12
126.31
216.19
127.88
221.45
126.12
212.19
131.35
214.58
123.04
208.91
123.28
210.47
3
2.91
5.58
2.97
5.73
2.93
5.61
3.23
6.31
2.84
5.43
2.84
5.40
6
3.40
6.74
3.57
7.11
3.43
6.79
3.90
7.75
3.29
6.56
3.30
6.51
12
4.13
8.25
4.53
8.97
4.19
8.32
4.99
9.74
3.91
8.11
3.93
7.90
3
1.42
2.99
1.43
3.06
1.42
2.98
1.61
3.61
1.38
2.89
1.39
2.91
6
1.87
4.19
1.91
4.36
1.82
4.11
2.11
4.82
1.77
3.93
1.81
4.03
12
2.43
5.54
2.59
5.93
2.34
5.29
2.83
6.39
2.19
4.98
2.33
5.21
blocks to model dynamic correlations. Auto-DSTSGN [32] employs an automated dilated graph framework with neural architecture search to construct adaptive adjacency matrices. STPGNN [33] introduces a pivotal graph convolution module to prioritize sensors with complex dependencies within a parallel feature extraction framework. Implementation Details We follow the original configurations for all backbone models. For RRN, we stack two invertible residual blocks with a hidden size of 32. To guarantee invertibility (Lemma 1), we set the Center Normalization scaling α = 0.9 and constrain the GCN weights to a spectral norm of 0.9. Models are trained using the Adam optimizer with a learning rate of 5 × 10−3 and a batch size of 64 for up to 100 epochs with early stopping. Experiments are conducted on NVIDIA L40S GPUs, and we report the average performance over five random seeds.
5.2
Overall Performance
The experimental results, summarized in Tables 1 and 2, highlight the robust capability of our Reversible Residual Normalization (RRN) framework in mitigating spatio-temporal distribution shifts. Performance Gains Across Backbones. As shown in Table 1, RRN consistently enhances forecasting accuracy across all five backbone models and three datasets. The improvements are most significant on the SDWPF dataset, which features high volatility due to meteorological changes; for instance, GWavenet’s MAE decreases by approximately 7.7% and 5.4% for the 3-step and 6-step horizons, respectively. These universal gains across diffusionbased, attention-based, and automated-search architectures demonstrate RRN’s effectiveness as a model-agnostic plugin for stabilizing latent representations.
9
Baseline
RRN (2 Blocks)
DCRNN: Hardware Efficiency
RRN (3 Blocks)
RRN (4 Blocks)
RRN (5 Blocks)
GWavenet: Hardware Efficiency
104
62
54
46
19
40
19
19
06
103
19
19
98
92
84
11
78
11
11
11
11
58
104
5
5
5
5
16
16
16
4 16
102
16
103
GPU Memory (MiB)
# Params (K)
101
Testing (batch/s)
GPU Memory (MiB)
# Params (K)
.48
.13
.7
.08 13
12
.81
Training (batch/s)
14
16
.98
26
26
28
.9
33
.72
57 39 .01
.15 24
.02 27
26
.53
.6
42 .62
.81 11
101
10
12
14
.07
19
.54
28
.52
.7
38
.5
37
.2
37
37
36
.7
.77
.05
102
Training (batch/s)
Testing (batch/s)
Figure 3: Hardware efficiency comparison between baseline and RRN models (2–5 blocks) across GPU memory, parameter counts, and throughput for DCRNN and Graph WaveNet Impact of RRN Depth: DCRNN
162.5 162.0 161.5
86.5
2
3
Number of Blocks
4
150.5 80.5
150.0 149.5 149.0
79.5
160.5
1
151.0
80.0
161.0
5
152.0 151.5
81.0
163.0
87.0
MAE RMSE
81.5
163.5
87.5
86.0
164.0
RMSE MAE
MAE
88.0
Impact of RRN Depth: GWavenet
82.0
164.5
RMSE
MAE RMSE
88.5
1
2
3
Number of Blocks
4
5
148.5
Figure 4: Impact of the number of RRN residual blocks on forecasting accuracy, illustrated through MAE and RMSE metrics for DCRNN and Graph WaveNet architectures Comparison with Normalization Strategies. We evaluate RRN against state-of-theart normalization methods using DCRNN in Table 2. First, RRN outperforms other invertible approaches (RevIN, Dish-TS, SAN) by integrating graph-aware operations, whereas these baselines typically treat nodes as independent time series. Second, the non-invertible method ST-Norm shows competitive performance, occasionally matching RRN on the SDWPF dataset. ST-Norm’s success in filtering high-frequency noise highlights the value of spatio-temporal decoupling, suggesting a promising direction for integrating multi-frequency decomposition into our invertible framework in future work.
5.3
Extra Experiments
Performance Analysis We evaluate the computational efficiency of RRN across varying depths, as shown in Figure 3. The framework demonstrates high memory efficiency, maintaining stable GPU usage and constant parameter counts regardless of the number of stacked blocks. However, inference throughput decreases linearly as depth increases. This latency is primarily attributed to the iterative fixed-point calculation required for the inverse transformation, highlighting a trade-off between structural depth and computational speed. Impact of Residual Blocks
We further investigate the relationship between RRN depth 10
Table 3: Ablation study of RRN components on the SDWPE dataset. We evaluate the impact of removing Center Normalization (CN) and the invertible GNN module (iGNN) across two backbone architectures. Best results are highlighted in bold. DCRNN (MAE)
Variant
GWavenet (MAE)
3 Horizon
6 Horizon
12 Horizon
3 Horizon
6 Horizon
12 Horizon
64.06 63.04 61.99 60.85
90.13 88.55 87.61 85.91
126.31 125.77 124.28 123.28
58.72 54.69 56.81 54.20
83.34 79.40 81.68 78.81
114.95 114.79 113.17 112.59
Baseline (Original) w/o CN w/o iGNN w/ CN + iGNN (RRN)
and forecasting accuracy, as illustrated in Figure 4. Increasing the number of residual blocks generally reduces prediction error, enhancing the model’s ability to map complex distributions. However, we observe diminishing marginal returns with greater depth. Our results indicate that a configuration of 2 to 3 blocks achieves the optimal balance between model expressivity and operational efficiency. Ablation Study We validate the contribution of each component within RRN using the SDWPF dataset, as detailed in Table 3. The results demonstrate that removing either Center Normalization (CN) or the invertible GNN module (iGNN) leads to clear performance degradation. Specifically, CN is essential for mitigating temporal non-stationarity, while iGNN is critical for capturing spatial heterogeneity. The complete RRN framework consistently yields the lowest errors, confirming that the synergy between temporal normalization and spatial structure preservation is required for robust forecasting. Node 0 Probability Density
Node 1 Original Transformed
8
6
4
2
2
-1
0
1
2
Value
3
4
0
Node 24 Original Transformed
6
6
4
0
Node 25 8
Original Transformed
8
-1
0
1
2
3
6
4
4
2
2
0
4
Value
Original Transformed
8
-1
0
1
2
Value
3
4
0
-1
0
1
2
3
Value
Figure 5: Comparison of probability density distributions for representative nodes before (Original) and after (Transformed) RRN transformation, illustrating the alignment of heterogeneous data into a normalized space. Distribution Smoothing Effect To better understand how RRN helps the backbone model, we visualize the probability density of the data before and after the transformation. As shown in Figure 5, the original data for different locations, such as Node 0, Node 1, Node 24, and Node 25 have very different and sometimes sharp shapes. After passing through the RRN framework, these original distributions are transformed into a much smoother and more aligned shape. This change is very important because when the data is finally input into the backbone model, it exists in a much more stable and unified space. This smooth distribution makes it much easier for the base model to learn clear patterns and make accurate predictions.
11
5.4
Reversibility and Sensitivity Analysis
In this section, we empirically validate the invertibility of our proposed framework and analyze its sensitivity to the Lipschitz constraints. Table 4: Sensitivity analysis and reversibility test of RRN. We fix the Center Normalization scaling α = 0.9 and vary the Frobenius norm constraint c. Configuration
GNN Frobenius Norm (c)
(Fixed α = 0.9)
Prediction (MAE)
0.8
0.9
1.0
5
10
15
3-horizon
60.45
61.27
220.04
1815.85
7.99 × 106
9.70 × 105
6-horizon
86.62
86.41
464.58
1836.87
7.63 × 106
1.01 × 106
12-horizon
122.70
122.48
308.33
1836.38
7.63 × 106
9.14 × 105
Iteration
Reconstruction Error
Reconstruction Error
5
3.78 × 10−3
2.46 × 10−3
3.15 × 10−2
8.89 × 10−2
2.58 × 100
1.30 × 101
10
2.14 × 10−3
1.18 × 10−3
3.16 × 10−2
8.44 × 10−2
1.06 × 101
3.36 × 102
20
1.23 × 10−3
4.89 × 10−4
3.25 × 10−2
1.15 × 10−1
7.36 × 101
1.49 × 105
50
2.94 × 10−4
5.92 × 10−5
3.30 × 10−2
2.91 × 10−1
1.07 × 105
7.18 × 1013
Table 5: Sensitivity analysis and reversibility test of RRN. We fix the GNN Frobenius norm constraint c = 0.9 and vary the Center Normalization scaling α. Configuration
Center Normalization Scaling (α)
(Fixed c = 0.9)
Prediction (MAE)
0.8
0.9
1.0
5
10
15
3-horizon
62.64
61.33
258.54
72.97
1.27 × 107
1.57 × 106
6-horizon
88.33
86.93
258.56
106.45
9.56 × 105
1.47 × 106
12-horizon
124.80
123.22
258.47
134.29
9.02 × 105
1.14 × 106
Iteration
Reconstruction Error
Reconstruction
5
1.72 × 10−3
2.52 × 10−3
1.37 × 10−1
2.01 × 10−1
3.63 × 101
4.72 × 100
10
3.42 × 10−4
9.39 × 10−4
1.64 × 10−1
2.75 × 10−1
9.06 × 103
6.04 × 101
20
4.59 × 10−5
3.17 × 10−4
1.70 × 10−1
3.11 × 10−1
7.39 × 107
7.04 × 104
50
7.96 × 10−7
3.43 × 10−5
1.73 × 10−1
3.26 × 10−1
5.92 × 1020
2.87 × 1015
We validate the invertibility condition from Lemma 1 by analyzing forecasting accuracy (MAE) and reconstruction error. As shown in Tables 4 and 5, increasing the GNN constraint c or Center Normalization scaling α gradually violates the contraction mapping condition. When parameters remain strictly below 1, the model maintains low forecasting error and negligible reconstruction error (10−3 –10−4 ). However, as c or α reaches or exceeds 1, fore12
casting precision deteriorates significantly, and the reconstruction error explodes (> 106 ) as the fixed-point iteration diverges. These results empirically confirm that strictly enforcing the Lipschitz constant (< 1) is essential for both training stability and the mathematical validity of the inverse transformation.
6
Limitations
While the proposed framework demonstrates superior forecasting accuracy, two primary limitations exist. First, the inverse transformation relies on fixed-point iteration, which introduces sequential computational overhead. Although early termination can accelerate this process, the cumulative latency across stacked blocks remains higher than standard non-invertible normalization. Second, to strict maintain Lipschitz continuity for invertibility, our Center Normalization module removes mean shifts but intentionally omits the division by standard deviation used in Instance Normalization. Since variance scaling is a critical component of distribution shift, this design may limit the model’s ability to handle complex scaling variations. Future work will explore incorporating variance normalization while preserving theoretical invertibility.
7
Conclusion
This study addresses the critical challenge of distribution shift in spatio-temporal forecasting, characterized by simultaneous temporal non-stationarity and spatial heterogeneity. To mitigate this, we propose Reversible Residual Normalization (RRN), a novel framework that integrates Lipschitz-continuous Center Normalization with spectral-constrained graph convolutions. By strictly enforcing invertibility, RRN projects complex spatio-temporal data into a stationary latent space while preserving topological dependencies. Extensive experiments demonstrate that our model-agnostic approach significantly enhances the robustness and generalization of deep forecasting models by explicitly accounting for coupled spatial and temporal distribution shifts.
References [1]
Shengnan Guo et al. “Attention Based Spatial-Temporal Graph Convolutional Networks for Traffic Flow Forecasting”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 33. 2019, pp. 922–929.
[2]
Xiaoming Shi et al. “A Spatial–Temporal Attention Approach for Traffic Prediction”. In: IEEE Transactions on Intelligent Transportation Systems (2020).
[3]
Shen Fang et al. “GSTNet: Global Spatial-Temporal Network for Traffic Flow Prediction”. In: Proceedings of the International Joint Conference on Artificial Intelligence (IJCAI). 2019.
[4]
Yuxuan Liang et al. “Airformer: Predicting Nationwide Air Quality in China with Transformers”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. 2023, pp. 14329–14337.
13
[5]
Keyong Hu et al. “Air Quality Prediction Using Spatio-Temporal Deep Learning”. In: Atmospheric Pollution Research 13.10 (2022), p. 101543.
[6]
Guangyin Jin et al. “Spatio-Temporal Graph Neural Networks for Predictive Learning in Urban Computing: A Survey”. In: IEEE Transactions on Knowledge and Data Engineering 36.10 (2023), pp. 5388–5408.
[7]
Xuanming Hu et al. “Boosting Urban Prediction via Addressing Spatial-Temporal Distribution Shift”. In: IEEE International Conference on Data Mining (ICDM). 2023, pp. 160–169.
[8]
Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky. “Instance Normalization: The Missing Ingredient for Fast Stylization”. In: arXiv preprint arXiv:1607.08022 (2016).
[9]
Xianbiao Qi et al. “Lipsformer: Introducing Lipschitz Continuity to Vision Transformers”. In: arXiv preprint arXiv:2304.09856 (2023).
[10]
Jiajun Zha et al. “Invertible Attention”. In: arXiv preprint arXiv:2106.09003 (2021).
[11]
Jens Behrmann et al. “Invertible Residual Networks”. In: International Conference on Machine Learning (ICML). 2019, pp. 573–582.
[12]
Moonjeong Park, Jaeseung Heo, and Dongwoo Kim. “Mitigating Oversmoothing Through Reverse Process of GNNs for Heterophilic Graphs”. In: International Conference on Machine Learning (ICML). 2024, pp. 39667–39681.
[13]
Takeru Miyato et al. “Spectral Normalization for Generative Adversarial Networks”. In: International Conference on Learning Representations (ICLR). 2018.
[14]
Taesung Kim et al. “Reversible Instance Normalization for Accurate Time-Series Forecasting against Distribution Shift”. In: International Conference on Learning Representations (ICLR). 2021.
[15]
Zhiding Liu et al. “Adaptive Normalization for Non-stationary Time Series Forecasting: A Temporal Slice Perspective”. In: Advances in Neural Information Processing Systems (NeurIPS) 36 (2023), pp. 14273–14292.
[16]
Wei Fan et al. “Dish-ts: A General Paradigm for Alleviating Distribution Shift in Time Series Forecasting”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. 2023, pp. 7522–7529.
[17]
Wei Fan et al. “IN-Flow: Instance Normalization Flow for Non-stationary Time Series Forecasting”. In: Proceedings of the 31st ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2025, pp. 295–306.
[18]
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. “Density Estimation using Real NVP”. In: International Conference on Learning Representations (ICLR). 2017.
[19]
Shuangfei Zhai et al. “Normalizing Flows are Capable Generative Models”. In: International Conference on Machine Learning (ICML). 2025.
[20]
Jenny Liu et al. “Graph Normalizing Flows”. In: Advances in Neural Information Processing Systems (NeurIPS) 32 (2019).
[21]
Durk P. Kingma and Prafulla Dhariwal. “Glow: Generative Flow with Invertible 1x1 Convolutions”. In: Advances in Neural Information Processing Systems (NeurIPS) 31 (2018). 14
[22]
Laurent Dinh, David Krueger, and Yoshua Bengio. “NICE: Non-linear Independent Components Estimation”. In: International Conference on Learning Representations (ICLR). 2015.
[23]
Durk P. Kingma et al. “Improved Variational Inference with Inverse Autoregressive Flow”. In: Advances in Neural Information Processing Systems (NeurIPS) 29 (2016).
[24]
Kaiming He et al. “Deep Residual Learning for Image Recognition”. In: Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition (CVPR). 2016, pp. 770–778.
[25]
Henry Gouk et al. “Regularisation of Neural Networks by Enforcing Lipschitz Continuity”. In: Machine Learning 110.2 (2021), pp. 393–416.
[26]
Thomas N. Kipf and Max Welling. “Semi-Supervised Classification with Graph Convolutional Networks”. In: International Conference on Learning Representations (ICLR). 2017.
[27]
Jingbo Zhou et al. “SDWPF: A Dataset for Spatial Dynamic Wind Power Forecasting over a Large Turbine Array”. In: Scientific Data 11.1 (2024), p. 649.
[28]
Jinliang Deng et al. “ST-Norm: Spatial and Temporal Normalization for Multi-Variate Time Series Forecasting”. In: Proceedings of the 27th ACM SIGKDD Conference on Knowledge Discovery and Data Mining. 2021, pp. 269–278.
[29]
Yaguang Li et al. “Diffusion Convolutional Recurrent Neural Network: Data-Driven Traffic Forecasting”. In: International Conference on Learning Representations (ICLR). 2018.
[30]
Zonghan Wu et al. “Graph WaveNet for Deep Spatial-Temporal Graph Modeling”. In: Proceedings of the 28th International Joint Conference on Artificial Intelligence (IJCAI). 2019.
[31]
Chuanpan Zheng et al. “GMAN: A Graph Multi-Attention Network for Traffic Prediction”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 34. 2020, pp. 1234–1241.
[32]
Guangyin Jin et al. “Automated Dilated Spatio-Temporal Synchronous Graph Modeling for Traffic Prediction”. In: IEEE Transactions on Intelligent Transportation Systems 24.8 (2022), pp. 8820–8830.
[33]
Weiyang Kong, Ziyu Guo, and Yubao Liu. “Spatio-Temporal Pivotal Graph Neural Networks for Traffic Flow Forecasting”. In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. 2024, pp. 8627–8635.
15