UniTS: A Universal Time Series Analysis Framework Powered by Self-Supervised Representation Learning

We first introduce the unified formulation of time series analysis, which serves as the basis of our framework.

In machine learning methods, the mapping function $f(\boldsymbol{x}_{i})$ is learned from a training dataset $\boldsymbol{X}=[\boldsymbol{x}_{1},\ldots,\boldsymbol{x}_{N}]\in\mathbb{R}^{N\times D\times T}$ and an optional label set $\boldsymbol{Y}\in\mathbb{R}^{N\times*}$. Unlike the traditional approaches that learn $f_{{}_{\mathcal{T}}}(\boldsymbol{x}_{i})$ from scratch for each task $\mathcal{T}$, our UniTS first pre-trains one or more task-independent encoders $h_{m}$ using only $\boldsymbol{X}$ to map the time series to unified representations, as $\boldsymbol{z}_{i,m}=h_{m}(\boldsymbol{x}_{i})\in\mathbb{R}^{K}$ ($m=1,\ldots,M$). Then, given a task $\mathcal{T}$, it fuses the features $\boldsymbol{z}_{i,1},\ldots,\boldsymbol{z}_{i,M}$ to one vector $\boldsymbol{z}_{i}^{\prime}\in\mathbb{R}^{K^{\prime}}$ to build a task-specific model based on the encoders, with $f_{{}_{\mathcal{T}}}(\boldsymbol{x}_{i})=g_{{}_{\mathcal{T}}}(\boldsymbol{z}_{i}^{\prime})$ where $g_{{}_{\mathcal{T}}}$ is an output model. The design goal is to take advantage of various self-supervised pre-training methods to easily achieve universal performance improvement for different analysis tasks and to address the challenges of partial labeling and domain shift. For the description, we denote $\boldsymbol{Z}_{m}=[\boldsymbol{z}_{1,m},\ldots,\boldsymbol{z}_{N,m}]\in\mathbb{R}^{N\times K}$.

To achieve universal time series analysis via self-supervised pre-training, UniTS is designed with three main modules, including the Pre-training Module, the Feature Fusion Module, and the Analysis Task Module, as illustrated in Figure 1.

First of all, UniTS creates one or more instances of the pre-training templates with their hyper-parameters, where each template is a self-supervised learning method. UniTS has implemented various types of template, as discussed in Section 3.1. During pre-training, each instance separately learns its encoder $h_{m}$ ($m=1,\ldots,M$), while all encoders $h_{1},\ldots,h_{M}$ are jointly used for the analysis tasks. The variety of pre-trained representations can be complementary with each other to achieve better performance.

After pre-training, the feature fusion module combines all learned representations of each sample $\boldsymbol{x}_{i}$ into one embedding, denoted as $\boldsymbol{z}_{i}^{\prime}\in\mathbb{R}^{K^{\prime}}$. The goal of this module is to automatically fuse the information from different pre-training instances to better facilitate the tasks. The details of this module are shown in Section 3.2.

Any analysis task $\mathcal{T}$ can be performed on top of $\boldsymbol{z}_{i}^{\prime}$ by using a task-specific function $g_{{}_{\mathcal{T}}}$ to map $\boldsymbol{z}_{i}^{\prime}$ to the prediction $\hat{\boldsymbol{y}}_{i}$ (or $\hat{y}_{i}$) of the target, and then fine-tuning the models (including the encoders, the learnable feature fusion model and the task-specific layers) by minimizing a loss function of the task. Presently, UniTS supports the five mainstream tasks illustrated in Table 1, while other tasks can be seamlessly integrated using the sklearn-style APIs. The analysis task module is discussed in detail in Section 3.3.

Discussion. From the above description, any time series analysis task can be carried out with UniTS in a unified way. At the training stage, the model $f_{{}_{\mathcal{T}}}$ is built following the above pipeline. During the inference stage, the learned $f_{{}_{\mathcal{T}}}$ is used to map the input series to the predictions. To deal with the problems of partial labeling, only a small size of the labeled data is required for fine-tuning, while the pre-training stage does not rely on any labels (Figure 2(a)). Similarly, when facing the problem of domain shift, the user can pre-train the encoders using the data from an available source domain to get the transferable representations, and then fine-tune $f_{{}_{\mathcal{T}}}$ using a small data set from the target domain (Figure 2(a)). In any case, the pre-trained encoders can be directly used without re-training, and the data size and the number of iterations for fine-tuning can be much smaller than the training from scratch (Liang et al., 2023) while guaranteeing competitive performance (Figure 3).

UniTS provides various self-supervised representation learning methods that can be used as pre-training templates. Below, we explain the currently implemented pre-training templates, and discuss how to easily add new templates via the sklearn-like APIs.

In general, UniTS has integrated the following pre-training methods, which are divided into three types based on different pre-training objective functions (see the references for more details).

Contrastive Learning. UniTS currently supports time series contrastive learning at three levels, i.e. the whole-series level (Liang et al., 2023), the sub-sequence level (Franceschi et al., 2019), and the timestamp level (Yue et al., 2022). These methods encourage the representations generated from the similar inputs to be closer than those of dissimilar inputs, which has shown effective in extracting informative features.

Autoregression. The autoregression-based pre-training algorithm (Zerveas et al., 2021) learns the representations by masking some observations of the input series (e.g., set to 0) and then predicting the masked values using the unmasked data. It is inspired by the masked language model in natural language processing, since both time series and sentences share the same nature of sequential dependencies.

Hybrid. This branch of approach (Eldele et al., 2021) optimizes a hybrid objective of the two types above, which may be better than the individual objectives in some cases, but not always.

The pre-training template is designed using a unified sklearn-like interface to allow flexible extension. A new algorithm can be seamlessly integrated as long as it is wrapped in a Python class with two methods: fit which takes the input $\boldsymbol{X}$ for pre-training, and transform which maps $\boldsymbol{X}$ to the representations $\boldsymbol{Z}$.

This module fuses the representations of a variety of pre-training instances. The module is also designed using sklearn-style interfaces. A feature fusion class contains a transform function that maps a set of representations $(\boldsymbol{Z}_{1},\ldots,\boldsymbol{Z}_{M})$ into one unified $\boldsymbol{Z}^{\prime}\in\mathbb{R}^{N\times K^{\prime}}$, and $\boldsymbol{Z}^{\prime}$ is used for the analysis tasks. Note that the parameters of a feature fusion instance can be optimized at the fine-tuning stage if they are learnable, which allows automatically extracting information from all pre-trained representations. UniTS now supports two basic feature fusion methods as described below.

Concatenation. The simplest but most popular way of feature fusion is directly concatenating the features of each sample, i.e., $\boldsymbol{z}_{i}^{\prime}=\boldsymbol{z}_{i,1}\oplus\ldots\oplus\boldsymbol{z}_{i,M}$, where $\oplus$ is the concatenation operation.

Projection. One can also use a learnable model $p$ to project the concatenated features into another latent space. This is especially effective in some cases such as clustering where dimension reduction is usually required. Formally, we have $\boldsymbol{z}_{i}^{\prime}=p(\boldsymbol{z}_{i,1}\oplus\ldots\oplus\boldsymbol{z}_{i,M})$.

The analysis task module is designed as a template wrapped with the sklearn-like APIs, where each task is an instance. The template contains two major components: a fit method which performs the task-specific fine-tuning, and a predict method which outputs the final prediction. UniTS currently supports the five important tasks described in Table 1. Below, we briefly explain the technical details.

Classification. This task is performed in a standard manner. The output model projects the representations into $C$ classes, and the softmax function is used to generate the class distribution. The standard cross-entropy loss is used for fine-tuning, and the class that gives the maximum probability is determined as the prediction.

Clustering. The clustering task can be directly performed by running a classical clustering algorithm (e.g., $k$-Means) on top of the representations. Moreover, one can also fine-tune the models for better performance. The fine-tuning is designed based on the $k$-Means loss. At each epoch, the $k$-Means algorithm is run on the representations to obtain the $C$ centroids. Then, the sum of the L2 norms between each representation vector and its centroid is added to the pre-training loss as a regularization term, which encourages the clustering structure of the representations. Note that we do not directly minimize the k-means loss to avoid the trivial solution, i.e, all representations become equal to their centroids.

Forecasting. Forecasting is performed in the standard way. A decoder is used to transform the representations into predictions. Then, the forecasting loss such as the mean square error (MSE) or the mean absolute error (MAE) is minimized for fine-tuning. In the inference stage, the decoder outputs are the forecasted values.

Anomaly detection. We employ a popular reconstruction-based framework (Schmidl et al., 2022) for anomaly detection. A decoder is added to project the representations to the predicted inputs $\hat{\boldsymbol{x}}_{i}$. The objective is to minimize the reconstruction loss as $||\boldsymbol{x}_{i}-\hat{\boldsymbol{x}}_{i}||$. For detection, an anomaly score $s_{t}$ is computed at each time step $t$ as $|\hat{\boldsymbol{x}}_{i,t}-\boldsymbol{x}_{i,t}|$. The observation with the score larger than a threshold $\tau$ is determined as an anomaly, i.e., $\hat{\boldsymbol{y}}_{i,t}=1$ if $s_{t}>\tau$ and 0 otherwise.

Missing value imputation. This task is performed using a modern structure named denoising autoencoder (DAE) (Vincent et al., 2008). During fine-tuning, we generate a random binary mask $\boldsymbol{m}_{i}\in\{0,1\}^{D\times T}$ for each $\boldsymbol{x}_{i}$. The masked sample $\boldsymbol{x}_{i}\otimes\boldsymbol{m}_{i}$ is input to the pre-trained encoders, where $\otimes$ denotes the element-wise multiplication. Then, a decoder is used on top of the representations to reconstruct the entire input $\boldsymbol{x}_{i}$ as $\hat{\boldsymbol{x}}_{i}$. DAE is learned by minimizing the reconstruction loss $||\boldsymbol{x}_{i}-\hat{\boldsymbol{x}}_{i}||$. At the inference stage, the missing values of an input $\boldsymbol{x}_{i}$ are first replaced by 0, i.e., $\boldsymbol{x}_{i,j,k}=0,\forall(j,k)\in P_{i}$. Next, $\boldsymbol{x}_{i}$ is input into the fine-tuned model to predict $\hat{\boldsymbol{x}}_{i}$. To this end, the missing value at position $(j,k)\in P_{i}$ is imputed using $\hat{\boldsymbol{x}}_{i,j,k}$.