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

arXiv: https://arxiv.org/abs/2601.22586 · HTML: https://mlsanity.com/paper/2601.22586/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 8 of its 28 references that are in the map.
- First posted on arXiv 30 Jan 2026.

## Builds on (most cited first, 8 of 8)

- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Autoformer: Decomposition Transformers with Auto-Correlation for Long-Term Series Forecasting](https://mlsanity.com/paper/2106.13008/) (2021)
- [Graph WaveNet for Deep Spatial-Temporal Graph Modeling](https://mlsanity.com/paper/1906.00121/) (2019)
- [Adaptive Graph Convolutional Recurrent Network for Traffic Forecasting](https://mlsanity.com/paper/2007.02842/) (2020)
- [Discrete Graph Structure Learning for Forecasting Multiple Time Series](https://mlsanity.com/paper/2101.06861/) (2021)
- [Invariance, Causality and Robustness](https://mlsanity.com/paper/1812.08233/) (2018)
- [Learning Invariant Molecular Representation in Latent Discrete Space](https://mlsanity.com/paper/2310.14170/) (2023)
- [Event-Aware Multimodal Mobility Nowcasting](https://mlsanity.com/paper/2112.08443/) (2021)

## Built on by (most cited first, 2 of 2)

- [HypeMed: Enhancing Medication Recommendations with Hypergraph-Based Patient Relationships](https://mlsanity.com/paper/2603.18459/) (2026)
- [Aggregate-then-Calibrate for Human-centered Assessment with Theoretical Guarantees](https://mlsanity.com/paper/2608.02455/) (2026)

## Abstract (arXiv)

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.