Balancing Imbalance: Data-Scarce Urban Flow Prediction via Spatio-Temporal Balanced Transfer Learning

1 Introduction

Accurate prediction of urban spatio-temporal data is crucial for smart city development, serving as an important topic in the field of urban computing that supports various tasks [Lv et al., 2014]. Although deep learning has significantly advanced spatio-temporal prediction tasks like predicting traffic flow [Qu et al., 2023; Wang et al., 2024a] and taxi demand [Geng et al., 2019], it faces challenges in scenarios where training data is scarce due to privacy concerns and varying urban developments [Wang et al., 2018]. These challenges exist in many cities such as Hong Kong [HKGov, 2025] and Liverpool [Kono, 2022], and they could potentially hindering the developments of smart city.

Inspired by transfer learning, researchers have proposed spatio-temporal transfer learning that aims to extract transferable knowledge from data-rich (source) domains to address urban computing tasks in data-scarce (target) domains, such as traffic prediction, air quality prediction [Wei et al., 2016], point of interest (POI) recommendation [Ding et al., 2020], and trajectory simulation [Wang et al., 2024b]. The strategies of existing work on learning transferable knowledge can be divided into two categories. (1) Source-target joint pre-training: Utilizing data from both source and target cities in pre-training phase, managing to learn knowledge beneficial to specific target by using some matching strategies [Wang et al., 2021; Tang et al., 2022; Jin et al., 2022]. (2) Source-only pre-training: Employing large-scale source data to train a prediction model by decoupling target from pretraining procedure, aiming to be used to any other city [Wang et al., 2018; Yao et al., 2019; Lu et al., 2022; Liu et al., 2023; Yuan et al., 2024].

∗ Corresponding author.

An ideal cross-city prediction model should be widely applicable to numerous unseen cities (i.e., universality) while also demonstrating high prediction accuracy in the specific cities where it is applied (i.e., precision). Existing research in this field often prioritizes either universality or precision in model design, posing challenges when attempting to achieve both simultaneously. Source-target joint pre-training methods tend to predict precisely on the single selected city due to the introduction of specific distribution biases from that target during pre-training, but they struggle to generalize extensively to other unselected cities and remain narrowly adaptability, as shown in Figure 1(a). Source-only pre-training methods often require multiple datasets and higher training costs to exhibit universality of unseen domains, yet they absorb distribution biases of the source cities inevitably, which may lead to sub-optimal target performances even with finetuning [Jin et al., 2022], as shown in Figure 1(b). Consequently, how to reduce the impact of harmful distribution biases from specific domains during pre-training, thereby achieving both universality and precision in generalizing to unseen cities, remains a challenging problem.

Harmful distribution biases in specific cities arise from inherent disparities between them and unseen cities. The disparities between cities are rooted in the unique characteristics of each city [Yuan et al., 2024]. Accordingly, the adverse effects of distribution biases stem from, and can be alleviated by removing this uniqueness of the data source cities. Despite the intricate factors contributing to city uniqueness, the uneven distribution of regional traffic patterns, i.e., regional imbalance, serves as its fundamental cause. For example, Washington has more cultural centers than transportation hubs, whereas Chicago is the opposite; New York exhibits a higher proportion of areas with significantly high or low traffic flow, in contrast to Chicago. Therefore, the key to building a predictive model that encapsulates universality and precision lies in identifying and calibrating regional imbalances for balanced source-only pre-training, capable of achieving the effect of "pre-train once, adapt everywhere", as illustrated in Figure 1(c).

To this end, this paper introduces STBaT, spatio-temporal balanced transfer learning, for data-scarce urban flow prediction. To quantify the regional imbalance of different cities, we propose a novel Regional Imbalance Acquisition Module (RIAM). It effectively obtains regional pattern density and inter-regional pattern similarity, serving as indicative information for balanced pre-training. After that, we devise a Spatio-Temporal Balanced Learning Module (STBLM) to train the prediction network. It facilitates knowledge transfer between similar regions and re-weights the intake of predictive supervision information to enhance the model's ability of generalization. To achieve information flow between RIAM and STBLM, we design a bidirectional complementary iterative learning algorithm to pre-train our spatio-temporal prediction module, promoting the reliable regional imbalance awareness and balanced predictive learning. In summary, we make the following contributions in this paper:

• Novel Study Perspective: To our knowledge, we are the first to study balanced cross-city transfer learning, offering a new perspective to achieve universality and precision in data-scarce urban flow prediction.
• Advanced Methodology: We propose STBaT, a novel spatial-temporal balanced transfer learning framework for urban flow prediction. By linking a RIAM and a STBLM through a bidirectional complementary iterative learning algorithm, it identifies and corrects regional imbalances, enabling the model's generalization capabilities to unseen cities.
• Strong empirical evidence: We conduct cross-city experiments across a multitude of tasks on three real world datasets. The results show the superiority of STBaT in addressing the data scarcity of urban flow prediction. The code is made publicly available at https://github.com/ShaunHao/stbat.

2 RelatedWork

Deep learning-based traffic prediction techniques can extract complex spatio-temporal relationships and accurately predict various traffic data, such as urban flow [Zhang et al., 2017; Gong et al., 2022; Chen et al., 2022; Liang et al., 2021; Ji et al., 2023; Xia et al., 2023; Wen et al., 2023] and traffic speed [Yu et al., 2017; Li et al., 2018; Guo et al., 2019; Zheng et al., 2020; Song et al., 2020; Cirstea et al., 2022; Jia et al., 2023], playing a crucial role in intelligent transportation systems. However, a common issue is data scarcity, leading to decreased accuracy in deep models. To address this, cross-city transfer learning methods have been proposed, encompassing twomain categories: source-target joint pre-training method and source-only pre-training method. The former leverages sufficient source data along with a small amount of target data for pre-training, using techniques such as domain adaptation [Wang et al., 2021; Fang et al., 2022], domain adversarial networks [Tang et al., 2022; Ouyang et al., 2024], selective meta-learning [Jin et al., 2022] and knowledge distillation [Jin et al., 2023] to alleviate data scarcity of specific target city. In contrast, the latter pre-trains models on rich data from source cities, employing strategies such as region matching [Wang et al., 2018], meta-learning [Yao et al., 2018; Lu et al., 2022], pattern bank construction [Liu et al., 2023], parameter generation [Yuan et al., 2023] and prompt learning [Yuan et al., 2024] to acquire knowledge that enables the model to adapt to more cities.

However, these two types of work focus solely on performance in specific city or generalization across a wide range of cities, lacking a combination of both. This paper differs from existing works by aiming to learn transferable knowledge that is independent of the target cities while removing distribution biases from the source cities, thereby achieving universality and precision in predictions for unseen cities.

3 Preliminaries

Definition 1 (Region). We divide city 𝑐 into a grid map of size 𝑊𝑐 × 𝐻𝑐 based on latitude and longitude, which contains 𝑊𝑐 rows and 𝐻𝑐 columns. Each grid is defined as a cell region 𝑟𝑐, and all the grids form a set of cell regions 𝑅𝑐.

Definition 2 (Spatio-temporal Urban Flow Series). In city 𝑐, we represent the time range as a set 𝑇𝑐 = {𝑡𝑐 - |𝑇𝑐 | + 1, . . . , 𝑡𝑐}, consisting of |𝑇𝑐 | evenly split nonoverlapping time intervals, where 𝑡𝑐 represents the latest timestamp of city 𝑐. Then, the spatio-temporal urban flow series in city 𝑐 is represented as 𝑋𝑐 = � 𝑥 (𝑡 ) 𝑟𝑐 |𝑟𝑐 ∈ 𝑅𝑐, 𝑡 ∈ 𝑇𝑐 � ,

where 𝑥 (𝑡 ) 𝑟𝑐 denotes the flow data of region 𝑟𝑐 at time 𝑡, e.g., the number of taxi pickups or drop-offs.

Definition 3 (Few-Shot Urban Flow Prediction). Given the urban flow series data 𝑋𝑆1 , 𝑋𝑆2 , . . . , 𝑋𝑆𝑁 from source cities 𝑆1, 𝑆2, . . . , 𝑆𝑁 and 𝑋𝑇 from a data-scarce target city 𝑇, our goal is to train a prediction model, such that it can minimize the prediction error on the test data 𝑋 ′ 𝑇 of 𝑇.

Definition 4 (Zero-Shot Urban Flow Prediction). Given 𝑋𝑆1 , 𝑋𝑆2 , . . . , 𝑋𝑆𝑁 from 𝑆1, 𝑆2, . . . , 𝑆𝑁, our goal is to train a prediction model that minimizes the prediction error on the unseen target city test data, even with no training data of it.

4 Methodology

In this section, we provide a detailed explanation of our proposed method, STBaT, and how it addresses the problem of spatio-temporal balanced transfer learning. The overall method, as shown in Figure 2, consists of two major components: the Regional Imbalance Acquisition Module (RIAM), and the Spatio-Temporal Balanced Learning Module (STBLM). First, we introduce our novel RIAM and explain how it comprehensively captures the regional imbalance characteristics of source cities from both functionality and flow perspectives. Next, we present our newly designed STBLM and describe how it assists the learning process of the spatio-temporal prediction network to acquire balanced and generalizable prediction knowledge. Finally, we outline the training algorithm for the above two components, which ensures the reliability of the acquired regional imbalance and the generalization of the learned balanced knowledge through a bidirectional complementary iterative learning algorithm.

4.1 Regional Imbalance AcquisitionModule

Capturing the imbalance inherent in regional traffic patterns is a prerequisite for eliminating distribution biases in cities. In diverse urban environments, traffic patterns are influenced by regional functionality and traffic flow factors, thus displaying distinctive imbalances. To concretize these imbalances, an intuitive approach involves modeling and accessing the distribution information of the regional traffic patterns. To implement the above intuitions, a Regional Imbalance Acquisition Module (RIAM) is proposed. RIAM utilizes a regional pattern autoencoder to extract region embeddings of source cities, representing their regional traffic patterns. A density estimator and a similarity measure are used to obtain regional pattern density and inter-regional pattern similarity, which then serve as imbalance indicators.

Regional Pattern Autoencoder

As aforementioned, the traffic pattern is influenced by both functional and flow-related factors. To simultaneously capture these two factors, a regional pattern autoencoder consisting of a regional pattern encoder and a decoder has been developed. Given normalized POI distributional tensor 𝑃𝑝𝑜𝑖 = � 𝑝𝑟𝑠 | 𝑟𝑠 ∈ 𝑅𝑠 � ∈ R |𝑅𝑠 |×𝐾 with 𝐾 categories of a city 𝑠, the encoder takes it into a feed-forward network to obtain the region embedding tensor Φ𝑠 = � 𝜙𝑟𝑠 | 𝑟𝑠 ∈ 𝑅𝑠 � ∈ R |𝑅𝑠 |×𝐷 that reflect traffic patterns, in which embedding of each region as

where 𝜃enc denotes the learnable parameters of encoder, 𝐹𝐶(·) denotes a fully connected layer, and 𝑅𝑒𝐿𝑈(·) denotes a ReLU activation layer. Then, 𝜙𝑟𝑠 is parallelly fed into a POI classifier and a spatio-temporal feature distribution predictor, which jointly form the decoder. The POI classifier outputs the prediction of 𝑝𝑟𝑠 as ˆ 𝑝𝑟𝑠 ∈ R 𝐾 , which is optimized through cross-entropy loss, aiming at encoding regional functionality into 𝜙𝑟𝑠 . The spatio-temporal feature distribution predictor predicts the mean ˆ𝜇𝑟𝑠 ∈ R 𝐽 and standard deviation ˆ𝜎𝑟𝑠 ∈ R 𝐽 of the 𝐽-dimensional regional flow features, which are refined through Kullback-Leibler Divergence loss LKLD to align with the actual feature statistics 𝜇𝑟𝑠 and 𝜎𝑟𝑠 , thereby integrating flow pattern into 𝜙𝑟𝑠 . The details of 𝜇𝑟𝑠 and 𝜎𝑟𝑠 are introduced in Section 4.2. Formally, the decoder can be defined as

where 𝜃dec denotes the learnable parameters. The total loss for the regional pattern autoencoder can be expressed as

𝛽 is a balancing hyper-parameter.

Density Estimator

As a fundamental imbalance indicator, it is essential to obtain the density of regional traffic patterns of the source cities. The density estimator derives the density of the traffic pattern distribution of a city 𝑠 based on its region embeddings 𝜙𝑟𝑠 extracted from the regional pattern autoencoder.

Since the region embedding space is a continuous space, each region's pattern is not independent of others. The more neighboring embeddings a region has in the embedding space, the higher the density of regional traffic patterns corresponding to it. Based on this intuition, the density estimator employs kernel density estimation [Rosenblatt, 1956; Parzen, 1962] to estimate the regional pattern density smoothly. It takes into account each 𝜙𝑟𝑠 and uses a specific kernel as a weight to infer the probability density function of the traffic pattern. For any given region 𝑟𝑠 ∈ 𝑅𝑠, its regional pattern density can be inferred as

where 𝐾(·) is the normal kernel function, and ℎ is the bandwidth parameter. The 𝜌𝑟𝑠 will further support balanced learning of the STBLM. For specific implementation details, please refer to Section 4.2.

Similarity Measure

Considering the continuous nature of traffic patterns 𝜙𝑟𝑠 , proximity and distance concepts are essential characteristics. Explicitly measuring inter-regional pattern similarity enhances the understanding of shared knowledge among regions, empowering balanced learning within the subsequent STBLM. The similarity measure calculates how close any two region embeddings are to obtain the inter-regional pattern similarity. Given any two regions 𝑖, 𝑗 ∈ 𝑅𝑠, the similarity measure obtains the similarity of their traffic patterns as

where 𝑠𝑖𝑚(·) denotes the cosine similarity metric. For further details on how 𝛿(𝜙𝑖, 𝜙𝑗 ) is used in STBLM, please refer to Section 4.2.

4.2 Spatio-Temporal Balanced LearningModule

Under the influence of distribution biases, training on unbalanced source data absorbs detrimental knowledge for generalizing to unseen cities. For instance, imbalanced learning may lead to a skewed distribution of extracted spatiotemporal features in regions with lower traffic pattern density, impacting the efficacy of predictive knowledge acquisition in these areas. To address these challenges, we introduce a Spatio-Temporal Balanced Learning Module (STBLM) with two functions: facilitating knowledge transfer among similar regions to rectify imbalances in feature extraction and reweighting label information to promote balanced predictive learning.

Spatio-Temporal Prediction Network

The spatio-temporal prediction network, comprising a spatiotemporal feature extractor and a feedforward predictor, can be formally expressed as

where 𝑧𝑟𝑠 is the extracted spatio-temporal feature, ˆ 𝑥 (𝑡 ) 𝑟𝑠 is the future traffic prediction, 𝜃feat, 𝜃pred denotes the learnable parameters of spatio-temporal feature extractor and feedforward predictor. Notice that STBaT has no specific structural assumptions on the spatio-temporal feature extractor, as long as it does not alter the spatial dimensions of the input data. Thus, various advanced spatio-temporal prediction models can be compatible with STBaT.

Balanced Learning for Spatio-Temporal Feature

Extraction

Regions with similar traffic patterns share common knowledge, which should be learned by the spatio-temporal feature extractor. By adjusting regional feature distributions based on inter-regional pattern similarity, knowledge can be transferred between similar regions, thus calibrating feature learning in low-density regions.

First, STBLM computes the mean 𝜇𝑟𝑠 and standard deviation 𝜎𝑟𝑠 of the spatio-temporal features 𝑧𝑟𝑠 of region 𝑟𝑠, serving as regional spatio-temporal feature statistics

where 𝐵 is number of samples in a training batch.

Then, a feedforward projector is utilized to fuse the 𝜇𝑟𝑠 and 𝜎𝑟𝑠 and extract the representation of the regional feature distribution 𝑑𝑟𝑠 , formalized as

where 𝜃proj denotes the learnable parameters of feedforward predictor, and | | is the concatenation operation.

Next, a newly developed Similarity-aware Distribution Contrastive Learning (SimDCL) is employed to adjust regional feature distributions based on inter-regional pattern similarity, thereby optimizing the spatio-temporal feature extractor. The core idea is, the feature distributions 𝑑𝑖 and 𝑑 𝑗 of two regions 𝑖 and 𝑗 should be closer when there is a higher similarity 𝛿(𝑖, 𝑗) between them. For a city 𝑠 with region set 𝑅𝑠, the SimDCL loss can be defined as

where L (𝑖,𝑘 ) dcl = - log exp(𝑠𝑖𝑚(𝑑𝑖 , 𝑑𝑘)/𝜏dcl) �|𝑅𝑠 | 𝑗 𝟙𝑗≠𝑖, 𝛿(𝜙𝑖 ,𝜙𝑘 )≥𝛿(𝜙𝑖 ,𝜙𝑗 ) exp(𝑠𝑖𝑚(𝑑𝑖, 𝑑 𝑗 )/𝜏dcl) ,

𝜏dcl is the temperature hyper-parameter, and 𝟙· is the indicator function.

Recalling from Section 4.1, the feature statistics are also utilized to optimize the decoder, in order to encode the regional flow patterns into 𝜙𝑟𝑠 . Thus, we feed the 𝜇𝑟𝑠 and 𝜎𝑟𝑠 obtained from Eq.10 into the RIAM to optimize Eq.5, promoting RIAM's comprehensive modeling of regional imbalances.

Balanced Learning for Spatio-Temporal Traffic Prediction

In existing work, the equal weighting of prediction losses across regions results in an imbalance in the model's attention to various traffic patterns, leading to a greater focus on accurately predicting high-density traffic patterns rather than low-density ones. By assigning equal weights to each traffic pattern within the loss function, rather than to each region, a more balanced focus can be achieved across diverse traffic patterns.

To balance the influence of regions with varying densities on the overall prediction loss while ensuring numerical stability, STBLM calculate the multiplicative inverse of the regional pattern density 𝜌𝑟𝑠 and scale it to transform it into the regional weight

Next, the 𝜔𝑟𝑠 is utilized to reweight the prediction loss, resulting in a density-aware weighted prediction loss, defined as

where ˆ 𝑥 (𝑡 ) 𝑟𝑠 is the prediction for the future traffic state, Lpred(·) is the typical prediction loss, e.g., mean squared error loss, etc.

4.3 Training Process

To streamline training for RIAM and STBLM without additional pre-training, synchronizing imbalance information acquisition with balanced prediction knowledge learning, we devised a bidirectional complementary iterative learning algorithm with the following steps:

1. Optimize STBLM and Obtain the Regional SpatioTemporal Feature Statistics. Using the current STBLM with 𝜃STBLM = {𝜃feat, 𝜃pred, 𝜃proj} and 𝜚 = � 𝜌𝑟𝑠 | 𝑟𝑠 ∈ 𝑅𝑠 � , Δ = � 𝛿(𝜙𝑖, 𝜙𝑗 ) | 𝑖, 𝑗 ∈ 𝑅𝑠 � derived from the current RIAM, the losses Ldcl and � Lpred are computed to optimize 𝜃STBLM. Then, the N = � 𝜇𝑟𝑠 | 𝑟𝑠 ∈ 𝑅𝑠 � , Σ = � 𝜎𝑟𝑠 | 𝑟𝑠 ∈ 𝑅𝑠 � are obtained to provide guidance for the subsequent optimization of RIAM.
2. Optimize RIAM and Derive the Regional Imbalance Information. Using the current RIAM with 𝜃RIAM = {𝜃enc, 𝜃dec} and the obtained N, Σ from the current STBLM, the loss Lemb is calculated to optimize 𝜃RIAM. After this, 𝜚 and Δ are derived from the current RIAM to aid in the next optimization of STBLM.

By iteratively executing the above two steps, a bidirectional connection is established between RIAM and STBLM, contributing to a balanced pre-training process of spatio-temporal prediction network, i.e., 𝜃feat and 𝜃pred. As a result, the pretrained 𝜃feat and 𝜃pred can be adapted to urban flow prediction task in new cities with limited or no training data.

5 Experiments

5.1 Experiment Settings

Datasets

Following previous work, we take urban flow prediction as an example task, and evaluate our proposed framework on real-world public datasets of three cities: New York (NY), Chicago (CHI), and Washington (DC), which contain vehicle pickup and drop-off records of bike and taxi. Each dataset covers a time range of one year with time intervals of one hour. The detailed statistics of datasets are shown in Table 1. Additionally, we also use public POI data of each city. All data are collected and opened by [Jin et al., 2022].

Task Settings

We use a similar few-shot traffic prediction setting to [Wang et al., 2021; Jin et al., 2022]. We choose each of three cities as the target city, while using one of (New York or Chicago) or all of the remaining cities as the source cities. This creates a total of eight tasks of cross-city transfer. For each source city, we divided the training and validation sets in a 2:1 ratio. For each target city, we allocate the last 2 months for testing, the preceding 2 months for validation, and the last 3 days or 0 days before the validation data for training, forming two data-scarce scenarios: few-shot and zero-shot. We first pretrain our framework on source training data, then fine-tune it on target training data. After that, we evaluate it on target test data. Min-max normalization is applied for data preprocessing.

Baselines

We compare the performance of STBaT and a number of baselines for urban flow prediction. Based on whether requires source data, the baselines can be classified into nontransfer baselines and transfer baselines.

• Non-transfer Methods: The non-transfer baselines consist of three methods, including ARIMA [Box and Jenkins, 1968], a statistical method, and two deep learning methods, i.e., ST-net [Yao et al., 2019] and PDFormer [Jiang et al., 2023].
• Transfer Methods: For transfer baselines, we select two types of learning-based approaches: four source-only pre-training methods, including vanilla finetuning (Finetuned), MAML [Finn et al., 2017] and MetaST [Yao et al., 2019]; and two source-target joint pre-training methods, consisting of ST-DAAN [Wang et al., 2021] and CrossTReS [Jin et al., 2022].

For ARIMA models, we employ 6 autoregressive steps, 1 moving average step, and 1 integration step. Following previous work, we use the official code and hyper-parameters reported by the original papers of PDFormer, MAML, MetaST, ST-DAAN, RegionTrans and CrossTReS.

Implementation Details

We implement STBaT using PyTorch. For the dimensions of FCs in the regional pattern autoencoder, we set them to (16, 32) for the encoder, (32, 14) for the POI classifier and (32, 256) for the spatio-temporal feature distribution predictor. We choose the ST-net model as the spatio-temporal feature extractor, stacking three residual blocks with 64 channels and a single-layer LSTM with 128 hidden units. For the STBLM, the dimensions of the FC layers of the feedforward predictor are set to 256 and 1, and the dimensions of the feedforward projector are set to 256 and 256. To capture temporal dependencies, we set the horizon 𝜏 to 6, meaning that the input data consisted of observations from the previous 6 intervals, to predict the urban flow in next time step. We train the model until the validation error does not decrease for 5 consecutive epochs on the source data, and then select the model with the lowest source validation error as the initialization for target adaptation. We evaluate the target performance using root mean squared error (RMSE) and mean absolute error (MAE). The mean error for both the pickup and drop-off predictions is reported.

5.2 Evaluations

Performance Comparison of Few-Shot and Zero-Shot Prediction

We evaluate each method in both few-shot and zero-shot scenarios, and report the average results of 5 independent runs in cross-city tasks in Table 2.

In few-shot prediction scenario, deep learning methods based solely on target data i.e., ST-net and PDFormer, exhibit unsatisfactory performance. Among the transfer baselines, source-target joint pre-training methods (e.g., ST-DAAN and CrossTReS) generally outperform source-only pre-training methods (e.g., MetaST and RegionTrans). This advantage stems from the utilization of few-shot target data during pretraining to help adapting to specific target city. Notably, STBaT does not utilize any target data during pre-training, whereas it reduces RMSE by up to 6.3% and MAE by up to 11.6% compared to the best source-target joint pre-training method. This improvement can be attributed to the balanced learning strategy, which acquires a favorable initialization that internalizes knowledge beneficial for generalization, thus facilitating precise transfer to the target domain.

In zero-shot prediction scenario, the non-transfer baselines and some transfer baselines that utilize target for pre-training are not applicable. Among the source-only pre-training methods, MetaST yields lower error compared to other baselines, resulting in suboptimal performance. This may be attributed to its model design, which includes memory units that can capture transferable patterns from the source cities. Thanks to the capability to identify and calibrate distribution biases arising from imbalance in the source cities, STBaT once again achieves optimal performance in zero-shot scenarios. It enhances RMSE by 5.8% and MAE by 9.4% compared to the best source-only pre-training method, showcasing the superiority of our proposed framework in situations of severe data scarcity.

Across all baselines, although source-only pre-training methods can fit to more data-scarce scenarios, they fail to outperform source-target joint pre-training methods across all tasks. Consequently, no single baseline achieves both universality and precision, except for STBaT, which consistently outperforms other methods in all tasks and scenarios.

Impact of Components

To assess the efficacy of each design within STBaT, we create two variants of our framework.

• (a) Impact of model structure

(b) Impact of training loss

Table 3: Training memory and time costs of different approaches

• Variants of model structure: Including removing the POI classifier (w/o POI) and removing the spatiotemporal feature distribution predictor (w/o FDP).
• Variants of training loss: Including removing the similarity-aware distribution contrastive learning loss (w/o SimDCL) and removing the density-aware weighted prediction loss (w/o WLoss).

Figure 3 depicts the RMSE of STBaT variants in few-shot taxi flow prediction task of DC. For the model structure, STBaT w/o POI and w/o FDP both lead to significant performance drop, highlighting the importance of capturing regional functionalities and flow characteristics for modeling urban regional imbalances. For the training loss, both STBat w/o SimDCL and w/o WLoss result in accuracy degradation, attributed to unbalanced learning in feature extraction and traffic prediction, respectively.

Impact of Hyper-parameters

To assess the impact of the hyper-parameters on STBaT's performance, we conduct analyses in few-shot taxi flow prediction tasks of DC. The results are shown in Figure 4. (1) We vary the dimension 𝐷 of the region embedding, and set it to 32. This allows us to effectively capture the traffic patterns. (2) We tune the parameter 𝛽 of embedding loss and find that a value of 0.01 consistently yields the best performance in different prediction tasks. (3) We adjust the bandwidth ℎ of the RIAM's density estimator, and set it to 0.85 to achieve the best performance. (4) We set the temperature 𝜏dcl to 1 to better learn balanced spatio-temporal feature extraction, adapting effectively to different target city prediction tasks.

Impact of Data Scarcity

We conducted an experiment based on taxi flow prediction tasks to investigate the performance of our method under different scarcities of target data. Figure 5 shows that as the amount of available data in the target city increases, the performance on different cross-city tasks gradually improves and eventually approaches saturation. The significant improvement from 0 days to 3 days of data indicates that model can adapt well to the target city, even with minimal data.

Figure 5: Performance with different scarcities of data

Efficiency Study

We studied the training efficiency of STBaT against several representative baselines, with detailed comparisons shown in Table 3. The results demonstrate that while vanilla finetuning achieves the lowest training costs, its performance lags significantly behind competitors; STBaT substantially outperforms MetaST and CrossTReS in both memory and time consumptions through its effective learning algorithms.

6 Conclusion

This paper introduces a spatio-temporal balanced transfer learning framework (STBaT), a novel paradigm to enhance the universality and precision of pre-training a model for datascarce urban flow prediction. By leveraging a Regional Imbalance Acquisition Module and a Spatio-Temporal Balanced Learning Module, STBaT demonstrates its superiority in fewshot and zero-shot scenarios of diverse prediction tasks. For future work, integrating the idea of STBaT into various urban computing tasks, such as graph-based predictions, shows promise. Our study paves an encouraging path for future spatio-temporal generalization learning.

Ethical Statement

There are no ethical issues.

Acknowledgments

This work was supported by the National Natural Science Foundation of China (No. 62272033).