UniFlow: Unified Urban Flow Prediction

• UniFlow is a unified model that converts heterogeneous grid and graph urban flow data into a common sequential token space.
• It employs a shared encoder-decoder transformer enhanced by spatio-temporal memory retrieval to capture recurring urban patterns.
• The framework enables joint training across diverse datasets, improving performance in limited-data and cross-city scenarios.

UniFlow is a foundation model for unified urban spatio-temporal flow prediction that treats grid-based and graph-based urban flow data within a single modeling framework rather than maintaining separate CNN-style and GNN-style families. Its defining premise is that structurally different urban flow datasets still share spatio-temporal laws such as locality, periodicity, temporal continuity, and recurring urban activity patterns, and that these commonalities can be exploited by converting both data types into a common sequential token space, processing them with a shared spatio-temporal transformer, and augmenting prediction with learnable retrieval memory (Yuan et al., 2024).

1. Problem formulation and conceptual scope

UniFlow addresses urban spatio-temporal flow prediction in the generic form

$X \in \mathbb{R}^{T \times N \times C},$

where $T$ is the temporal length, $N$ is the number of spatial locations, and $C$ is the number of variables or channels. The forecasting task is defined as

$\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$

with $H$ historical observations mapped to a prediction horizon of length $P$ (Yuan et al., 2024).

The paper places under this formulation several urban data types, including crowd flows, taxi flows, traffic speeds, and dynamic population flows. It distinguishes two standard spatial organizations. In the grid-based case, the city is represented as

$$X \in \mathbb{R}^{T \times H \times W \times C}, \quad H \times W = N,$$

which preserves regular Euclidean locality. In the graph-based case, the city is represented as

$X \in \mathbb{R}^{T \times N \times C},$

with topology

$$\mathcal{G} = (\mathcal{V}, \mathcal{E}, \mathcal{A}), \qquad \mathcal{A} \in \mathbb{R}^{N \times N},$$

which captures irregular connectivity such as roads and intersections (Yuan et al., 2024).

A central motivation is that prior work usually separated these two formats into different paradigms: CNN / ResNet / video-prediction style models for grids, and GNN / graph recurrent / graph transformer models for graphs. UniFlow instead proposes a single “one-for-all” model. Its stated goal is to accept both data types, convert them into a unified token sequence format, learn shared spatio-temporal patterns across datasets and data types, and improve performance especially under limited-data regimes through cross-learning. This suggests that the paper uses “foundation model” in the sense of joint multi-dataset, multi-structure training rather than a separate pretraining/fine-tuning pipeline.

2. Unified representation of grids and graphs

A key architectural step is the conversion of heterogeneous spatial structures into a common sequential representation

$T$0

where $T$1 is the number of patches and $T$2 is the embedding dimension. UniFlow adopts channel independence following PatchTST-style practice, so each channel is treated as a separate $T$3 spatio-temporal field before patching (Yuan et al., 2024).

The paper uses different patch extractors for grids and graphs, but both produce the same tokenized output type.

Data type	Input form	Patching mechanism
Grid	$T$4	3D convolution with kernel and stride $T$5
Graph	$T$6, with $T$7	Conv1D temporal extraction plus METIS subgraph partitioning

For grid data, UniFlow applies a 3D convolution: $T$8 followed by flattening,

$T$9

The number of tokens is

$N$0

This creates spatio-temporal cuboid patches ordered in scan-line order. The mechanism is explicitly local in both time and 2D space (Yuan et al., 2024).

For graph data, UniFlow uses a two-stage process: $N$1

$N$2

$N$3

$N$4

Here each token corresponds to a subgraph patch rather than a Euclidean cuboid. The use of METIS is central because graphs lack regular Euclidean neighborhoods and therefore cannot use the same patch geometry as grids (Yuan et al., 2024).

The paper calls this patching mechanism “multi-view” partly because the broader system extracts patterns from multiple complementary perspectives—time-domain, frequency-domain, time-derived spatial, and frequency-derived spatial—and partly because distinct structure-aware patchers are unified into one sequential format. A common misconception is that “multi-view” here primarily refers to multimodal input; the paper instead defines it through complementary spatio-temporal perspectives and structure-specific tokenization.

3. Transformer backbone and end-to-end forecasting pipeline

After patchification, both grid and graph data are processed by the same encoder-decoder transformer. The end-to-end pipeline contains three modules: spatio-temporal patching, sequential modeling with a transformer, and Spatio-Temporal Memory Retrieval Augmentation. The operational flow is raw grid or graph data to patch tokens $N$5, then masking of historical versus future positions, encoder processing of history, decoder prediction of future tokens, and projection back to flow values (Yuan et al., 2024).

The encoder-decoder process is written as

$N$6

$N$7

$N$8

$N$9

where $C$0 is a mask token embedding for future positions. The transformer is described as a stack of $C$1 identical layers, each containing multi-head self-attention and a feed-forward neural network. The paper does not describe a factorized spatial/temporal attention design; instead, attention is applied jointly over the token sequence induced by the patching stage (Yuan et al., 2024).

Local spatio-temporal structure is therefore injected before attention, not through separate spatial and temporal operators inside the backbone. The paper’s own interpretation is explicit: local structure comes from patching, global spatio-temporal dependency comes from self-attention, and reusable pattern priors come from memory retrieval. This is one of UniFlow’s main design commitments.

The implementation settings reported in the experimental configuration are concise: 4 encoder layers, 4 decoder layers, hidden size 256, and 512 embeddings in each memory. Optimization uses Adam with learning rate $C$2, decayed to $C$3, for at most 200 epochs with early stopping, on an A100 GPU (Yuan et al., 2024).

The paper does not provide a separate positional encoding equation. It explicitly notes that transformer embeddings and mask token embedding are used, but the main text does not spell out a standalone positional encoding formula. This is important because readers might expect a standard ViT-style positional design to be central; the paper leaves that part implicit.

4. Spatio-Temporal Memory Retrieval Augmentation

The most distinctive mechanism in UniFlow is Spatio-Temporal Memory Retrieval Augmentation, abbreviated ST-MRA. Its purpose is to store and retrieve reusable spatio-temporal pattern prototypes shared across datasets, cities, and data types. The paper’s motivating examples include rush hours, periodicity, central business district dynamics, and residential regularity (Yuan et al., 2024).

ST-MRA constructs four query views from the historical token sequence $C$4:

1. time-domain,
2. frequency-domain,
3. time-derived spatial,
4. frequency-derived spatial.

The temporal queries are

$C$5

The spatial queries are derived from adaptive topologies: $C$6 then

$C$7

The paper emphasizes that these are not purely spatial representations, but entangled spatio-temporal ones (Yuan et al., 2024).

Memory is organized into four external banks: time-domain, frequency-domain, time-spatial, and frequency-spatial. Each stores learnable key-value pairs $C$8, with shape $C$9, where here $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$0 denotes the number of memory units rather than the number of spatial locations. In the experiments, each memory contains 512 embeddings (Yuan et al., 2024).

Retrieval is fully differentiable: $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$1

$\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$2

$\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$3

The retrieved prompt $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$4 is then added to the decoder input embedding for augmentation. The paper is explicit that this is not classical non-parametric retrieval-augmented generation. The memory is fully learnable and end-to-end trainable. A common misconception is to treat ST-MRA as a static external database; the paper instead presents it as a trainable memory-augmented prediction module.

The empirical interpretation offered in the case study is that central urban areas in Beijing and Shanghai retrieve similar prompts, as do residential regions in both cities. This suggests that the memory is meant to encode cross-city pattern reuse rather than merely dataset-specific templates.

5. Training regime, datasets, and empirical behavior

UniFlow is not trained through a separate pretraining/fine-tuning pipeline for its main results. Instead, it is trained from scratch, using all datasets simultaneously, in an end-to-end way. For each batch, one dataset is randomly sampled so that the model can jointly train across different data types. This detail is central to the paper’s foundation-model framing (Yuan et al., 2024).

All datasets are normalized to $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$5, while reported prediction metrics are computed on denormalized values. Each dataset is split into 60% train, 20% validation, and 20% test. The loss is MSE over the prediction horizon only: $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$6 The paper does not describe auxiliary losses, contrastive objectives, or explicit regularization terms (Yuan et al., 2024).

To handle dataset size variation, batch size is adjusted according to

$\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$7

where $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$8 is the number of training samples in a dataset and $\hat{X}_{H+1:H+P} = \mathcal{F}(X_{1:H}),$9 is a predetermined number of iterations. This is intended to keep iterations per epoch roughly similar across datasets while respecting GPU memory constraints.

The evaluation uses nine real-world datasets. The grid-based datasets are TaxiBJ, TaxiNYC, CrowdNJ, CrowdBJ, FlowSH, and PopSH. The graph-based datasets are TrafficBJ, TrafficSH, and TrafficNJ. The paper states that these datasets were deliberately chosen from cities where both graph and grid data are available, so that cross-type and cross-city learning can be meaningful (Yuan et al., 2024).

The central empirical claim is that UniFlow outperforms existing models on both grid-based and graph-based flow prediction, with especially strong behavior in limited-data scenarios. The paper further presents cross-learning as particularly beneficial in few-shot settings, zero-shot transfer, noisy settings, and large graph scenarios. Because the detailed metric tables are not fully reproduced in the provided text, the strongest encyclopedic statement is qualitative: the reported advantage is broad, and the paper identifies data scarcity and cross-domain transfer as the settings where the unified design is most consequential.

6. Interpretation, scope, and common misconceptions

UniFlow’s contribution is not simply “using a transformer for urban flow prediction.” Its distinctive claim is that a common token sequence representation can mediate between grid-based and graph-based urban data, after which a shared transformer and a learnable memory can exploit recurring spatio-temporal regularities across cities and spatial structures. This is the paper’s central unification thesis (Yuan et al., 2024).

Several points are easily misunderstood. First, UniFlow is not presented as a pretrain-then-finetune foundation model; its main results come from joint end-to-end training across datasets. Second, ST-MRA is not a frozen retrieval system but a learnable memory of prototypes. Third, the transformer stage does not separately factorize spatial and temporal operators; the separation of roles is instead distributed across patching, self-attention, and memory retrieval. Fourth, “multi-view” refers primarily to multiple spatio-temporal perspectives rather than multiple modalities.

The method’s scope is urban spatio-temporal forecasting across grids and graphs, not universal spatio-temporal modeling. The data formulation, patching operators, and memory views are all specialized to urban flow settings. A plausible implication is that the model’s strongest transfer should occur when different datasets share the periodicity, locality, and recurrent activity signatures that the memory banks are designed to store. The paper’s Beijing–Shanghai case study is consistent with that reading.

Within the literature of urban flow prediction, UniFlow can therefore be characterized as a unified tokenization-and-retrieval framework: it standardizes heterogeneous spatial structures into a common sequence space, processes them through a shared encoder-decoder transformer, and augments decoding with four-view spatio-temporal memory. Its significance lies less in any single module than in this combined claim that shared urban laws can be operationalized across grid and graph data inside one model (Yuan et al., 2024).