统一时序模型（Unified Time Series Model）

项目主页 | 论文链接 (NeurIPS 2024)

UniTS 是一个统一的时序模型（time series model），能够使用共享参数处理多个领域中的各种任务，且不包含任何任务特定模块（task-specific modules）。

作者：Shanghua Gao Teddy Koker Owen Queen Thomas Hartvigsen Theodoros Tsiligkaridis Marinka Zitnik

概述

基础模型（Foundation models），尤其是大语言模型（LLMs），正在深刻改变深度学习。我们不再需要为每个任务单独训练模型，而是可以通过少样本提示（few-shot prompting）或微调（fine-tuning）将一个预训练模型适配到多种任务上。然而，当前的基础模型主要适用于序列数据，尚不能直接应用于时序数据。时序数据具有独特挑战：其内在的多样性、跨领域的数据集差异、预测（forecasting）、分类（classification）等任务规范的显著不同，以及对任务专用模型的明显依赖。

我们开发了 UniTS——一个统一的时序模型，支持通用任务规范（universal task specification），可同时处理分类、预测、插补（imputation）和异常检测（anomaly detection）任务。这是通过一种新颖的统一网络主干（unified network backbone）实现的，该主干结合了序列注意力（sequence attention）与变量注意力（variable attention），并引入动态线性算子（dynamic linear operator），以统一模型的方式进行训练。

在 38 个跨领域数据集上的实验表明，UniTS 的性能优于任务专用模型和经过改造的基于自然语言的 LLMs。在新数据领域和新任务上的评估中，UniTS 展现出卓越的零样本（zero-shot）、少样本（few-shot）和提示学习（prompt learning）能力。

环境配置

1. 依赖项

安装 PyTorch 2.0+ 及所需包：

pip install -r requirements.txt

2. 准备数据

bash download_data_all.sh

data_provider 文件夹中的 .yaml 文件展示了不同多任务设置下的数据集配置。

默认情况下，所有实验均采用多任务设置，即单个 UniTS 模型在多个数据集上联合训练。

3. 训练与评估模型

1. 预测与分类任务的多任务学习：

• 预训练 + 提示学习

bash ./scripts/pretrain_prompt_learning/UniTS_pretrain_x128.sh

• 监督学习

bash ./scripts/supervised_learning/UniTS_supervised.sh

2. 新预测与分类任务的少样本迁移学习：

注意：请先按照以下训练脚本中的说明获取预训练检查点（pretrained ckpt）。

• 微调（Finetuning）

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_newdata/UniTS_finetune_few_shot_newdata_pct20.sh

• 提示调优（Prompt tuning）

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_newdata/UniTS_prompt_tuning_few_shot_newdata_pct20.sh

3. 异常检测任务的少样本迁移学习：

• 微调

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_anomaly_detection/UniTS_finetune_few_shot_anomaly_detection.sh

• 提示调优

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_anomaly_detection/UniTS_prompt_tuning_few_shot_anomaly_detection.sh

4. 插补任务的少样本迁移学习：

• 微调

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_imputation/UniTS_finetune_few_shot_imputation_mask050.sh

• 提示调优

# 请在脚本中设置预训练模型路径。
bash ./scripts/few_shot_imputation/UniTS_prompt_tuning_few_shot_imputation_mask050.sh

5. 新预测长度的零样本学习：

# 请在脚本中设置预训练模型路径。
bash ./scripts/zero_shot/UniTS_forecast_new_length_unify.sh

6. 新预测数据集的零样本学习：

# 此设置需训练一个特殊版本的 UniTS，其使用共享的提示/掩码 token。
bash ./scripts/zero_shot/UniTS_zeroshot_newdata.sh

在您自己的数据上使用 UniTS

UniTS 是一个高度灵活的统一时序模型，通过单一共享模型和共享权重即可支持预测、分类、插补和异常检测等任务。我们提供了一份 教程（Tutorial），帮助您将自有数据与 UniTS 结合使用。

预训练权重

我们在 checkpoints 中提供了上述模型的预训练权重。

引用

@article{gao2024building,
  title={UniTS: Building a Unified Time Series Model},
  author={Gao, Shanghua and Koker, Teddy and Queen, Owen and Hartvigsen, Thomas and Tsiligkaridis, Theodoros and Zitnik, Marinka},
  journal={arXiv},
  url={https://arxiv.org/pdf/2403.00131.pdf},
  year={2024}
}

致谢

本代码库基于 Time-Series-Library 构建。感谢！

免责声明

分发声明：已批准公开发布，分发不受限制。

本材料基于美国空军合同 FA8702-15-D-0001 下由国防部负责研究与工程的副部长资助的工作。本材料中表达的任何观点、发现、结论或建议均为作者个人观点，不一定反映国防部负责研究与工程的副部长的观点。

© 2024 麻省理工学院。

受 FAR52.227-11 专利权条款约束 — 承包商拥有所有权（2014 年 5 月）

本软件/固件按“现状”提供。

根据 DFARS 第 252.227-7013 或 7014 条款（2014 年 2 月版）定义，本作品已授予美国政府无限权利。尽管存在任何版权声明，美国政府在本作品中的权利仍由上述 DFARS 252.227-7013 或 DFARS 252.227-7014 条款详细规定。未经美国政府明确授权而使用本作品，可能侵犯本作品中存在的任何版权。

UniTS 快速上手指南

UniTS 是一个统一的时间序列模型，支持预测、分类、插补和异常检测等多种任务，使用共享参数而无需任务专用模块。

环境准备

• 操作系统：Linux / macOS（推荐）
• Python ≥ 3.8
• PyTorch ≥ 2.0（建议使用 清华镜像源 加速安装）
• GPU（可选但推荐）

安装步骤

1. 克隆代码仓库：

git clone https://github.com/mims-harvard/UniTS.git
cd UniTS

2. 安装依赖（建议使用虚拟环境）：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

3. 下载官方数据集（用于示例训练）：

bash download_data_all.sh

基本使用

示例 1：多任务预训练 + Prompt 学习（预测与分类）

bash ./scripts/pretrain_prompt_learning/UniTS_pretrain_x128.sh

示例 2：监督学习训练

bash ./scripts/supervised_learning/UniTS_supervised.sh

提示：如需在自定义数据上使用 UniTS，请参考项目中的 Tutorial.md 文件。预训练权重可在 GitHub Releases 下载。