🤖 AI Summary
To address weak out-of-distribution (OOD) generalization, compromised in-distribution (ID) performance, and strong assumption dependence in existing causal approaches under distributional shift, this paper proposes Distributionally Robust Invariant Gradients (DRIG), a gradient-based distributionally robust method. DRIG models linear causal structures via generalized additive interventions and—uniquely—integrates data-dependent noise interventions with gradient invariance, establishing finite-sample robustness under class shifts through theoretical analysis. It unifies and extends the anchor regression framework to accommodate broader perturbation classes and further generalizes to semi-supervised domain adaptation. Experiments on synthetic data, single-cell transcriptomics, and intensive care datasets demonstrate that DRIG achieves superior trade-offs between OOD robustness and ID accuracy, consistently outperforming standard empirical risk minimization (ERM), distributionally robust optimization (DRO), and state-of-the-art causal baselines.
📝 Abstract
Since distribution shifts are common in real-world applications, there is a pressing need to develop prediction models that are robust against such shifts. Existing frameworks, such as empirical risk minimization or distributionally robust optimization, either lack generalizability for unseen distributions or rely on postulated distance measures. Alternatively, causality offers a data-driven and structural perspective to robust predictions. However, the assumptions necessary for causal inference can be overly stringent, and the robustness offered by such causal models often lacks flexibility. In this paper, we focus on causality-oriented robustness and propose Distributional Robustness via Invariant Gradients (DRIG), a method that exploits general noise interventions in training data for robust predictions against unseen interventions, and naturally interpolates between in-distribution prediction and causality. In a linear setting, we prove that DRIG yields predictions that are robust among a data-dependent class of distribution shifts. Furthermore, we show that our framework includes anchor regression as a special case, and that it yields prediction models that protect against more diverse perturbations. We establish finite-sample results and extend our approach to semi-supervised domain adaptation to further improve prediction performance. Finally, we empirically validate our methods on synthetic simulations and on single-cell and intensive health care datasets.
