Improving Generalizability of Graph Machine Learning Models

The School of EECS is hosting the following PhD Progress Review 1 Confirmation Seminar:

Improving Generalizability of Graph Machine Learning Models

Speaker: Haiyang Jiang
Host: A/Prof.  Sen Wang

Abstract: Urban flow prediction is a classic spatial-temporal forecasting task that estimates the amount of future traffic flow for a given location. Though models represented by Spatial-Temporal Graph Neural Networks (STGNNs) have established themselves as capable predictors, they tend to suffer from distribution shifts that are common with the urban flow data due to the dynamics and unpredictability of spatial-temporal events. Thus, we propose a novel framework named Memory-enhanced Invariant Prompt learning (MIP) for urban flow prediction under constant distribution shifts. Specifically, MIP is equipped with a learnable memory bank that is trained to memorize the causal features within the spatial-temporal graph. By querying a trainable memory bank that stores the causal features, we adaptively extract invariant and variant prompts (i.e., patterns) for a given location at every time step. Then, instead of intervening in the raw data based on simulated environments, we directly perform an intervention on variant prompts across space and time. With the intervened variant prompts in place, we use invariant learning to minimize the variance of predictions, so as to ensure that the predictions are only made with invariant features. With extensive comparative experiments on two public urban flow datasets, we thoroughly demonstrate the robustness of  MIP against OOD data.

Bio: Haiyang Jiang is a second-year Ph.D. student at the University of Queensland, Australia. Prior to his candidacy, he earned a Bachelor of Engineering from Dalian University of Technology and Master of data science from Hong Kong Polytechnic University. Haiyang is currently working on the research topic of distribution shift on graph data, advised by Dr. Rocky Tong Chen and Prof. Hongzhi Yin.