Robust Machine Learning for Regulatory Sequence Modeling under Biological and Technical Distribution Shifts

TL;DR

This paper evaluates the robustness of deep learning models in regulatory genomics under various biological and technical distribution shifts, highlighting robustness gaps and proposing methods to improve reliability.

Contribution

It introduces a comprehensive robustness framework combining simulation and real data analysis to assess and enhance model reliability under genomic distribution shifts.

Findings

Models are robust to mild GC content shifts but fail under motif rewiring and noise.

Adding biological priors improves robustness but has limited effect against high noise.

Uncertainty-aware prediction helps recover low-risk predictions under distribution shifts.

Abstract

Robust machine learning for regulatory genomics is studied under biologically and technically induced distribution shifts. Deep convolutional and attention based models achieve strong in distribution performance on DNA regulatory sequence prediction tasks but are usually evaluated under i.i.d. assumptions, even though real applications involve cell type specific programs, evolutionary turnover, assay protocol changes, and sequencing artifacts. We introduce a robustness framework that combines a mechanistic simulation benchmark with real data analysis on a massively parallel reporter assay (MPRA) dataset to quantify performance degradation, calibration failures, and uncertainty based reliability. In simulation, motif driven regulatory outputs are generated with cell type specific programs, PWM perturbations, GC bias, depth variation, batch effects, and heteroscedastic noise, and CNN, BiLSTM, and transformer models are evaluated. Models remain accurate and reasonably calibrated under mild GC content shifts but show higher error, severe variance miscalibration, and coverage collapse under motif effect rewiring and noise dominated regimes, revealing robustness gaps invisible to standard i.i.d. evaluation. Adding simple biological structural priors motif derived features in simulation and global GC content in MPRA improves in distribution error and yields consistent robustness gains under biologically meaningful genomic shifts, while providing only limited protection against strong assay noise. Uncertainty-aware selective prediction offers an additional safety layer that risk coverage analyses on simulated and MPRA data show that filtering low confidence inputs recovers low risk subsets, including under GC-based out-of-distribution conditions, although reliability gains diminish when noise dominates.