UniST: A Prompt-Empowered Universal Model for Urban Spatio-Temporal Prediction

In plain words

Cities produce many grid-shaped streams of numbers over time: taxi pickups, crowd counts, bike rentals, phone-network load, road speeds. Usually a team trains one forecasting model per stream, from scratch, and the model does not carry over to a new city or a new kind of data. UniST tries to build one model that covers many such streams.

It is trained in two stages. First, a Transformer learns to fill in hidden parts of many city datasets at once, the way a masked word model learns language. Second, a small add-on network looks at the recent history and the same time on earlier days, and produces a short "hint" vector that is fed into the model. The hints are stored in learned memory slots, so the model can pick up patterns it has seen before. The idea is that a new dataset may share some local pattern with the training data, and the hint lets the model reuse it, even with little or no data from the new place.

TL;DR

• One Transformer encoder-decoder, pre-trained with four masking strategies on many city grids (traffic speed, taxi, bike, crowd, cellular) §3.3-§3.4.
• A memory-based prompt network, built from spatial closeness and hierarchy and temporal closeness and period, adds a prompt to the decoder §3.5.
• Author-reported: best short-term RMSE and MAE on the seven datasets shown tab. 2; largest gains in few-shot and zero-shot settings §4.4-§4.5.
• The official code has no license file. Pseudocode in this card is written from the code, which is cited, not quoted.

Problem

Urban spatio-temporal data comes in varied formats: different numbers of variables, time spans, spatial sizes and sampling rates §1. Existing predictors are built for one scenario and need a lot of target-domain data; the paper argues this blocks a universal model and makes few-shot and zero-shot use hard §1, §2. The paper assumes that fine-grained patterns in a new dataset resemble some patterns seen in training data, and that distinct patterns need distinct prompts §3.5.1.

Method

Input and grid. Each dataset is a tensor of time steps, variables and an H x W grid of equal cells §3.1. Each channel is processed independently as a 3D tensor (time, height, width) §3.3. A 3D convolution with kernel size equal to stride (l, h, w) turns it into a token sequence §3.3. The paper does not give the patch size; the code uses 2 spatially and 2 in time [code: src/main.py#L52-L53], Embed.py.

Positional information. The paper uses fixed sine and cosine encodings, applied separately to space and time §3.3. The code's default is the same (SinCos) and freezes them main.py, [code: src/model.py#L403-L426].

Base model. An MAE-style encoder-decoder §3.3. The encoder sees only the unmasked tokens; the decoder receives the encoder output plus a learned mask token at each hidden position. Unlike MAE, the decoder is full-sized and is used in both stages §3.3. The code's default size is called middle: encoder width 128 with 6 layers and 8 heads, decoder width 128 with 4 layers and 4 heads, MLP ratio 2 model.py. The paper's model-size table lists this same configuration as "10M Params" app. D.2; the cost table reports 6.71M for UniST app. E.3. See Pitfalls.

Stage 1: pre-training §3.2, §3.4, alg. 1. For each step, pick one dataset and a mini-batch, pick one of four masks, mask the input, reconstruct, and take a mean squared error step.

• Random masking: hide tokens at random §3.4.
• Tube masking: hide the same spatial cells at all time steps, which imitates sensors that are offline §3.4.
• Block masking: hide a whole block of spatial cells at all time steps §3.4.
• Temporal masking: hide the future steps and predict them from the past §3.4.

# Pre-training on many datasets  [alg. 1] [code: src/train.py#L152-L159]
# Input:  datasets D_1..D_M, base model F, MSE loss
# Output: pre-trained parameters theta
def pretrain(batches, model, optimizer):
    # One mini-batch comes from one dataset [code: src/train.py#L178]
    for dataset_name, batch in batches:
        # Choose a masking strategy uniformly; the ratio is always 0.5
        # [§3.4] [code: src/train.py#L30, L157-L158]
        strategy = random_choice(["random", "temporal", "tube", "block"])
        ratio = 0.5
        # Mask, then reconstruct  [§3.3] [code: src/model.py#L729-L735]
        predicted, target, hidden = model(batch, ratio, strategy)
        # MSE on hidden tokens only  [alg. 1] [code: src/model.py#L720-L724]
        loss = mean_squared_error(predicted[hidden], target[hidden])
        # Clip grad norm to 0.05, then step  [code: src/train.py#L33-L39]
        clip_gradient_norm(model, max_norm=0.05)
        optimizer.step()

Masking ratio. The paper names a masking percentage r but does not give its value §3.4. The code fixes r = 0.5 for all four strategies [code: src/train.py#L30]. The block-masking code only accepts r in {0.25, 0.5, 0.75} mask_strategy.py.

Stage 2: spatio-temporal knowledge-guided prompt learning §3.2, §3.5, alg. 2. A prompt network is trained on top of the pre-trained base model, sampling from the same dataset collection §3.2, alg. 2. The network has four knowledge inputs §3.5.2:

• spatial closeness: a 2D convolution with kernel 3 over the history §3.5.2;
• spatial hierarchy: convolutions with larger kernels, 2^i+1 for i in {2, 3, 4} §3.5.2;
• temporal closeness: an attention network over the last M steps §3.5.2;
• temporal period: an attention network over the same time slots from the previous N days §3.5.2.

Each knowledge feature queries a learned key-value memory, and the retrieved memory vector is the prompt §3.5.3. Memory pools are learned: one spatial pool and one temporal pool, each with N slots of key and value §3.5.3.

# Memory lookup that turns one knowledge feature into a prompt  [§3.5.3]
# features: (N, C) knowledge features; keys, values: (M, C) learned memory
def memory_prompt(features, keys, values):
    # Project the query and squash it  [code: src/Prompt_network.py#L59]
    query = tanh(linear(features))                  # (N, C)
    # Dot each query with every key, normalise over the M slots
    # [code: src/Prompt_network.py#L61-L63]
    weights = softmax(query @ keys.T, axis=-1)      # (N, M)
    # The prompt is the weighted sum of the memory values
    # [code: src/Prompt_network.py#L65]
    return weights @ values                         # (N, C)

Insertion. The paper writes the prompt as an addition to the input, X_p = X + P alg. 2. The code adds the summed prompts to the decoder tokens of the future (predicted) positions only, after the positional and time embeddings model.py. The default prompt_content is s_p_c, which sums the spatial, temporal-period and temporal-closeness prompts model.py, main.py.

Training in stage 2 alg. 2, app. D.3. The paper's loss is MSE plus an L2 penalty with weight lambda, which is not stated. The code uses MSE on the hidden (future) tokens and AdamW weight decay of 1e-6 [code: src/train.py#L22], main.py. See Paper vs code.

Hyperparameters

Results

Author-reported. Setup: 6-step input and 6-step horizon, one dedicated model per baseline, UniST trained on the dataset collection and evaluated on each dataset §4.2, app. D.5. The paper does not say which checkpoint produced the UniST row. In the code, the pre-training evaluation reports a 'temporal' mask entry, which hides the same six future steps; this is the closest analogue to the forecast, and it is an inference from the code train.py. Errors are in the paper's table units; the appendix says Crowd and Cellular values should be scaled by 10^3 app. A.1. The baseline column is the paper's underlined second-best value, which is the best baseline for each metric.

Short-term tab. 2, the seven datasets in the main table; the text says the full results are in the appendix §4.2, app. E:

Long-term, 64-step input and horizon tab. 3; the paper says UniST's average improvement over the best baseline on each dataset is 10.1% §4.3:

Few-shot and zero-shot §4.4-§4.5, app. E, Crowd and BikeNYC tables. The paper's few-shot tables use 10%, 5% and 1% of the training data. Zero-shot uses no target training data. Baselines are the best value in each column of the same table, by RMSE and by MAE separately. Crowd, UniST RMSE / MAE:

BikeNYC, UniST RMSE / MAE:

These are author-reported app. E. On both datasets the zero-shot error is higher than the 1% few-shot error, and on BikeNYC the zero-shot RMSE is above the best 1% few-shot baseline's RMSE (8.50).

Ablations §5.1, fig. 5, app. E.5: removing the spatial-closeness-and- hierarchy, temporal-period or temporal-closeness prompt each lowers performance on the four traffic-speed datasets §5.1, fig. 5. Masking ablation on TrafficCD, columns prediction / imputation / spatial extrapolation; the metric is not labelled in the table app. E.5, mask table:

Noise app. E.2, noise table, sine-cosine encoding. The paper says UniST "consistently outperforms" baselines under noise. At 10% Gaussian noise, the TaxiBJ row is 28.76 against a best baseline of 27.36, so the claim does not hold for that dataset. The other four are below their baselines (Crowd 3.29 vs 3.85; Cellular 14.91 vs 16.48; BikeNYC 3.695 vs 3.93; TrafficSH 0.6877 vs 0.742). With learnable positional encodings the 10% row is much worse (TaxiBJ 43.80) app. E.2.

Efficiency app. E.3, cost table: UniST 6.71M parameters, 2875 MB memory, training 1.4 min per epoch over 20+ datasets, inference 0.034 min. PatchTST: 2.59M, 2859 MB. iTransformer: 25.27M, 2935 MB.

Scaling §5.3, fig. 8: training loss and masked-reconstruction RMSE fall with more steps for every model size. Larger models converge faster, with diminishing returns; the paper attributes the flattening to limited data diversity.

Paper vs code

The official code is at tsinghua-fib-lab/UniST at code_commit. The repository has no license file, so no code is quoted; all points below are cited by path and line.

• Prompt lookup: the paper writes a dot product of the query with the keys, then a weighted sum of values §3.5.3. The code applies a linear layer and tanh to the query, then a softmax over the slots Prompt_network.py.
• Spatial kernels: the paper uses kernel 3 for closeness and 5, 9, 17 for hierarchy §3.5.2. The code uses kernels 1, 3 and 5 in one stack Prompt_network.py.
• Temporal inputs: the paper's closeness uses the previous M steps §3.5.2. The code's closeness uses the decoder's token stream, which includes the encoder output for history and mask tokens for the future. Spatial prompts use the raw history frames, and period uses raw period values passed through the patch convolution model.py, model.py.
• Prompt insertion: Algorithm 2 adds the prompt to the input, X_p = X + P alg. 2. The code adds it to the future decoder tokens after positional and time embeddings model.py.
• Freezing: Algorithm 2 says the attention and feed-forward layers are fixed, and then updates theta as well as the prompt parameters alg. 2. The code does not freeze any of these layers; every parameter with requires_grad is in the optimizer train.py.
• Block masking: the paper's block is a spatial sub-block over both height and width §3.4. The code takes one half of the flattened spatial index, which is a contiguous set of rows mask_strategy.py.
• Masking ratio: the paper leaves r open §3.4. The code fixes r = 0.5 train.py.
• Few-shot data: the paper trains on 1%, 5% and 10% of the target training set §4.4. The subsampling code checks args.mode == 'few-shot' [code: src/DataLoader.py#L96-L97], but main sets the mode only after data is loaded main.py and sets it to prompting or testing main.py. The subsampling never runs.
• Zero-shot path: with few_ratio = 0.0, main sets testing, and run_loop evaluates the validation split, not the test split train.py.
• Validation interval: the paper evaluates every ten epochs app. D.2. The code evaluates every 20 epochs, at epoch 10 and at the last epoch train.py.
• Checkpoint choice: the paper picks the best model on validation app. D.2. The code compares the mean normalised masked MSE over datasets and masks, although the argument is named rmse train.py, train.py.
• Evaluation dropout: the code never calls model.eval(), so the 0.1 embedding dropout stays active during validation and testing Embed.py.
• Prediction clamp: the code clamps predictions to [-1, 1] before computing RMSE and MAE train.py.
• Scaler: the min and max are taken over train, validation and test DataLoader.py.
• Model size: the code's middle default matches the paper's "10M" row app. D.2 but not the 6.71M in the cost table app. E.3. The parameter count was not computed here; check by instantiating the model.
• Learning rates: pre-training 3e-4 matches the paper app. D.2 and the script pretrain.sh. The prompt-tuning script uses 3e-4 prompt_tuning.sh, not the 5e-5 the paper reports.
• Pre-training data: the released pre-training script lists 21 datasets, including Crowd pretrain.sh. The few-shot and zero-shot scripts use Crowd as target few_shot.sh, zero_shot.sh. The paper says the zero-shot dataset has no training data §4.5; the scripts do not show how that was enforced.
• Checkpoint for zero-shot: the zero-shot script loads a pretrained path with strict=False main.py. Pre-training runs do not build the prompt network model.py, so the prompt weights come from a prompt-tuned checkpoint. The paper does not say which checkpoint was used.
• Grid sizes: the code computes batch_size from H + W, with branches for sums below 64 only DataLoader.py. TaxiBJ and TDrive are 32 x 32 in the paper's dataset table app. A.1, so H + W = 64 matches no branch. Not run; the dataset files are not in the repository.
• Embedding variant: TDrive and BikeNYC2 use a separate 24-slot hour embedding for both the encoder and the period input model.py, model.py. The paper does not mention this.

Pitfalls

• Algorithm 2 has an undefined symbol. It updates gamma and theta, but gamma is never defined and the algorithm never says which parameters are fixed alg. 2. The code trains every parameter (see Paper vs code).
• Few-shot results from this code are not 1/5/10% results. The subsampling line never runs; the few-shot scripts train on the full training split (see Paper vs code).
• Zero-shot from the released scripts is validation-split numbers. The testing branch evaluates the validation split and prints it as validation output train.py.
• Dropout is on during evaluation. Without model.eval(), the 0.1 embedding dropout also affects validation and test predictions Embed.py.
• Evaluation reseeds the global RNG. random_masking_evaluate calls torch.manual_seed(111) each time it runs mask_strategy.py. Other evaluation masks use np.random for the half-block choice mask_strategy.py. Results depend on call order.
• Early stopping ends the process. After five evaluations without improvement, the code writes an "Early stop!" line to the results files and calls exit() train.py.
• Normalisation sees the test split. The min-max scaler includes test and validation values DataLoader.py. Denormalised errors depend on that range.
• The prompt network needs a pre-trained base. Prompt weights do not exist in a pre-training checkpoint model.py. Loading one with strict=False leaves them at random initialisation without an error main.py.
• Training loss covers hidden tokens only. The loss is on hidden tokens model.py. For temporal masking those are the future steps, so errors on the visible history are never trained or scored.
• Units differ across tables. Crowd and Cellular values are scaled by 10^3 in appendix notes app. A.1; the few-shot and zero-shot Crowd values need the same care. The masking-ablation table does not label its metric app. E.5.
• Paper text overstates the noise result. "Consistently outperforms" is not true for TaxiBJ at 10% noise app. E.2.
• Missing details that change results. Batch size, number of epochs, patch size, masking ratio, learning-rate schedule, gradient clip and seed are not in the paper. The code fixes them; reproduce from the code and cite it.
• Cost. Pre-training across all datasets takes longer than per-dataset baselines app. E.3. Memory is about 2.9 GB in the cost table, similar to the baselines. The table's "1.4 min" is one epoch over 20+ datasets, not a full run.

Lineage

• Builds on: masked autoencoders (MAE, He et al., 2022) for the masked-patch base model §3.3; the memory-network prompt idea from Sukhbaatar et al. (2015) §3.5.3. Neither has a card on this wiki.
• Compared against: STResNet, ACFM, MC-STL, STGSP, STNorm, STID and PromptST for urban prediction; PredRNN, MAU, MIM, SimVP and TAU for video prediction; PatchTST and iTransformer for time series; MAML and MetaST for meta-learning §4.1. None has a card on this wiki.
• Superseded by: not known from wiki content.