Graph-Relational Distributionally Robust Optimization
Abstract
Out-of-distribution (OOD) generalization is a challenging machine learning problem yet highly desirable in many high-stake applications.
Distributionally robustoptimization (DRO) is a promising learning paradigm to tackle this challenge butsuffers from several limitations.
To address this challenge, we propose graphrelational distributionally robust optimization that trains OOD-resilient machinelearning models by exploiting the graph structure of data distributions.
Our approach can uniformly handle both fully-known and partially-known graph structures.
Empirical results on both synthetic and real-world datasets demonstrate theeffectiveness and flexibility of our method.
Introduction
Recent years have witnessed a surge of applying machine learning (ML) in high-stake and safetycritical applications, such as health diagnoses and self-driving cars.
Such applications pose anunprecedented out-of-distribution (OOD) generalization challenge [16]: ML models are constantlyexposed to unseen distributions that lie outside their training space.
Despite well-documentedsuccess for interpolation , modern ML models ( e.g. , deep neural networks) are notoriously weak forextrapolation ; a highly accurate model on average can fail catastrophically when presented with rareor unseen distributions [1].
Without addressing this challenge, ML models cannot be safely deployed.
A promising solution for out-of-distribution generalization is to conduct distributionally robustoptimization (DRO) [13, 21, 11].
Different from empirical risk minimization (ERM) [24] thatminimizes the average loss, DRO aims to optimize the worst-case generalization risk over a set oftraining groups.
For instance, data with a similar distribution can compose a group [18].
However,it suffers from critical limitations.
(1) DRO recklessly prioritizes the worst-case groups withoutassessing the risk that they might be outliers [27]; optimizing over outliers would fundamentallydamage OOD generalization.
(2) The worst-case groups are not necessarily the influential ones thatare truly connected to unseen distributions; optimizing over the worst-case rather than influentialgroups would yield mediocre generalization performance.
To address these challenges, we propose a novel method for graph-relational distributionally robustoptimization.
Instead of the worst-case distributions, our key idea is to minimize the generalizationrisks over influential groups.
Such influential groups can be identified by analyzing the graph of datadistributions.
Graph structures widely exist in the real world and can usually be represented by agraph.
For instance, to capture the similarity of weather events in the U.S. [26], one can construct agraph where each state (group) realizes a node, and the physical adjacency between two states resultsin an edge.
A significant merit of our approach is that it can uniformly handle various scenarioswhen the graph structure is either fully or partially available.
Empirical results on both synthetic andreal-world datasets demonstrate the superior performance of our method over SOTA.
Related Work
Distributionally Robust Optimization .
In the context of distributionally robust optimization (DRO),[3] and [20] argued that minimizing the maximal loss over a set of possible distributions can providebetter generalization performance than minimizing the average loss.
The robustness guarantee of
DRO heavily relies on the quality of the uncertainty set which is typically constructed by momentconstraints [2], f -divergence [13] or Wasserstein distance [19].
To avoid yielding overly pessimisticmodels [5], group DRO [8, 18] is proposed to leverage pre-defined data groups to formulate theuncertainty set as the mixture of these groups.
However, none of these methods incorporate thephysical prior that widely exists in real-world applications.
Out-of-Distribution Generalization .
The goal of OOD generalization is to generalize models fromsource distributions to unseen target distributions.
There are mainly two branches of methods to tackle
OOD generalization: domain-invariant learning [1, 9, 12] and distributionally robust optimization.
The goal of domain-invariant learning is to exploit the causally invariant correlations across multipledistributions.
Invariant Risk Minimization (IRM) is one of the most representative methods whichlearns the optimal classifier across source distributions.
However, recent work [17] shows IRM canfail catastrophically unless the test data are sufficiently similar to the training distribution.
Problem Formulation and Preliminary Works
Problem Formulation .
We focus on the problem of out-of-distribution (OOD) generalization.
Let
X be the input space and Y be the target space.
( X, Y ) are random variables defined over samples( x, y ) ∈
X × Y and the joint distribution P ( X, Y ) .
Since we cannot sample directly from P ( X, Y ) ,we usually assume data are drawn from a set of groups E all , where each group e ∈ E all is sampled froma distinct distribution P ( X e , Y e ) , e.g. , the distribution of medical images varies at different hospitalsdue to equipment or demographic differences.
Let E train ⊊ E all be a finite subset of training groups,and assume that for each e ∈ E train , we have access to a dataset D e := (cid:8)(cid:0) x
ej , y ej (cid:1)(cid:9) n e j =1 sampled i.i.d.from P ( X e , Y e ) .
Given a function class F and a loss function ℓ , our goal is to learn a predictorf ∈ F using the data from D e that minimizes the worst-case risk over the entire family of E all :
[Equation]
It is challenging to learn a predictor f ∈ F that generalizes from the finite set of training domains
E train to perform well on the set of all domains E all since we do not have access to data from anyunseen group e ∈ E test , where E test =
E all \ E train .
Empirical Risk Minimization (ERM) [24].
ERM minimizes the average loss over the distributionof all training groups E train :
[Equation]
where m = |E train | is the number of training groups.
Models trained via ERM heavily rely on spuriouscorrelations that do not always hold under distributional drifts [1].
Distributionally Robust Optimization (DRO) [18].
Instead of minimizing the average loss, DROminimizes the worst-combination loss of different training groups:
[Equation]
where q is the mixture vector of E train and ∆ m = { q ∈ R m | (cid:80) mk
=
1 q k = 1; ∀ k, q k ≥ 0 } .
We empirically found that DRO blindly prioritizes the worst-case groups that incur higher losses than others.
However, favoring the worst-case groups would inevitably ignore the influential ones that are trulyconnected to unseen distributions; optimizing over the worst-case rather than influential groups wouldyield compromised OOD resilience
[Table]
[Figure]
Approach
4.1 Graph-Relational Distributionally Robust Optimization
Generalizing ML models to arbitrary unseen distributions without any prior knowledge or structuralassumption is impossible [6].
Fortunately, the graph structures of E all are often available as priorknowledge and can be captured by a graph G =
( V, E ) , where the nodes V =
∪ e ∈E all X e symbolizethe groups and the edges E represent interactions among groups.
We assume the graph capturescovariate shift ( P e ( X ) ̸
= P e ′
( X ) ) rather that concept shift ( P e ( Y | X ) ̸
= P e ′ ( Y | X ) ) . Given thegraph G , we can identify influential groups and incorporate them as a physical prior p (see Sec. 4.2)to constrain the optimization in Eq.
2: D ( q ∥ p ) ≤ τ , where D ( · ) is a distance metric over the spaceof distributions.
τ is a fixed margin the controls the extent to which we enforce the prior constraint.
4.2 Implementation of Physical Prior
Motivated by centrality analysis [14] in social networks, we propose to assess the group centrality toidentify influential groups that are truly connected to unseen distributions, which can be calculatedusing graph measurements [23] such as degree, betweenness, and closeness.
In this paper, wecalculate the betweenness centrality of each node in G as a physical prior p to identify influentialgroups.
Betweenness centrality measures how often a node is on the shortest path between two othernodes in the graph.
[4] indicates that nodes with higher betweenness centrality would have morecontrol over the graph as more information will pass through them.
We consider two scenarios: graphstructure is fully known and partially known.
Fully-known structure denotes the graph structure of all groups E all is available.
Let s ∈
E trainand t ∈ E test be the start and end of a path in G .
We define the centrality of group e as the fractionof shortest paths that pass through it: c full e = (cid:80) s ∈E train ,t ∈E test σ ( s,t | e ) σ ( s,t ) , where σ ( s, t ) is the number ofshortest paths between groups s and t in the graph (( s, t ) -paths), and σ ( s, t | e ) is the number of( s, t ) -paths that go through group e .
Partially-known structure denotes only the graph structures of training groups E train is available.
The underlying assumption is that the unobserved part of the graph should not be very differentfrom the observed part and training groups with high centrality also exert strong influence on unseengroups.
Instead of sampling groups pairs from two separate sets, we sample ( s, t ) from E train .
Wedefine the centrality of group e as: c partial e
[Equation]
We use softmax function to normalize c e and the prior probability for group e is: p e =exp ( c e ) / (cid:80) me =1 exp ( c e ) .
In Sec. 5, we empirically found that the proposed method with c partial estill outperforms other baselines by a large margin and is only slightly worse than that with c full e .
[Table] [Table]
Experiments
Datasets .
(1) {+DG-
 [26]+} is a synthetic binary classification dataset with 15 groups.
Each groupcontains 100 data points.
In this dataset, adjacent groups have similar decision boundaries.
Following[26], we use six connected groups as the training groups
[Table]
and use others as test groups.
(2) DG-
[26]is another synthetic dataset generated using the same procedure as DG-15 , except that it contains 60groups, with 6,000 data points in total.
We randomly select six groups as the training groups, anduse others as test groups.
Visualization of DG-15 and DG-60 are shown in Fig.
(3) TPT-48 [25]contains the monthly average temperature for the 48 contiguous states in the US from 2008 to 2019.
We focus on the regression task to predict the next 6 months’ temperature based on the previous first6 months’ temperature.
We consider two generalization tasks: E(24) → W(24): we use the 24 easternstates as training groups and the 24 western states as test groups; N(24) → S(24): we use the 24northern states as training groups, the 24 southern states as test groups.
Baselines .
We compare our method with the following methods: (1) Empirical Risk Minimization(ERM) [24]; (2) Group distributionally robust optimization (DRO) [18]; (3) Invariant Risk Minimization (IRM) [1]; (4) Risk Extrapolation (REx) [10]; (5) Spectral Decoupling (SD) [15].
Following [7],we perform model selection based on a validation set constructed from training groups only.
Results .
Results of DG-15 and DG-60 are summarized in Tab. 1.
As seen, in both datasets, ourmethod achieves the best performance.
In DG-15 , all other baselines are inferior or ERM while oursoutperforms ERM by 9.56%.
We show the results for task N (24) → S (24) on TPT-48 in Tab. 2.
Asobserved, our method achieves the lowest average MSE.
We also report the average MSE of Hop-1,
Hop-2, and Hop-3 test groups.
Although REx achieves the lowest error on Hop-1 and Hop-2 groups,it yields the highest prediction error on Hop-3 groups.
Our method achieves the best performance on
Hop-3 groups, indicating its generalization capability under large distributional drifts.
Ablation Study .
We evaluate our method with partially-known graph structure.
In this scenario,we assume only the graph structure of training groups are available.
We report the results in Tab. 3.
As seen, in all datasets, ours (partial) is only slightly worse than ours (full), indicating the strongeffectiveness and flexibility of our method.
Conclusion
In this paper, we proposed Graph-Relational Distributionally Robust Optimization.
We integrategraph information into distributionally robust optimization to develop OOD resilience.
Our methodstrikes a good balance between the worst-case and influential groups, preventing the model fromoverfitting to worst-case groups and rationally improving generalization performance.
Empiricalresults on both synthetic and real-world datasets demonstrate the effectiveness of our method.