Does Vector Quantization Fail in Spatio-Temporal Forecasting? Exploring a Differentiable Sparse Soft-Vector Quantization Approach

## Abstract

Abstract Spatio-temporal forecasting, pivotal in numerous fields, hinges on the delicate equilibrium between isolating nuanced patterns and sifting out noise. To tackle this, we introduce Sparse Regression-based Vector Quantization (SVQ), a novel technique that leverages sparse regression for succinct representation, an approach theoretically and practically favored over classical clustering-based vector quantization methods. This approach preserves critical details from the original vectors using a regression model while filtering out noise via sparse design. Moreover, we approximate the sparse regression process using a blend of a two-layer MLP and an extensive codebook. This approach not only substantially cuts down on computational costs but also grants SVQ differentiability and training simplicity, resulting in a notable enhancement of performance. Our empirical studies on five spatial-temporal benchmark datasets demonstrate that SVQ achieves state-of-the-art results. Specifically, on the WeatherBench-S temperature dataset, SVQ improves the top baseline by 7.9%. In video prediction benchmarks—Human3.6M, KTH, and KittiCaltech—it reduces MAE by an average of 9.4% and improves image quality by 17.3% (LPIPS). Code is publicly available at https://anonymous.4open.science/r/SVQ-Forecasting .

## 1 Introduction

Spatio-temporal forecasting (Wang et al., ) has board applications in many domains including environment, climate, transportation, etc. Many traditional methods relying on feature engineering have been studied in the past several decades. Recently, deep learning models are becoming more popular (Wang et al., ). However, as more applications (e.g., weather forecasting) arise, it calls for deep models to accurately predict future dynamics with precise spatial and temporal resolutions.

Despite the various forms of spatio-temporal forecasting approaches, such as the use of graph neural network (GNN) models for unevenly resolved data, our work concentrates on grid-based datasets, one of the most widely adopted structures in the field. Within this context, we address a primary challenge: how to finely balance the preservation of important details against the need to filter out irrelevant noise, thereby enhancing the generalized performance of forecasting models. Although vector quantization (VQ) is effective in reducing noises by constraining latent patterns to a clustered codebook, it is usually limited in preserving the diversity of input patterns, making it suboptimal for spatial-temporal forecasting. We performed a comprehensive suite of experiments, the results of which indicate that, contrary to expectations, nearly all sota VQ methods we examined did not enhance the accuracy of a basic weather forecasting model as shown in Table 4. This outcome may provide insight into the infrequent application of VQ techniques in spatial-temporal prediction tasks.

Figure: Figure 1: Codebook visualization using t-SNE. Left: VQ updated by k-means clustering and exponential moving average as described in (van den Oord et al., ). Right: SVQ with randomly initialized, frozen or learnable codebook. Two methods are trained on WeatherBench-S temperature dataset with the same codebook size (1024).
Refer to caption: x1.png

To address this limitation of VQ, we propose a Sparse regression-based Vector Quantization (SVQ) method. By associating each vector to multiple codes from a large codebook, it is able to cover diverse visual patterns; by introducing a sparse regularization mechanism in the regression model, we effectively filter out noise and thus ensure a good generalized performance. As shown in Figure 1, our proposed SVQ method exhibits a more uniform distribution of quantized vectors compared to VQ, indicating that SVQ is able to preserve more diverse and fine-grained information from the original inputs than VQ.

Despite employing a series of vectors for representation instead of a single codebook vector, regression-based vector quantization remains true to the foundational aspects of the technique. It still quantizes input vectors into a finite set of discrete levels, utilizing a dynamic codebook-like structure. The key process of mapping input data to representative vectors is maintained, but with the added sophistication of regression models that adaptively combine vectors for a more accurate representation. This method preserves the essence of vector quantization—efficiently compressing and approximating data—while enhancing its ability to capture complex patterns within the input space.

To address the high computational cost of sparse regression, which
scales quadratically with codebook size, we approximate the sparse regression model with a MLP model and a large codebook. It not only dramatically simplifies the computation, but more importantly, makes SVQ differentiable and thus easy to train. We emphasize the importance of differentiability of SVQ for its success in spatial-temporal forecasting. Our experiments reveal that lookup-based VQ methods can substantially increase prediction error, despite some employing a residual design that theoretically should perform on par with or better than non-VQ cases. This suggests the “straight-through” estimator, a non-differentiable component of VQ, may be a major factor behind VQ’s underperformance. In particular, the non-differentiable operator such as argmin or lookup introduces errors into gradient estimation of all modules preceding the quantization module, consequently undermining prediction accuracy. In contrast, our proposed SVQ is differentiable and easy to compute accurate gradients for training. We finally note that SVQ also differs from traditional VQ methods by not relying on intricate models to construct the prior distribution of codes, such as the GPT-style model in VQGAN (Esser et al., ) or the BERT-style model in MaskGIT (Chang et al., ).

Besides differentiability, there are two interesting properties of the proposed SVQ scheme. First, according to our empirical study, we find that with large enough codebook size, we can freeze the randomly initialized codebook without sacrificing any performance, making SVQ potentially more efficient for training. Secondly, unlike lookup-based methods, SVQ allows for both positive and negative weights, enhancing the ability to add or subtract patterns to improve representational power.

Finally, despite its simplification, our experiments across multiple real-world benchmark datasets—including weather forecasting, traffic flow prediction, and video forecasting, demonstrate the superior performance of our SVQ method. In WeatherBench-S temperature forecasting benchmark, it improves the prior SOTA method by 7.9%. For video prediction across renowned datasets—Human3.6M, KTH, and KittiCaltech—SVQ not only achieves a consistent 9.4% reduction in MAE but also enhances perceptual quality, evidenced by a 17.3% lower LPIPS score. These results highlight SVQ’s versatility and superior accuracy in diverse spatial-temporal prediction scenarios.

## 2 Related Work

### 2.1 Recurrent-based forecasting model

The majority of spatio-temporal forecasting models leverage techniques such as Conv2D (Xu et al., ), Conv3D (Wang et al., ), and attention (Liu et al., ) for spatial modeling. Distinctions among these models primarily arise from how they incorporate temporal information. Recurrent-based models, exemplified by ConvLSTM (SHI et al., ), have been widely used to capture motion dynamics by iteratively processing multi-frame predictions. Variants like PredRNN (Wang et al., ) introduce the Spatio-Temporal LSTM (ST-LSTM) unit, integrating spatial appearances and temporal variations within a single memory pool. Further advancements include PredRNN++ (Wang et al., ) and PredRNNV2 (Wang et al., ), which deepen the model and expand the receptive field through a cascading LSTM mechanism and a memory decoupling strategy, respectively.

### 2.2 Non-recurrent forecasting model

Despite the effectiveness of recurrent-based methods, they suffer from high computational cost caused by their inherent unparallelizable architecture. Recent efforts in spatio-temporal forecasting have shifted towards non-recurrent models that decouple the forecasting task from autoregressive processes. Notably, SimVPv1 (Gao et al., ) and SimVPv2 (Tan et al., ) separate spatial and temporal learning into distinct phases within an encoder-translator-decoder structure, consistently surpassing recurrent counterparts in video prediction tasks. Due to their similar architectures, we will collectively refer to them as the SimVP model in the remainder of our paper. TAU (Tan et al., ) refines the architecture by incorporating a visual attention mechanism into the translator. OpenSTL (Tan et al., ) further enhances the translator by MetaFormers.

### 2.3 Vector quantization and sparse coding

Instead of using continuous latent, VQVAE (van den Oord et al., ), a seminal work, integrates vector quantization (VQ) in learning discrete latent representations, typically assigning each vector to the nearest code in a codebook. Subsequent enhancements include Residual VQ (Zeghidour et al., ), which quantizes the residuals recursively, and Multi-headed VQ (Mama et al., ), which adopts multiple heads for each vector. While effective, these methods often rely on a relatively small number of codes for representing continuous vectors. Addressing this, SCVAE (Xiao et al., ) employs sparse coding, allowing vectors to be represented through sparse linear combinations of multiple codes and achieving end-to-end training via the Learnable Iterative Shrinkage Thresholding Algorithm (LISTA) (Gregor & LeCun, ). However, a significant drawback of the sparse coding method using LISTA (dubbed SVQ-raw here) is its high computational complexity, which scales quadratically with codebook size.

Our work is closely related to a simultaneous research work (Tschannen et al., ), which employs an infinite cookbook with a linear layer for continuous vector quantization applied in image generation. Despite the similarity, our approach is largely inspired by sparse regression, which is clearly evidenced by our analysis. In particular, we provide a better theoretical foundation and address the challenges arising from spatio-temporal forecasting well.

## 3 Sparse Vector Quantization (SVQ)

In this section, we will first describe the mathematical foundation for sparse vector quantization, followed by the description of detailed implementation.

### 3.1 Vector quantization by sparse regression

Let $\{z_{i}\in\mathbb{R}^{d}\}_{i=1}^{m}$ be the set of codes. A typical vector quantization method assigns a data point $x\in\mathbb{R}^{d}$ to the nearest code in $\{z_{i}\}_{i=1}^{m}$. The main problem with such an approach is that a significant part of the information in $x$ will be lost due to quantization. Sparse regression turns the code assignment problem into an optimization problem

$$ $\displaystyle w=\mathop{\arg\min}_{w\in\mathbb{R}_{+}^{m}}\frac{1}{2}\left|x- \sum_{i=1}^{m}w_{i}z_{i}\right|^{2}+\lambda|w|_{1},$ (1) $$

where $w=(w_{1},\ldots,w_{m})\in\mathbb{R}_{+}^{m}$ is the weight for combining codes $\{z_{i}\}_{i=1}^{m}$ to approximate $x$. By introducing $L_{1}$ regularizer in the optimization problem, we effectively enforce $x$ to be associated with a small number of codes. Compared to classic VQ methods where codes have to be learned through clustering, according to (Chiu et al., ), it is sufficient to use randomly sampled vectors as codes as long as its number is large enough, thus avoiding the need of computing and adjusting codes. The theoretical guarantee of sparse regression is closely related to the property of subspace clustering, as revealed in Theorem 4.1.

As shown in Figure 2, the obvious downside of sparse regression for VQ is its high computational cost, as it needs to solve the optimization problem in (1) for EVERY data point. Below, we will show that sparse regression can be approximated by a two-layer MLP and a randomly fixed or learnable matrix, making it computationally attractive.

To solve the optimization problem (1), we consider the composite optimization method whose iteration is given as follows

$$ $\displaystyle w^{\prime}_{t+1}$ $\displaystyle=$ $\displaystyle w_{t}-\eta Z^{\top}(Zw_{t}-x),$ $\displaystyle\left[w_{t+1}\right]_{i}$ $\displaystyle=$ $\displaystyle\mbox{sgn}\left(\left[w^{\prime}_{t+1}\right]_{i}\right)\left( \left|\left[w^{\prime}_{t+1}\right]_{i}\right|-\lambda\eta\right)_{+},$ $$

where $Z=(z_{1},\ldots,z_{m})$, $[z]_{i}$ is the $i$th element of vector $z$, and $(a)_{+}$ outputs 0 0 if $a<0$ and $a$ otherwise. We consider the first step of the iteration where $w_{0}=0$ and have $w=\eta\mbox{sgn}\left(Z^{\top}x-\lambda\mathbf{1})\right[Z^{\top}x-\lambda
\mathbf{1}]_{+}=\eta\mbox{sgn}\left(Z^{T}x-\lambda\right)\sigma\left(Z^{T}x-
\lambda\right)$ and the resulting output for $x$ is given as

$$ $x^{\prime}=\sum_{i=1}^{m}w_{i}z_{i}=\eta Zsgn\left(Z^{T}x-\lambda\right)\sigma \left(Z^{T}x-\lambda\right).$ $$

By generalizing $\eta Z$ into another matrix $B$, we have output vector $x^{\prime}$ exactly expressed as a matrix and a two-layer MLP over $x$. We finally note that although it is convenient to form the codebook by randomly sampling vectors, we found empirically that tuning codebook does bring slight additional gains in some cases.

Figure: Figure 2: Effect of SVQ approximation: Computational cost (FLOPS) and prediction error (MSE) on WeatherBench-S temperature dataset with SVQ-raw and SVQ. The computational complexity of SVQ-raw increases quadratically with the size of codebook, making it suffer from out-of-memory (OOM) issue when scaling codebook size up to $2^{12}$.
Refer to caption: extracted/5396914/graphs/compare_lista_mlp.png

Figure: Figure 3: Top: Architecture of backbone model and the proposed quantization module. The encoder, translator, decoder are inherited from SimVP. A quantization module is added between the encoder and translator to effectively ensure a good generalized performance. Bottom: Quantization process of classic VQ (Left) and our proposed SVQ (Right). In contrast, SVQ select multiple codes (red dots) from a huge codebook (gray dots), and the codebook can be either learnable or frozen.
Refer to caption: x2.png

### 3.2 Detailed implementation

Architecture of backbone model. As depicted in Figure 3, we adopt SimVP as our backbone model, which encompasses an encoder for spatial feature extraction, a translator for temporal dependency learning, and a decoder for frame reconstruction. We integrate the quantization module between the encoder and translator. The input data is a 4D tensor $X\in\mathbb{R}^{H*W*T*C}$, representing height ($H$), width ($W$), time steps ($T$), and channels ($C$). The encoder ${\rm En}$ condenses $X$ into downsampled latent representation ${\rm En}(X)\in\mathbb{R}^{H^{\prime}*W^{\prime}*T*C^{\prime}}$, maintaining temporal dimensionality while altering spatial and channel dimensions. This latent space, composed of $H^{\prime}*W^{\prime}*T$ tokens, each a $C^{\prime}$-dimensional vector, undergoes vector quantization.

Quantization module. SVQ comprises a two-layer MLP and a huge codebook. The codebook is a randomly initialized matrix $\mathcal{M}\in\mathbb{R}^{N*C^{\prime}}$, where $N$ denotes the size of codebook. To achieve automatic selection of codes, we generate a weight matrix $\mathcal{W}\in\mathbb{R}^{H^{\prime}*W^{\prime}*T*N}$ by nonlinear projection from latent representation $\rm{En}(X)$. This projection is formally expressed as $\mathcal{W}={\rm MLP}(\rm{En}(X))$, wherein the $\rm MLP$ comprises two linear layers and an intermediate ReLU activation function. Subsequently, the quantized output $\mathcal{Q}$ is obtained by computing the dot product of weight matrix $\mathcal{W}$ and codebook matrix $\mathcal{M}$, a process that can be conceptualized as a selection operation as shown in Figure 3. To encourage sparsity within the generated weight matrix, we apply a Mean Absolute Error (MAE) loss to the output as a surrogate form of regularization.

Despite the potential feasibility of placing the quantization module either before or after the Translator, we note that post-Translator placement induces significant training instability in classic VQ, attributable to the non-differentiable nature of the straight-through estimator. Conversely, our SVQ module demonstrates consistent stability during training. For consistency, the quantization module is placed before the Translator in our primary experiments, with a comparative assessment of its alternative placement provided in Section B.1.

## 4 Efficient Utilize of Cookbook Using Sparse Regression

To understand the difference between sparse regression-based quantization scheme and clustering-based quantization scheme, we measure the number of codes that we need in order to approximate any vector within an unit ball $\mathcal{B}$ with error less than $\delta$. We denote this number by $T(\mathcal{B},\delta)$. As indicated by the theorem below, using spare regression, $T(\mathcal{B},\delta)$ can be reduced from $O(1/\delta^{d})$ to $O(1/\delta^{p})$, with $p\ll d$ for high dimensional vectors.

###### Theorem 4.1 .

$T(\mathcal{B},\delta)$ for the clustering-based method is at least $1/\delta^{d}$, while $T(\mathcal{B},\delta)$ for sparse regression can be $(4d/\delta)^{p}$, with

$$ $p\geq\max\left(3,\frac{\log(4/\delta)}{\log\log(2d/\varepsilon)}\right),$ $$

as long as the non-zero elements used by sparse regression is at least $\frac{4d}{\delta\left(\log C+p\log(4d)-(p+1)\log\delta\right)}.$

###### Proof.

To bound $T(\mathcal{B},\delta)$ for the clustering method, we use the covering number for a unit ball $\mathcal{B}$, requiring at least $1/\delta^{d}$ code vectors to approximate any vector within error $\delta$. Considering $Ug$, with $U=(u_{1},\ldots,u_{m})$, $u_{k}\sim\mathcal{N}(0,I_{d}/m)$, and $g\in\Delta_{s}$ an $s$-sparse unit vector, we have:

$$ $\Pr\left(\|UU^{\top}-I\|_{2}\geq\gamma\right)\leq 2d\exp\left(-\frac{m\gamma^{ 2}}{3d}\right),$ $$

implying

$$ $\|UU^{\top}-I\|_{2}\leq\Delta:=\sqrt{\frac{d}{m}\log\frac{2d}{\varepsilon}}$ $$

with probability at least $1-\varepsilon$. Thus, $\|g^{\prime}-g\|_{2}\geq(1+\Delta)^{-1}\|Ug-Ug^{\prime}\|_{2}$. With the $s$-sparse unit vector covering number bounded by $(Cm/s\delta)^{s}$, we obtain:

$$ $\left(\frac{Cm}{s\delta}\right)^{s}\geq\left(1+\frac{2}{\delta}\right)^{d}(1+ \Delta)^{d},$ $$

By choosing $m=(4d/\delta)^{p}$, we have

$$ $(1+\Delta)^{d}\leq\exp\left(d\Delta\right)\leq e,$ $$

yielding $s\log(C^{\prime}m/\delta)\geq 2d/\delta+s\log s$, given $C^{\prime}=Ce$.

As $s\geq 4d/[\delta(\log C+p\log(4d)-(p+1)\log\delta)]$, it is easy to verify that $s\log s\leq 2d/\delta$, leading to $m\geq(4d/\delta)^{p}$.
∎

###### Theorem 4.1 .

###### Proof.

## 5 Experiments

We extensively evaluate SVQ on a wide range of real-world spatio-temporal forecasting datasets under the unified framework of OpenSTL (Tan et al., ). Note that since SimVP holds leading performance across almost all benchmarks, it is adopted as our primary baseline. We begin by comparing the SimVP+SVQ model with state-of-the-art (SOTA) counterparts. Subsequently, we assess the capability of SVQ to enhance various backbones. Finally, we investigate the delicate balance between detail preservation and noise reduction, and conduct experiments with artificial noise injection to show the effect of SVQ on mitigating noise.

Moreover, we provide a comprehensive comparison between SVQ and existing VQ methods in Appendix B.3, evaluating their downstream prediction performance, codebook utilization, and computational efficiency.

Dataset.
To evaluate the proposed SVQ, we conduct extensive experiments on five real-world spatio-temporal forecasting tasks, including weather (WeatherBench), traffic flow (TaxiBJ), human pose dynamics (Human3.6M), driving scenes (KittiCaltech), and human actions (KTH Action). More details about datasets are described in Appendix
A.1.

Experimental details. In deployment, we found SVQ to be quite robust to codebook size, as its performance remains consistently strong when using a sufficiently large codebook. Therefore, we fix the codebook size at 10,000 for WeatherBench, TaxiBJ, and Human3.6M datasets, and at 6,000 for KittiCaltech and KTH datasets. The hidden dimension of nonlinear projection layer is fixed at 128. Our experiments are conducted on either 1 or 4 NVIDIA V100 32GB GPUs, with a total batch size of 16. More details about backbone architectures, VQ parameters, computational costs, and metrics are described in Appendix A.2, A.3, B.2, and A.5 respectively.

### 5.1 Benchmarks on various tasks

**Table 1: WeatherBench results: Performance comparison for SVQ module and baseline models on WeatherBench. WeatherBench-S is single-variable, one-hour interval forecasting setup trained on data from 2010-2015, validated on 2016, and tested on 2017-2018. WeatherBench-M targets broader application, which is multi-variable, six-hour interval forecasting setup trained on data from 1979-2015, validated on 2016, and tested on 2017-2018. The top two results are highlighted by bold or underlined.**
| Dataset | Variable | Temperature | Humidity | Wind Component | Total Cloud Cover |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | MSE$\downarrow$ | MAE$\downarrow$ | MSE$\downarrow$ | MAE$\downarrow$ | MSE$\downarrow$ | MAE$\downarrow$ | MSE$\downarrow$ | MAE$\downarrow$ |  |
| WeatherBench-S | ConvLSTM(SHI et al., ) | 1.521 | 0.7949 | 35.146 | 4.012 | 1.8976 | 0.9215 | 0.0494 | 0.1542 |
| E3D-LSTM(Wang et al., ) | 1.592 | 0.8059 | 36.534 | 4.100 | 2.4111 | 1.0342 | 0.0573 | 0.1529 |  |
| PredRNN(Wang et al., ) | 1.331 | 0.7246 | 37.611 | 4.096 | 1.8810 | 0.9068 | 0.0550 | 0.1588 |  |
| MIM(Wang et al., ) | 1.784 | 0.8716 | 36.534 | 4.100 | 3.1399 | 1.1837 | 0.0573 | 0.1529 |  |
| MAU(Chang et al., ) | 1.251 | 0.7036 | 34.529 | 4.004 | 1.9001 | 0.9194 | 0.0496 | 0.1516 |  |
| PredRNN++(Wang et al., ) | 1.634 | 0.7883 | 35.146 | 4.012 | 1.8727 | 0.9019 | 0.0547 | 0.1543 |  |
| PredRNN.V2(Wang et al., ) | 1.545 | 0.7986 | 36.508 | 4.087 | 2.0072 | 0.9413 | 0.0505 | 0.1587 |  |
| TAU(Tan et al., ) | 1.162 | 0.6707 | 31.831 | 3.818 | 1.5925 | 0.8426 | 0.0472 | 0.1460 |  |
| SimVP (w/o VQ)(Tan et al., ) | 1.105 | 0.6567 | 31.332 | 3.776 | 1.4996 | 0.8145 | 0.0466 | 0.1469 |  |
| SimVP+SVQ (Frozen) | 1.023 | 0.6131 | 30.863 | 3.661 | 1.4337 | 0.7861 | 0.0456 | 0.1456 |  |
|  | SimVP+SVQ (Learnable) | 1.018 | 0.6109 | 30.611 | 3.657 | 1.4186 | 0.7858 | 0.0458 | 0.1463 |
|  | Improvement | $\uparrow$7.9% | $\uparrow$7.0% | $\uparrow$2.3% | $\uparrow$3.2% | $\uparrow$5.4% | $\uparrow$3.5% | $\uparrow$2.1% | $\uparrow$0.9% |
| WeatherBench-M | Variable | Temperature | Humidity | Wind U Component | Wind V Component |  |  |  |  |
| ConvLSTM(SHI et al., ) | 6.303 | 1.7695 | 368.15 | 13.490 | 30.002 | 3.8923 | 30.789 | 3.8238 |  |
| PredRNN(Wang et al., ) | 5.596 | 1.6411 | 354.57 | 13.169 | 27.484 | 3.6776 | 28.973 | 3.6617 |  |
| MIM(Wang et al., ) | 7.515 | 1.9650 | 408.24 | 14.658 | 35.586 | 4.2842 | 36.464 | 4.2066 |  |
| MAU(Chang et al., ) | 5.628 | 1.6810 | 363.36 | 13.503 | 27.582 | 3.7409 | 27.929 | 3.6700 |  |
| PredRNN++(Wang et al., ) | 5.647 | 1.6433 | 363.15 | 13.246 | 28.396 | 3.7322 | 29.872 | 3.7067 |  |
| PredRNN.V2(Wang et al., ) | 6.307 | 1.7770 | 368.52 | 13.594 | 29.833 | 3.8870 | 31.119 | 3.8406 |  |
| TAU(Tan et al., ) | 4.904 | 1.5341 | 342.63 | 12.801 | 24.719 | 3.5060 | 25.456 | 3.4723 |  |
| SimVP (w/o VQ)(Tan et al., ) | 4.833 | 1.5246 | 340.06 | 12.738 | 24.535 | 3.4882 | 25.232 | 3.4509 |  |
| SimVP+SVQ (Frozen) | 4.427 | 1.4160 | 360.15 | 12.445 | 23.915 | 3.4078 | 24.968 | 3.4117 |  |
|  | SimVP+SVQ (Learnable) | 4.433 | 1.4164 | 360.53 | 12.449 | 23.908 | 3.4060 | 24.983 | 3.4095 |
|  | Improvement | $\uparrow$8.4% | $\uparrow$7.1% | $\downarrow$5.9% | $\uparrow$2.3% | $\uparrow$2.6% | $\uparrow$2.4% | $\uparrow$1.0% | $\uparrow$1.2% |

**Table 2: Video prediction results: Performance comparison for SVQ module and baseline models on Human3.6M, KTH, and KittiCaltech. The top two results are highlighted by bold or underlined.**
| Dataset | Human3.6M | KittiCaltech | KTH |  |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Metric | MAE$\downarrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ | LPIPS$\downarrow$ | MAE$\downarrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ | LPIPS$\downarrow$ | MAE$\downarrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ | LPIPS$\downarrow$ |
| ConvLSTM(SHI et al., ) | 1583.3 | 0.9813 | 33.40 | 0.03557 | 1583.3 | 0.9345 | 27.46 | 0.08575 | 445.5 | 0.8977 | 26.99 | 0.26686 |
| E3D-LSTM(Wang et al., ) | 1442.5 | 0.9803 | 32.52 | 0.04133 | 1946.2 | 0.9047 | 25.45 | 0.12602 | 892.7 | 0.8153 | 21.78 | 0.48358 |
| PredNet(Lotter et al., ) | 1625.3 | 0.9786 | 31.76 | 0.03264 | 1568.9 | 0.9286 | 27.21 | 0.11289 | 783.1 | 0.8094 | 22.45 | 0.32159 |
| PhyDNet(Guen & Thome, ) | 1614.7 | 0.9804 | 39.84 | 0.03709 | 2754.8 | 0.8615 | 23.26 | 0.32194 | 765.6 | 0.8322 | 23.41 | 0.50155 |
| MAU(Chang et al., ) | 1577.0 | 0.9812 | 33.33 | 0.03561 | 1800.4 | 0.9176 | 26.14 | 0.09673 | 471.2 | 0.8945 | 26.73 | 0.25442 |
| MIM(Wang et al., ) | 1467.1 | 0.9829 | 33.97 | 0.03338 | 1464.0 | 0.9409 | 28.10 | 0.06353 | 380.8 | 0.9025 | 27.78 | 0.18808 |
| PredRNN(Wang et al., ) | 1458.3 | 0.9831 | 33.94 | 0.03245 | 1525.5 | 0.9374 | 27.81 | 0.07395 | 380.6 | 0.9097 | 27.95 | 0.21892 |
| PredRNN++(Wang et al., ) | 1452.2 | 0.9832 | 34.02 | 0.03196 | 1453.2 | 0.9433 | 28.02 | 0.13210 | 370.4 | 0.9124 | 28.13 | 0.19871 |
| PredRNN.V2(Wang et al., ) | 1484.7 | 0.9827 | 33.84 | 0.03334 | 1610.5 | 0.9330 | 27.12 | 0.08920 | 368.8 | 0.9099 | 28.01 | 0.21478 |
| TAU(Tan et al., ) | 1390.7 | 0.9839 | 34.03 | 0.02783 | 1507.8 | 0.9456 | 27.83 | 0.05494 | 421.7 | 0.9086 | 27.10 | 0.22856 |
| DMVFN (Hu et al., ) | - | - | - | - | 1531.1 | 0.9314 | 26.95 | 0.04942 | 413.2 | 0.8976 | 26.65 | 0.12842 |
| SimVP (w/o VQ)(Tan et al., ) | 1441.0 | 0.9834 | 34.08 | 0.03224 | 1507.7 | 0.9453 | 27.89 | 0.05740 | 397.1 | 0.9065 | 27.46 | 0.26496 |
| SimVP+SVQ (Frozen) | 1264.9 | 0.9851 | 34.07 | 0.02380 | 1408.6 | 0.9469 | 28.10 | 0.05535 | 364.6 | 0.9109 | 27.28 | 0.20988 |
| SimVP+SVQ (Learnable) | 1265.1 | 0.9851 | 34.06 | 0.02367 | 1414.9 | 0.9458 | 28.10 | 0.05776 | 360.2 | 0.9116 | 27.37 | 0.20658 |
| Improvement | $\uparrow$12.2% | $\uparrow$0.2% | $\downarrow$0.0% | $\uparrow$26.2% | $\uparrow$6.6% | $\uparrow$0.2% | $\uparrow$0.8% | $\uparrow$3.6% | $\uparrow$9.3% | $\uparrow$0.6% | $\downarrow$0.3% | $\uparrow$22.0% |

We explore both fixed (frozen) and learnable versions of SVQ on various forecasting tasks. Interestingly, we find that when codebook size is large, the performance of frozen, randomly-initialized codebook is on par with carefully learned codebook. Our comparison baselines consist of two categories: 1) Non-recurrent models including SimVP (Tan et al., ) and TAU (Tan et al., ); 2) Recurrent-based models including ConvLSTM (SHI et al., ), PredNet (Lotter et al., ), PredRNN (Wang et al., ), PredRNN++ (Wang et al., ), MIM (Wang et al., ), E3D-LSTM (Wang et al., ), PhyDNet (Guen & Thome, ), MAU (Chang et al., ), PredRNNv2 (Wang et al., ), and DMVFN (Hu et al., ). Baseline results are copied from the original OpenSTL paper (Tan et al., ). To preclude ambiguity, we select the best MetaFormer of SimVP for each dataset, detailed in Appendix A.2.

We present benchmark results of WeatherBench and three video prediction datasets (Human3.6M, KTH, and KittiCaltech) in Tables 1 and 2, respectively. Due to page limit, the benchmark results of TaxiBJ dataset are provided in Appendix B.4. These datasets have different characteristics. WeatherBench and TaxiBJ are macro forecasting tasks with low-frequency collection (30min or 1-6h). Human3.6M features subtle differences between frames as low-frequency signals. KittiCaltech poses challenges due to its rapidly changing backgrounds and limited training data. The KTH dataset, on the other hand, tests the capacity for long-horizon forecasting, requiring the prediction of 20 future frames based on 10 observed frames.

However, despite the distinct characteristics among datasets, a common thread is the need for improved noise reduction coupled with enhanced representational capabilities, which can universally benefit their respective forecasting tasks. Notably, the SimVP+SVQ model achieves either the best or comparable performance across all datasets. For instance, on the WeatherBench-S temperature dataset, SVQ significantly improves the best baseline by 7.9% (1.105 → 1.018). On these three popular video prediction tasks, SVQ not only delivers a reduction in forecasting errors (average 9.4% decrease in MAE), but also significantly improves subjective image quality (average 17.3% decrease in LPIPS). Additional visualizations of forecasting samples can be found in Appendix D.

### 5.2 Boosting performance as a versatile plug-in

In this section, SVQ serves as a versatile plug-in module applicable to various MetaFormers (Yu et al., ). The adopted MetaFormers can be categorized into three types. 1) CNN-based: SimVPv1(IncepU) (Gao et al., ), SimVPv2(gSTA) (Tan et al., ), ConvMixer (Trockman & Kolter, ), ConvNeXt (Liu et al., ), HorNet (Rao et al., ), and MogaNet (Li et al., ). 2) Transformer-based: ViT (Dosovitskiy et al., ), Swin Transformer (Liu et al., ), Uniformer (Li et al., ), Poolformer (Yu et al., ), and VAN (Guo et al., ). 3) MLP-based: MLPMixer (Tolstikhin et al., ). We conduct experiments on WeatherBench-S temperature dataset because it is lightweight and fast for training.

As shown in Table 3, SVQ consistently improves the performance across all MetaFormers, showcasing its universality across diverse backbone types. We observe an average reduction in MSE and MAE by 4.8% and 6.0%, respectively. In detail, SVQ leads to an average MSE reduction of 4.1% for CNN-based backbones, 5.1% for transformer-based, and 10.7% for MLP-based. The more pronounced enhancement in transformer-based and MLP-based models indicates that our approach is especially effective with architectures that prioritize global interactions. Notably, SimVPv2(gSTA) is the best backbone, while our SVQ further improves it by 7.9%. These findings also aligns with our motivation that mitigating noise in the learning process significantly benefits spatio-temporal forecasting, irrespective of model architecture. By integrating SVQ to constrain the diversity of predicted patterns and cut out noise, researchers can focus on crafting high-quality and general base models.

**Table 3: Boosting performance: The effect of SVQ for various MetaFormers on WeatherBench-S temperature dataset.**
| MetaFormer | MSE | MAE |  |  |
| --- | --- | --- | --- | --- |
| w/o SVQ | w SVQ | w/o SVQ | w SVQ |  |
| SimVPv1(IncepU)(Gao et al., ) | 1.238 | 1.216 | 0.7037 | 0.6831 |
| SimVPv2(gSTA)(Tan et al., ) | 1.105 | 1.018 | 0.6567 | 0.6109 |
| ConvMixer(Trockman & Kolter, ) | 1.267 | 1.257 | 0.7073 | 0.6780 |
| ConvNeXt(Liu et al., ) | 1.277 | 1.159 | 0.7220 | 0.6568 |
| HorNet(Rao et al., ) | 1.201 | 1.130 | 0.6906 | 0.6472 |
| MogaNet(Li et al., ) | 1.152 | 1.067 | 0.6665 | 0.6271 |
| ViT(Dosovitskiy et al., ) | 1.146 | 1.111 | 0.6712 | 0.6375 |
| Swin(Liu et al., ) | 1.143 | 1.088 | 0.6735 | 0.6320 |
| Uniformer(Li et al., ) | 1.204 | 1.110 | 0.6885 | 0.6400 |
| Poolformer(Yu et al., ) | 1.156 | 1.097 | 0.6715 | 0.6297 |
| VAN(Guo et al., ) | 1.150 | 1.083 | 0.6803 | 0.6342 |
| MLP-Mixer(Tolstikhin et al., ) | 1.255 | 1.120 | 0.7011 | 0.6455 |
| Average improvement | $\uparrow$ 4.8% | $\uparrow$ 6.0% |  |  |

### 5.3 Delicate balance between detail preservation and noise reduction

Investigating vector quantization’s (VQ) role in spatio-temporal forecasting, we evaluated several cutting-edge VQ methods akin to the SVQ framework, implemented as plug-in modules coordinated with the primary model. Table 4 shows that SVQ markedly improves forecasting as a plug-in, whereas other VQ approaches generally lead to heightened prediction errors. Enhanced detail retention within the representational capacity is linked to lower forecasting errors. Classic VQ methods suffer from notable information losses, as evidenced by a higher MSE of 1.854. In contrast, both residual VQ and grouped residual VQ outperform traditional VQ with lower MSEs of 1.213 and 1.174, respectively, affirming their ability to preserve intricate details.

It is commonly understood that the size of the codebook in clustering-based VQ is critical: larger codebooks capture more details, whereas smaller ones enhance noise reduction. To explore this trade-off, we analyzed how the codebook size influences the prediction MSE in grouped residual VQ. As Figure 4 indicates, MSE initially decreases but increases with overly large codebooks, echoing findings from (Yu et al., ) that an excessively large codebook may degrade image generation performance. This underscores the necessity for dataset-specific tuning in clustering-based VQ approaches. In contrast, SVQ naturally balances detail preservation and noise reduction using sparse regression, negating the need for such fine-tuning.The codebook size in our approach exhibits a low-maintenance profile: performance remains solid across various codebook sizes, suggesting that utilizing a default large codebook can yield robust results without the need for extensive tuning. We are not suggesting that our SVQ outperforms other VQ methods in general image generation tasks, as this is beyond the scope of our current objective. Rather, we emphasize SVQ’s effectiveness as a noise reduction tool that directly enhances real-world spatio-temporal forecasting tasks, while the application to general image generation remains a topic for future exploration.

**Table 4: Comparison of vector quantization methods: All methods share identical backbone, with the recommended setting as described in Appendix A.3. The performance enhancements over baseline are highlighted in bold.**
| Method | MSE$\downarrow$ | MAE$\downarrow$ |
| --- | --- | --- |
| Baseline (SimVP w/o VQ) | 1.105 | 0.6567 |
| VQ(van den Oord et al., ) | 1.854 | 0.8963 |
| Residual VQ(Zeghidour et al., ) | 1.213 | 0.6910 |
| Grouped Residual VQ(Yang et al., ) | 1.174 | 0.6747 |
| Multi-headed VQ(Mama et al., ) | 1.211 | 0.6994 |
| Residual VQ (Stochastic)(Lee et al., ) | 1.888 | 0.9237 |
| Residual Finite Scalar Quantization(Mama et al., ) | 1.319 | 0.7505 |
| Lookup Free Quantization (LFQ)(Yu et al., ) | 2.988 | 1.1103 |
| Residual LFQ(Yu et al., ) | 1.281 | 0.7281 |
| SVQ-raw(Xiao et al., ) | 1.123 | 0.6456 |
| SVQ | 1.018 | 0.6109 |

Figure: Figure 4: Predition MSE curves on WeatherBench-S temperature dataset with Grouped Residual VQ and SVQ.
Refer to caption: extracted/5396914/graphs/compare_balance_vq.png

### 5.4 Robustness to artificial noise injection

To clearly demonstrate the noise reduction effect of SVQ, we conduct an experiment by adding artificial noise to the training data. Table 5 shows that the SVQ-equiped model experiences a much lower rise in MSE and MAE relative to the model without SVQ, regardless of the proportion of injected noise. For instance, when the proportion of injected noise is 10%, the MSE of the model without SVQ rises by 25.4%, whereas the model with SVQ shows a modest increase of just 3.2%. These results explicitly confirm that our method effectively mitigates noise by constraining latent patterns via quantization.

**Table 5: Noise injection analysis: The proportion of injected noise is indicated by $\eta$. We present MSE and MAE, and their percentage increase over baseline without artificial noise.**
| Noise proportion | MSE | MAE |  |  |
| --- | --- | --- | --- | --- |
| w/o SVQ | w SVQ | w/o SVQ | w SVQ |  |
| $\eta$=0 | 1.105 | 1.018 | 0.6567 | 0.6109 |
| $\eta$=10% | 1.386(+25.4%) | 1.051(+3.2%) | 0.7702(+17.3%) | 0.6269(+2.6%) |
| $\eta$=20% | 1.554(+40.7%) | 1.196(+17.5%) | 0.8282(+26.1%) | 0.6710(+9.8%) |
| $\eta$=30% | 1.750(+58.4%) | 1.255(+23.2%) | 0.8821(+34.3%) | 0.6953(+13.8%) |
| $\eta$=40% | 2.081(+88.3%) | 1.568(+54.0%) | 1.0031(+52.7%) | 0.7973(+30.5%) |
| $\eta$=50% | 2.646(+139.4%) | 1.529(+50.2%) | 1.1339(+72.7%) | 0.7881(+29.0%) |

Figure: Figure 5: Distribution of regression weight $\mathcal{W}$ and codebook $\mathcal{M}$: Higher kurtosis represents more compact and concentrate distribution near zero, as well as sparser regression weights. Left: Learnable SVQ with MAE loss. Middle: Learnable SVQ with MSE loss. Right: Frozen SVQ with MAE loss. Learnable setting and MAE loss encourage sparser weights and a more structured codebook.
Refer to caption: x3.png

### 5.5 Ablation study

We conduct a series of ablation studies on WeatherBench-S temperature dataset to understand the contribution of important designs based on the default setting: SVQ with 10000 codebook size, learnable codebook, and MAE loss.

Codebook size and learnability. We compare the effect of codebook size when codebook is learnable or frozen (Table 6). The results indicate that that increasing codebook size gradually improves performance. When significantly increasing codebook size to a large value, such as 10000, the performance difference between frozen and learnable codebook narrows, evidenced by a mere 0.5% (1.023 → 1.018). We observe that enlarging the codebook ensures comprehensive latent space coverage through random codes, negating the need for meticulous learning and allowing random initializations to rival learned counterparts. Additionally, our tailored SVQ structure surpasses alternative designs, including single-layer, bucket-shaped, and post-ReLU variants, while maintaining a similar parameter count.

**Table 6: Ablation of model structure.**
| Learnability | Projection dim | Codebook size | MSE$\downarrow$ | MAE$\downarrow$ |
| --- | --- | --- | --- | --- |
| Frozen | 128 | 10 | 1.070 | 0.6227 |
| 128 | 1000 | 1.044 | 0.6198 |  |
| 128 | 10000 | 1.023 | 0.6131 |  |
| Learnable | 128 | 10 | 1.060 | 0.6194 |
| 128 | 1000 | 1.048 | 0.6182 |  |
| 128 | 10000 | 1.018 | 0.6109 |  |
| 1280(Bucket-shape) | 1280 | 1.035 | 0.6149 |  |
| None(One-layer) | 10000 | 1.043 | 0.6144 |  |
|  | 128(Post-ReLU) | 10000 | 1.032 | 0.6136 |

Self-learned sparse regression structure. The original SimVP model adopts MSE as prediction loss. We individually replace it with MAE loss and add the SVQ module. The results show that the joint use of SVQ and MAE loss is crucial for significantly improving the model’s performance (Table 7). We suggest that the sparsity of the weight matrix $\mathcal{W}$ impacts vector representation learning and use kurtosis to quantify this after normalizing $\mathcal{W}$. Figure 5 demonstrates that both a learnable codebook and MAE loss contribute to increased sparsity. Analyzing four codebook initialization methods, in both learnable and fixed settings (Table 8), we find that a learnable codebook promotes sparsity irrespective of the initialization, indicating that sparsity is a self-learned property that enhances intermediate representation learning.

**Table 7: Module ablation.**
| Module | MSE$\downarrow$ | MAE$\downarrow$ |
| --- | --- | --- |
| SimVP (MSE loss) | 1.105 | 0.6567 |
| SimVP (MAE loss) | 1.126 | 0.6509 |
| SimVP+SVQ (Learnable, MSE loss) | 1.099 | 0.6527 |
| SimVP+SVQ (Learnable, MAE loss) | 1.018 | 0.6109 |

**Table 8: Ablation of codebook initialization method.**
| Initialization | Learnability | MSE$\downarrow$ | MAE$\downarrow$ | Kurtosis |
| --- | --- | --- | --- | --- |
| kaiming uniform | Frozen | 1.023 | 0.6131 | 1.596 |
| Learnable | 1.018 | 0.6109 | 7.213 |  |
| sparse(sparsity=0.9) | Frozen | 1.050 | 0.6183 | 4.165 |
| Learnable | 1.034 | 0.6160 | 41.558 |  |
| trunc normal | Frozen | 1.049 | 0.6166 | 1.582 |
| Learnable | 1.031 | 0.6161 | 4.236 |  |
| orthogonal | Frozen | 1.034 | 0.6170 | 1.561 |
| Learnable | 1.030 | 0.6131 | 35.774 |  |

## 6 Conclusions

In this work, we present Sparse regression-based Vector Quantization (SVQ), a concise yet effective method for spatiotemporal forecasting enhancement. SVQ elegantly balances detail preservation and noise reduction using sparse regression, realized via a two-layer MLP and a codebook, offering computational efficiency and full differentiability. Tested across diverse benchmarks, from weather to traffic and video prediction, SVQ consistently outperforms existing models, setting new performance standards without complex priors. Its differentiability and seamless integration with baseline models highlight SVQ as a significant advancement for efficient and effective spatio-temporal forecasting.

## Impact Statements

This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

## Appendix A Implementation details

### A.1 Dataset details

The details of the benchmark datasets are summarized as follows: WeatherBench (Rasp et al., ) and TaxiBJ (Zhang et al., ) are two macro forecasting tasks. These two datasets are collected at low frequencies (30min or 1-6h) and exhibit stable states, which is especially challenging for the model to capture subtle changes. Human3.6M (Ionescu et al., ), KittiCaltech (Geiger et al., ; Dollár et al., ), and KTH Action (Schüldt et al., ) are three popular video prediction tasks. Due to the diversity and complexity of real-world videos, these datasets have totally different characteristics. A summary of dataset statistics is provided in Table 9.

**Table 9: The detailed statistics of benchmark datasets.**
| Dataset | Size | Seq. Len. | Img. Shape | Interval |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| train | test | in | out | H $\times$ W $\times$ C |  |  |
| WeatherBench-S | 2,167 | 706 | 12 | 12 | 32 $\times$ 64 $\times$ 1 | 1 hour |
| WeatherBench-M | 54,019 | 2,883 | 4 | 4 | 32 $\times$ 64 $\times$ 4 | 6 hour |
| TaxiBJ | 20,461 | 500 | 4 | 4 | 32 $\times$ 32 $\times$ 2 | 30 min |
| KittiCaltech | 3,160 | 3,095 | 10 | 1 | 128 $\times$ 160 $\times$ 3 | Frame |
| Human3.6M | 73,404 | 8,582 | 4 | 4 | 256 $\times$ 256 $\times$ 3 | Frame |
| KTH Action | 4,940 | 3,030 | 10 | 20 | 128 $\times$ 128 $\times$ 1 | Frame |

### A.2 Architecture configuration of SimVP

Table 10 reports the architectures of SimVP on all datasets. We select the best MetaFormer to replace the translator module based on OpenSTL benchmarks(^1^11[https://openstl.readthedocs.io/en/latest/model_zoos/video_benchmarks.html](https://openstl.readthedocs.io/en/latest/model_zoos/video_benchmarks.html))(^2^22[https://openstl.readthedocs.io/en/latest/model_zoos/weather_benchmarks.html](https://openstl.readthedocs.io/en/latest/model_zoos/weather_benchmarks.html))(^3^33[https://openstl.readthedocs.io/en/latest/model_zoos/traffic_benchmarks.html](https://openstl.readthedocs.io/en/latest/model_zoos/traffic_benchmarks.html)). The parameters remain unchanged, following the original configurations. It should be noted that, due to the reproducibility issue of ConvNeXt on the TaxiBJ dataset, we have opted to utilize gSTA as our backbone model.

**Table 10: Detailed configuration of SimVP backbone.**
| Dataset | MetaFormer (Translator) | spatio_kernel | hid_S | hid_T | N_T | N_S | drop_path | LR scheduler |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| WeatherBench-S temperature | gSTA | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.1 | cosine |
| WeatherBench-S humidity | Swin | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.2 | cosine |
| WeatherBench-S wind component | Swin | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.2 | cosine |
| WeatherBench-S total cloud cover | gSTA | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.1 | cosine |
| WeatherBench-M | MogaNet | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.1 | cosine |
| TaxiBJ | gSTA | enc=3, dec=3 | 32 | 256 | 8 | 2 | 0.1 | cosine |
| Human3.6M | gSTA | enc=3, dec=3 | 64 | 512 | 6 | 4 | 0.1 | cosine |
| KTH | IncepU | enc=3, dec=3 | 64 | 256 | 6 | 2 | 0.1 | onecycle |
| KittiCaltech | gSTA | enc=3, dec=3 | 64 | 256 | 6 | 2 | 0.2 | onecycle |

### A.3 Parameters of compared VQ methods

In Table 4, we compared SVQ with a series of well-known VQ methods. These methods are reproduced using the source code provided on the GitHub repository(^4^44[https://github.com/lucidrains/vector-quantize-pytorch/tree/master](https://github.com/lucidrains/vector-quantize-pytorch/tree/master)). We maintain the parameters consistent with the recommended settings to guarantee performance, as detailed in Table 11. Importantly, to maintain comparable complexity among the methods, we enable a shared codebook for all quantizers in Residual VQ, as well as across the heads in Multi-headed VQ. Moreover, we also conduct an extensive additional experiment for VQ methods using the same codebook size in Section B.3.

**Table 11: Parameters of the compared VQ methods.**
| Vector quantization method | codebook_size | num_quantizers | groups | heads | shared_codebook | Specific parameters |
| --- | --- | --- | --- | --- | --- | --- |
| VQ | 512 | - | - | - | - | - |
| Residual VQ | 1024 | 8 | - | - | $\checkmark$ | - |
| Grouped Residual VQ | 1024 | 8 | 2 | - | $\checkmark$ | - |
| Multi-headed VQ | 1024 | - | - | 8 | $\checkmark$ | - |
| Residual VQ (Stochastic) | 1024 | 8 | - | - | $\checkmark$ | stochastic_sample_codes=True |
| Residual Finite Scalar Quantization | - | 8 | - | - | - | levels=[8, 5, 5, 3] |
| Lookup Free Quantization (LFQ) | 8192 | - | - | - | - | entropy_loss_weight=0.1 |
| Residual LFQ | 256 | 8 | - | - | - | - |

### A.4 Evaluation of perplexity

In contrast to other VQ methods that rely on a single code, our SVQ generates multiple regression weights to merge several codes. To evaluate its perplexity, we first normalize the regression weights and then convert them into binary form using a threshold set at $\theta$ times the standard deviation, where $\theta$ serves as the threshold value. We utilize two thresholds (2 and 3) to obtain reasonable perplexity.

### A.5 Metrics

Forecasting accuracy is evaluated using mean squared error (MSE) and mean absolute error (MAE), while the image quality of predicted frames is assessed using structural similarity index measure (SSIM) (Wang et al., ), peak signal-to-noise ratio (PSNR), and learned perceptual image patch similarity (LPIPS) (Zhang et al., ). The training process is early stopped with a patience of 10, and the models with the minimal loss are saved for subsequent evaluation.

## Appendix B Additional quantitative results

### B.1 Position of quantization module

To illustrate how the position of quantization module affects representation learning, we consider two designs as shown in Figure 6. In our main experiments, we adopt the first design where quantization is performed before the translator. In this section, we also investigate an alternative design where quantization is performed after the translator. Interestingly, we found that the classic VQ method (van den Oord et al., ) suffers pronounced instability and codebook collapse issues under the latter design, as shown in Figure 7. It is essential to highlight that the backbones without VQ maintain their MSE within the acceptable range of approximately 1 to 2 (refer to Table 3). Yet, integrating classic VQ causes a substantial rise in MSE values, exceeding 10 for different backbones—a level considered excessively high. We hypothesize that this instability is attributed to the non-differentiability of the straight-through estimator, which introduces errors into the gradient flow for preceding modules. In contrast, our SVQ module never encounters such issues and remains highly stable throughout training. To maintain the integrity of all VQ methods, we opt for the first design in our principal manuscript, wherein quantization is executed preceding the translator module.

Figure: Figure 6: Comparison of two quantization designs with different positions. Top: Quantization before translator. Bottom: Quantization after translator.
Refer to caption: x4.png

Figure: Figure 7: VQ (Top) and SVQ (Bottom) training curves. We perform post-translator quantization on various backbones, with the same codebook size (1024), employing early stopping (patience of 10) on the WeatherBench-S temperature dataset. Perplexity for SVQ is averaged over different $\theta$ values, detailed in Appendix A.4.
Refer to caption: x5.png

### B.2 Computational cost

Table 12 presents the computational costs of SVQ module and forecasting models. It shows that recurrent-based models have significantly higher FLOPS requirements, while non-recurrent models are more efficient. The proposed SVQ module is not only effective but also computationally cheap. Across all datasets, SVQ only slightly adds the number of parameters and FLOPS. The computational burden of SimVP+SVQ remains significantly smaller than recurrent-based models.

**Table 12: Number of parameters and computing performances for all forecasting models.**
| Model type | Dataset | Human3.6M | KTH | KittiCaltech | WeatherBench-S | TaxiBJ |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Model | Params | FLOPS | Params | FLOPS | Params | FLOPS | Params | FLOPS | Params | FLOPS |  |
| Recurrent-based | ConvLSTM | 15.5M | 347.0G | 14.9M | 1368.0G | 15.0M | 595.0G | 14.98M | 136G | 14.98M | 20.74G |
| E3D-LSTM | 60.9M | 542.0G | 53.5M | 217.0G | 54.9M | 1004G | 51.09M | 169G | 50.99M | 98.19G |  |
| PredNet | 12.5M | 13.7G | 12.5M | 3.4G | 12.5M | 12.5M | - | - | 12.5M | 0.85G |  |
| PhyDNet | 4.2M | 19.1G | 3.1M | 93.6G | 3.1M | 40.4G | 3.09M | 36.8G | 3.09M | 5.60G |  |
| MAU | 20.2M | 105.0G | 20.1M | 399.0G | 24.3M | 172.0G | 5.46M | 39.6G | 4.41M | 6.02G |  |
| MIM | 47.6M | 1051.0G | 39.8M | 1099.0G | 49.2M | 1858G | 37.75M | 109G | 37.86M | 64.10G |  |
| PredRNN | 24.6M | 704.0G | 23.6M | 2800.0G | 23.7M | 1216G | 23.57M | 278G | 23.66M | 42.40G |  |
| PredRNN++ | 39.3M | 1033.0G | 38.3M | 4162.0G | 38.5M | 1803G | 38.31M | 413G | 38.40M | 62.95G |  |
| PredRNN.V2 | 24.6M | 708.0G | 23.6M | 2815.0G | 23.8M | 1223G | 23.59M | 279G | 23.67M | 42.63G |  |
|  | DMVFN | - | - | 3.5M | 0.88G | 3.6M | 1.2G | - | - | 3.54M | 0.057G |
| Non-recurrent | TAU | 37.6M | 182.0G | 15.0M | 73.8G | 44.7M | 80.0G | 12.22M | 6.70G | 9.55M | 2.49G |
| SimVP (w/o VQ) | 28.8M | 146.0G | 12.2M | 62.8G | 15.6M | 96.3G | 12.76M | 7.01G | 7.84M | 2.08G |  |
| SimVP+SVQ | 30.7M | 178.0G | 13.3M | 110.0G | 16.8M | 156G | 14.37M | 16.8G | 9.45M | 3.72G |  |

### B.3 Comprehensive comparison of VQ methods using the same codebook size

To make a fair comparison, we extend Table 4 by setting the codebook size to the same value (1024) for compared VQ methods. They are comprehensively evaluated from different aspects including downstream performance (prediction MSE), codebook usage (perplexity), and computational complexity (FLOPS, inference FPS, and training time per epoch). The quantitative results are shown in Table 13. The convergence performance is shown in Figure 8, where SVQ quickly converges to the lowest prediction error and satisfactory utilization of the codebook. Residual VQ with stochastic sampling has the highest codebook usage. However, its prediction MSE is worse than residual VQ without stochastic sampling. This demonstrates that codebook usage does not guarantee better downstream performance. SVQ generally outperforms the other VQ methods in computational efficiency, due to the simplified approximation described in Section 3.1.

**Table 13: Quantitative comparison of VQ methods with the same codebook size (1024) on WeatherBench-S temperature dataset.**
| Vector quantization method | Prediction MSE$\downarrow$ | Perplexity$\uparrow$ | FLOPS$\downarrow$ | Inference FPS$\uparrow$ | Training time per epoch(min)$\downarrow$ |
| --- | --- | --- | --- | --- | --- |
| VQ | 1.8544 | 51.95 | 7.207G | 21.1 | 7.11 |
| Residual VQ | 1.2131 | 142.47 | 8.616G | 7.7 | 13.25 |
| Residual VQ (Stochastic) | 1.8882 | 817.91 | 8.616G | 8.1 | 17.27 |
| Grouped Residual VQ | 1.1737 | 132.57 | 8.616G | 4.8 | 19.98 |
| Multi-headed VQ | 1.2113 | 16.36 | 8.717G | 6.2 | 13.15 |
| SVQ (freeze) | 1.0393 | 335.72($\theta$=3)/438.39($\theta$=2) | 8.037G | 24.6 | 7.27 |
| SVQ (Learnable) | 1.0403 | 246.44($\theta$=3)/331.41($\theta$=2) | 8.037G | 24.9 | 7.30 |

Figure: Figure 8: Prediction error and codebook usage of different VQ methods during the training process. All methods adopt the same codebook size (1024) and are early stopped with a patience of 10 on WeatherBench-S temperature dataset. For simplicity, the perplexity of SVQ is averaged on different $\theta$.
Refer to caption: extracted/5396914/graphs/compare_vq.png

### B.4 Benchmark on TaxiBJ dataset

**Table 14: Performance comparison for SVQ and baseline models on TaxiBJ.**
| Model | MSE$\downarrow$ | MAE$\downarrow$ | SSIM$\uparrow$ | PSNR$\uparrow$ |
| --- | --- | --- | --- | --- |
| ConvLSTM(SHI et al., ) | 0.3358 | 15.32 | 0.9836 | 39.45 |
| E3D-LSTM(Wang et al., ) | 0.3427 | 14.98 | 0.9842 | 39.64 |
| PhyDNet(Guen & Thome, ) | 0.3622 | 15.53 | 0.9828 | 39.46 |
| PredNet(Lotter et al., ) | 0.3516 | 15.91 | 0.9828 | 39.29 |
| PredRNN(Wang et al., ) | 0.3194 | 15.31 | 0.9838 | 39.51 |
| MIM(Wang et al., ) | 0.3110 | 14.96 | 0.9847 | 39.65 |
| MAU(Chang et al., ) | 0.3268 | 15.26 | 0.9834 | 39.52 |
| DMVFN(Hu et al., ) | 3.3954 | 45.52 | 0.8321 | 31.14 |
| PredRNN++(Wang et al., ) | 0.3348 | 15.37 | 0.9834 | 39.47 |
| PredRNN.V2(Wang et al., ) | 0.3834 | 15.55 | 0.9826 | 39.49 |
| TAU(Tan et al., ) | 0.3108 | 14.93 | 0.9848 | 39.74 |
| SimVP (w/o VQ)(Tan et al., ) | 0.3246 | 15.03 | 0.9844 | 39.71 |
| SimVP+SVQ (Frozen) | 0.3171 | 14.68 | 0.9848 | 39.83 |
| SimVP+SVQ (Learnable) | 0.3191 | 14.64 | 0.9849 | 39.86 |
| Improvement | $\uparrow$1.7% | $\uparrow$2.6% | $\uparrow$0.1% | $\uparrow$0.4% |

### B.5 Ablation of frozen module

The SVQ module consists of a two-layer MLP and a large codebook. The MLP can be seen as a projection from the input vector to the regression weights. In Section 5.5, we have already examined the impact of freezing the codebook on forecasting performance. To further investigate the effect of freezing the two-layer MLP, we conducted an ablation study in this section. The results, presented in Table 15, show that freezing the codebook only has a slight impact on forecasting performance, while freezing the MLP significantly impairs the performance.

**Table 15: Ablation of frozen modules on WeatherBench-S temperature dataset.**
| Frozen module | None (All learnable) | Codebook | Two-layer MLP projection | Both |
| --- | --- | --- | --- | --- |
| MSE | 1.018 | 1.023 | 1.060 | 1.093 |
| MAE | 0.6109 | 0.6131 | 0.6194 | 0.6387 |

## Appendix C Effect of SVQ on latent representation

We investigated the impact of SVQ on the sparsity of regression weights in Figure 5. To further investigate its effect on the latent representation, we compare the distribution of batch tensors before and after applying SVQ. We transform the tensors into normalized vectors and estimate their density distributions. As depicted in Figure 9, the representation after SVQ demonstrates a more compact distribution, indicating improved robustness to noise. The results further prove that SVQ can enhance forecasting performance by effectively handling noise in the data.

Figures 10, 11, 13, 14, and 12 present the comparison of latent feature maps before and after applying SVQ. These figures illustrate that the difference between foreground and background in the feature maps increases after SVQ. For example, in the KittiCaltech dataset, a clear distinction is observed between road conditions and sky (Figure 10). Similarly, in the WeatherBench-S temperature dataset, distinctive regions are identified between high and low latitudes (Figure 11). These findings suggest that SVQ helps in enhancing the discriminative power of the latent representations, which in turn contributes to improved downstream forecasting performance.

Figure: Figure 9: Distribution of latent vector on WeatherBench-S temperature dataset.
Refer to caption: extracted/5396914/graphs/svq_effect_weather_density.png

Figure: Figure 10: Latent feature map on the KittiCaltech dataset: (a) feature map before SVQ, and (b) feature map after SVQ.
Refer to caption: x6.png

Figure: Figure 11: Latent feature map on the WeatherBench-S temperature dataset: (a) feature map before SVQ, and (b) feature map after SVQ.
Refer to caption: x7.png

Figure: Figure 12: Latent feature map on the TaxiBJ dataset: (a) feature map before SVQ, and (b) feature map after SVQ.
Refer to caption: x8.png

Figure: Figure 13: Latent feature map on the KTH dataset: (a) feature map before SVQ, and (b) feature map after SVQ.
Refer to caption: x9.png

Figure: Figure 14: Latent feature map on the Human3.6M dataset: (a) feature map before SVQ, and (b) feature map after SVQ.
Refer to caption: x10.png

## Appendix D Additional qualitative results

### D.1 Forecasting errors on WeatherBench and TaxiBJ datasets

Figure: Figure 15: The qualitative forecasting errors on WeatherBench-S temperature dataset.
Refer to caption: x11.png

Figure: Figure 16: The qualitative forecasting errors on TaxiBJ dataset.
Refer to caption: x12.png

### D.2 Additional forecasting samples

Figure: Figure 17: Forecasting samples of SimVP+SVQ model on the test set of TaxiBJ (32$\times$32), WeatherBench (64$\times$32), KittiCaltech (160$\times$128), KTH (128$\times$128), and Human3.6M (256$\times$256). Zoom in for details. Our model produces accurate predicted frames for different tasks.
Refer to caption: x13.png

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 1
questions: 1. What is the motivation of using vector quantization into spatiotemporal prediction?
2. What is the significance of theoretical analysis in Chapter 4? Is this theoretical analysis related to video prediction?
3. What are the ''improvement'' in Table 1.2.3. refer to, SimVP?
4. Can the method proposed in the paper be compared with diffusion based models? For example, ExtDM: Distribution Extrapolarization Diffusion Model for Video Prediction, CVPR2024. What are the differences between these two methods, e.g., their application scenarios or efficiency?
5. Why are there different types of comparison results between WeatherBench-S and WeatherBench-M in Tab. 1 (Total Cloud Cover in WeatherBench-S and Wind UV in WeatherBench-M)? Why not compare the same subjects? What are the differences in SVQ performance across different physical quantities and data scales?
soundness: 2
strengths: The paper shows originality through the development of SVQ, a novel combination of sparse regression and soft vector quantization for spatio-temporal forecasting with theoretical analysis. Empirical validation is extensive, with experiments on multiple real - world datasets and comparisons to existing methods, achieving state-of-the-art results and validating the method's effectiveness and quality. It has potential applications in various domains and can inspire future research, opening new avenues for exploration and providing insights for model development.
summary: This paper identifies the limited performance of traditional vector quantization (VQ) in spatiotemporal forecasting due to non-differentiability and limited representation power. It proposes Differentiable Sparse Soft - Vector Quantization (SVQ), which approximates sparse regression with a two-layer MLP for differentiability and uses Soft-VQ with sparse regression to capture patterns and filter noise. Empirical results on multiple datasets show SVQ achieves state-of-the-art performance, is versatile as a plug-in, and effectively balances detail and noise reduction. Ablation studies and additional analyses confirm its key components’ importance and its robustness.
weaknesses: 1. The novelty of the approach is limited, and the performance improvement appears marginal. Furthermore, there is no discussion on the additional computational overhead that these slight performance enhancements (such as those observed in wind speed and humidity in Table 1, and in the KTH dataset in Table 2, as well as in Table 18) will incur. A more detailed analysis of the computational overhead introduced by SVQ, compared to baseline methods, is needed, especially for cases where the performance gains are smaller.

2. The motivation is unclear. It is not surprising to use huge codebook and sparse representation to improve the effect, because huge codebook itself brings a lot of extra parameter redundancy. So why use vector quantization? Because in other fields (e.g., video compression, video generation), VQ is to compress redundant information, not to add redundant information. You could further explain their rationale for using vector quantization in this context, given its typical use for compression in other fields.

3. The theoretical analysis provided seems unrelated to the content of the article. Furthermore, the article fails to discuss the relationship between the information or features extracted after compression using SVQ and the original spatiotemporal data. Consequently, there is a notable lack of corresponding theoretical discussion. A more in-depth exploration of how the features learned through SVQ are related to the original spatiotemporal data would greatly aid in fully elucidating the mechanism underlying SVQ.

4. The ablation study conducted is insufficient. There is no doubt that setting the code size to 10000 will yield better performance compared to 1000. A more detailed discussion of the trade-offs involved (such as efficiency, convergence, etc.) with larger code sizes would be helpful.

5. The introduction of redundant over-complete codebooks and additional computational overhead has resulted in a lack of discussion on computational efficiency, speed, and complexity, among other factors. Empirical measurements of training and inference times, memory usage, and computational complexity, as a function of codebook size, would provide a more comprehensive illustration of the advantages of SVQ.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1.Will the SVQ module added as a plug-in to the model have similar performance improvements for other tasks (such as image generation or natural language processing)? What are the applicable application scenarios?
soundness: 3
strengths: 1. The manuscript proposes a differentiable sparse soft vector quantization (SVQ) method, which is the first vector quantization method applied to spatiotemporal prediction and shows significant improvement.
2. The SVQ method proposed in the manuscript has achieved leading performance in multiple real-world spatiotemporal prediction tasks, significantly reducing errors on multiple benchmark datasets, such as reducing errors by 7.9% on the WeatherBench dataset.
3. The SVQ proposed in the manuscript can be seamlessly integrated into different types of spatiotemporal prediction models as a plug-in, and has improved performance in various architectures, demonstrating the versatility of the method.
summary: Vector quantization (VQ) is insufficient in improving the accuracy of spatiotemporal prediction. This paper introduces differentiable sparse soft vector quantization (SVQ) that can strike a balance between detail preservation and noise reduction, providing a solid foundation for full differentiability and sparse regression. Empirical studies on five spatiotemporal benchmark datasets show that SVQ achieves the best results, including a 7.9% improvement on the WeatherBench-S temperature dataset, a 9.4% average MAE reduction in video prediction benchmarks (Human3.6M, KTH, and KittiCaltech), and a 17.3% improvement in image quality (LPIPS).
weaknesses: 1.The SVQ method proposed in the manuscript still requires a lot of computing resources, especially in the case of high-dimensional data and large-scale codebooks.
2.The comparison methods cited by the author in Tables 1 and 2 are only up to date in 2022, and lack comparisons of the latest methods in the past two years.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: See weakness
soundness: 2
strengths: The proposed Differentiable Sparse Soft-Vector Quantization (SVQ) method represents a novel advancement in vector quantization techniques.
summary: The paper introduces a novel Differentiable Sparse Soft-Vector Quantization (SVQ) method, which integrates sparse regression with differentiability to tackle optimization challenges in vector quantization. This approach aims to enhance representation capacity in spatio-temporal forecasting tasks, marking a significant advancement in the field. While the SVQ method presents an innovative approach to vector quantization, the paper would benefit from clearer connections between theory and practice, updated baselines, deeper integration insights, and improved mathematical clarity to strengthen its contributions to the field.
Soundness:2
weaknesses: The method does not adequately demonstrate how the theoretical advantages of sparse regression are translated into tangible improvements in quantization performance. While the authors discuss optimization strategies, they fail to provide a clear connection between these theoretical claims and the practical outcomes. A more comprehensive explanation of how these optimization techniques directly enhance quantization would bolster the credibility of their approach.
	In the experimental section, I noticed that the baseline models employed are relatively outdated. Given the recent advancements in spatio-temporal forecasting, particularly the emergence of various diffusion-based methods that have demonstrated significant improvements in predictive performance, it would be beneficial for the authors to consider incorporating these state-of-the-art models as baselines. This would provide a more comprehensive evaluation of the proposed method's effectiveness and advantages.
	The explanation of the quantization module's implementation lacks depth regarding its integration with the overall spatio-temporal forecasting model. While the authors outline the architecture and components involved, they do not provide sufficient details on how the quantization process interacts with other model elements or influences the final forecasting results. A more thorough exploration of these interactions would enhance the clarity and applicability of their proposed method.
	In Section 4, the mathematical proof lacks clarity in the notation used, which may hinder readers' understanding. For example, what’s the meaning of g' after Eq. (8)? Additionally, the proof does not establish a strong connection to the problem being addressed. I recommend revising this section to improve the clarity of the symbols and to explicitly link the proof to the main objectives of the paper.

---

confidence: 5
contribution: 4
initial_rating: 6
presentation: 4
questions: 1. The paper emphasizes the advantages of the proposed method in spatio-temporal forecasting, but VQ is also widely used in generative tasks (e.g., VQ-VAE). Could this method be applied to such tasks?
soundness: 3
strengths: 1. Innovative Approach: The paper effectively combines sparse regression with differentiable quantization, addressing the non-differentiability and limited representational power of traditional VQ. Using MLP to approximate sparse regression allows the model to capture complex patterns efficiently.  
2. Simplicity and Effectiveness: The proposed method is intuitive and easy to implement, with straightforward derivations and motivations. It demonstrates significant improvements across multiple tasks and models.  
3. Comprehensive Experiments: The paper provides detailed evaluations of the proposed quantization mechanism, including ablation studies and supplementary materials that address key questions. The well-designed visualizations offer excellent insights into the behavior and strengths of SVQ.
summary: This paper addresses the limitations of traditional Vector Quantization (VQ) and demonstrates impressive performance in spatio-temporal forecasting. SVQ uses a two-layer MLP to approximate sparse regression, reducing computational complexity while maintaining the flexibility to map each input vector to multiple codebook vectors. This soft quantization captures the complex dynamics of spatio-temporal data, preserving essential information while minimizing noise. The experiments confirm that SVQ is an efficient and expressive quantization mechanism applicable to various forecasting tasks. The visualizations provide valuable insights into the behavior and advantages of SVQ.
weaknesses: 1. Visual Layout: Perhaps due to space constraints, the layout of the figures and tables could be more aesthetically pleasing.

Meta-review: The paper was previously submitted to NeruIPS 2024, with title “A Differentiable Sparse Soft-Vector Quantization (SVQ) for Spatio-Temporal Forecasting”, but got rejected due to some weaknesses. The authors didn’t make material improvements except chancing the title into “Does Vector Quantization Fail in Spatio-Temporal Forecasting? Exploring a Differentiable Sparse Soft-Vector Quantization Approach”. Hence, I would recommend rejecting the paper. 

Some comments for the authors to consider:

1.  The proposed Sparse Soft-Vector Quantization (SVQ) is not specific to spatio-temporal forecasting, but instead a general compression method that can be used for various tasks. The authors consider only spatio-temporal forecasting for experimental evaluation though. Since spatio-temporal forecasting is a comprehensive task that can be affected by many components, the presented results are not convincing to demonstrated the superiority of SVQ over its competitors. On the other hand, while evaluating the paper from the viewpoint of spatio-temporal forecasting, the link between SVQ and the problem space is unclear, as pointed out by Reviewer PA7x and Reviewer AkZ2.

2. It seems that the theorems presented in Section 4 are specific to the traditional SVQ, not the proposed differentiable SVQ. 

3. What do you mean by $R_+^m$? Is it to enforce all components being nonnegative? If so, the deductions in Eq(2) and Eq(3) are wrong.

Decision: Reject