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

AI summary

Overview

Research area: Urban spatio-temporal traffic forecasting, specifically weather-aware prediction of city taxi flow, combined with causal inference techniques for out-of-distribution (OOD) generalization.

Technical level: Advanced. The paper assumes familiarity with Transformer attention, spatio-temporal graph neural networks, structural causal models, do-calculus, and domain-adversarial training.

Scope: The paper introduces WED-Net, a dual-branch Transformer that disentangles intrinsic traffic patterns from weather-induced ones and uses a causality-driven data augmentation strategy to improve taxi-flow prediction under extreme weather across three U.S. cities.

What This Paper Is About

Existing urban flow models either ignore weather or treat it as a coarse-grained, globally shared auxiliary input, and because extreme weather is rare in training data, they learn spurious correlations that collapse when conditions shift. The paper's goal is to build a forecasting model that explicitly separates the stable, intrinsic structure of urban mobility from the transient effects of weather, and that generalizes better when rare extreme events (such as heavy rain) actually occur.

Key Contributions

1. A weather-aware disentangled architecture. WED-Net uses two branches: an Intrinsic Spatio-Temporal Dependency Encoder (I-STEnc) built on self-attention, and a Weather-Induced Spatio-Temporal Dependency Encoder (W-STEnc) built on cross-attention where traffic features are queries and weather features are keys and values. This models fine-grained, spatially heterogeneous weather effects rather than coarse weather labels.

2. Memory augmentation plus adversarial weather discrimination. Each branch has its own learnable spatio-temporal memory bank that retrieves representative historical patterns, and a Weather Discriminator with a Gradient Reversal Layer pushes the intrinsic branch toward weather-invariant representations. An adaptive gating mechanism fuses the two branches with data-dependent weights.

3. A spatio-temporal causal augmentation strategy. Attention maps from a pre-experimental reconstruction identify causal parcels and time steps per target; non-causal traffic and weather values are replaced with values from a reference sample matched on day type (weekday or weekend) and hour index. This is grounded in an explicit structural causal model and do-calculus formulation, and unlike prior work it considers temporal and spatial causality jointly.

4. Evaluation on three cities under split weather conditions. Experiments on taxi data from New York City, Chicago, and Washington DC, with test sets partitioned into normal and extreme weather, and with the augmentation applied to seven baselines as well as WED-Net.

Main Findings

• Best overall accuracy on most benchmarks. WED-Net achieves the lowest MAE and RMSE across most cities and weather conditions in Table 2. Without augmentation, on NYC extreme weather it records MAE 0.1281 and RMSE 0.2330 versus the second-best model MTGNN at 0.1344 and 0.2494; on CHI extreme it records 0.0386 and 0.1558 versus MTGNN's 0.0398 and 0.1639; on DC extreme it records 0.0871 and 0.2583 versus MTGNN's 0.0881 and 0.2598.

• The advantage is largest under extreme weather. The paper reports improvements of 14% and 9% over the second-best model, MTGNN, in NYC and DC respectively under extreme conditions.

• One exception is reported. Under normal weather in DC, MTGNN's numbers (MAE 0.0848, RMSE 0.2392) are lower than WED-Net's without augmentation (MAE 0.0887, RMSE 0.2556). With causal augmentation, WED-Net reaches MAE 0.0825 and RMSE 0.2297 there.

• Causal augmentation benefits most baselines. In the lower half of Table 2, DCRNN improves by 25%, GTS by 8%, and STGCN, GWNet, STAEformer, and WED-Net by around 5%, while WED-Net stays best throughout. With augmentation WED-Net improves to MAE 0.1172 / RMSE 0.2087 on NYC extreme, 0.0372 / 0.1541 on CHI extreme, and 0.0863 / 0.2506 on DC extreme.

• Ablations isolate each component. Removing the weather modality entirely (Ours-w) causes the largest performance drop, especially under extreme weather; replacing cross-attention with self-attention (Ours-wca) also degrades performance; removing the ST memory (Ours-mem) or the weather discrimination loss (Ours-ddl) consistently hurts, with the latter causing especially large drops on DC; degrading spatial causal localization to adjacency-matrix neighbors (CaAu-s) or temporal localization to all time steps (CaAu-t) also degrades performance. The ablation covers CHI and DC under extreme and normal weather.

• Intrinsic causal neighbors are stable; weather-effect neighbors are not. For target parcel v37 at 10:00 on Fridays under precipitation of 1.0, 0.3, and 0.0 inches per hour, the intrinsic causal neighbors stay largely consistent across weather conditions while the weather-effect causal neighbors vary markedly with precipitation.

• Latent representations confirm the separation. PCA on hidden features shows I-STEnc features forming clusters largely invariant to weather, while W-STEnc features exhibit clear separation across weather conditions.

• Qualitative prediction behavior. Under normal weather WED-Net and STAEformer perform similarly, but under heavy rain STAEformer shows clear deviations during intense precipitation while WED-Net better tracks the true flow trend.

• The motivating example. On three Tuesdays (March 7, 14, and 21) in NYC, the March 14 rainstorm produces a sustained drop in total taxi flow versus the two rain-free Tuesdays, changes the spatial distribution during the morning peak, and causes parcel 43's flow to drop sharply to roughly the level of parcel 50.

Methodology in Plain English

The authors start from the observation that urban traffic at any moment is a mix of two things: routine structure (commuting rhythms, which neighborhoods connect to which) and disruption (rain, storms). Their model takes in both past taxi flow and past weather for each land parcel and routes them through two separate processing streams. One stream looks only at traffic and learns what is stable about it using attention over time and space. The other stream asks how traffic relates to weather by letting traffic features query weather features, so rainfall at a specific place and hour can modulate the traffic representation there rather than being smeared globally across the whole city.

Each stream also consults a learned memory bank of representative historical patterns, retrieving the closest matches to the current situation — useful when the current situation is rare. An adversarial discriminator is attached to the intrinsic stream: it tries to guess the weather condition from the intrinsic representation, while a gradient reversal layer forces that representation to become harder to classify, so the intrinsic stream stops encoding weather. A small gating network decides per-situation how much to trust the intrinsic stream versus the weather stream, and the weighted result goes to an MLP that outputs the forecast. Training minimizes MAE on flow plus a weighted cross-entropy term from the discriminator.

The second half of the method addresses data scarcity. Rare rainstorms give the model too few examples, so the authors generate more. They first run the model on the training data and read out its attention maps to decide, for each target parcel and each timestep, which other parcels and timesteps actually influence it (the causal ones) and which do not. They then take an extreme-weather sample and overwrite the non-causal parcels and timesteps with the corresponding values from a normal-weather sample that matches the same weekday/weekend type and hour. The causal parts stay untouched. The authors describe this as a do-calculus intervention that blocks the backdoor path from non-causal variables to the prediction, forcing the model to depend on stable relationships. The augmentation is applied both to WED-Net and to the baseline models.

Setup details: they use taxi records from NYC, Chicago, and Washington DC with meteorological data from weather stations, interpolated to each parcel by inverse-distance weighting. Splits are 50%/25%/25% in chronological order. Samples with average precipitation above 0.1 inches per hour are labeled extreme rainstorm cases. The model uses 12 hidden features for the base embedding, 18 each for the spatial- and time-adaptive embeddings, and 12 each for the time-of-day and day-of-week embeddings, with 4 Transformer blocks and 4-head attention in each encoder, batch size 128, learning rate 1e-3, OneCycleLR, and AdamW with weight decay 5e-4, on NVIDIA A40 GPUs. Each model predicts 12 future steps from 12 past steps.

Table 1 reports 66 parcels for NYC (03/2017–10/2017), 77 for CHI (03/2017–10/2017), and 69 for DC (01/2017–08/2017). Normal-weather counts are NYC 3505/1044/1056, CHI 3505/1085/997, DC 3476/1077/1038; the extreme-weather rows list two counts each (NYC 365 and 290, CHI 282 and 329, DC 353 and 384).

Why This Matters

Research impact. The paper argues that prior causal approaches to OOD generalization in spatio-temporal learning either overlook temporal causality, use spatio-temporal mechanisms that can distort causal structures, or depend on fixed confounder stratification that does not adapt to changing conditions. WED-Net's contribution is to combine fine-grained continuous weather modeling with a disentangled architecture and a causal augmentation that operates jointly over space and time, and to demonstrate that the augmentation transfers to seven other models. It also shows that coarse weather labels used by earlier systems such as MemeSTN, EAST-Net, and MT-C2G miss localized, continuously varying effects.

Real-world applications:

• Emergency response and disaster management, where accurate flow forecasts during storms guide where resources and responders are needed.
• Traffic management and congestion mitigation, including route planning and travel time estimation under adverse conditions.
• Taxi and ride-hailing operations that need demand forecasts that do not collapse during rain or storms.
• Urban planning and resilience analysis, giving quantitative insight into how mobility redistributes spatially during disruptions.

Industry relevance. Any operator building mobility, logistics, or smart-city services on spatio-temporal forecasts faces the same failure mode: a model tuned on ordinary days degrades exactly when conditions are abnormal, which is when forecasts matter most. The weather-invariance/discrimination mechanism and the augmentation procedure are architecture-agnostic in the sense that the augmentation is shown to improve graph-based and Transformer-based baselines alike, and the authors have released code at the linked public repository.

Future Directions

• The paper does not include an explicit future work section; the following questions follow from its scope.
• The threshold used to label extreme rainstorm cases (average precipitation above 0.1 inches per hour) is set from the standard definition of rainfall intensity and the empirical distributions of the three cities. Whether a single threshold transfers to other climates and cities, and whether other precipitation bands behave differently, is open.
• The paper focuses on precipitation and notes that rainfall has a more direct impact on road conditions and travel demand than temperature or wind, although temperature and wind velocity are collected as meteorological attributes. How other extreme events mentioned in the introduction — hurricanes, earthquakes, large public events, traffic accidents — interact with the disentanglement and augmentation machinery is not evaluated.
• The evaluation is limited to three U.S. cities and taxi flow. Generalization to other mobility modes, other countries, and other spatial granularities is untested.
• The causal identification step depends on attention maps from a pre-experimental reconstruction and on a top-proportion selection rule for causal neighbors; the sensitivity of results to that proportion, and to errors in the discovered causal sets, is not reported in the available content.
• The computational cost of a dual-branch Transformer with two memory banks plus an adversarial discriminator plus repeated augmented samples is not reported, which matters for real-time deployment.

Target Audience

Researchers and graduate students working on spatio-temporal forecasting, urban computing, and causal representation learning; practitioners at transportation agencies, mobility and ride-hailing platforms, and disaster-response organizations who need forecasts that hold up under extreme weather; and anyone interested in how causal inference tools and domain-adversarial training can be combined with Transformer architectures for out-of-distribution robustness. Readers should be comfortable with attention mechanisms and causal graphical models.

Authors’ 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.