PaperlayerView exact revision on arXiv

Source-linked AI summary

UniTS: A Unified Multi-Task Time Series Model

Shanghua Gao, Teddy Koker, Owen Queen, Thomas Hartvigsen, Theodoros Tsiligkaridis, Marinka Zitnik

arXiv:2403.00131v3cs.LGcs.AI

TL;DR

Time-series models often specialize by task, and unifying predictive and generative tasks remains challenging across heterogeneous datasets. UNITS combines task tokenization with a shared time-series architecture and masked-reconstruction pre-training. Across multi-task benchmarks, it reports high average performance and strong few-shot and prompt-based results, while its dataset coverage and zero-shot generalization remain areas for future work.

  • Problem

    Unifying predictive and generative time-series tasks is challenging because datasets vary in domains, temporal scales, sampling rates, and dynamic patterns, while many models remain task-specific.

    Source: §1 · pp. 1–2

  • Method

    UNITS uses task tokenization, a shared architecture for heterogeneous series, and masked-reconstruction pre-training to train one model across predictive and generative tasks.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    Across 38 tasks, shared-weight UNITS achieves the highest average performance and excels on 27 tasks; in single-task settings, it outperforms task-specialized models and repurposed LLMs across four task types.

    Source: §1 · p. 2

  • Takeaways & Limitations

    UNITS demonstrates prompt-based and few-shot performance and can generalize to new domains and tasks within the evaluated scope.

    Source: §6 · p. 10

    The collected datasets do not cover all available time-series datasets, and zero-shot generalization remains for future exploration.

    Source: M Limitations and Future Directions · p. 34

Abstract

from arXiv · show

Although pre-trained transformers and reprogrammed text-based LLMs have shown strong performance on time series tasks, the best-performing architectures vary widely across tasks, with most models narrowly focused on specific areas, such as time series forecasting. Unifying predictive and generative time series tasks within a single model remains challenging. We introduce UniTS, a unified multi-task time series model that utilizes task tokenization to integrate predictive and generative tasks into a single framework. UniTS employs a modified transformer block to capture universal time series representations, enabling transferability from a heterogeneous, multi-domain pre-training dataset-characterized by diverse dynamic patterns, sampling rates, and temporal scales-to a wide range of downstream datasets with varied task specifications and data domains. Tested on 38 datasets across human activity sensors, healthcare, engineering, and finance, UniTS achieves superior performance compared to 12 forecasting models, 20 classification models, 18 anomaly detection models, and 16 imputation models, including adapted text-based LLMs. UniTS also demonstrates strong few-shot and prompt capabilities when applied to new domains and tasks. In single-task settings, UniTS outperforms competitive task-specialized time series models. Code and datasets are available at https://github.com/mims-harvard/UniTS.

1 Introduction

Time-series models remain largely task-specific despite demand for a model that handles heterogeneous data and both predictive and generative tasks. UNITS combines task tokenization with a shared architecture and reports strong benchmark, few-shot, and prompt-based performance.

1 Introduction

  • Existing time-series models often specialize by task, and prior unified or LLM-adapted approaches still rely on extensive fine-tuning or task- and dataset-specific modules.Forecasting-focused generative pre-training does not address other task types.Source: §1 · p. 2
  • Task tokenization, a unified architecture, and masked-reconstruction pre-training let UNITS share weights across predictive and generative tasks.The architecture handles varying variable counts and sequence lengths; self-attention spans time and variable dimensions, while dynamic linear and gating modules model temporal relationships and reduce feature interference.Source: §1 · p. 2
  • UNITS achieves the highest average performance across 38 multi-task tasks and excels on 27 of 38 tasks.In single-task settings, it also outperforms task-specialized models and repurposed LLMs across forecasting, classification, anomaly detection, and imputation.Source: §1 · pp. 2–3
  • Few-shot transfer improves on the strongest baseline by 12.4% in MSE on imputation and 2.3% in F1-score on anomaly detection.UNITS also addresses out-of-domain forecasting and classification without specialized modules.Source: §1 · p. 3

2 Related Work

Prior time-series work includes specialized models for individual tasks and broader foundation-model efforts. UNITS is positioned as a time-series-only approach that supports both predictive and generative tasks without extensive task-specific adjustments.

Traditional time series modeling

  • Traditional time-series methods typically train task-specific models on individual datasets, often requiring specialized heads and modules for fixed class counts or variable counts.UNITS instead aims to use one network architecture and universal task specification across tasks.Source: §2 · p. 3

General time series modeling

  • Time-series foundation-model efforts include diverse-signal architectures, LLM reprogramming, and multi-domain pre-training, but often still require task-specific modules and tuning.Lag-Llama [90] focuses on forecasting; UNITS is trained exclusively on time-series data and supports predictive and generative tasks without extensive task-specific adjustments.Source: §2 · p. 3

Prompt learning

  • Prior prompt methods adapt LLMs for forecasting, whereas UNITS uses time-series-only training and frozen-model prompt adaptation for tasks including classification and imputation.Examples of prior methods include PromptCast [118], LLMTime [81], TEMPO [10], and GPT4MTS.Source: §2 · p. 3

3 Problem Formulation

The formulation treats each dataset as samples paired with task specifications, with variation in sample lengths, variable counts, and task instances. A unified model should handle heterogeneous inputs, specify tasks universally, and share weights across tasks.

3 Problem Formulation

  • Each dataset is represented as samples X_i paired with task specifications Y_i, and a sample x has time length t and variable count v.The task set includes forecasting, classification, anomaly detection, and imputation, with multiple possible horizons or class counts.Source: §3 · p. 3
  • A unified model is defined as one weight-shared model that accepts heterogeneous series and uses a universal task specification across task types.Its input handling must accommodate varying lengths and variable counts.Source: §3 · p. 4
  • UNITS supports multi-task training, frozen-model prompt learning across tasks, and single-task training using the setup of existing models.Prompt tokens provide task and dataset context; the figure illustrates GEN tokens for forecasting and CLS tokens for classification.Source: §3 · p. 4, Figure 2 · p. 4

4 UNITS Model

UniTS encodes samples, prompts, and task types as tokens, then uses a modified transformer and shared task towers to support predictive and generative tasks. Its training combines unified masked-reconstruction pretraining with distinct prompt-tuning, supervised multi-task, and single-task regimes.

4.1 Prompting UNITS with Unified Time Series Data Tokens

  • Sample, prompt, and GEN/CLS task tokens form a shared representation for time-series inputs, task context, and generative or predictive outputs.Prompt tokens provide dataset and task context; task tokens are converted into predictions.Source: §4 · p. 4, §4.1 · p. 5
  • Forecasting and classification use task-specific token outputs, while GEN tokens also support imputation and reconstruction-based anomaly detection.Forecasting GEN tokens are replicated to the requested length; classification compares the output CLS token with class embeddings.Source: §4.1 · p. 5

4.2 Unified Network Architecture in UNITS

  • Two-way attention models both time and variable dimensions, while DyLinear interpolates weights to accommodate varying sequence lengths and capture token dependencies.The authors report that DyLinear improves generative-task performance.Source: §4.2 · pp. 5–6
  • A gating module rescales layer features to stabilize latent representations, and shared GEN and CLS towers serve generative and predictive tasks.These shared towers replace standalone task-specific heads for individual datasets.Source: §4.2 · p. 6

4.3 UNITS Model Training

  • Unified masked-reconstruction pretraining trains the backbone and GEN/CLS towers to reconstruct full samples using both prompt and CLS-token semantics.The pretraining loss combines reconstruction through the GEN tower with a CLS-related reconstruction path.Source: §4.3 · p. 6
  • Experiments distinguish frozen-pretrained token tuning (UNITS-PMT), supervised multi-task training from scratch (UNITS-SUP), and per-dataset single-task training (UNITS-ST).The single-task comparison covers forecasting, classification, anomaly detection, and imputation.Source: §4.3 · p. 6, Table 1 · p. 7

5 Experiments

Experiments evaluate UniTS across 38 multi-task datasets and separate single-task benchmarks, then test prompt-based transfer, direct forecasting at new lengths, and few-shot adaptation. Across these settings, the results report strong performance against task-specialized and multi-task baselines.

5.2 Benchmarking UNITS for Multi-Task Learning

  • A shared UniTS model achieves the highest average performance across 38 tasks and excels on 27 of 38, including 17/20 forecasting and 10/18 classification tasks.On forecasting, UNITS-SUP scores 0.439 vs. 0.466 MSE and 0.381 vs. 0.394 MAE against iTransformer; classification accuracy is 81.6% vs. 80.9% for TimesNet.Source: Table 2 · p. 8, §5.2 · p. 8

5.1 Benchmarking UNITS on Single-Task Learning

  • On single-task benchmarks, UNITS-ST leads on 28/32 forecasting datasets for MSE, 27/32 for MAE, and all 16 imputation datasets for MSE and MAE.It also gains 1.4% average classification accuracy over TimesNet and 3.95% F1 over TimesNet on anomaly detection.Source: §5.1 · p. 8

5.2 Benchmarking UNITS for Multi-Task Learning

  • UNITS-PMT achieves performance comparable to its fully supervised counterpart; separately, UNITS-SUP beats GPT4TS on forecasting by 2.2% (0.439 vs. 0.449; MSE), despite being 48× smaller.UNITS-PMT also has a lower forecasting MAE than UNITS-SUP (0.379 vs. 0.381).Source: §5.2 · p. 9

5.3 UNITS for Direct Multi-Step Forecasting

  • At the +384 forecasting extension, direct multi-step inference reduces MSE by 10.5% (0.442 vs. 0.494) against iTransformer and yields a 3× speedup.It reduces average inference steps from 3.66 to 1.Source: §5.3 · p. 9, Figure 3 · p. 9

5.4 UNITS for Few-Shot Learning on New Datasets and Tasks

  • Few-shot UNITS improves on iTransformer across data ratios and outperforms it on imputation, while UNITS-FT scores 86.3 F1 versus PathTST’s 84.3 on anomaly detection.At 20% data, UNITS-FT gains 8.8% classification accuracy and reduces forecasting MSE by 5.7%; on 25%-masked imputation it improves MSE by 12.4% and MAE by 7.9%.Source: §5.4 · pp. 9–10, Table 3 · p. 9, Table 4 · p. 10

6 Conclusion

UniTS uses a universal task specification to handle heterogeneous time series and multiple task types. The conclusion reports strong performance across multi-domain and multi-task datasets, along with few-shot and prompt-based generalization, while identifying alternative backbones as future work.

6 Conclusion

  • UniTS unifies heterogeneous time series tasks with a universal task specification and outperforms task-specific models and reprogrammed LLMs across 38 multi-domain, multi-task datasets.The authors report strong few-shot and prompt-based performance and generalization to new domains and tasks; they plan to explore MLP-based blocks and Mamba as alternative backbones.Source: §6 · p. 10

A Extended Related Work

  • A unified model must handle varying variable counts and time lengths while supporting generative and predictive tasks in one model.PatchTST handles variable variation but still requires task-specific heads for fixed-length forecasting and classification with predetermined classes.Source: Table 6 · p. 19, A Extended Related Work · p. 19

B Datasets

  • The dataset collection spans multiple sources and task settings; its combined training set contains over 35 million timesteps and over 6,000 variables.The appendix describes separate collections for multi-task, few-shot, zero-shot forecasting, imputation, and anomaly detection.Source: B Datasets · p. 19

C.1 All learning settings supported by UNITS

  • UniTS supports multi-task, prompt, few-shot, zero-shot, and single-task learning; its zero-shot examples include new forecasting lengths and out-of-domain datasets with new variable counts.Few-shot evaluation covers forecasting and classification on new out-of-domain datasets, plus imputation and anomaly detection as new task types.Source: §C.1 · pp. 19–20

C.2 Generalizing Task Tokens to Various Tasks

  • For imputation, GEN tokens replace missing sample tokens before the augmented sequence is combined with prompt tokens; anomaly detection reconstructs samples and uses reconstruction error as its criterion.The shared GEN and CLS towers transform task tokens into generative and predictive results, respectively.Source: §C.2 · pp. 21–22, Figure 4 · p. 22

C.3 Implementation of UNITS Network Architecture

The appendix describes UniTS as a stack of blocks with shared generative and classification towers, using time- and variable-wise attention, a dynamic feed-forward module, and gates. The default multi-task model uses three blocks and shared task towers.

C.3 Implementation of UNITS Network Architecture

  • The network combines time and variable multi-head self-attention; variable attention averages queries and keys over time to share an attention map across time points.The authors report evaluating both attention components in Table 22.Source: §C.3 · p. 22
  • Dynamic FFN replaces the standard first linear layer with a 3-kernel temporal convolution and inserts DyLinear before the unchanged output linear layer.Splitting the intermediate embeddings into two routes produces a scale-combination effect, while the CLS token is skipped by DyLinear.Source: §C.3 · p. 23
  • A gate after each block component uses a sigmoid-scaled factor to rescale its input features between 0 and 1.The GEN tower is shared across forecasting, imputation, and anomaly detection; the CLS tower is shared across classification tasks.Source: §C.3 · p. 23

D.1 Model Details

  • The default multi-task model uses three UniTS blocks, one GEN tower, and one CLS tower, with embedding dimension 64 for UNITS-SUP and 128 for UNITS-PMT.Forecasting tasks from the same data source but with different forecast lengths share prompt and GEN tokens.Source: §D.1 · p. 23

D.2 Training Details

The appendix specifies joint training protocols for supervised multi-task learning and self-supervised pre-training, evaluation metrics, and how baseline architectures are adapted for unified comparison.

D.2 Training Details

  • Supervised multi-task training runs for five epochs with effective batch size 1,024; self-supervised pre-training runs for ten epochs with effective batch size 4,096.Experiments use A100-40G GPUs, one or two per experiment, with maximum runtime under 48 hours.Source: §D.2 · p. 24
  • Evaluation averages performance by task type, using MSE and MAE for forecasting and imputation, accuracy for classification, and precision, recall, and F1 for anomaly detection.The authors report no task-specific hyperparameter tuning for UniTS, and baselines use the same training setting.Source: §D.2 · p. 24

D.3 Further Information on Pre-training

  • Pre-training randomly truncates samples to 50%–100% length, then applies random-position or right-side masking at 70%–80% with equal probability.Right masking is designed to enhance forecasting ability.Source: §D.3 · p. 24

D.4 Implementation Details of Baselines

  • Baselines receive data-specific input heads and task-specific output heads to support unified training, but these separate modules make generalization to new datasets and tasks difficult.For the proposed method and patch-based baselines, patch size and stride are fixed at 16.Source: §D.4 · p. 25

E Additional Results: Prompt Learning and Pre-training

Prompt learning benefits from model scaling and task-specific prompt tokens, while pre-training outcomes depend on the reconstruction losses, training scale, task mix, and domain mix. The results show that forecasting-only pre-training transfers to classification, whereas classification-only pre-training performs worse across both task types.

Prompt learning with model scaling

Effect of prompt tokens

  • 81.6% classification accuracy, 0.439 average MSE, and 0.381 average MAE are reported with prompt tokens, versus 81.0%, 0.460, and 0.391 without them.Using 10 rather than 5 prompt tokens gives greater forecasting gains and a 0.1% classification-accuracy improvement; sharing prompts across tasks reduces performance but still beats using none.Source: Table 15 · p. 25, Table 16 · p. 25, E Additional Results: Prompt Learning… · p. 26

Unified pre-training

  • 78.0% to 33.1% classification accuracy follows removal of the CLS-token reconstruction loss, while removing prompt-token loss worsens forecasting MSE from 0.471 to 0.967.The two losses are components of the unified mask-reconstruction pre-training objective.Source: E Additional Results: Prompt Learning… · p. 26, Table 17 · p. 26

Pre-training with scaled numbers of epochs and data sizes

Cross-task pre-training

  • Forecasting-only pre-training achieves similar performance to joint forecasting-and-classification pre-training, and performs well on classification despite seeing no classification data.Classification-only pre-training instead reduces performance on both task types; the passage suggests the larger forecasting-data volume may matter more than task type.Source: E Additional Results: Prompt Learning… · pp. 26–27, Table 20 · p. 27

Cross-domain pre-training

  • Joint Weather-and-Traffic pre-training outperforms single-domain pre-training; Weather-only pre-training performs better across both domains than Traffic-only pre-training.On Weather datasets, MSE changes from 0.253 with joint pre-training to 0.259 with Weather-only pre-training, and to 0.289 when pre-trained on Traffic.Source: E Additional Results: Prompt Learning… · p. 27, Table 21 · p. 27

F Additional Results: Ablation Studies of UNITS

Ablations report contributions from both attention modules, the Dynamic FFN, and the gate module; UniTS also outperforms a Transformer baseline when both use unified tokenization and co-training.

Effect of time and variable MHSA

Effect of Dynamic FFN

Effect of gate module

Comparison with Transformer

  • UniTS outperforms the original Transformer in classification and forecasting when both use unified tokenization and co-training.The comparison suggests that transformer structure alone is insufficient for robust multi-task time-series performance.Source: F Additional Results: Ablation Studies… · p. 28

G Additional Results: UNITS for Zero-Shot Forecasting on New Datasets

The additional results cover zero-shot forecasting on new datasets, relationships among prompt tokens, per-dataset classification results, and inference-step costs for direct multi-step forecasting.

G Additional Results: UNITS for Zero-Shot Forecasting on New Datasets

  • 45.2% MSE improvement over LLMTime (0.456 vs. 0.832) on River; UniTS is approximately 106 times faster in inference.The zero-shot evaluation uses one sample from each of five tasks with new forecasting lengths and variable counts; Time-LLM is excluded because it requires those dimensions to match training.Source: G Additional Results: UNITS for… · pp. 28–29

H Additional Results: Relation among Prompt Tokens

  • After training, prompt tokens from the same datasets cluster, while some tokens from different datasets remain close, indicating shared information across domains.Datasets within the same class, including two EEG datasets, also have higher prompt-token similarity.Source: H Additional Results: Relation among… · p. 29

I Additional Results: Classification Performance Stratified by Datasets

J Additional Results: Direct Multi-step Forecasting on New Forecasting Lengths

  • For maximum extra inference length 384, sliding-window inference requires an average of 3.66 times more steps than direct multi-step inference, which completes in one step.Figure 6 compares the average inference steps for the two approaches.Source: J Additional Results: Direct… · p. 30

K Additional Results: Benchmarking in the Single-Task Regime

The single-task benchmark trains separate models with tuned hyperparameters across forecasting, classification, imputation, and anomaly detection; a same-hyperparameter comparison also reports advantages for multi-task learning.

K Additional Results: Benchmarking in the Single-Task Regime

  • The single-task benchmark separately trains a model per dataset, using 36 forecasting, 10 classification, 4 imputation, and 5 anomaly-detection datasets.Forecasting uses horizons of 96, 192, 336, and 720 with fixed and optimal look-back windows; imputation masks 12.5%, 25%, 37.5%, or 50% of time points.Source: K Additional Results: Benchmarking in… · pp. 30–31, Table 30 · p. 31, Table 31 · p. 32
  • For fair comparisons, Anomaly Transformer is evaluated using reconstruction error only, and TimesNet results are reproduced from the Time-Series-Library.Full benchmark results are provided for forecasting, classification, imputation, and anomaly detection.Source: K Additional Results: Benchmarking in… · pp. 31–32

L Additional Results: Multi-task versus Single-task Learning

M Limitations and Future Directions

The collected datasets do not cover all available time series, and generalization to entirely new data through zero-shot learning remains future work. The impact statement identifies no societal impacts requiring specific emphasis.

M Limitations and Future Directions

  • The collected datasets omit some available series, including univariate UCR datasets and more physiologic PhysioNet signals.The authors plan to explore larger dataset collections.Source: M Limitations and Future Directions · p. 34
  • Generalization to entirely new data through zero-shot learning remains for future exploration, distinct from adapting to new time series data.The paper demonstrates adaptability through prompt learning and few-shot learning.Source: M Limitations and Future Directions · p. 34

N Impact Statement

Loading 2403.00131v3…