TIMBA: Time series Imputation with Bi-directional Mamba Blocks and Diffusion models

## Abstract

Abstract The problem of imputing multivariate time series spans a wide range of fields, from clinical healthcare to multi-sensor systems. Initially, Recurrent Neural Networks (RNNs) were employed for this task; however, their error accumulation issues led to the adoption of Transformers, leveraging attention mechanisms to mitigate these problems. Concurrently, the promising results of diffusion models in capturing original distributions have positioned them at the forefront of current research, often in conjunction with Transformers. In this paper, we propose replacing time-oriented Transformers with State-Space Models (SSM), which are better suited for temporal data modeling. Specifically, we utilize the latest SSM variant, S6, which incorporates attention-like mechanisms. By embedding S6 within Mamba blocks, we develop a model that integrates SSM, Graph Neural Networks, and node-oriented Transformers to achieve enhanced spatiotemporal representations. Implementing these architectural modifications, previously unexplored in this field, we present Time series Imputation with Bi-directional mamba blocks and diffusion models (TIMBA). TIMBA achieves superior performance in almost all benchmark scenarios and performs comparably in others across a diverse range of missing value situations and three real-world datasets. We also evaluate how the performance of our model varies with different amounts of missing values and analyse its performance on downstream tasks. In addition, we provide the original code to replicate the results.

## 1 Introduction

The issue of missing values is widely recognized and can arise from various causes, including difficulties in data collection, problems with recording mechanisms, or human errors. In the field of Multivariate Time Series (MTS) analysis, where multiple variables are recorded, sometimes with irregular frequency, this problem is even more pronounced. It affects many fields, from defective sensor records that transmit erroneous data (Wu et al., ), to the Internet of Things (IoT) where these failures are even more common (Ahmed et al., ), and even clinical contexts (Moor et al., ). Properly filling these gaps is crucial, as failure to do so can distort the data distribution and reduce model performance (Solís-García et al., ).

In recent years, this problem has been tackled using Deep Learning (DL) techniques, employing Recurrent Neural Networks (RNNs) to capture temporal representations and Graph Neural Networks (GNNs) for spatial ones (Cini et al., ). However, the error accumulation problem in RNNs has led to their replacement by Transformers (Tashiro et al., ), which use attention mechanisms to avoid these issues. Furthermore, the latest state-of-the-art results have been achieved using diffusion models (Tashiro et al., ; Liu et al., ), which excel at modeling multiple distributions in the original data. Additionally, the use of State-Space Models (SSMs) has expanded in recent years, thanks to their discretization, as presented in S4 (Gu et al., ). These models, reminiscent of RNNs, have shown a high capacity to handle time series and long sequences, such as those found in speech processing, text, and even DNA sequencing (Gu et al., ). This architecture has been briefly explored in the context of multivariate time series imputation (MTSI) (Alcaraz & Strodthoff, ). Recently, S6, an improvement over S4 that incorporates an attention-like mechanism, was developed to potentially eliminate error accumulation issues, as demonstrated in selective copy tasks and speech processing (Gu & Dao, ). However, the S6 and Mamba blocks have not yet been tested in the context of MTSI.

Starting from the state-of-the-art models for MTSI, CSDI and PriSTI (Tashiro et al., ; Liu et al., ), in this article we replaced time-oriented Transformers in their architecture with S6 layers embedded within bi-directional Mamba blocks (Gu & Dao, ), carefully maintaining a similar parameter count for fair comparability. This combination results in our proposed method, Time series Imputation with bi-directional Mamba Blocks and Diffusion models (TIMBA).

Our main contributions are as follows: 1) We propose TIMBA, a model that replaces the time-oriented transformers in state-of-the-art diffusion models with bi-directional Mamba blocks, an approach not previously explored to the best of our knowledge. 2) Through an extensive benchmark using three real-world datasets, we demonstrate that TIMBA consistently achieves superior performance in almost all scenarios and performs comparably in others when compared to state-of-the-art models 3) Additionally, we expand the model’s analysis by presenting an ablation study, evaluating its sensitivity to different missing rates, and using it for data imputation followed by an evaluation on a downstream task. The rest of the paper is organized as follows: In Section 2, we review related work in the literature. Section 3 introduces the mathematical background supporting our approach. Section 4 presents TIMBA. In Section 5, we discuss the benchmark, experimental setup, and results obtained. Finally, Section 6 provides the conclusion. Additional information about the code and data can be found in Appendix A.

## 2 Related works

##### Multivariate time series imputation

MTSI has been explored through various approaches. Traditional solutions such as mean, zero, or linear imputation have been utilized to address missing values (Moor et al., ). However, these methods often distort the original data distribution, thus degrading the data quality. In contrast, Machine Learning (ML) offers a range of simple to sophisticated techniques, including k nearest neighbors (Beretta & Santaniello, ), linear predictors (Seaman et al., ), matrix factorization (MF) (Cichocki & Phan, ), vector autoregressive model-imputation (VAR) (Bashir & Wei, ), and multivariate Gaussian processes (MGP) (Li & Marlin, ), which have improved outcomes in many cases.

With the advent of DL, MTSI applications were soon discovered. Initially, Recurrent Neural Networks (RNNs) (Suo et al., ; Lipton et al., ) excelled due to their robust capacity to extract accurate temporal representations, crucial for effective imputation. Advancements in RNNs led to the adoption of architectures such as Bi-directional RNNs (BiRNNs), exemplified by BRITS (Cao et al., ), which analyze time series data both forwards and backwards, thereby enhancing imputation capabilities. However, recent years have seen the rise of Transformers as a dominant architecture, highlighting the limitations of RNNs. The attention mechanism in Transformers is particularly effective in identifying significant or real samples in time series, addressing the error accumulation issues prevalent in RNNs, as evidenced by several studies (Tashiro et al., ; Liu et al., ).

Lastly, the role of GNNs, which are based on graph theory, must be noted for their profound ability to extract spatial relationships among variables. Many researchers argue that for accurate imputation, it is essential to derive robust spatio-temporal representations, hence the integration of RNNs or Transformers with GNNs has become a focal point (Cini et al., ; Liu et al., ). To further enhance the effectiveness of MTSI, generative models and State-Space Models have emerged as powerful tools, offering unique advantages in handling complex data distributions and dynamics.

##### Generative models for imputation

Recently, there has been growing interest in applying generative methods to MTSI due to their ability to learn the original data distribution, making them highly suitable for addressing this issue. Variational Autoencoders (VAE) have been used in this context for several years (Fortuin et al., ). However, interest in generative techniques for imputation exploded following the introduction of Generative Adversarial Networks (GANs), with the groundbreaking work of Generative Adversarial Imputation Networks (GAIN) (Yoon et al., ) marking a significant milestone. Following this development, numerous GAN applications emerged in MTSI, often incorporating RNN layers to analyze temporal sequences (Luo et al., ; Miao et al., ).

Nevertheless, the issue of model collapse in GANs (Thanh-Tung & Tran, ), coupled with the emergence of Denoising Diffusion Probabilistic Models (DDPMs) and their superior ability to capture diverse and high-quality distributions, has led to a shift away from GANs. Research such as CSDI by Tashiro et al. () demonstrated how DDPMs, with the right conditional information, could efficiently address MTSI challenges. Further studies, such as those by Yun et al. (), confirmed that unconditional DDPMs might fail, whereas Liu et al. () enhanced DDPMs’ ability to generate better conditional information through specialized layers, furthering their effectiveness in MTSI applications.

##### State-Space Models

State-Space Models (SSMs) (Hangos et al., ), derived from control theory to model system dynamics, have been successfully adapted to the Deep Learning (DL) realm as Structured State Space Models (S4) (Gu et al., ). These models, similar to RNNs, have recently demonstrated their proficiency in modeling diverse sequences, leading to their application in text, audio, and even video-related problems (Gu et al., ). Like RNNs, however, they lack an attention mechanism, which is crucial for avoiding error accumulation during tasks such as imputation.

Nevertheless, the latest iteration, Selective State Space Models (S6), has incorporated an attention mechanism into the SSM structure, sparking a surge in applications. These models have been employed in selective copy tasks, language modeling, DNA modeling, and audio modeling and generation (Gu & Dao, ), occasionally outperforming traditional transformers. Specifically, in the field of time series imputation with DDPMs, the application of S4 has so far been limited to the study by (Alcaraz & Strodthoff, ), while S6 remains unexplored.

## 3 Background

### 3.1 Multivariate time series imputation

A multivariate time series (MTS) consists of $N_{t}$ variables or channels recorded at different time instants denoted by $t$. Consequently, these data can be represented as ${\bm{X}}_{t}\in{\mathbb{R}}^{N_{t}\times d}$, where each row $i$ contains the $d$-dimensional vector $x^{i}_{t}\in{\mathbb{R}}^{d}$, associating the $i$th variable with the time instant $t$. Additionally, a matrix ${\bm{M}}_{t}\in\{0,1\}^{N_{t}\times d}$ is used, where each entry contains a 0 if the corresponding value in ${\bm{X}}_{t}$ is missing, and a 1 if it is an observed data point.

Considering this, we define $\widetilde{{\bm{X}}}_{t}$ as the unknown ground truth variable-measure matrix, i.e., the complete time series without missing data; $\widehat{{\bm{X}}}_{t}$ denotes the time series imputed by the model, and $\mathcal{X}$ the series imputed by linear interpolation technique.

Finally, given that our work will focus on graph-based models, we model the time series as a sequence of graphs following the approach established by Cini et al. (). In this approach, each instant $t$ is defined as a graph $\mathcal{G}_{t}$ with $N_{t}$ nodes at each time instant. Moreover, the graph is defined as $\mathcal{G}_{t}=\langle{\bm{X}}_{t},\mathcal{A}_{t}\rangle$, where ${\bm{X}}_{t}$ represents the data matrix and $\mathcal{A}_{t}$ represents the adjacency matrix at instant $t$. However, this paper only considers the approach in which the graph topology does not change over time, thus $\mathcal{A}$ is constant for each $t$.

### 3.2 Denoising Diffusion Probabilistic Models

Denoising Diffusion Probabilistic Models (DDPM) rely on a Markov chain with time steps $t=1,\ldots,T$. Beginning with an initial state $x_{T}\sim\mathcal{N}(0,1)$, these models aim to reverse the sequence of latent variables $x_{t}$ in the same space as $x_{0}$, to learn $p_{\theta}(x_{0})$ that approximates the original data distribution $q(x_{0})$ (Ho et al., ).

The forward process or diffusion process involves progressively corrupting an initial sample $x_{0}$ by gradually adding Gaussian noise until reaching $x_{T}\sim\mathcal{N}(0,1)$. Noise addition is moderated by a variance scheduler $\beta_{t}=\beta_{1},\ldots,\beta_{T}$ to ensure distribution variance remains controlled and converges appropriately. This process is formalized by:

$$ $q(x_{1:T}|x_{0}):=\prod^{T}_{t=1}{q(x_{t}|x_{t-1})},\quad q(x_{t}|x_{t-1}):= \mathcal{N}(x_{t};\sqrt{1-\beta_{t}}x_{t-1},\beta_{t}I)$ (1) $$

Reparameterizing the parameters from Equation 1, we define $\overline{\alpha}_{t}=1-\beta_{t}$ and $\alpha_{t}:=\prod_{t=1}^{T}\overline{\alpha}_{t}$. Utilizing the transformation $\mathcal{N}(\mu,\sigma^{2})=\mu+\sigma\cdot\epsilon$, where $\epsilon\sim\mathcal{N}(0,1)$, allows us to derive a simplified model: $q(x_{t}|x_{0})=\sqrt{\overline{\alpha_{t}}}x_{0}+\sqrt{1-\overline{\alpha_{t}}}\epsilon$.

The reverse process, to recover the initial sample $x_{0}$, is defined by the following Markov chain:

$$ $\displaystyle p_{\theta}(x_{0:T}):=p_{\theta}(x_{T})\prod^{T}_{t=1}{p_{\theta} (x_{t-1}|x_{t})},\quad x_{T}\sim\mathcal{N}(0,1)$ (2) $\displaystyle p_{\theta}(x_{t-1}|x_{t}):=\mathcal{N}(x_{t-1};\mu_{\theta}(x_{t },t),\sigma_{\theta}(x_{t},t)I)$ $$

As in (Ho et al., ), we do not predict $\sigma_{\theta}$ as it remains fixed, controlled by the scheduler. $\mu_{\theta}$ is recalculated as:

$$ $\mu_{\theta}(x_{t},t)=\frac{1}{\sqrt{\alpha_{t}}}(x_{t}-\frac{\beta_{t}}{\sqrt {1-\overline{\alpha_{t}}}}\epsilon_{\theta}(x_{t},t))$ (3) $$

With the above considerations, in order to learn to model the original distribution, for each step of the chain $t$, and an $x_{t}$ generated according to $q(x_{t}|x_{0})$, our model must solve the following optimization problem defined by $\mathcal{L}(\theta):=\mathbb{E}_{t,x_{o},\epsilon}[||\epsilon-\epsilon_{\theta
}(x_{t},t)||^{2}]$.

### 3.3 State-Space Models and Mamba

State-Space Models (SSMs) are based on the theory of continuous control systems and process an input signal as defined in Equation 4 (Brogan, ).

$$ $\displaystyle\dot{h}(t)={\bm{A}}h(t)+{\bm{B}}x(t)$ (4) $\displaystyle y(t)={\bm{C}}h(t)+{\bm{D}}x(t)$ $$

An input signal $x(t)$ is mapped to a latent state $h(t)$ and then projected to an output signal $y(t)$. The matrices ${\bm{A}}$, ${\bm{B}}$, ${\bm{C}}$, and ${\bm{D}}$ determine the transformations involved in this process. Structured State Space Models (S4) (Gu et al., ) adapt SSMs by discretizing them. This involves modeling ${\bm{D}}$ as a skip-connection layer, defining a step size $\Delta=t_{n+1}-t_{n}$, and discretizing all matrices, where $\overline{{\bm{C}}}={\bm{C}}$, $\overline{{\bm{D}}}={\bm{D}}$, $\overline{{\bm{A}}}=({\bm{I}}-\Delta/2\cdot{\bm{A}})^{-1}({\bm{I}}+\Delta/2
\cdot{\bm{A}})$, and $\overline{{\bm{B}}}=({\bm{I}}-\Delta/2\cdot{\bm{A}})^{-1}\Delta{\bm{B}}$. This results in Equation 5, noting that ${\bm{D}}$ is not represented as it is modeled as a skip-connection.

$$ $\displaystyle h_{t}=\overline{{\bm{A}}}h_{t-1}+\overline{{\bm{B}}}x_{t}$ (5) $\displaystyle y_{t}=\overline{{\bm{C}}}h_{t}$ $$

Following S4, Selective State Space Models (S6) were introduced (Gu & Dao, ), improving on S4 by making $\Delta,{\bm{B}}$, and ${\bm{C}}$ dependent on a selection mechanism such that $\Delta\leftarrow\tau_{\Delta}(\text{Parameter}+s_{\Delta}(x))$, ${\bm{B}}\leftarrow s_{B}(x)$ and ${\bm{C}}\leftarrow s_{C}(x)$, where $s_{B}(x)$, $s_{C}(x)$ and $s_{\Delta}(x)$ are implemented as parameterized linear projections focusing on the hidden state of each $t$ in the original sequence $x$.

Finally, in Gu & Dao (), S6 is embedded within a Mamba block, which essentially consists of Multi-layer Perceptron (MLP) layers that serve to increase the hidden state size and generate two distinct information channels. In the first channel, a convolutional layer and S6 are applied, while the second channel applies a non-linear activation function to create a gate controlling the retention of information. Lastly, a MLP layer compresses everything back to its original dimensions.

## 4 TIMBA

Our approach builds upon the architecture established by Tashiro et al. () and Liu et al. (), which is centered on a diffusion model incorporating various blocks for refining data imputation.

### 4.1 Training DDPM for imputation

Although the general theory behind DDPMs was discussed in Section 3.2, these models need conditional information to perform effectively for general imputation, as proposed by Tashiro et al. () and demonstrated by Yun et al. ().

To train a DDPM for MTSI, key modifications include adapting the reverse process to incorporate this conditional information. Thus, we revise Equation 2 to include our conditional variables, which are $\mathcal{{\bm{X}}}$ and $\mathcal{{\bm{A}}}$:

$$ $\displaystyle p_{\theta}(x_{0:T}):=p_{\theta}(x_{T})\prod^{T}_{t=1}{p_{\theta} (x_{t-1}|x_{t},\mathcal{{\bm{X}}},\mathcal{{\bm{A}}})}$ (6) $\displaystyle p_{\theta}(x_{t-1}|x_{t},\mathcal{{\bm{X}}},\mathcal{{\bm{A}}}): =\mathcal{N}(x_{t-1};\mu_{\theta}(x_{t},\mathcal{{\bm{X}}},\mathcal{{\bm{A}}}, t),\beta_{t}I)$ $$

Following this, the optimization objective will be similar to that described in Section 3.2. However, as in Tashiro et al. (), we need to generate synthetic missing values (generate imputation targets $x_{0}\in M^{ta}$) where the imputation error is calculated. The model generates imputations for all missing values, but we filter the results and retain only those specifically marked as targets for error calculation. The optimization function then changes to that defined by Equation 7. Finally, it should be noted that there are many strategies for generating ${\bm{M}}^{ta}$, which are detailed further in Section 5.3.

$$ $\mathcal{L}(\theta):=\mathbb{E}_{t,x_{o},\epsilon}[{\bm{M}}^{ta}\cdot|| \epsilon-\epsilon_{\theta}(x_{t},\mathcal{{\bm{X}}},\mathcal{{\bm{A}}},t)||^{2}]$ (7) $$

### 4.2 Model architecture

Our architecture builds upon the foundational design by Tashiro et al. () and incorporates enhancements introduced by Liu et al. (). Referring to Figure 1, and following the information flow from left to right, we begin with the Conditional Feature Extraction Module (CFEM) (Liu et al., ). This module is responsible for generating conditional information that aids other components in improving imputation quality. The CFEM takes $\mathcal{{\bm{X}}}$ and $\mathcal{{\bm{A}}}$ as inputs and processes temporal data with bidirectional Mamba blocks, replacing the original transformers—a modification we propose and will justify subsequently. Spatial data is processed using transformer and Message Passing Neural Network (MPNN) blocks, inspired by Wu et al. (). This information is then refined with an MLP to produce $H^{pri}$, as illustrated in the CFEM bubble of Figure 1. The operation is defined as: $H^{pri}=\text{CFEM}(\mathcal{{\bm{X}}},\mathcal{{\bm{A}}})$.

Figure: Figure 1: Architecture of the TIMBA model
Refer to caption: extracted/5909928/red.png

Afterwards, a series of Noise Estimation Modules (NEM) (Tashiro et al., ) is employed to estimate the final noise. These modules receive $H^{pri}$, the diffusion time embedding $t$, $\mathcal{{\bm{A}}}$, and $H^{in}$. For the initial NEM, $H^{in}$ is a concatenation of noisy ${\bm{X}}_{t}$ and $\mathcal{{\bm{X}}}$. The output generates two information flows: one as input for the next NEM block (acting as $H^{in}$) and another as $H^{out}$, serving as the initial noise estimation. Outputs from each NEM are aggregated and processed by a convolutional layer to finalize the noise estimation. The internal workings of these blocks are akin to CFEM, as depicted in the NEM bubble in Figure 1, and described by: $H^{out}=\text{NEM}(H^{in},H^{pri},\mathcal{{\bm{A}}},t)$.

Now, we will discuss our decision to replace the transformer blocks, which focus on the temporal dimension, with Mamba blocks. Originally, this architecture utilized transformers for the temporal processing of information. However, despite their known advantages in MTSI, transformers lack an intrinsic inductive bias for temporal data. In contrast, Mamba blocks provide this bias while incorporating attention mechanisms, making them more suitable for accurately capturing temporal relationships. Therefore, we propose replacing time-focused transformers with Mamba blocks.

Furthermore, the mechanism within Mamba blocks’ S6 layers resembles that of RNNs. This similarity allows us to adopt a BiRNN-like approach, where each Mamba block processes time series bidirectionally. Inspired by the Vision Mamba block (Zhu et al., ), which processes image patches in both directions, we have modified this approach to work effectively with NEM and CFEM blocks.

Figure: Figure 2: Implementation of the Mamba block within the NEM module capable of bi-directional time series analysis.
Refer to caption: extracted/5909928/bimamba.png

Our bidirectional block, depicted in Figure 2, processes a data tensor. If it is a NEM block, it concatenates $H^{pri}$ and applies layer normalization. Dropout is employed for regularization, and MLP layers expand the hidden state to create two information channels, similar to the original Mamba block. The first channel is modified to establish two additional paths that analyze the time series forwards and backwards, each with its convolutional layer and S6 layer. The second channel remains unchanged, maintaining its activation and gating mechanism. Outputs are aggregated and compressed back to their original dimensions via another MLP, followed by layer normalization. Finally, a residual connection incorporates previously processed information, producing the final output of this module.

## 5 Experiments

### 5.1 Datasets

We employed datasets identical to those used in Cini et al. (), which established benchmarks later used as comparison metrics (Tashiro et al., ; Liu et al., ). The AQI-36 dataset (Yi et al., ) comprises air quality data from 36 stations in Beijing, recorded hourly over a year, with 13.24% missing values. We also utilized the METR-LA and PEMS-BAY datasets (Li et al., ), documenting traffic in Los Angeles and the San Francisco Bay Area, respectively. METR-LA includes data from 207 sensors over four months with an 8.10% missing rate, while PEMS-BAY covers 325 sensors over six months with 0.02% missing values, both sampled every five minutes.

Consistent with the Cini et al. () benchmark, we replicated the adjacency matrices using thresholded Gaussian Kernels on geographical distances, as suggested in (Li et al., ) and (Wu et al., ). The same training, validation, and testing splits were used: 70%, 10%, and 20% of the data, respectively, with identical seeds. For AQI-36, testing was conducted in February, May, August, and November, as originally designated, with 10% validation samples using the same seed.

Evaluation of these datasets includes three scenarios with synthetic missing data. In the ”Block missing” scenario, 5% of the data was randomly masked, and 1-4 hour blocks were masked at each sensor with a 0.15% probability; in the ”Point missing” scenario, 25% of the values were randomly masked; in the last scenario, missing values are simulated with the same distribution as the original data.

### 5.2 Baselines

We compared our proposed method against established benchmark methods originally presented in the GRIN article (Cini et al., ), and subsequently extended by CSDI (Tashiro et al., ) and PriSTI (Liu et al., ). To ensure the accuracy of our comparison, we replicated the experiments of CSDI and PriSTI using the parameters reported in their original papers. The only modification made was the adjustment of small differences found in the training and validation splits of the AQI-36 dataset between CSDI, PriSTI and GRIN to maintain a fair comparison. The evaluated methods are:

1) Statistical techniques: MEAN (using historical mean values), DA (daily averages), and KNN (geographical proximity). 2) Machine learning algorithms: Lin-ITP (linear interpolation), MICE (White et al., ) (multiple imputation), VAR (vector autoregressive model), and KF (Kalman Filter). 3) Matrix factorization strategies: TRMF (Yu et al., ) (temporal regularized matrix factorization) and BATF (Chen et al., ) (Bayesian augmented tensor factorization). 4) Autoregressive models: BRITS (Cao et al., ) (bidirectional RNN), MPGRU (one-step-ahead GNN-based predictor similar to DCRNN (Li et al., )), GRIN (Cini et al., ). 5) Generative models: CSDI (Tashiro et al., ), PriSTI (Liu et al., ), V-RIN (Mulyadi et al., ) (VAE with quantified uncertainty), GP-VAE (Fortuin et al., ) (VAE combined with Gaussian process), and rGAIN (Yoon et al., ) (GAN-based imputation with a recurrent structure).

### 5.3 Experimental settings

For DDPMs, 100 imputations were generated per missing value, with the median serving as the final imputation. Additionally, all experiments have been repeated 3 times using different seeds. During training, four synthetic missing data generation techniques were used:

1) Point strategy: A random [0, 100]% of data was masked per batch. 2) Block strategy: Sequences of missing values of length [L/2, L] were generated with a [0, 15]% probability, plus an additional 5% missing. 3) Historical strategy: Real imputation masks already present in the data were used. 4) Hybrid strategy: Each training sample was masked using the point strategy as a first option, or block or historical strategies as a second option with a 50% probability.

For the ”Point missing” scenarios described in Section 5.1, the point strategy was used. For ”Block missing” scenarios, the hybrid strategy was applied with the block strategy as the secondary option. For AQI-36, the hybrid strategy with the historical strategy as the secondary option was employed.

It is important to highlight that for the benchmark comparison in Section 5.4.1, the CSDI, PriSTI, and TIMBA models were trained for the number of epochs specified in their original papers: 200 epochs for the AQI-36 dataset and 300 epochs for the Traffic datasets. However, for the remaining experiments discussed in Section 5.4, the training epochs were limited to 50 due to time constraints during the execution of the experiments.

Finally, as in the original work by Tashiro et al. (), the model was trained with a learning rate of $10^{-3}$, which was reduced to $10^{-4}$ after 75% of the training epochs, and again to $10^{-5}$ after 90%. Consistent with Tashiro et al. (), a quadratic scheduler was used for the noise scheduling. Detailed hyperparameter settings, as well as information on training and inference times, and code reproducibility, are provided in Appendix A.

Aiming for precise parameter alignment in model comparisons, when we developed TIMBA, we replaced PriSTI’s dedicated temporal dimension transformer with a Mamba block with the most similar parameter count possible, while maintaining the approach described in Gu & Dao () of using a factor of 2 expansion on the original tensor’s hidden space. Using the METR-LA dataset as a reference, TIMBA has 876,765 parameters, while PriSTI and CSDI have 797,533 and 416,305 parameters, respectively. This results in a proportionally smaller increase in capacity, as PriSTI increases its parameters by 91.5% more compared to CSDI, whereas TIMBA only increases by 9.93% compared to PriSTI.

### 5.4 Results and discussion

To evaluate the imputation performance of TIMBA, we used the Mean Absolute Error (MAE) and Mean Square Error (MSE) metrics, following the benchmark established by Cini et al. ().

#### 5.4.1 Benchmark results

Table 1 presents the evaluation results of TIMBA compared to the previously defined benchmark. Our method generally outperforms previous results, consistently achieving better or at least comparable outcomes to PriSTI.

**Table 1: Results obtained after testing TIMBA against the benchmark established in the literature (Cini et al., ). The results are shown in terms of MAE and MSE.**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | AQI-36 | METR-LA | PEMS-BAY |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Simulated failure (24.6%) | Block-missing (16.6%) | Point-missing (31.1%) | Block-missing (9.2%) | Point-missing (25.0%) |  |  |  |  |  |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell |
| Mean\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 53.48 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4578.08 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.48 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 139.54 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.56 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 142.22 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.46 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 87.56 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.42 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 86.59 $\pm$ 0.00\endcollectcell |
| DA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 50.51 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4416.10 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.53 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 445.08 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.57 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 448.66 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.30 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 43.76 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.35 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 44.50 $\pm$ 0.00\endcollectcell |
| KNN\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 30.21 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2892.31 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.79 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 124.61 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.88 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 129.29 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.30 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 49.90 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.30 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 49.80 $\pm$ 0.00\endcollectcell |
| Lin-ITP\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.46 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 673.92 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.26 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 33.76 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.43 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.75 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.54 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.14 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.76 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.74 $\pm$ 0.00\endcollectcell |
| KF\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 54.09 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4942.26 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.75 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 534.69 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.66 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 529.96 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.64 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 93.19 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.68 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 93.32 $\pm$ 0.00\endcollectcell |
| MICE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 30.37 $\pm$ 0.09\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2594.06 $\pm$ 7.17\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.22 $\pm$ 0.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 51.07 $\pm$ 1.25\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.42 $\pm$ 0.07\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 55.07 $\pm$ 1.46\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.94 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 28.28 $\pm$ 0.37\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.09 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 31.43 $\pm$ 0.41\endcollectcell |
| VAR\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.64 $\pm$ 0.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 833.46 $\pm$ 13.85\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.11 $\pm$ 0.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 28.00 $\pm$ 0.76\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.69 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 21.10 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.09 $\pm$ 0.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.06 $\pm$ 0.73\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.30 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.52 $\pm$ 0.01\endcollectcell |
| TRMF\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.46 $\pm$ 0.06\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1379.05 $\pm$ 34.83\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.96 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 22.65 $\pm$ 0.13\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.86 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 20.39 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.95 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.21 $\pm$ 0.06\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.85 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.03 $\pm$ 0.00\endcollectcell |
| BATF\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.21 $\pm$ 0.27\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 662.87 $\pm$ 29.55\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.56 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 35.39 $\pm$ 0.03\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.58 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 36.05 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.05 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.48 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.05 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.90 $\pm$ 0.06\endcollectcell |
| V-RIN\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.00 $\pm$ 0.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 838.05 $\pm$ 24.74\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.84 $\pm$ 0.17\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 150.08 $\pm$ 6.13\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.96 $\pm$ 0.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 49.98 $\pm$ 1.30\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.49 $\pm$ 0.04\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 36.12 $\pm$ 0.66\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.21 $\pm$ 0.03\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.08 $\pm$ 0.29\endcollectcell |
| GP-VAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 25.71 $\pm$ 0.30\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2589.53 $\pm$ 59.14\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.55 $\pm$ 0.09\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 122.33 $\pm$ 2.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.57 $\pm$ 0.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 127.26 $\pm$ 3.97\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.86 $\pm$ 0.15\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 26.80 $\pm$ 2.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.41 $\pm$ 0.23\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 38.95 $\pm$ 4.16\endcollectcell |
| rGAIN\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.37 $\pm$ 0.26\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 641.92 $\pm$ 33.89\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.90 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 21.67 $\pm$ 0.15\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.83 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 20.03 $\pm$ 0.09\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.18 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 13.96 $\pm$ 0.20\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.88 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.37 $\pm$ 0.20\endcollectcell |
| MPGRU\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.79 $\pm$ 0.52\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1103.04 $\pm$ 106.83\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.57 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 25.15 $\pm$ 0.17\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.44 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 22.17 $\pm$ 0.03\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.59 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.19 $\pm$ 0.11\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.11 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.59 $\pm$ 0.02\endcollectcell |
| BRITS\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 14.50 $\pm$ 0.35\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 622.36 $\pm$ 65.16\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.34 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 17.00 $\pm$ 0.14\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.34 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.46 $\pm$ 0.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.70 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.50 $\pm$ 0.07\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.47 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.94 $\pm$ 0.03\endcollectcell |
| GRIN\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.08 $\pm$ 0.47\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 523.14 $\pm$ 57.17\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.03 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 13.26 $\pm$ 0.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.91 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.41 $\pm$ 0.03\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.14 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.60 $\pm$ 0.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.67 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.55 $\pm$ 0.01\endcollectcell |
| CSDI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.74 $\pm$ 0.16\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 388.37 $\pm$ 11.42\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.90 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.27 $\pm$ 0.18\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.77 $\pm$ 0.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.42 $\pm$ 0.47\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.84 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.06 $\pm$ 0.04\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.58 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.30 $\pm$ 0.04\endcollectcell |
| PriSTI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.84 $\pm$ 0.11\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 376.11 $\pm$ 10.62\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.78 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.64 $\pm$ 0.13\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.70 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.47 $\pm$ 0.04\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.87 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.64 $\pm$ 0.21\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.59 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.61 $\pm$ 0.03\endcollectcell |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.56 $\pm$ 0.4\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 352.29 $\pm$ 5.33\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.76 $\pm$ 0.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.36 $\pm$ 0.34\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.69 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.36 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.84 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.57 $\pm$ 0.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.58 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.63 $\pm$ 0.08\endcollectcell |

However, a more detailed analysis reveals that our method does not perform as well in the PEMS-BAY point-missing scenario, achieving results comparable to PriSTI. Interestingly, in this specific scenario, CSDI outperforms both methods. Given that the primary difference between PriSTI and TIMBA with CSDI lies in the hyperparameters of the noise scheduler, a finer adjustment of these values might further improve our model’s performance.

#### 5.4.2 Ablation analysis

This section presents the results of an ablation study comparing the performance of TIMBA using both bidirectional and unidirectional Mamba blocks. The objective of this experiment was to determine to what extent bidirectional information processing is beneficial for Mamba blocks. The results are shown in Table 2.

**Table 2: Ablation results for TIMBA. TIMBA refers to the configuration utilizing bidirectional blocks, while TIMBA-Uni represents the model using only unidirectional Mamba blocks. This experiment was conducted with training limited to 50 epochs, as described in Section 5.3**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | AQI-36 | METR-LA | PEMS-BAY |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Simulated failure (24.6%) | Block-missing (16.6%) | Point-missing (31.1%) | Block-missing (9.2%) | Point-missing (25.0%) |  |  |  |  |  |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell |
| TIMBA-Uni\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.18 $\pm$ 0.06\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 402.62 $\pm$ 6.57\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.88 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.18 $\pm$ 0.09\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.79 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.59 $\pm$ 0.05\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.92 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.01 $\pm$ 0.14\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.64 $\pm$ 0.00\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.99 $\pm$ 0.06\endcollectcell |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.66 $\pm$ 0.32\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 360.66 $\pm$ 25.02\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.79 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.73 $\pm$ 0.24\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.71 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.56 $\pm$ 0.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.86 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.80 $\pm$ 0.19\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.59 $\pm$ 0.01\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.65 $\pm$ 0.14\endcollectcell |

The analysis indicates that TIMBA achieves superior performance across all scenarios when utilizing bidirectional blocks. This finding supports the necessity of these blocks for effectively extracting improved temporal representations from the data.

#### 5.4.3 Missing rate sensitivity analysis

In the following experiment, we analyze the sensitivity of our model to varying levels of missing values. We evaluate CSDI, PriSTI, and TIMBA using the METR-LA dataset under the Point-missing scenario with different rates of missing data. For each model, we used the best-performing weights obtained during a previous training of 50 epochs, consistent with the description in Section 5.3. The results for the three models are presented in terms of MAE and MSE in Tables 3 and 4, respectively.

**Table 3: Sensitivity analysis for different levels of missing values in the METR-LA dataset under the Point missing scenario, presented in terms of MAE.**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE - METR-LA (P) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 20%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 30%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 40%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 50%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 60%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 70%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 80%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 90%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| CSDI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.77\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.82\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.88\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.96\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.07\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.20\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.39\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.70\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.29\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| PriSTI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.64\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.67\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.71\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.76\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.81\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.89\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.99\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.14\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.43\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.64\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.66\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.70\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.75\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.80\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.88\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.98\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.13\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.41\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |

**Table 4: Sensitivity analysis for different levels of missing values in the METR-LA dataset under the Point missing scenario, presented in terms of MSE.**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE - METR-LA (P) |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 20%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 30%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 40%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 50%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 60%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 70%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 80%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 90%\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| CSDI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.50\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.10\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.92\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.88\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.14\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 13.86\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 16.54\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 21.75\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 32.54\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| PriSTI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.71\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.07\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.61\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.29\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 10.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.24\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.88\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.88\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 22.08\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.60\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.91\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.48\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.15\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 9.96\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.09\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.65\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 15.55\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 21.57\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |

Analyzing the results, we find that TIMBA consistently delivers the best performance, particularly in terms of MSE. This demonstrates that our model effectively manages higher rates of missing values in the input data, further highlighting the advantages of SSMs over classical Transformers for these time series tasks.

#### 5.4.4 Downstream task analysis

Evaluating imputation models requires assessing their effectiveness on downstream tasks (Wang et al., ). This section outlines a specific task using imputed data to evaluate CSDI, PriSTI, and TIMBA. The downstream task designed for this experiment involves predicting the value of a node at time $t$ using the values of other nodes at the same time $t$.

For this section, we performed the node forecasting task using two different nodes. We conducted all experiments twice, once targeting node 14 and once targeting node 31, to ensure a more informative comparison. These nodes correspond to stations with the highest and lowest connectivity in the graph, following a similar approach to Cini et al. () in selecting nodes for experiments.

We use the AQI-36 dataset, imputing missing values with the best weights obtained from the results in Table 1. For this task, we only impute validation and test data to ensure that the imputation models work with data they have never seen. These new imputed data are split into a new training set (80%) and test set (20%), normalized using a MinMax scaler.

We employ an MLP with one hidden layer of 100 neurons and an output layer of 1. This network is trained for 500 epochs to minimize the MSE between the actual and predicted values, using the same five seeds for the imputed data with CSDI, PriSTI, and TIMBA. Final test results are measured using MSE and MAE and only consider actual values. Imputed target values are excluded from the final test results for the calculation of both metrics.

**Table 5: Results of the downstream task for node value prediction using data imputed by CSDI, PriSTI, and TIMBA.**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Sensor 14 | Sensor 31 |  |  |
| --- | --- | --- | --- | --- |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MAE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | MSE\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| CSDI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.51 $\pm$ 0.69\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 96.99 $\pm$ 21.25\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.99 $\pm$ 1.92\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 376.40 $\pm$ 148.06\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| PriSTI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.46 $\pm$ 0.71\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 92.70 $\pm$ 20.27\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.70 $\pm$ 1.80\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 361.19 $\pm$ 132.99\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 6.45 $\pm$ 0.69\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 91.90 $\pm$ 20.33\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 11.68 $\pm$ 1.77\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 359.80 $\pm$ 131.74\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |

As shown in Table 5, TIMBA achieves the best results for both nodes and both metrics. This indicates that the quality of imputation provided by our method is advantageous for use as a preprocessing step to improve subsequent tasks.

#### 5.4.5 Limitations

Analyzing the limitations of our model, it should be noted that our method does not assume any distribution in the missing values beyond their occurrence through a stationary process. Throughout the paper, we mainly focus on the missing at random (MAR) scenario (Rubin, ). Beyond this, we make no assumptions about the amount or length of the missing data.

## 6 Conclusions

In this paper, we presented TIMBA, a model that replaces the time-oriented transformer layers in state-of-the-art diffusion models for multivariate time series imputation with bidirectional Mamba blocks. Our extensive benchmark with three real-world datasets demonstrates that TIMBA either surpasses or performs comparably to the current state-of-the-art. Additionally, we showed that TIMBA can scale effectively with longer temporal sequences, generally achieving better results as the number of time steps per sample increases.

For future work, it would be valuable to explore ways to further reduce training and inference time, potentially through the application of latent diffusion models or denoising diffusion implicit models. Additionally, it would be interesting to apply our architecture to time series forecasting problems.

## Appendix A Implementation details

### A.1 Complete hyperparams description

Table 6 lists all the hyperparameters used to obtain the results presented in this paper. Most of these hyperparameters are directly taken from the PriSTI GitHub repository to ensure a fair comparison.

**Table 6: The hyperparameters of TIMBA for all datasets.**
| Description | AQI-36 | METR-LA | PEMS-BAY |
| --- | --- | --- | --- |
| Batch size | 16 | 4 | 4 |
| Time length $L$ | 36 | 24 | 24 |
| Epochs | 200 | 300 | 300 |
| Learning rate | 0.001 | 0.001 | 0.001 |
| Layers of noise estimation | 4 | 4 | 4 |
| Channel size $d$ | 64 | 64 | 64 |
| Number of attention heads | 8 | 8 | 8 |
| Minimum noise level $\beta_{1}$ | 0.0001 | 0.0001 | 0.0001 |
| Maximum noise level $\beta_{T}$ | 0.2 | 0.2 | 0.2 |
| Diffusion steps $T$ | 100 | 50 | 50 |
| Number of virtual nodes $k$ | 16 | 64 | 64 |
| Mamba block dropout | 0.1 | 0.1 | 0.1 |
| SSM state expansion factor | 16 | 16 | 16 |
| Mamba local convolution width | 4 | 4 | 4 |
| Mamba block expansion factor | 2 | 2 | 2 |

To expand on this information, it is important to note that the virtual nodes
$k$ were introduced by Liu et al. (). These nodes help reduce the complexity of transformers in handling spatial information by compressing real nodes into a defined number of virtual nodes $k$.

### A.2 Code reproducibility

This project has been conducted with a focus on facilitating code reproducibility and easy access to the datasets used. The code is available on GitHub: [Hidden for anonymization – check supplementary material during review].

The implementation is in Python (Van Rossum & Drake, ), utilizing the following open-source libraries:

- •
Pytorch (Paszke et al., ).
- •
Pytorch Lightning (Falcon & The PyTorch Lightning team, )
- •
Numpy (Harris et al., ).
- •
Torch spatio-temporal (Cini & Marisca, ).
- •
Pandas (pandas development team, ; Wes McKinney, ).
- •
Hydra (Yadan, ).

Additionally, to simplify execution, a Docker image and container (Merkel, ) have been developed, along with scripts to create and run them easily.

All experiments were conducted on a computer with the following specifications: Ubuntu 22.04.2 LTS, AMD Ryzen Threadripper PRO 3955WX 16-Cores CPU, NVIDIA RTX A5000 24 GB GPU, and 8X16 GB (128GB) DDR4 RAM. The runtime results obtained using this setup are presented in Table 7.

**Table 7: Training and inference times measured in hours.**
| \endcollectcell\eqmakebox[CFNetvsNeuMF][r] | AQI-36 | METR-LA | PEMS-BAY |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Models\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Training\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Inference\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Training\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Inference\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Training\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | Inference\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| CSDI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.29\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.22\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 7.28\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 1.74\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 19.11\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 4.62\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| PriSTI\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.40\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.33\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.21\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 2.44\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 19.72\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 5.99\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |
| TIMBA\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.56\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 0.44\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 12.83\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 3.65\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 32.74\endcollectcell\eqmakebox[CFNetvsNeuMF][r] | 8.71\endcollectcell\eqmakebox[CFNetvsNeuMF][r] |

### A.3 Data availability

The datasets used in this paper are publicly and freely accessible. Specifically, the Torch Spatio-Temporal library (Cini & Marisca, ) includes code for downloading and preparing these datasets, so there is no need to obtain them separately before running our code.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: Please see my comments in the Weaknesses section.
soundness: 3
strengths: - Overall, the clarity of the paper is good. It is well-written and easy to follow.
- By and large, the proposed method is technically sound.
- Empirical evaluation shows improved performance against existing methods.
summary: This paper introduces a model called Time series Imputation with bi-directional Mamba Blocks and Diffusion models (TIMBA) to address the time series imputation. The model combines a bi-directional Mamba block and a diffusion-based generative model to do imputation for time-series data. Empirical evaluation using three benchmark datasets show improved performance against various baseline methods for time series imputation.
weaknesses: - The major weakness of the paper is the limited novelty. Using diffusion model for time series imputation has been explored in a few existing work. Particularly, MSTI (Alcaraz & Strodthoff, 2023) explored using Mamba as the backbone of diffusion for time-series data. Given the existing work, it seems that the main technical contribution is to replace S4 with another SSM variant, S6. The technical contribution looks incremental to me.
- Most experiments limit the training epochs to 50, which might not be sufficient for models to converge, especially for generative models.
- In the limitations section, it is argued that the paper addresses the missing-at-random (MAR) problem. Please clarify if the paper addresses MAR or missing-completely-at-random (MCAR). How is the probability of missingness in MAR being considered?

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 2
questions: Same as above.
soundness: 2
strengths: 1. It explores some interesting technical combinations of existing models, such as bidirectional Mamba and S6, which could be beneficial for the imputation task as it requires bidirectional context.

2. The GNN model is used to capture spatial relationships between variables, which could complement temporal dependencies.

3. It examines the sensitivity to various levels of missing values in the imputation tasks.
summary: The paper presents the TIMBA model for multivariate time series imputation (MTSI). This work builds on the idea of using diffusion models to learn the latent distribution of time series data, enabling the generation of complete sequences from incomplete data. It also replaces the time-oriented Transformer model with a bidirectional Mamba model containing S6 blocks. A GNN-based model is used to represent the time series data, which serves as conditional variables for the diffusion models and the inputs for the Conditional Feature Extraction Module (CFEM). The Noise Estimation Module (NEM) is used to model the final noise. The paper claims that this approach enhances the spatiotemporal representation for the imputation task.

Overall, the paper highlights various techniques widely used in the time-series and machine learning communities. Although it attempts to combine these techniques to solve the imputation task, it is difficult to understand the utility of each component and the relationships between them. I believe a few components could be combined to perform imputation more effectively. Therefore, this study proposes an over-designed architecture that may lack sufficient rationale to justify each component.
weaknesses: 1. This paper introduces a variety of existing techniques, combining them into a new model, TIMBA. However, it is primarily a combination of existing techniques and offer little in terms of new insights or techniques.

2. The technical contribution is limited. For instance, the second and the third contribution are benchmark experiments and ablation studies, which are generally not considered as contribution.

3. Although this paper focuses on diffusion and bidirectional Mamba, many components appear unnecessary. In diffusion-based time series models, a typical setup includes only a backbone model and a diffusion model. For example, the foundational work TimeGrad uses RNN as the backbone model [1]. In the imputation task, the foundational work, CSDI, employs a single-layer transformer encoder as the backbone and demonstrates that bidirectional models generally outperform unidirectional ones. Therefore, this study would likely only need bidirectional Mamba as the backbone for imputation. 

    [1] Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting.

    [2] CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation

4. Althought this paper proposes ablation studies as a contribution, however, the benefits of various blocks are not known. For example, the study believes that the bidirectional Mamba with S6 is a better design, while we do not know whether the S6 can improve performance. A complete ablation studies should be conducted for all involved components, helping justify the necessity of each component. For example, for your proposed Mamba architecture, it should include the ablation on Mamba with and without S6, unidirectional vs bidirectional.

5. Many blocks are not discusses and do not have clear rationale. For example, the Conditional Feature Extraction Module (CFEM) and Noise Estimation Modules (NEM) are not discussed sufficiently. It does not show the motivation behind these blocks. We also do not know why these models should be combined with other blocks.  The authors may address the following questions, but are not limited to these: what is the specific purpose of the CFEM in this architecture? How does the NEM contribute to the overall performance of the model?


6. I find the use of GNN-based representation for time series confusing, as its necessity is unclear.Could you elaborate on why GNN-based representation is necessary for this task, and what specific advantages it provides over alternative approaches?


7. Although it is positioned as a diffusion model for imputation tasks, it only compares with two baselines in this category, CSDI and PriSTI. Most of the baselines are from traditional ML methods, which may not be suitable for comparison.

---

confidence: 3
contribution: 1
initial_rating: 3
presentation: 2
questions: 1  Why Choose SSM over Time-Oriented Transformers?: Why does the model replace time-oriented transformers with State-Space Models (SSMs)? What makes SSMs better suited for time modeling compared to transformers?

2  Overfitting Concerns: In the results from the PEMS-Bay dataset, the CSDI model using only the diffusion mechanism slightly outperforms the proposed model, which includes additional techniques such as GNN. Does this suggest that the proposed model might be overfitting to the spatio-temporal datasets?

3  Ablation Study: The ablation study focuses only on TIMBA-Uni and TIMBA. What about the contributions of other components of the model? How do they impact performance individually?

4  Latest Works in Time Series Imputation: In 2024, many papers on time series imputation have been published. I recommend incorporating more recent works in the introduction to provide a more comprehensive and up-to-date background.
ImputeFormer: Low rankness-induced transformers for generalizable spatiotemporal imputation; SADI: Similarity-Aware Diffusion Model-Based Imputation for Incomplete Temporal EHR Data; Higher-order Spatio-temporal Physics-incorporated Graph Neural Network for Multivariate Time Series Imputation; and other recent works.

5  Training Strategy and Result Variations: Could the differences in results from prior works (e.g., PriSTI) be due to different training strategies, or are there other factors that could explain these variations?
soundness: 2
strengths: In this paper, the authors propose TIMBA, a novel model that replaces time-oriented transformers in diffusion models with bi-directional Mamba blocks. Through extensive benchmarking on three real-world datasets, TIMBA consistently outperforms or matches state-of-the-art models. The paper further assesses the model's robustness through an ablation study, sensitivity analysis to varying missing rates, and downstream task evaluation after data imputation.
summary: In this paper, the authors propose the TIMBA model for multivariate time series imputation, integrating the Diffusion Probabilistic Model (DDPM) with State-Space Models (SSM), Graph Neural Networks (GNN), and other advanced techniques. By replacing time-oriented Transformers with SSMs and embedding the Selective State Space Models (S6) architecture within Mamba blocks, TIMBA effectively combines SSMs, GNNs, and node-based Transformers to enhance spatiotemporal representations.
weaknesses: 1  Unclear Motivation: The paper combines several machine learning methods (e.g., Transformers, SSMs, diffusion models, GNNs), but the motivation behind each choice is unclear. There is little to no theoretical analysis to support the use of these models or to explain why this specific combination of techniques is beneficial.

2  Lack of Comparison to Relevant Baselines: The experiments do not compare the proposed model to other state-of-the-art diffusion-based models like SSSD or diffusion with RNNs. This omission leaves uncertainty about whether the proposed TIMBA models truly offer better performance than existing methods.

3 Results Discrepancies: Some of the results reported in the paper differ from those in previous works (e.g., PriSTI). However, the authors do not provide an explanation for these discrepancies. This raises concerns about whether the differences arise from the training strategy, model implementation, or some other factor.

4 Inadequate Coverage of Time Series Imputation: While the paper claims to focus on time series imputation, the experiments primarily involve spatio-temporal datasets. There are no experiments on non-spatio-temporal datasets like clinical healthcare, which limits the generalizability of the model. This mismatch between the title and the scope of experiments weakens the overall contribution.

5  Format Issues: Some tables in the paper are out of range, indicating formatting problems that need to be addressed to improve readability and presentation quality.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 3
questions: Please comment on the highlighted weaknesses.
soundness: 3
strengths: * **Relevant problem**: Integrating SSMs into time series models is an important direction for future work.
* **Somewhat strong empirical results**: Good performance on relevant benchmarks, though improvements are limited, especially considering the reported standard deviations. Furthermore, the ablation analysis is limited (see weaknesses).
summary: The paper introduces TIMBA, an architecture for time series imputation. The main idea behind the proposed approach is to replace the Transformer layers of an existing architecture (PriSTI) with bidirectional Mamba blocks. Empirical performance matches or even surpasses the state of the art.
weaknesses: **Major weaknesses**

* **No clear motivation**: The authors argue that the proposed Mamba blocks should be used instead of Transformers as they are more tailored to processing time series given their inductive biases. However, the same could be said for models like RNNs and TCNs and for specialized Transformer architectures for time series. While SSMs have proven particularly useful for processing long sequences, the sequences processed here are a maximum of 36 time steps. It's unclear what specific benefits the proposed architectural modification offers. I suspect that the performance gap (often not particularly significant given the reported deviations) relative to existing models could be eliminated with better hyperparameter selection.

* **Limited novelty**: The paper suggests replacing the sequence modeling part of an existing architecture with another (existing) sequence modeling layer. The contributions, therefore, appear limited by ICLR standards.

* **Limited analysis**: Since the contribution is primarily an architectural modification, this should be the focus of additional ablation studies. Simply showing that the bidirectional version of the model achieves better results is not sufficient. One possibility would be to replace the S6 layers with RNNs/TCNs and/or replace the Mamba block entirely with other sequence modeling blocks.

**Minor comments**
* Some relevant missing references, e.g.,
    - [1] SPIN, an attention-based graph neural network for imputation, which mitigates the error accumulation seen in GRIN.
    - [2] SAITS, a popular attention-based imputation architecture.

#### References

[1] Marisca et al., "Learning to Reconstruct Missing Data from Spatiotemporal Graphs with Sparse Observations", NeurIPS 2022\
[2] Du et al., “SAITS: Self-Attention-based Imputation for Time Series”, Expert Systems with Applications, 2022