Causal Structure-guided Distributionally Robust Optimization under Domain Shifts

Abstract

Domain generalization (DG) aims to learn predictive models from multiple source domains that maintain performance on unseen target domains. Distributionally robust optimization (DRO) addresses distribution shift by minimizing the worst-case risk over an uncertainty set of plausible test distributions. However, if this uncertainty set is overly large, it may include unrealistic shifts, leading to low-confidence predictions. To address this issue, we propose Causal Structure-guided DRO (CS-DRO), which estimates a directed acyclic graph (DAG) that encodes the predictive relationships between representations and labels, serving as a proxy for causal structure shared across source domains. Using gradient-based signals, we quantify how well the candidate distributions explored during optimization maintain the estimated structure and impose a structure-preserving constraint on the uncertainty set. The resulting DRO objective is reformulated via Lagrangian relaxation into a tractable primal–dual learning problem. Experiments on standard DG benchmarks show that our method achieves competitive performance compared to state-of-the-art methods while improving robustness to unseen target domains. Our Github repository including code and supplementary material is here: https://github.com/gyeomo/CS-DRO.