Model or dataset
WenjieDu/SAITS avatar
WenjieDu/SAITS

SAITS: a self-attention imputation model for incomplete time series

The official PyTorch implementation of the paper "SAITS: Self-Attention-based Imputation for Time Series". A fast and state-of-the-art (SOTA) deep-learning neural network model for efficient time-series imputation (impute multivariate incomplete time series containing NaN missing data/values with machine learning). https://arxiv.org/abs/2202.08516

514 stars70 forksPythonMIT

At a glance

What is it?
SAITS is the official PyTorch implementation of a self-attention model for filling NaN values in multivariate time series. It suits researchers who want the paper's training recipe in their own pipeline, not teams needing a maintained library with a stable API.
Who is it for?
Adopt SAITS if you are reproducing the paper, comparing imputation methods under a controlled setup, or willing to read run_models.py and the modeling/ package before trusting the output.
Can I use it commercially?
Yes. MIT is a permissive licence: you can use, modify and sell software built on it, as long as you keep its copyright and licence notices.
Is it still maintained?
Yes. The repository last received commits 36 days ago.
What is it written in?
Mainly Python, according to GitHub's language statistics.

Answers come from the project's GitHub data, last synced on September 29, 2026, and from our analysis. They are not legal advice.

Editorial analysis

What SAITS is for, and who ends up using it

The problem is narrow and concrete: you have a multivariate time series where some entries are NaN, and you want those entries filled with values that respect the temporal and cross-variable structure rather than a column mean. SAITS is the official PyTorch implementation of the paper "SAITS: Self-Attention-based Imputation for Time Series", published in Expert Systems with Applications. The README describes it as the first work applying pure self-attention without any recursive design to general time series imputation.

That last phrase is the useful part. Recurrent imputation models carry hidden state forward step by step, which makes long gaps expensive to model and training hard to parallelize. SAITS drops the recursion and lets attention operate across the sequence directly. The README also positions the repository more broadly than the paper title suggests: "More generally, you can use it for sequence imputation." If your data is a sequence with holes, not strictly a sensor panel, the same code applies.

The audience is therefore research-shaped. The README invites modification under the MIT licence for research and domain applications, and it says outright that the code "probably needs a bit of modification in the model structure or loss functions for specific scenarios or data input." That sentence is an honest description of scope, not a marketing hedge. If you need a supported library with versioned releases, this repository is not the right entry point; the README points to PyPOTS for that instead.

How the self-attention imputation actually runs end to end

The repository layout tells you the data flow without needing to read the paper. run_models.py is the entry point. Global_Config.py holds shared settings. configs/ holds experiment configuration, and modeling/ holds the model code. dataset_generating_scripts/ produces the datasets, and NNI_tuning/ contains the hyperparameter search setup. Paper_SAITS.pdf is checked in alongside the code.

The README's PyPOTS example shows the shape contract that the original code follows too. Data arrives as a three-dimensional array, printed in the example as (n_samples, n_steps, n_features). The README notes that train and validation sets have different sample counts but the same sequence length and feature dimension. The training set contains only the incomplete series. The validation set carries an extra key, X_ori, described as the ground truth needed for evaluation and for picking the best model checkpoint. The test set again holds only incomplete series, with a separate ground-truth array used for scoring.

Two design choices deserve attention. First, checkpoint selection depends on ground truth in validation, which means you must hold back real observations and then delete them yourself to build the task. Second, the README's example constructs an indicating mask from the test ground truth with np.isnan, so the evaluation is explicit about which entries were removed. There is also Simple_RNN_on_imputed_data.py in the repository root, which the filename presents as a downstream check: impute first, then run a simple recurrent model on the filled series. That is a sanity path rather than part of the imputation model itself.

Installing SAITS and running your first imputation

There is no package on an index for this repository. The README gives no pip install line for SAITS itself. What it does provide is conda_env_dependencies.yml at the repository root, which is the environment definition, and the PyPOTS route for users who want the model wrapped. Clone the repository, then create the environment from that file:

bash
git clone https://github.com/WenjieDu/SAITS.git
cd SAITS
conda env create -f conda_env_dependencies.yml
conda activate saits

Once the environment is active, run the training entry point. The README does not document its command-line flags, so read run_models.py and configs/ to see which arguments it accepts before launching a long run:

bash
python run_models.py

If you prefer the wrapped path, the README shows the PyPOTS example for PhysioNet-2012. The preprocessing call downloads and extracts the dataset, and the missing rate is a parameter:

python
import numpy as np
from sklearn.preprocessing import StandardScaler
from pygrinder import mcar, calc_missing_rate
from benchpots.datasets import preprocess_physionet2012
data = preprocess_physionet2012(subset='set-a', rate=0.1)
train_X, val_X, test_X = data["train_X"], data["val_X"], data["test_X"]
print(train_X.shape)  # (n_samples, n_steps, n_features)
print(f"We have {calc_missing_rate(train_X):.1%} values missing in train_X")

The validation set in that same example is built as a dictionary with X and X_ori, and the test set as a dictionary with X only. Expect the printed shape to be three-dimensional and the missing-rate line to reflect the rate you passed. If the shape is two-dimensional, your loader has collapsed the feature axis and the model will not receive what it expects.

Where SAITS is the wrong tool

The clearest limitation is packaging. The README does not document an installable release, a version number, or a rollback procedure. There is no changelog to tell you whether a checkpoint trained last month still loads after a pull. For a research reproduction that is acceptable; for a service that must rebuild an artifact on demand, it is a real operational risk that you would have to absorb by pinning a commit yourself.

The second limitation is scope of validation. The README says the model "probably needs a bit of modification in the model structure or loss functions for specific scenarios or data input." That is a warning that the defaults encode assumptions about missingness and about what a good reconstruction looks like. If your missingness is not random, if it correlates with the value that is missing, the training signal changes meaning. The repository does not claim to solve that; it gives you the scripts to generate missingness patterns and leaves the mapping to your problem to you.

The third is that SAITS is a point-imputation model. The README's own survey update describes a taxonomy of deep-learning imputation methods "based on uncertainty and model architecture", and the related searches around this project include CSDI, a conditional score-based diffusion model for probabilistic imputation. If your downstream decision needs a distribution over plausible fills rather than a single estimate, a deterministic self-attention reconstruction is the wrong family, and no amount of tuning fixes that. The README also notes the code would need modification for specific scenarios, so treat the out-of-the-box configuration as a baseline, not a finished product.

SAITS against BRITS, CSDI and ImputeFormer

The honest alternative for most users is PyPOTS, which the README describes as a Python toolbox for data mining on partially observed time series and which now contains SAITS. The difference is not the model, it is the surrounding contract. PyPOTS gives you a consistent dataset interface, the preprocessing helpers shown in the example, and a shared training loop across many models, so switching from SAITS to another imputer is a configuration change rather than a rewrite. The cost is that you are one layer removed from the paper code, and the README notes that PyPOTS adapted more than 20 forecasting models by applying SAITS embedding and training strategies, which means the wrapper is opinionated about how those models are trained.

BRITS is the recurrent counterpart. It carries hidden state through the sequence, so it is the natural comparison when you want to know whether dropping recursion actually helps on your data. The README's claim that SAITS is the first pure self-attention approach without recursive design is exactly the axis on which these two differ.

CSDI takes a different route again, modelling imputation as a conditional diffusion problem and producing samples rather than one fill. ImputeFormer is a transformer variant that the related searches pair with low-rank structure for spatiotemporal data, which matters if your series have a spatial or graph component that a plain per-series attention block does not encode. None of these is strictly better; they encode different assumptions about what missingness is and what the output should be. The TSI-Bench work the README mentions, with nearly 35,000 experiments across 28 imputation methods and 3 missing patterns, is the place to look for a controlled comparison rather than trusting any single repository's framing.

Maintenance, upgrade cost and the MIT licence

The repository is not archived, and the last push was on 2026-08-25. That is recent enough that the codebase is not abandoned, but the README does not describe a release cadence, and no recent releases were retrieved, so there is no version boundary to upgrade across. In practice this means upgrades are git pulls. You should expect to read diffs in modeling/ and run_models.py before merging them, because nothing in the repository states which changes are breaking.

The upgrade cost is dominated by checkpoints. Since the README describes validation-time checkpoint selection using X_ori, a model you trained is tied to the code that produced it. If the model definition in modeling/ changes, retraining is the safe assumption unless you have verified otherwise. Budget for that rather than assuming forward compatibility.

The licence is MIT, which is permissive and places few constraints on reuse, including in commercial settings. That is a statement about the licence text, not legal advice; if you are embedding the code in a product, have your own counsel review how you distribute it and how you attribute the paper. The README asks that you cite SAITS in publications if it helps your work, and CITATION.cff is checked in at the repository root for that purpose. Note that citation is a request, not a licence condition, and the two should not be conflated when you write your own compliance notes.

Editorial conclusion

Adopt SAITS if you are reproducing the paper, comparing imputation methods under a controlled setup, or willing to read run_models.py and the modeling/ package before trusting the output. Do not adopt it as a drop-in production dependency: the repository is a paper artifact with no released package, the README documents no rollback or checkpoint-compatibility guarantee, and the README itself says the model structure or loss functions probably need modification for specific scenarios. Verify first that your missingness pattern matches what you intend to simulate (the dataset_generating_scripts/ directory is where the synthetic masks come from), that your validation set carries ground truth because the README's example requires X_ori for checkpoint selection, and that the PyPOTS route covers the model you actually want, since the README states PyPOTS has adapted more than 20 forecasting models to the imputation task using SAITS embedding and training strategies.

Frequently asked questions

Is SAITS available as a pip package, or do I have to clone the repository?

The README gives no pip install line for SAITS itself. The repository provides conda_env_dependencies.yml for the environment, and the README points to PyPOTS as the wrapped way to train SAITS.

What input shape does SAITS expect for time series imputation?

The README's PyPOTS example prints train_X.shape as (n_samples, n_steps, n_features), and notes that train and validation sets can differ in sample count while sharing sequence length and feature dimension. Training sets hold only the incomplete series.

Why does the SAITS validation set need X_ori?

The README states that in the validation set you need ground truth for evaluation and for picking the best model checkpoint. X_ori is the key that carries it, and the test set uses a separate ground-truth array for scoring.

Can SAITS handle missing data that is not randomly missing?

The README does not make that claim. It says the code probably needs modification in the model structure or loss functions for specific scenarios or data input, so non-random missingness is something you would have to address yourself.

What licence is SAITS released under?

The repository is MIT licensed, and the README states the code is open source under the MIT licence, inviting modification for research and domain applications. The README also asks that you cite SAITS in publications.

Official sources

  1. Issues
  2. License: MIT
  3. Project website
  4. README
  5. WenjieDu/SAITS on GitHub
Add this badge to your README

If you maintain this project, the badge below links readers to this analysis and shows its maintenance status from the daily GitHub snapshot. Paste the markdown into your README; add ?metric=license or ?metric=stars to the image URL for a different field.

Add this badge to your README

markdown
[![Hysen Labs](https://hysenlabs.com/badge/wenjiedu-saits.svg)](https://hysenlabs.com/projects/wenjiedu-saits)