T OPOLOGY - AWARE R OBUST O PTIMIZATION FOR O UT - OF -D ISTRIBUTION G ENERALIZATION
A BSTRACT
Out-of-distribution (OOD) generalization is a challenging machine learning problem yet highly desirable in many high-stake applications.
Existing methods suffer from overly pessimistic modeling with low generalization confidence.
As generalizing to arbitrary test distributions is impossible, we hypothesize that further structure on the topology of distributions is crucial in developing strong OOD resilience.
To this end, we propose topology-aware robust optimization (TRO) that seamlessly integrates distributional topology in a principled optimization framework.
More specifically, TRO solves two optimization objectives: (1) Topology Learning which explores data manifold to uncover the distributional topology; (2) Learning on
Topology which exploits the topology to constrain robust optimization for tightlybounded generalization risks.
We theoretically demonstrate the effectiveness of our approach and empirically show that it significantly outperforms the state of the arts in a wide range of tasks including classification, regression and semantic segmentation.
Moreover, we empirically find the data-driven distributional topology is consistent with domain knowledge, enhancing the explainability of our approach
I NTRODUCTION
Recent years have witnessed a surge of applying machine learning (ML) in high-stake and safetycritical applications.
Such applications pose an unprecedented out-of-distribution (OOD) generalization challenge : ML models are constantly exposed to unseen distributions that lie outside their training space.
Despite well-documented success for interpolation , modern ML models ( e.g. , deep neural networks) are notoriously weak for extrapolation ; a highly accurate model on average can fail catastrophically when presented with rare or unseen distributions (Arjovsky et al., 2019).
For example, a flood predictor, trained with data of all 89 major flood events in the U.S. from 2000 to 2020, would erroneously predict on event “Hurricane Ida” in 2021.
Without addressing this challenge, it is unclear when and where a model can be applied and how much risk is associated with its use.
A promising solution for out-of-distribution generalization is to conduct distributionally robust optimization (DRO) (Namkoong & Duchi, 2016; Staib & Jegelka, 2019; Levy et al., 2020).
DRO minimizes the worst-case expected risk over an uncertainty set of potential test distributions.
The uncertainty set is typically formulated as a divergence ball surrounding the training distribution endowed with a certain distance metric such as f -divergence (Namkoong & Duchi, 2016) and Wasserstein distance (Shafieezadeh Abadeh et al., 2018).
Compared to empirical risk minimization (ERM) (Vapnik, 1998) that minimizes the average loss, DRO is more robust against distributional drifts from spurious correlations, adversarial attacks, subpopulations, or naturally-occurring variation (Robey et al., 2021).
However, it is non-trivial to build a realistic uncertainty set that truly approximates unseen distributions.
On the one hand, to confer robustness against extensive distributional drifts, the uncertainty set has to be sufficiently large, which increases the risks of conferring implausible distributions, e.g. , outliers, and thus yielding overly pessimistic models with low prediction confidence (Hu et al., 2018; Frogner et al., 2021).
On the other hand, the worst-case distributions are not necessarily the influential ones that are truly connected to unseen distributions; optimizing over worst-case rather than influential distributions would yield compromised OOD resilience.
As generalizing to arbitrary test distributions is impossible, we hypothesize further structure on the topology of distributions is crucial in constructing a realistic uncertainty set.
More specifically, we propose topology-aware robust optimization (TRO) by integrating two optimization objectives:
(1) Topology learning : We model the data distributions as many discrete groups lying on a common low-dimensional manifold, where we can explore the distributional topology by either using physical priors or measuring multiscale Earth Mover’s Distance (EMD) among distributions.
(2) Learning on topology : The acquired distributional topology is then exploited to construct a realistic uncertainty set, where robust optimization is constrained to bound the generalization risk within a topology graph, rather than blindly generalizing to unseen distributions.
Our contributions include: 1.
A new principled optimization method that seamlessly integrates topological information to develop strong OOD resilience.
Theoretical analysis that proves our method enjoys fast convergence for both convex and non-convex loss functions while the generalization risk is tightly bounded.
Empirical results in a wide range of tasks including classification, regression, and semantic segmentation that demonstrate the superior performance of our method over SOTA.
Data-driven distributional topology that is consistent with domain knowledge and facilitates the explainability of our approach.
P ROBLEM F ORMULATION AND P RELIMINARY W ORKS
The problem of out-of-distribution (OOD) generalization is defined by a pair of random variables ( X, Y ) over instances x ∈ X ⊆ R d and corresponding labels y ∈ Y , following an unknown joint probability distribution P ( X, Y ) .
The objective is to learn a predictor f ∈ F such that f ( x ) →
y for any ( x, y ) ∼ P ( X, Y ) .
Here F is a function class that is model-agnostic for a prediction task.
However, unlike typical supervised learning, the OOD generalization is complicated since one cannot sample directly from P ( X, Y ) .
Instead, it is assumed that we can only measure ( X, Y ) under different environmental conditions e so that data is drawn from a set of groups E all such that ( x, y ) ∼ P e ( X, Y ) .
For example, in flood prediction, these environmental conditions denote the latent factors ( e.g. , stressors, precipitation, terrain, etc) that underlie different flood events.
Let E train ⊊ E all be a finite subset of training groups (distributions), given the loss function ℓ , an OOD-resilient model f can be learned by solving a minimax optimization:
[Equation]
Intuitively, Eq. 1 aims to learn a model that minimizes the worst-case risk over the entire family E all .
It is nontrivial since we do not have access to data from any unseen distributions E test =
E all \E train .
Empirical Risk Minimization (ERM).
Typically, classic supervised learning employs ERM (Vapnik, 1998) to find a model f that minimizes the average risk under the training distribution P tr :
[Equation]
Though proved to be effective in i.i.d.
settings, models trained via ERM heavily rely on spurious correlations that do not always hold under distributional drifts (Arjovsky et al., 2019).
Distributionally Robust Optimization (DRO).
To develop OOD resilience, DRO (Namkoong &
Duchi, 2016) minimizes the worst-case risk over an uncertainty set Q by solving:
[Equation]
Here the uncertainty set Q approximates potential test distributions.
It is usually formulated as a divergence ball with a radius of ρ surrounding the training distribution P
( P tr )
= { Q : D ( Q, P tr ) ≤ ρ } endowed with a certain distance metric D ( · , · ) such as f -divergence (Namkoong & Duchi, 2016) or Wasserstein distance (Shafieezadeh Abadeh et al., 2018).
To construct a realistic uncertainty set without being overly conservative, Group DRO is further developed to formulate the uncertainty set as the mixture of training groups (Hu et al., 2018; Sagawa et al., 2019).
Despite the well-documented success, existing DRO methods suffer from critical limitations.
(1) To endow robustness against a wide range of potential test distributions, the radius of the divergence ball has to be sufficiently large with high risks of containing implausible distributions; optimizing for implausible distributions would fundamentally damage the OOD resilience by yielding overlypessimistic models with low prediction confidence.
(2) The worst-case groups are not necessarily the influential ones that are truly connected to unseen distributions; optimizing over worst-case rather than influential groups would yield compromised OOD resilience.
T OPOLOGY - AWARE ROBUST OPTIMIZATION
We propose a new principled optimization method (TRO) to develop OOD resilience, which integrates topology and optimization via a two-phase scheme: Topology Learning and Learning on Topology .
3. T OPOLOGY LEARNING : E XPLORE THE DISTRIBUTIONAL TOPOLOGY
[Figure]
[Table] [Table]
[Table]
Distributional Topology
[Table]
Data
[Table]
Data Label
EVENT
Data Label
EVENT 2 Data Label
EVENT
Data ?
EVENT N <latexit sha1_base64=
[Table]
=
[Table]
[Figure]
We model the data groups E all as many discrete distributions lying on a common low-dimensional manifold in a high-dimensional data measurement space.
In such case their structure, i.e. distributional topology , can be naturally captured by a graph G =
( V, E ) , where the entities V =
∪ e ∈E all X e symbolize the groups and the edges E represent interactions among groups.
The topology graph is constructed by: (1) identifying entity : we assume the entities are defined by the given group identities; and (2) uncovering interactions : we consider two scenarios to measure the connectivity between discrete distributions as illustrated in Fig 1.
Physical-based distributional topology .
In the scenario where the distributional adjacency information is available, we can instantly acquire the topology G physic by simply imposing the predefined neighborhood information.
For example, to capture the similarity of weather events in the U.S., one can construct a graph where each state realizes an entity, and the physical adjacency between two states results in an edge (see Fig. 1).
In this case, G physic functions as a physical prior to constrain the robust optimization introduced in Sec. 3.2.
We empirically find G physic yields an improvement of 9.56% over the state of the art regarding OOD generalization reported in Sec. 5.1.
Data-driven distributional topology .
In the absence of G physic , we propose a data-driven approach to learn the topology G data from training data.
Specifically, we embed the individual groups onto a shared data graph based on an affinity matrix of the combined data.
Inspired by Leeb & Coifman (2016), such a data graph can be viewed as a discretization of an underlying Riemann closed manifold.
By simulating a time-dependent diffusion process over the graph, we will obtain density estimates at multiple scales for each group, which will be used to calculate ℓ 1 distances between two groups.
Such multiscale ℓ 1 distance has been proved to be topologically equivalent to the Earth Mover’s Distance (EMD) on the manifold geodesic, but cutting down the computational complexity from O (cid:0) m 2 n 3 (cid:1) to ˜ O
( mn ) between m distributions over n data points (Tong et al., 2021).
We obtain the data-driven topology through three steps: (1) Data graph construction : we construct a data graph through an affinity matrix K of the combined data.
K can be implemented through kernel functions ( e.g. , RBF kernel) which capture the similarity of data.
Instead of calculating the similarity between raw data, we calculate the similarity between features extracted from an ERM-trained model as it captures spurious correlations which preserve group identity (Creager et al., 2021).
Specifically, we define the affinity matrix as: K i,j = exp (cid:16) −∥
f
( x i ) − f ( x j ) ∥ 2 /σ 2 (cid:17) , where σ 2 is the kernel scale.
(2) Multiscale diffusion density estimation : to simulate the diffusion process over the graph, we obtain a Markov diffusion operator P from K .
Following Coifman & Lafon (2006), we normalize the affinity matrix:
M = Q − 1 KQ − 1 , where Q is a diagonal matrix and Q i,i = (cid:80)
j K i,j .
The diffusion operator is defined as P = D − 1 M , where D is a diagonal matrix and D i,i = (cid:80) j M i,j .
The operator P will be used to approximate the multiscale density estimates µ e for each data group
X e : µ te = 1 n e P t 1 X e , where t is the diffusion time, P t denotes the t -th power of P , and 1 X e is the indicator function for group e .
Intuitively, P ti,j sums the probabilities of all possible paths of length t between x i and x j .
By taking multiple powers of P , µ e reveals the topological structure of X e at multiple scales.
(3) Diffusion EMD measurement : we follow Tong et al.
(2021) to measure the geodesic distance W α,K
( X e , X e ′ ) between X e and X e ′ by aggregating the ℓ 1 distances between the multiscale density estimates:
[Equation]
where α is used to balance long- and short-range distances, K is the maximum scale, and
[Equation]
Although G data is computationally more expensive than G physic , our experimental results in Sec.
5.2 indicate that optimizing with
G data can yield improved OOD resilience.
Besides, the ablation study in
Sec.
5.3 also indicates that G data is consistent with domain knowledge and enhances the explainability of TRO.
Last but not least, the data-driven method is fully differentiable, making it amenable to jointly conducting topology learning and learning on topology in an end-to-end manner.
We leave this as future work.
3. L EARNING ON TOPOLOGY : E XPLOIT TOPOLOGY FOR ROBUST OPTIMIZATION
[Table] f
Topology Learning : if G physic exists then
G ← G physic else
Obtain the affinity matrix K from data
[Equation]
M ij (cid:17) P ← D
[Equation]
1 M
Obtain G data via Eq.
G ← G data end
Learning on Topology :
Calculate topological prior p from G while not converged do
Sample ( x, y ) ∼ P e ( X, Y ) ∀ e ∈ E train Calculate
R ( f, q ) via Eq.
Update θ and q via Eq.end
Next, we propose a principled method that integrates distributional topology to develop TRO.
The key challenge is how to leverage G to construct a uncertainty set which can approximate unseen distributions with bounded generalization risk.
Our main idea is to assess the group centrality of training distributions.
Graph centrality is widely used in social network analysis (Newman, 2005) to measure how much information is propagated through each entity.
Here we introduce group centrality to identify influential groups that are truly connected to unseen distributions, which can be calculated using graph measurements (Tian et al., 2019) such as degree, betweenness, and closeness.
More specifically, we first calculate the centrality of each entity in
G as a topological prior p to identify influential groups.
Then, we construct the uncertainty set as an arbitrary mixture of training groups
[Equation]
} where q e de- notes the weight of group e , P e is the distribution of group e , and ∆ m is a ( m − 1) − dimensional probability simplex.
Finally, we use the prior p to constrain the uncertainty set Q by solving the minimax optimization problem as:
[Equation]
Intuitively, groups with high training loss and centrality will be assigned with large weights; this can tightly bound the OOD generalization risk within a topological graph.
D is an arbitrary distributional distance metric.
We use ℓ 2 distance to implement D due to its strong convexity and simplicity.
However, solving Eq. 4 often leads to a non-convex problem, wherein methods such as stochastic gradient descent (SGD) cannot guarantee constraint satisfaction (Robey et al., 2021).
To address this issue, we leverage Karush–Kuhn–Tucker conditions (Boyd et al., 2004) and introduce a Lagrange multiplier to convert the constrained problem into its unconstrained counterpart:
[Equation]
where λ is the dual variable.
Let θ ∈ Θ be the model parameters of f , we can solve the primal-dual problem effectively by alternatively updating:
[Equation]
where η tθ ( η t q ) is gradient descent (ascent) step size.
P ∆ m ( q ) projects q onto simplex ∆ m for regularization.
The overall algorithm of TRO is shown in Alg.
In Sec. 4, we show TRO enjoys fast convergence for both convex and non-convex loss functions, while the generalization risk is tightly bounded with topological constraints.
We empirically demonstrate TRO achieves strong OOD resilience by striking a good balance between the worst-case and influential groups (see Sec. 5.2).
Calculation of group centrality .
We use betweenness centrality to measure the centrality of groups.
Betweenness centrality measures how often an entity is on the shortest path between two other entities in the topology.
Freeman (1977) reveals that entities with higher betweenness centrality would have more control over the topology as more information will pass through them.
For physical-based topology G physic , we define the centrality of group e by computing the fraction of shortest paths that pass through it: c physice = (cid:80) s ∈E train ,t ∈E test σ ( s,t | e ) σ ( s,t ) , where σ ( s, t ) is the number of shortest 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 .
Intuitively, c physice measures how much information is propagated through e from the start (training) to the end (test).
For data-driven topology G data , the underlying assumption is that training groups with high centrality also exert strong influence on unseen groups.
Instead of sampling group pairs from two separate sets, we sample ( s, t ) from E train .
The centrality is modified as: c datae = (cid:80) s,t ∈E train σ ( s,t | e ) σ ( s,t ) .
We use softmax function to normalize c e and the prior probability for group e ∈ E train is defined as: p e = exp ( c e ) / (cid:80) e ∈E train exp ( c e ) .
T HEORETICAL A NALYSIS 4. C ONVERGENCE A NALYSIS
In this section, we show that by choosing appropriate step sizes η tθ and η tq , TRO yields fast convergence rates for both convex and non-convex loss functions.
We first state the assumptions of the theorems.
Next, we give the convergence rate for convex loss functions in Theorem 1 and the convergence rate for non-convex loss functions in Theorem 2.
Definition 1.
(Lipschitz continuity)
A mapping f : X → R m is G -Lipschitz continuous if for any x, y ∈ X ,
∥ f ( x ) − f ( y ) ∥ ≤ G ∥ x − y ∥ .
Definition 2.
(Smoothness)
A function f :
X → R is L -smooth if it is differentiable on X and the gradient ∇ f is L-Lipschitz continuous, i.e. , ∥∇ f ( x ) − ∇ f ( y ) ∥
≤ L ∥ x − y ∥ for all x, y ∈ X .
Assumption 1.
We make the following assumptions throughout the paper:
Given θ , the loss function ℓ ( f θ ( x ) , y ) is G-Lipschitz continuous and L-smooth with respect to x .
Convex Loss.
The expected number of stochastic gradient computations is utilized to estimate the convergence rate.
To reach a duality gap of ϵ (Nemirovski et al., 2009), the optimal rate of convergence for solving the stochastic min-max problems is O (cid:0) 1 /ϵ 2 (cid:1) if it is convex-concave.
The duality gap of the pair (˜ θ, ˜ q ) is defined as max q ∈ ∆ m R (˜ θ, q ) − min θ ∈ Θ R ( θ, ˜ q ) .
In the case of strong duality, (˜ θ, ˜ q ) is optimal iif the duality gap is zero.
We show TRO achieves the optimal O (cid:0) 1 /ϵ 2 (cid:1) rate.
Theorem 1.
Consider the dual problem in Eq.
5 when the loss function is convex and Assumptionholds.
Let Θ bounded by R Θ , E (cid:104) ∥∇ θ R ( θ, q ) ∥ 22 (cid:105) ≤ ˆ G 2 θ , and E (cid:104) ∥∇ q R ( θ, q ) ∥ 22 (cid:105) ≤ ˆ G 2 q .
With step sizes η θ = 2 R Θ / (cid:16) ˆ G θ √ T (cid:17) and η q = 2 / (cid:16) ˆ G q √ T (cid:17) , the output of TRO satisfies:
[Equation]
Theorem 1 shows that our method requires T
= O (cid:0) 1 /ϵ 2 (cid:1) iterations to reach a duality gap within ϵ .
To derive the convergence rate for non-convex functions.
, we define ϵ -stationary points as follows:
Definition 3.
( ϵ -stationary point)
For a differentiable function f : X → R , a point x ∈ X is said to be
first-order ϵ -stationary if
[Equation]
Nonconvex Loss.
The loss function ℓ ( f θ ( x ) , y ) can be nonconvex and as a result, R ( θ, q ) is nonconvex in θ .
Following Collins et al.
(2020), we define (˜ θ, ˜ q ) is an ( ϵ, δ ) -stationary point of R if:
[Equation]
[Equation]
(cid:13)(cid:13)(cid:13) 2 ≤ ϵ and R (˜ θ, ˜ q ) ≥ max q ∈ ∆ m R (˜ θ, q ) − δ , where ϵ, δ > 0 .
Theorem 2.
If Assumption 1 holds and the loss function is bounded by B and is M-smooth, the output of Alg.
1 satisfies:
[Equation]
Theorem 2 shows that our method converges in expectation to an ( ϵ, δ ) -stationary point of R in
O (1 /ϵ 4 ) stochastic gradient evaluations.
4. G ENERALIZATION B OUNDS
In this section, we provide learning guarantees for TRO.
Compared to DRO, TRO achieves a lower upper bound on the generalization risks over unseen distributions with the topological constraint.
Let H denote the family of losses associated with a hypothesis set F : H = { ( x, y ) (cid:55)→ ℓ ( f ( x ) , y ) : f ∈ F} , and n = ( n 1 , . . . , n m ) denote the vector of sample sizes for all training groups.
Follow- ing Mohri et al.
(2019), we define weighted Rademacher complexity for any
F as:
[Equation]
where e denotes group index, S e a sample of size n e , P e the distribution of group e , and σ = ( σ e,i ) e ∈ [ m ] ,i ∈
[ n e ] a collection of Rademacher variables.
The minimax weighted Rademacher complexity for a subset Λ ⊆ ∆ m is defined by R n ( H , Λ)
= max q ∈ Λ R n ( H , q ) where n = (cid:80) me =1 n e .
Let P Λ be the distribution over the mixture of training groups and ˆ P Λ be its empirical counterpart.
Let the distribution of some test group be P .
The learning guarantee for P is shown in Theorem 3.
Theorem 3.
Assume the loss function ℓ is bounded by B .
For any ϵ ≥ 0 and δ > 0 , with probability at least 1 − δ , the following inequality holds for all f ∈ F :
[Equation]
The upper bound of the generalization risk on P is mainly determined by its distance to P Λ :
D ( P ∥ P Λ ) .
With the topological prior, risks on P can be tightly bounded by minimizing D ( P ∥ P Λ ) , as long as P falls into the convex hull of training groups.
We empirically verify the effectiveness of the topological prior in minimizing the generalization risks over unseen distributions (see Sec. 5).
E XPERIMENTS
We evaluate TRO in a wide range of tasks including classification, regression, and semantic segmentation.
We compare TRO with SOTA baselines on OOD generalization and conduct ablation study on the key components of TRO.
Following Gulrajani & Lopez-Paz (2021), we perform model selection based on a validation set constructed from training groups only.
We provide implementation details in Appendix 7.2 and results on DomainBed (Gulrajani & Lopez-Paz, 2021) in Appendix 7.3.
Baselines .
We compare TRO with the following methods: (1) Empirical Risk Minimization (ERM) (Vapnik, 1998); (2) Group distributionally robust optimization (DRO) (Sagawa et al., 2019); (3) Invariant Risk Minimization (IRM) (Arjovsky et al., 2019); (4) Risk Extrapolation (REx) (Krueger et al., 2021); (5) Spectral Decoupling (SD) (Pezeshki et al., 2021).
[Table]
[Figure] [Figure]
5. C LASSIFICATION
Datasets .
DG-15 (Xu et al., 2022) is a synthetic binary classification dataset with 15 groups.
Each group contains 100 data points.
In this dataset, adjacent groups have similar decision boundaries.
Following Xu et al.
, we use six connected groups as the training groups, and use others as test groups.
Note that, different from Xu et al.
(2022) which focuses on domain adaptation, the data of test groups are unseen in OOD generalization.
DG-60 (Xu et al., 2022) is another synthetic dataset generated using the same procedure as DG-15 , except that it contains 60 groups, with 6,000 data points in total.
We randomly select six groups as the training groups and use others as test groups.
Visualization of DG-15 and DG-60 are shown in Fig. 2 (a) and (b), respectively.
Results .
The results of DG-15 and DG-60 are summarized in Tab. 1.
In both datasets, our method yields the highest accuracy.
For DG-15 , we show the detailed results of all groups in Fig.
We visualize the decision boundary of DG-15 and DG-60 in Appendix 7.3.
Ablations study .
TRO significantly improves the generalization performance by discovering influential groups.
To investigate the reason why TRO outperforms DRO, we show group weights q of DRO and TRO on DG-15 in Fig.
DRO assigns the highest weight to group “ 1 ” which is the furthest group to test groups.
Instead, TRO prioritizes influential groups “ 2 ”, “ 5 ”, and “ 6 ” which are truly connected to the test ones, yielding improved performance on unseen distributions.
5. R EGRESSION
Datasets .
TPT-48 (Vose et al., 2014) 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 first 6 months’ temperature.
We consider two generalization tasks: (1) E(24) → W(24): we use the 24 eastern states as training groups and the 24 western states as test groups; (2) N(24) → S(24): we use the 24 northern states as training groups and the 24 southern states as test groups.
Test groups one hop away from the closest training group are defined as Hoptest groups, those two hops away are Hop-2 test groups, and the remaining groups are Hop-3 test groups.
The visualization of N(24) → S(24) on TPT-48 is shown in Fig.
4 (left).
TRO (data-driven topology) consistently outperforms TRO (physical-based topology) in both tasks, indicating the data-driven topology captures the distributional relation more accurately.
[Table]
[Figure]
generalization risks by not only prioritizing the worst-case groups but also the influential ones.
Results .
We show the results of TPT-48 in Tab. 2.
TRO yields the lowest average MSE on both tasks.
We also report the average MSE of Hop-1, Hop-2, and Hop-3 test groups for both tasks.
Although
REx yields the lowest error on Hop-1 and Hop-2 groups in N (24) → S (24), it yields the highest prediction error on Hop-3 groups.
The results indicate that REx may yield compromised performance under large distributional drifts.
TRO yields the best performance on Hop-3 groups, indicating its strong generalization capability under large distributional drifts.
Ablations study .
(1) Data-driven topology yields better performance than physical-based topology.
We show group centrality of both physical and data topology on the task of North → South in Fig.
“ PA ” is identified by TRO as the influential group in physical-based topology; “ NY ”, “ PA ”, and “ MA ” are identified by TRO as influential groups in data-driven topology.
The results prove that the influential groups in data topology are more effective in minimizing the generalization error
Ignoring either the worst-case (IW-ERM) or influential (DRO) groups would yield compromised performance.
[Figure]
[Table]
(2) Strong OOD resilience of TRO comes from the synergy of the worst-case and influential groups.
To investigate which components contribute to the superior performance of TRO.
We build a simple baseline based on ERM: we directly use the group importance acquired from the topology to weight training groups and the weights are fixed during the training.
We name this baseline as importance weighted ERM (IW-ERM).
We show the results of “N(24) → S(24)” on TPT-48 in Tab. 3.
The results of IW-ERM are inferior to ERM and DRO, possibly because IW-ERM merely considers influential
Data Label
Data ?
Data Label
Data Label 2019USA FLOOD 2018BOL FLOOD 2018SOM FLOOD 2016IND FLOOD USA
BOL PRY
IND ESP
GHA
LKA PAK SOM
KHM NGA
USA
PRY
IND ESP
GHA
LKA
PAK
SOM
KHM NGA
Data Label 2019USA FLOOD USA
Data ?
2018BOL
FLOOD BOL PRY
ESP
GHA NGA
Data Label 2018SOM FLOOD
Data Label 2016IND FLOOD IND
LKA
PAK
KHM SOM
USA
PRY
[Figure] ESP
GHA
LKA
PAK
SOM
[Figure] groups.
and in Fig generalization risks by not only prioritizing the worst-case groups but also the influential ones
5. S EMANTIC S EGMENTATION
Datasets .
Sen1Floods11 (Bonafilia et al., 2020) is a public dataset for flood mapping at the global scale.
The dataset provides global coverage of 4,831 chips of 512 x 512 10m satellite images across 11 distinct flood events, covering 120,406 km 2 .
Each image is associated with its pixel-wise label.
Locations of the 11 flood events are shown in Fig. 6 (left).
Flood events vary in boundary conditions, terrain and other latent factors, posing significant OOD challenges to existing models in terms of reliability and explainability.
Following Bonafilia et al.
(2020), event “ BOL ” is held out for testing, and data of other events are split into training and validation sets with a random 80-20 split.
[Figure]
[Table]
[Figure]
[Table]
TRO yields better performance than other baselines on unseen flood events.
[Figure]
Results .
We show the results of Sen1Floods11 in Tab. 4.
ERM achieves the highest IoU on the validation set while TRO achieves the highest IoU on the test set.
The results prove that TRO yields better performance than other baselines on unseen flood events.
Ablations study .
(1) Data-driven distributional topology is consistent with domain knowledge.
We visualize the distributional topology as well as group centrality in Fig.
6 (right).
The learned distributional topology is consistent with domain knowledge, enhancing the explainability of TRO.
(2) Ablation study on λ .
We report IoU under different λ on Sen1Floods11 in Fig.
IoU remains stable for a wide range of λ , and λ = 0 .
01 yields the best performance.
C ONCLUSION
In this paper, we proposed a new principled optimization method that seamlessly integrates topological information to develop strong OOD resilience.
Empirical results in a wide range of tasks including classification, regression, and semantic segmentation demonstrate the superior performance of our method over SOTA.
Moreover, the data-driven distributional topology is consistent with domain knowledge and facilitates the explainability of our approach.
A CKNOWLEDGEMENTS
This work is partially supported by National Science Foundation (NSF) CMMI-2039857, General
University Research (GUR), and University of Delaware Research Foundation (UDRF).
The authors would like to thank Kien X. Nguyen for helping with the experiments on DomainBed
A PPENDIX 7. R ELATED W ORK
Distributionally Robust Optimization .
In the context of distributionally robust optimization (DRO),
Duchi & Namkoong (2021) and Shalev-Shwartz & Wexler (2016) argued that minimizing the maximal loss over a set of possible distributions can provide better 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 moment constraints (Delage & Ye, 2010), f -divergence (Namkoong & Duchi, 2016) or Wasserstein distance (Shafieezadeh Abadeh et al., 2018).
To avoid yielding overly pessimistic models group DRO (Hu et al., 2018; Sagawa et al., 2019) is proposed to leverage pre-defined data groups to formulate the uncertainty set as the mixture of these groups.
Although the uncertainty set of Group DRO is of a wider radius while not being too conservative, our preliminary results show that Group DRO recklessly prioritizes the worst-case groups that incur higher losses than others.
Such worst-case groups are not necessarily the influential ones that are truly connected to unseen distributions; optimizing over the worst-case rather than influential groups would yield mediocre OOD generalization performance.
Out-of-Distribution Generalization .
The goal of OOD generalization is to generalize the model from source distributions to unseen target distributions.
There are mainly two branches of methods to tackle OOD generalization: group-invariant learning (Arjovsky et al., 2019; Koyama & Yamaguchi, 2020; Liu et al., 2021) and distributionally robust optimization.
The goal of group-invariant learning is to exploit the causally invariant correlations across multiple distributions.
Invariant Risk Minimization (IRM) is one of the most representative methods which learns the optimal classifier across source distributions.
However, recent work (Rosenfeld et al., 2021) shows that IRM methods can fail catastrophically unless the test data are sufficiently similar to the training distribution.
7. I MPLEMENTATION DETAILS
In Sec. 3.1, for all hyperparameters such as the kernel scale σ 2 and the maximum scale K , we use the default values from the official implementation 1 of Tong et al.
In Sec. 3.2, for learning rate of model parameters η θ , we use default values from Xu et al.
( DG-15/-60 and TPT-48 ) and Bonafilia et al.
(2020) ( Sen1Floods11 ).
Therefore, we only tune the learning rate of the mixture distribution η q and the dual variable λ .
All results are reported over 3 random seed runs, which is consistent with Koh et al.
(2021) and Shi et al.
We select λ from {1e-3, 1e-2, 1e-1, 1, 10, 100} and select η q from {1e-4, 1e-3, 1e-2, 1e-1, 1}.
7. A DDITIONAL RESULTS
DG-15 and DG-60 .
We provide detailed classification results for each group.
The results are shown in
Fig. 8.
We can see that, compared to other baselines, TRO significantly improves the generalization performance on groups that are far from the training groups, such as group “ 13 ”.
We further visualize the decision boundary of DG-15 and DG-60 in Fig.
9 and Fig. 10, respectively
[Figure]
[Figure] [Figure]
[Figure]
Incognita
Art
Photo
Sketch Cartoon
[Figure]
SUN09 LabelMe
VOC
Art Photo
Sketch
Data-driven topology on PACS
[Table] [Table]
DomainBed .
Following the instructions of the official implementation of DomainBed Gulrajani &
Lopez-Paz (2021), we have conducted experiments on PACS (Li et al., 2017), Terra (Beery et al., 2018), and VLCS (Fang et al., 2013).
Image samples of the three datasets are shown in Fig.
11 (left).
(1) PACS is one of the most popular dataset for out-of-distribution generalization.
It consists of images from four groups: “ Art ”, “ Cartoon ”, “ Photo ” and “ Sketch ”.
Results on PACS are shown in Tab. 5.
Results of other baselines are from Appendix B.4 of Gulrajani & Lopez-Paz (2021).
In average accuracy, TRO outperforms the SOTA by 0.5%.
To further investigate the results, we visualize the learned topology in Fig.
11 (right).
As observed, when “ Cartoon ” is the test group, the topology is a chain graph consisting of three nodes where “ Art ” is the most influential group.
A possible explanation is that “ Art ” may contain more information than “ Photo ” and “ Sketch ” as “ Art ” can be viewed as the combination of photos and various kinds of styles.
Even though the topology is so simple, it enables our method to significantly outperforms ERM and DRO by 0.8% and 2.4% on average.
The results empirically demonstrate the strong explainability of our method when the number of training groups is quite limited, i.e. , 3.
We would like to point out that when the distributional shift across different groups is small (see explanation on the results of VLCS ), the influential group may not exist and all groups share the same centrality.
In this special case, TRO aims to strike a good balance between the average (ERM) risk and the worst-case (DRO) risk.
[Table]
(2) Terra consists of images of wild animals captured by camera traps under four locations.
Results on Terra are shown in Tab. 6.
Results of other baselines are from Appendix B.6 of Gulrajani &
Lopez-Paz (2021).
As observed, in average accuracy, TRO achieves comparable results with the
SOTA and and by 1.8% and 2.
(3) Results on VLCS are shown in Tab. 7.
Results of other baselines are from Appendix B.3 of
Gulrajani & Lopez-Paz (2021).
The average accuracy of DRO and TRO is the same.
We assume the reason is that the distributional shift across different groups is small (Li et al., 2017), and therefore the influential group may not exist and all groups share the same centrality.
In this special case, TRO aims to strike a good balance between the average (ERM) risk and the worst-case (DRO) risk.
The images of VLCS are all photos and the distributional shift is not as significant as PACS ( e.g. , Photo vs. Sketch).
As stated in Sec. 2.1 of Li et al.
(2017), “despite the famous analysis of dataset bias that motivated the creation of the VLCS benchmark, it was later shown that the domain shift is much smaller with recent deep features”, and PACS (Li et al., 2017) was proposed to address this limitation.
7. P ROOF OF T HEOREM
Proof .
By using the property of convex loss functions, we can obtain:
[Equation]
where ∀ t ≥ 1 :
R
[Equation]
By rearranging the terms in the above equation, we obtain:
[Equation]
Following Collins et al.
[Equation]
we will derive the combined bound by bounding the expectation of each term in the above equation
[Equation]
For the first term
[Equation]
by utilizing the telescoping sum
[Equation]
we can obtain:
[Equation]
Similarly, for the second term:
[Equation]
The third term and the last term are bounded by √ T ˜ σ q and R Θ √ T ˜ σ θ , respectively.
To this end, we can derive the overall bound as:
[Equation]
The above bound can be minimized by setting the step sizes η θ = 2 R Θ / (cid:16) ˆ G θ √ T (cid:17) and
η q = 2 /
[Equation]
ˆ G
[Equation]
√ T (cid:17) .
7. P ROOF OF T HEOREM
Proof .
Inspired by Qian et al.
and Collins et al.
(2020), we utilize the property of M-smooth to start the proof:
[Equation]
We rearrange these terms by:
[Equation]
The second term of the above equation can be bounded by:
[Equation]
By using the Law of Iterated Expectations, we can obtain:
[Equation]
Next
[Equation]
we investigate the convergence of q :
[Equation]
By aggregating the difference at all time steps, we obtain:
[Equation]
Since the above equation holds for all q ∈ ∆ m , we maximize the right hand side over q ∈ ∆ m :
[Equation]
Eqs.
9 and 10 show that TRO converges in expectation to an ( ϵ, δ ) -stationary point of R in O (1 /ϵ 4 ) stochastic gradient evaluations.