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

2026-01-30

AI Summary

This work addresses the challenge of spatiotemporal traffic flow forecasting under urban extreme weather events—such as heavy rainfall—which are rare yet highly complex, rendering existing methods inadequate in capturing the fine-grained impact of weather on traffic dynamics. To this end, we propose WED-Net, a dual-branch Transformer architecture that leverages self-attention to model intrinsic traffic patterns and cross-attention to capture weather-induced effects. The two branches are effectively disentangled through a memory bank, an adaptive gating mechanism, and an explicit weather discriminator. Furthermore, we introduce a causality-preserving data augmentation strategy to enhance generalization under rare extreme scenarios. Experiments on three real-world urban taxi datasets demonstrate that WED-Net significantly outperforms state-of-the-art approaches, offering robust support for resilient urban transportation and disaster response planning.

Technology Category

• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Machine Learning: Mixture of Experts (MoE)
• Computer Vision: Multi-modal Vision

Application Category

• Web Mining and Content Analysis: Large pretrained models with web data
• Search and Retrieval-Augmented AI: Multilingual and cross-lingual Web search
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

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.

Problem

• urban flow prediction
• extreme weather
• spatio-temporal dynamics
• out-of-distribution generalization
• weather effect

Innovation

• weather-effect disentanglement
• causal data augmentation
• dual-branch Transformer
• spatio-temporal prediction
• out-of-distribution generalization