[Paper Review] Shifts 2.0: Extending The Dataset of Real Distributional Shifts

Andrey Malinin, Ανδρέας Αθανασόπουλος+16 more|arXiv (Cornell University)|Jun 30, 2022

Advanced Glycation End Products research18 citations

TL;DR

This paper extends the Shifts dataset with two new industrial-scale benchmarks—3D MRI lesion segmentation for Multiple Sclerosis and marine vessel power consumption estimation—featuring real-world distributional shifts. The authors provide detailed data collection, preprocessing, and baseline results, enabling robust evaluation of model generalization and uncertainty estimation under realistic, high-stakes distributional shifts.

ABSTRACT

Distributional shift, or the mismatch between training and deployment data, is a significant obstacle to the usage of machine learning in high-stakes industrial applications, such as autonomous driving and medicine. This creates a need to be able to assess how robustly ML models generalize as well as the quality of their uncertainty estimates. Standard ML baseline datasets do not allow these properties to be assessed, as the training, validation and test data are often identically distributed. Recently, a range of dedicated benchmarks have appeared, featuring both distributionally matched and shifted data. Among these benchmarks, the Shifts dataset stands out in terms of the diversity of tasks as well as the data modalities it features. While most of the benchmarks are heavily dominated by 2D image classification tasks, Shifts contains tabular weather forecasting, machine translation, and vehicle motion prediction tasks. This enables the robustness properties of models to be assessed on a diverse set of industrial-scale tasks and either universal or directly applicable task-specific conclusions to be reached. In this paper, we extend the Shifts Dataset with two datasets sourced from industrial, high-risk applications of high societal importance. Specifically, we consider the tasks of segmentation of white matter Multiple Sclerosis lesions in 3D magnetic resonance brain images and the estimation of power consumption in marine cargo vessels. Both tasks feature ubiquitous distributional shifts and a strict safety requirement due to the high cost of errors. These new datasets will allow researchers to further explore robust generalization and uncertainty estimation in new situations. In this work, we provide a description of the dataset and baseline results for both tasks.

Motivation & Objective

• To address the critical gap in machine learning benchmarks that fail to capture real-world distributional shifts common in high-risk applications.
• To provide industrial-scale datasets with realistic distributional shifts in medical imaging and energy systems, where model failure has severe consequences.
• To enable systematic evaluation of model robustness and uncertainty estimation under actual deployment shifts, not synthetic or curated shifts.
• To extend the Shifts benchmark with two new, high-impact tasks: 3D white matter lesion segmentation in MRI and power consumption prediction in cargo vessels.
• To support the development of safer, more reliable ML systems through realistic evaluation of uncertainty and generalization under distributional shift.

Proposed method

• Collected and curated a new 3D MRI dataset of Multiple Sclerosis lesions from clinical sources, with diverse scanner protocols and acquisition parameters to induce real distributional shifts.
• Constructed a marine vessel power consumption dataset from operational telemetry of cargo ships, capturing shifts due to varying weather, load, and routing conditions.
• Designed canonical train/val/test splits for both datasets to enable evaluation under known distributional shifts, with in-domain and out-of-domain splits.
• Applied standard deep learning models (e.g., U-Net, MLP, ResNet) and uncertainty estimation techniques (MC dropout, variational inference, deep ensembles) as baselines.
• Evaluated performance using standard metrics: R-AUC, F1-AUC, F1@95%, and retention curves to assess uncertainty calibration under shift.
• Reported results with standard deviation across single seeds to ensure reproducibility and statistical rigor in benchmark evaluation.

Experimental results

Research questions

• RQ1 How do state-of-the-art deep learning models perform under real-world distributional shifts in 3D medical imaging and marine energy applications?
• RQ2 To what extent do uncertainty estimation methods (e.g., MC dropout, deep ensembles) maintain reliability when deployed on shifted data in high-stakes domains?
• RQ3 How do model performance and uncertainty calibration degrade when shifting from in-domain to out-of-domain data in real industrial datasets?
• RQ4 Can the new benchmarks support the development of models that generalize robustly and provide well-calibrated uncertainty estimates in safety-critical settings?
• RQ5 What are the key challenges in applying existing ML benchmarks to real-world, high-risk industrial applications with complex, non-i.i.d. data shifts?

Key findings

• The 3D MRI lesion segmentation task exhibits significant performance degradation under distributional shift, with F1-AUC dropping from 0.618 (in-domain) to 0.536 (out-of-domain) for MC dropout with deep ensemble.
• For marine vessel power prediction, the best-performing model (symbolic with ensemble) achieved an R-AUC of 17.51×10⁵ on the evaluation set, significantly outperforming DNN and single-model baselines.
• Uncertainty calibration (measured via R-AUC and retention curves) deteriorated substantially on out-of-domain data, with VI and MC dropout showing lower R-AUC on shifted data compared to in-domain.
• The symbolic model with ensemble training achieved the highest F1@95% on the evaluation set (0.509), indicating better robustness to uncertainty miscalibration under shift.
• Baseline models showed poor generalization under shift, with F1-AUC dropping by up to 20% on out-of-domain splits, highlighting the need for improved robustness and uncertainty estimation.
• The new datasets reveal that standard models and uncertainty methods fail to generalize reliably in real-world settings, underscoring the necessity for shift-aware training and evaluation.