Source-linked AI summary

RiskTraf: Risk-Extrapolated Residual Learning for Multi-Variate Traffic Flow Prediction

Guangyu Wang, Zhidan Liu

arXiv:2608.20656v1cs.LGcs.AI

TL;DR

Traffic forecasting lacks reliable benchmarks and robust ways to use raw flow, speed, and occupancy across traffic regimes. The paper introduces PEMSB-3V and RiskTraf, which applies risk-aware residual correction to frozen backbones; it consistently improves diverse models, achieving the best MAE and RMSE across four subsets.

  • Problem

    Traffic forecasting lacks reliable three-variable benchmarks and faces regime-dependent speed and occupancy correlations that can undermine robust flow prediction.

  • Method

    The paper constructs PEMSB-3V from validated raw measurements and adds RiskTraf, a frozen-backbone residual plug-in using historical speed and occupancy with risk extrapolation.

  • Results

    RiskTraf achieves the best MAE and RMSE on all four PEMSB-3V subsets, including an 8.75% MAE reduction on PEMS03-B versus its strongest standalone competitor.

  • Takeaways & Limitations

    The study supports standardized evaluation and risk-aware residual correction for using raw sensor-native variables across diverse traffic forecasting backbones.

  • Takeaways & Limitations

    Speed and occupancy are available only as historical auxiliary measurements.

Abstract

from arXiv · show

Traffic sensors commonly record flow, speed, and occupancy, but standard traffic flow forecasting benchmarks and models rarely exploit all three raw measurements reliably. Although speed and occupancy provide sensor-native traffic-state information beyond flow alone, existing releases often omit these variables, replace them with proxies, or contain logically inconsistent records. Moreover, direct empirical risk minimization over three-variable inputs may exploit regime-dependent shortcuts, as the relationships among flow, speed, and occupancy vary substantially between free-flow and congested states. We introduce \textbf{PEMSB-3V}, a public benchmark suite that preserves raw flow, speed, and occupancy measurements from PeMS detectors for flow prediction. We also propose \textbf{RiskTraf}, a model-agnostic risk-extrapolated residual plug-in. For each trained spatio-temporal backbone, RiskTraf freezes the selected checkpoint and learns a lightweight zero-start residual head from historical speed and occupancy. The residual head constructs ordered traffic-risk environments and optimizes horizon-wise flow corrections with a risk extrapolation objective, thereby mitigating regime-specific shortcut correlations without modifying the backbone. Extensive experiments demonstrate that RiskTraf consistently improves diverse forecasting backbones and outperforms debiasing and distribution-shift adaptation methods. Our code and benchmark are available at https://github.com/Guangyu4/RiskTraf.

1 Introduction

The section motivates multivariate traffic forecasting by identifying unreliable three-variable benchmarks and regime-dependent shortcuts in speed–occupancy correlations. It introduces PEMSB-3V for standardized evaluation and RiskTraf as a residual, risk-extrapolation plug-in for trained forecasting backbones.

  • Motivation: Architectural advances in traffic forecasting yield increasingly marginal gains while model complexity and computational costs continue to grow.This motivates exploiting richer contextual information for improved forecasting.
  • Challenges: Reliable use of flow, speed, and occupancy requires both valid measurements from the same sensors and learning methods that avoid regime-specific shortcuts.The section presents these as coupled requirements for multivariate forecasting.
  • Challenges: Existing releases may contain invalid or inconsistent flow, speed, and occupancy records because detectors require strict screening and can exhibit faults or configuration-sensitive measurements.Reported issues include stuck-off, stuck-on, and hanging-on detector behavior.
  • Challenges: Speed and occupancy have regime-dependent relationships with flow, so treating them as ordinary covariates can encourage shortcut learning instead of robust traffic-state modeling.Under congestion, slower vehicles can increase occupancy even when observed flow is similar.
  • Contributions: PEMSB-3V provides four district-aligned datasets preserving raw flow, speed, and occupancy measurements for standardized three-variable flow-prediction evaluation.The datasets are PEMS03-B, PEMS04-B, PEMS07-B, and PEMS08-B, with sensor-level validation.
  • Contributions: RiskTraf freezes a trained three-variable backbone and learns a lightweight zero-initialized residual head from historical speed and occupancy for horizon-wise flow correction.It is presented as a REx-based plug-in that improves trained backbones without replacing their architectures and is evaluated across multiple datasets and diverse spatio-temporal backbones.

2 PEMSB-3V Benchmark

PEMSB-3V is a four-dataset benchmark suite preserving raw flow, speed, and occupancy measurements from PeMS detectors. Its construction combines metadata-based sensor screening, cleaned multivariate records, road-topology adjacency, and geographically grounded district splits.

  • Dataset composition: PEMSB-3V comprises PEMS03-B, PEMS04-B, PEMS07-B, and PEMS08-B, constructed from California’s statewide PeMS freeway monitoring system.The source system contains nearly 40,000 detectors, and the benchmark retains native measurements rather than proxies or temporal codes.
  • Benchmark construction: The pipeline keeps detectors with raw flow, speed, and occupancy channels, filters low-completeness sensors, cleans five-minute records, and aggregates them into X ∈ R^T×N×3.Malformed timestamps are removed and only short gaps are interpolated.
  • Sensor screening: Metadata screening uses earliest-failed checks for ID coverage, spatial localization, static-attribute validity, and road-type consistency before final sensor selection.Mainline and HOV sensors are retained, while ramp and connector sensors are excluded from the mainline subset.
  • Sensor screening: 1,013 high-quality PEMS03-B sensors remain after removing 41 missing-metadata IDs, 2 invalid-geolocation records, and excluding 626 ramp or connector sensors.Of 1,665 geolocated records, 1,039 pass all metadata checks before duplicate-location merging.
  • Geographic organization: The four datasets correspond to Caltrans districts 03, 04, 07, and 08, preserving geographic provenance while maintaining comparability with widely used PeMS benchmarks.The split follows PeMS administrative organization rather than an arbitrary partition.
  • Benchmark audit: An audit finds that most representative benchmarks omit raw speed and occupancy, substitute proxy variables, or contain logically inconsistent measurements.PEMS03 and PEMS07 use temporal encodings, whereas PEMS04 and PEMS08 include physically implausible speed records.

3 Methodology

RiskTraf is a model-agnostic risk-extrapolation plug-in with a paired two-stage design that leaves the forecasting backbone unchanged. It first trains a spatio-temporal backbone on 12 historical steps of flow, speed, and occupancy, then freezes it while learning speed- and occupancy-driven residual flow corrections across traffic-risk environments.

  • Overall design: RiskTraf uses a paired two-stage design without modifying the forecasting backbone.Stage I trains the backbone, while Stage II adds the residual correction module.
  • Stage I: Backbone training: Stage I trains a standard spatio-temporal backbone with 12 historical steps of flow, speed, and occupancy to predict future flow.The backbone receives all three traffic measurements as inputs.
  • Stage II: Residual learning: Stage II freezes the trained backbone and attaches a lightweight residual head driven by historical speed and occupancy.The residual head provides corrections without retraining or altering the backbone.
  • Stage II: Residual learning: The residual head uses a risk extrapolation objective to learn robust flow corrections across traffic-risk environments.Its optimization targets flow corrections conditioned on historical speed and occupancy.

3.1 Problem Formulation and Paired Protocol

The protocol uses historical flow, speed, and occupancy across sensor nodes to predict future flow only, with speed and occupancy serving as auxiliary measurements. RiskTraf then freezes a validated backbone checkpoint and learns a speed- and occupancy-based residual correction in a paired two-stage protocol.

  • Problem Formulation: Historical inputs contain flow, speed, and occupancy observed over all sensor nodes, while the target is future flow over horizon H.Speed and occupancy are available only from the historical input and are not predicted as outputs.
  • Paired Protocol: RiskTraf is applied separately to each dataset-backbone pair using an arbitrary spatio-temporal backbone for the Stage-I baseline.The protocol first trains the backbone to produce a baseline flow prediction.
  • Paired Protocol: After selecting the checkpoint by validation MAE, RiskTraf freezes backbone parameters and learns a residual head from historical speed and occupancy.The residual head corrects the baseline prediction while the backbone remains fixed.
  • Paired Protocol: During Stage II, only residual parameters are updated, with a small residual scaling factor controlling the correction.This stage trains the lightweight residual component rather than modifying the backbone.

3.2 Stage I: Baseline Backbone Training

Stage I trains the traffic-forecasting backbone under standard empirical risk minimization using historical three-variable sequences to predict future flow. The checkpoint with the best validation MAE becomes the reference for the vanilla backbone and subsequent RiskTraf stage.

  • Stage I: Baseline Backbone Training: Stage I applies standard empirical risk minimization to train backbone B_ϕ from historical three-variable sequence X for future flow prediction Y_t.This follows the setting used by existing traffic forecasting models.
  • Stage I: Baseline Backbone Training: The best-validation-MAE checkpoint is selected as the reference checkpoint for both the vanilla backbone and subsequent RiskTraf stage.MAE is computed as masked mean absolute error on inverse-scaled flow values.

3.3 Stage II: Risk-Aware Residual Plug-in

Stage II freezes the trained backbone and learns a lightweight, model-agnostic residual correction from historical speed and occupancy. It orders samples by traffic risk and regularizes residual errors across regimes while retaining only validation-improving corrections.

  • Stage II: Risk-Aware Residual Plug-in: RiskTraf freezes the Stage-I backbone prediction and learns only a constrained horizon-wise residual correction, preserving the backbone’s spatio-temporal representations.The plug-in is model-agnostic and addresses regime-specific correlations that the backbone’s ERM objective may exploit.
  • Residual head: The residual head combines a trainable node-identity embedding with node-level historical summaries, including mean speed, mean occupancy, and occupancy-minus-speed.The embedding is learned from scratch without external sensor metadata, providing node-specific context for heterogeneous detector locations and configurations.
  • Conservative correction: RiskTraf initially reproduces the baseline through a zero-initialized final layer and retains the corrected checkpoint only when validation MAE improves over the frozen baseline.Otherwise, it rolls back to the Stage-I prediction, making the plug-in conservative.
  • Risk environments: RiskTraf computes a normalized risk score from historical speed and occupancy, with larger ρ indicating higher occupancy and lower speed along the free-flow–congestion axis.Samples are sorted by ρ and partitioned into ordered low-to-high risk environments without incident annotations or manually defined congestion labels.
  • Risk-extrapolated objective: The ordered-environment objective combines flow-prediction loss with penalties that prevent residual errors from increasing with risk and suppress excessive highest-risk loss.Ppair penalizes adjacent inversions, while Pextra compares the highest-risk loss against a detached lowest-risk reference; penalty weighting is warmed up during Stage II.

3.4 Inference

At inference, the selected Stage-I backbone generates a baseline forecast, and an approved Stage-II residual head computes a correction from historical traffic-state and node features. If validation rejects the residual head, RiskTraf returns the Stage-I baseline prediction.

  • Inference: The selected Stage-I backbone first produces the baseline forecast ˆ𝑌base.This baseline is the starting point for inference.
  • Inference: After passing the validation safeguard, the Stage-II residual head computes Δ ˆ𝑌 from historical speed and occupancy, node embeddings, and node-level summaries.The residual correction is then used to form the final prediction.
  • Inference: If the residual head is rolled back, the output remains the Stage-I baseline prediction.Rollback preserves the backbone-only forecast.

4 Experiments

Experiments on PEMSB-3V evaluate data validity, model-agnostic effectiveness, method comparisons, component choices, and practicality. RiskTraf consistently improves forecasting across diverse backbones, with traffic-state variables and risk environments providing interpretable residual corrections.

  • Data validity: Across four PEMSB-3V subsets, higher occupancy generally corresponds to lower speed, while identical flow levels span different occupancy states and speeds.These patterns support speed and occupancy as traffic-state indicators beyond flow alone.
  • Overall effectiveness: RiskTraf reduces MAE and RMSE in nearly all displayed backbone–dataset cases, including recurrent, graph-based, attention-based, normalization-based, and heterogeneity-aware architectures.The paired protocol freezes the same validation-selected vanilla backbone and trains only a lightweight speed-and-occupancy residual head.
  • Method comparison: 8.75%, 5.12%, 13.68%, and 6.09% MAE reductions are achieved on PEMS03-B, PEMS04-B, PEMS07-B, and PEMS08-B against each subset’s strongest standalone competitor.RMSE reductions are 7.99%, 3.51%, 7.87%, and 1.04% on the same subsets, respectively.
  • Component analysis: The full speed+occ variant performs best overall, while occupancy alone outperforms removing both auxiliary variables on most backbone–dataset pairs.Occupancy primarily reflects congestion intensity, whereas speed contributes complementary vehicle-movement information.
  • Component analysis: K = 3 yields the lowest MAE in the PEMS03-B sweep, but K = 4 is fixed for the main experiments to avoid setting-specific tuning.The non-monotonic curve indicates that finer risk partitioning is not uniformly better.
  • Practicality and interpretability: Test MAE decreases from 28.45 to 21.45 for GTS, from 19.65 to 17.43 for GWNet, and from 17.31 to 15.78 for STNorm.The gains accompany small overhead because only the residual head is optimized, supporting an accuracy–efficiency trade-off.

5 Related Work

Related forecasting work progressed from recurrent models to graph-based architectures, while invariant-learning and non-stationarity methods address robustness through different assumptions. RiskTraf differs by targeting regime-varying relationships among flow, speed, and occupancy rather than ordinary statistical or spectral shifts.

  • Traffic Forecasting: Traffic forecasting evolved from LSTM models to graph architectures using predefined topology or learned adaptive graphs for spatial dependencies.DCRNN and STGCN use predefined adjacency matrices, whereas Graph WaveNet learns adaptive graphs; the passage also introduces MegaCRN but truncates its description.
  • Invariant Learning: Invariant-learning methods seek stable predictive mechanisms across environments, including shared classifiers, risk-variance minimization, and worst-group optimization.IRM, IRM Games, REx, and Group DRO represent these approaches through shared optimality, Nash equilibrium, risk variance, and worst-group risk objectives.
  • Invariant Learning: Invariant-learning effectiveness depends on environment definitions, since partitions misaligned with relevant spurious correlations can yield limited gains over empirical risk minimization.DomainBed provides the cited evidence for this limitation.
  • Non-Stationarity: Normalization and transformation methods address non-stationarity in statistical or spectral series properties, whereas this work studies regime-varying relationships among traffic measurements.RevIN, Dish-TS, Non-stationary Transformers, and multi-order wavelet derivative transforms target shifts in statistics, attention stationarity, or wavelet-domain dynamics.

6 Conclusion and Future Works

The paper examines how to use speed and occupancy for traffic flow prediction without relying on regime-dependent correlations that harm generalization. RiskTraf addresses this by adding a lightweight REx-based residual correction to a frozen three-variable backbone.

  • Conclusion and Future Works: Direct training on flow, speed, and occupancy can exploit regime-dependent correlations and hurt generalization.The paper motivates separating backbone prediction from residual correction to reduce this risk.
  • Conclusion and Future Works: RiskTraf uses historical speed and occupancy to learn a lightweight residual correction under a risk extrapolation objective.The method is designed as a paired plug-in for traffic flow prediction.
  • Conclusion and Future Works: RiskTraf freezes a trained three-variable backbone before learning its residual correction.This separates the pretrained backbone from the subsequent correction module.
Loading 2608.20656v1…