WED-Net: A Weather-Effect Disentanglement Network with Causal Augmentation for Urban Flow Prediction

TL;DR

WED-Net is a novel dual-branch Transformer model that disentangles weather effects from urban traffic patterns using causal augmentation, improving prediction robustness under extreme weather conditions.

Contribution

The paper introduces WED-Net, a new architecture with causal augmentation and disentanglement mechanisms for better urban flow prediction during extreme weather events.

Findings

WED-Net outperforms existing models under extreme weather conditions.

Disentanglement improves prediction accuracy and robustness.

Causal augmentation enhances generalization to rare scenarios.

Abstract

Urban spatio-temporal prediction under extreme conditions (e.g., heavy rain) is challenging due to event rarity and dynamics. Existing data-driven approaches that incorporate weather as auxiliary input often rely on coarse-grained descriptors and lack dedicated mechanisms to capture fine-grained spatio-temporal effects. Although recent methods adopt causal techniques to improve out-of-distribution generalization, they typically overlook temporal dynamics or depend on fixed confounder stratification. To address these limitations, we propose WED-Net (Weather-Effect Disentanglement Network), a dual-branch Transformer architecture that separates intrinsic and weather-induced traffic patterns via self- and cross-attention, enhanced with memory banks and fused through adaptive gating. To further promote disentanglement, we introduce a discriminator that explicitly distinguishes weather conditions. Additionally, we design a causal data augmentation strategy that perturbs non-causal parts while preserving causal structures, enabling improved generalization under rare scenarios. Experiments on taxi-flow datasets from three cities demonstrate that WED-Net delivers robust performance under extreme weather conditions, highlighting its potential to support safer mobility, highlighting its potential to support safer mobility, disaster preparedness, and urban resilience in real-world settings. The code is publicly available at https://github.com/HQ-LV/WED-Net.