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

TL;DR

This work addresses the challenge of urban flow forecasting under extreme weather by explicitly disentangling intrinsic traffic dynamics from weather-induced effects. It proposes WED-Net, a dual-branch Transformer with memory banks, a Weather Discriminator, and an adaptive fusion strategy, augmented by a spatio-temporal causal augmentation framework to preserve causal structure during distribution shifts. Empirical results across three cities show robust gains, especially under extreme rain, and ablations confirm the importance of each architectural component and the causal strategy. The approach yields more reliable forecasts for safety-critical urban mobility, disaster preparedness, and resilience planning in real-world settings.

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.

Figures (7)

• Figure 1: Comparison of spatio-temporal patterns of NYC taxi flow on Mar. 7th, Mar. 14th (rainstorm), and Mar. 21st.
• Figure 2: The architecture of WED-Net. Embedded traffic and weather inputs are processed by I-STEnc and W-STEnc to disentangle intrinsic dynamics and weather-induced effects via attention. ST Memory enhances each branch with historical patterns, a Weather Discriminator enforces weather invariance, and Adaptive Fusion combines both branches for prediction.
• Figure 3: Spatio-temporal causal augmentation.
• Figure 4: Intrinsic and Weather-effect causal neighbors of parcel $v_{37}$ under different weather conditions.
• Figure 5: (a) Taxi flow prediction under varying weather; (b) Latent representation visualization of our model across weather scenarios.