FactoST-v2: Learning to Factorize and Adapt Toward Universal Spatio-Temporal Foundation Models

FactoST-v2 is a versatile spatio-temporal foundation model designed to decouple universal temporal learning from domain-specific spatial adaptation. It features a minimalist Encoder-Only Backbone trained via Randomized Sequence Masking to enable arbitrary-length generalization and Probabilistic Quantile Prediction.

For downstream adaptation, it employs a lightweight, factorized adapter that incorporates:

• Spatio-Temporal Metadata Fusion (STMF) for context injection.
• Spatio-Temporal Filtering (STF) for dynamic relation modeling.
• Domain-Specific Prompt Alignment (DSPA) for distribution alignment (replacing the legacy HDA module).
• Continual Memory Replay (CMR) for mitigating catastrophic forgetting.

This repository maintains backward compatibility with the original FactoST-v1 (NeurIPS) while providing the enhanced capabilities of the v2 architecture.

📋 Table of Contents

• 🚀 Quick Start
• 📦 Installation
• 💻 Basic Usage
• 📁 Project Structure
• 🔧 Implementation Details
• Citations

🚀 Quick Start

📦 Installation

# Clone the repository
git clone [https://github.com/CityMind-Lab/FactoST.git](https://github.com/CityMind-Lab/FactoST.git)
cd FactoST

# Install dependencies
# We recommend using a virtual environment (Python 3.9+)
pip install -r requirements.txt

💻 Basic Usage

# 1. Spatio-Temporal Adaptation (Stage 2: Fine-tuning)
# Adapts the frozen backbone to target domains using DSPA and STF
sh script/exp_factost_sta.sh

# 2. Inference after Fine-tuning
# Runs evaluation on test sets using the adapted model
sh script/inference_factost_sta.sh

# 3. Zero-shot Temporal Inference (Stage 1)
# Evaluates the backbone's universal temporal capabilities without adaptation
sh script/exp_factost_utp.sh

Note: The experiment scripts automatically select the appropriate pretrained model checkpoint (e.g., factost_utp2_tiny_4sta.pt) based on the configured model_size.

📁 Project Structure

FactoST/
├── script/                       # Experiment scripts
│   ├── exp_factost_sta.sh         # Stage 2: Spatio-temporal adaptation (runs FactoST_sta.py)
│   ├── inference_factost_sta.sh   # Inference entry (runs FactoST_sta.py in inference mode)
│   └── exp_factost_utp.sh         # Stage 1: Universal temporal entry (runs FactoST_utp.py)
│
├── src/                          # Core library
│   ├── models/                   # Model implementations
│   │   ├── FactoST.py             # Main FactoST-v2 Model (w/ DSPA, STF, STMF)
│   │   ├── FactoST_finetune.py    # UTP fine-tune wrapper
│   │   ├── FactoST_zeroshot.py    # UTP zero-shot wrapper
│   │   └── layers/               # Building blocks
│   │       ├── revin.py           # Reversible Instance Normalization
│   │       ├── utp.py             # Encoder-Only Backbone & Attention layers
│   │       └── utils.py           # Utilities
│   │
│   ├── data/                     # Data pipeline
│   │   ├── datamodule.py          # DataLoaders construction
│   │   ├── pred_dataset.py        # Dataset class
│   │   ├── timefeatures.py        # Temporal feature encoding
│   │   └── batch_scheduler.py     # CMR Batch sampling
│   │
│   ├── learner_FactoST_finetune.py # Trainer for Fine-tuning
│   ├── learner_FactoST_zeroshot.py # Trainer for Zero-shot
│   ├── metrics.py                 # Evaluation metrics (MAE, RMSE, Quantile Loss)
│   └── utils.py
│
├── data/                         # Datasets configuration (JSONs)
├── checkpoints/                  # Pretrained backbone weights
├── vis/                          # Visualization tools
│
├── FactoST_sta.py                # Entry point for Adaptation
├── FactoST_utp.py                # Entry point for UTP
└── requirements.txt

🔧 Implementation Details

Key building blocks (where to look)

1. Data pipeline

• Entry helper: datautils.py:get_dls (builds src/data/datamodule.py:DataLoaders)
• Dataset/time features: src/data/pred_dataset.py, src/data/timefeatures.py, src/data/func_time_feature_v2.py
• Batch sampling: src/data/batch_scheduler.py

2. Models

• Spatio-temporal FactoST: src/models/FactoST.py
• Unified temporal variants: src/models/FactoST_finetune.py, src/models/FactoST_zeroshot.py

3. Training / evaluation

• Learners: src/learner_FactoST_finetune.py, src/learner_FactoST_zeroshot.py
• Prune-then-finetune pipeline: src/learner_prune_then_finetune.py + src/pruning/

4. Callbacks & wrappers

• Patch creation (including spatio-temporal patching): src/callback/patch_mask.py
• RevIN wrapper (normalize only value channel): src/callback/transforms_wrapper.py:RevInCBWrapper

---

Note: This wrapper-based implementation maintains full backward compatibility while adding powerful spatio-temporal capabilities. All original functionality remains unchanged, and spatio-temporal features can be enabled/disabled as needed.

Citations

If you find this repository useful, please cite our work:

@inproceedings{zhonglearning,
  title={Learning to Factorize Spatio-Temporal Foundation Models},
  author={Zhong, Siru and Qiu, Junjie and Wu, Yangyu and Zou, Xingchen and Rao, Zhongwen and Yang, Bin and Guo, Chenjuan and Xu, Hao and Liang, Yuxuan},
  booktitle={The Thirty-ninth Annual Conference on Neural Information Processing Systems}
}