SAITS: A Self-Attention Imputer for Multivariate Time Series with NaN Gaps
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
At a glance
- What is it?
- SAITS is the reference PyTorch implementation of a self-attention imputation model published in Expert Systems with Applications. It is a research codebase first and a library second, and the README says as much about where you should expect to modify it.
- Who is it for?
- Adopt SAITS if you are reproducing or extending the ESWA paper, or if you want a single self-attention imputer you can read end to end and edit. Do not adopt it if you need a maintained pipeline with dataset loaders, missingness injectors and a uniform training API across many models; the README points that work at PyPOTS instead.
- 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 22 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 15, 2026, and from our analysis. They are not legal advice.
DEEP OPEN-SOURCE ANALYSIS
The gap SAITS fills: NaN blocks in multivariate sensor and clinical series
Most time series tooling assumes a rectangular array. Real recordings from clinical monitors, industrial sensors and wearable devices are not rectangular. Values go missing at single timestamps, in runs, or in whole blocks, and the downstream model either refuses the input or silently propagates the gap. SAITS targets that specific case: multivariate incomplete time series containing NaN, where the task is to produce a filled-in series rather than a forecast. The README frames the scope as general time series imputation and notes the model can be used for sequence imputation more broadly. The intended reader is someone with a partially observed array and a PyTorch environment, not someone looking for a drop-in preprocessing library. The repository is the official implementation accompanying the ESWA paper, so it is organized around the experiment in that paper rather than around a general-purpose API. That distinction shapes everything else in this review: file layout, config handling, and how much you are expected to write yourself.
Why the architecture drops recurrence: two diagonal-masked self-attention blocks
The README states that SAITS is the first work applying pure self-attention without any recursive design for general time series imputation. There is no RNN cell carrying state across timestamps, so the model does not have to walk the sequence in order to relate one observation to another. That is the mechanism the paper is built around, and it is also the reason the project describes itself as fast relative to recurrent imputers. The second structural idea visible in the README is the training strategy. The README refers to SAITS embedding and training strategies being applied to a long list of forecasting architectures (iTransformer, PatchTST, DLinear, Informer, Autoformer and others) inside PyPOTS to make them usable for imputation. That tells you where the reusable value sits: the pretraining-and-fine-tuning scheme and the embedding, not only the attention stack. The README does not lay out the full data flow, loss terms or the ORT/MT split in the text supplied here, so treat the paper as the specification and the code as the implementation of it. Anyone evaluating SAITS should read the arXiv preprint alongside the source rather than inferring the architecture from the README alone.
Getting it running: the training, validation and test dictionary contract
The README's example uses PyPOTS rather than the standalone repository, and it is worth reading closely because it documents the expected shapes and the dictionary keys. Data is loaded with preprocess_physionet2012(subset='set-a', rate=0.1) from benchpots.datasets, which the README says downloads and extracts the dataset automatically. Missingness is generated or measured with pygrinder, via mcar and calc_missing_rate. The shapes are printed as (n_samples, n_steps, n_features) for train_X and val_X. The dictionary contract is the part that trips people up. In the training set you put only the incomplete series: train_set = {"X": train_X}. In the validation set you supply ground truth as well, because it is used for evaluation and for picking the best checkpoint: val_set = {"X": val_X, "X_ori": data["val_X_ori"]}. The test set again carries only incomplete series, test_set = {"X": test_X}, with test_X_ori held separately for scoring, and the README shows an indicating_mask built from np.isnan on the test ground truth. StandardScaler is imported in the example, which implies scaling is your responsibility rather than the model's. If you work from the standalone SAITS repository instead of PyPOTS, expect to write the training loop, the batching and the checkpoint logic yourself; the README explicitly warns that the code may need modification in model structure or loss functions for specific scenarios or data input.
Where SAITS is the wrong tool: small samples, non-NaN missingness and production pipelines
The clearest limitation is stated by the authors rather than discovered by a user: the code probably needs a bit of modification in the model structure or loss functions for specific scenarios or data input. That is a research codebase being honest about itself, and it means there is no configuration key that will adapt SAITS to a new missingness mechanism for you. A second constraint follows from the input contract. SAITS expects missing values encoded as NaN in a floating-point array of shape (n_samples, n_steps, n_features). If your gaps are represented as sentinel values such as -1, 9999 or an empty string, you must convert them first, and a sentinel that is also a legal measurement will silently become training signal. Third, the self-attention design has no built-in notion of irregular timestamps. The repository topics include irregular sampling, but the tensor shape in the README is a fixed n_steps grid, so genuinely unevenly spaced observations need to be resampled onto a regular grid before they reach the model, and that resampling is a modelling decision the library does not make for you. Finally, if you need one API across many imputation models with shared preprocessing, SAITS standalone is not that. The README directs that use case to PyPOTS.
PyPOTS versus the standalone repository: same model, different amount of work
The README treats PyPOTS as the practical entry point and the standalone repository as the paper's artifact. The difference is not the model weights; it is everything around them. In PyPOTS you get the ecosystem libraries that the example imports: benchpots.datasets for preprocessing and downloading PhysioNet-2012, pygrinder for injecting MCAR missingness and computing missing rates, and a training interface where the model is selected by name and fed the three dictionaries shown above. In the standalone repository you get the SAITS implementation itself, which is what you want if you are changing the architecture, swapping the loss, or reproducing a table from the paper. The README also notes that SAITS embedding and training strategies were applied to more than twenty forecasting models inside PyPOTS to make them applicable to imputation. That is a meaningful signal about where the project's maintenance energy sits: the ecosystem repository is where new models land, while this repository is the origin of the method. If your goal is a working imputer this week, the PyPOTS path has fewer unknown steps. If your goal is to understand or alter the imputation mechanism, the standalone code is the shorter read.
Maintenance, licence and what the MIT terms let you do
The repository is MIT licensed and is not archived, with a last push timestamp of 2026-08-25 in the supplied metadata. MIT is permissive: it allows commercial use, modification and redistribution provided the copyright notice and permission notice are retained. That matters here because the README explicitly invites you to modify the SAITS code for your own research purpose and domain applications. If you fork it into a product, the obligation is the notice, not a copyleft release of your changes. This is not legal advice; check the LICENSE file in the repository and your own counsel for anything consequential. On upgrade cost, the material does not include a changelog or release notes, and no releases were retrieved, so there is no documented version history to reason about. The practical implication is that you should pin the commit you build against rather than tracking main, since a research repository with no release cadence can change shape without a version bump to warn you. The README's own framing reinforces this: it describes the code as something you adapt, which is a different contract from a library that promises backward compatibility.
Who should pick this up, and the first thing to check
Pick up SAITS if you have a multivariate array with NaN gaps, a GPU or the patience to train on CPU, and a reason to want the imputation model itself rather than a pipeline around it. That includes reproducing the ESWA results, extending the self-attention design, or reusing the training strategy on a different backbone, which is exactly what the PyPOTS integrations did for the forecasting models listed in the README. Skip it if your missingness is not NaN-encoded, if your timestamps are genuinely irregular and you have no plan to resample, or if you need a supported library with a version history. The first thing to verify in your own setup is the validation dictionary. The README's example builds val_set with both X and X_ori and states that the validation set is where ground truth is needed for evaluation and for picking the best model checkpoint, while test_set carries only X. If your validation split lacks X_ori, checkpoint selection has nothing to score against, and the model you ship will be whichever epoch happened to run last. Confirm that path before you trust any reported imputation metric.
Editorial conclusion
Adopt SAITS if you are reproducing or extending the ESWA paper, or if you want a single self-attention imputer you can read end to end and edit. Do not adopt it if you need a maintained pipeline with dataset loaders, missingness injectors and a uniform training API across many models; the README points that work at PyPOTS instead. Before committing, verify that the checkpoint selection path in your training script actually uses a validation set carrying X_ori ground truth, because the README's own example shows val_set built with both X and X_ori while test_set carries only X.
Community notes