T OPOLOGY - AWARE R OBUST O PTIMIZATION
Anonymous authors
Paper under double-blind review
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 that the learned topology is highly explainable and consistent with human knowledge and scientific plausibility.
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 that 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.
Topology learning methods that are orders of magnitude faster than Earth Mover’s distance to uncover distributional structure from massively collected datasets.
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.
Explainable distributional topology that is consistent with human knowledge and scientific plausibility.
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
[Table]
Data Graph <latexit sha1_base64="WZ4EL4dDpWzyj5/crhTqcPS1RzQ=">AAAB6nicbVBNS8NAEJ3Ur1q/qh69LBbBU0ikqBeh4MVjRfsBbSib7aZdutmE3YlQSn+CFw+KePUXefPfuG1z0NYHA4/3ZpiZF6ZSGPS8b6ewtr6xuVXcLu3s7u0flA+PmibJNOMNlshEt0NquBSKN1Cg5O1UcxqHkrfC0e3Mbz1xbUSiHnGc8iCmAyUiwSha6QFvvF654rneHGSV+DmpQI56r/zV7Scsi7lCJqkxHd9LMZhQjYJJPi11M8NTykZ0wDuWKhpzE0zmp07JmVX6JEq0LYVkrv6emNDYmHEc2s6Y4tAsezPxP6+TYXQdTIRKM+SKLRZFmSSYkNnfpC80ZyjHllCmhb2VsCHVlKFNp2RD8JdfXiXNC9e/dKv31UrNzeMowgmcwjn4cAU1uIM6NIDBAJ7hFd4c6bw4787HorXg5DPH8AfO5w/NRY1u</latexit> t = <latexit sha1_base64="o6X1JomQk5xBvKEipRNNSU2gQNk=
[Table]
/wiDM3fPLSa9SdWxnBrJM3IJUoUCjV/ny+gnLIh4jk1Trruuk6OdUoWCST8pepnlK2YgOeNfQmEZc+/ns5gk5NUqfhIkyFSOZqb8nchppPY4C0xlRHOpFbyr+53UzDG/8XMRphjxm80VhJgkmZBoA6QvFGcqxIZQpYW4lbEgVZWhiKpsQ3MWXl0nrwnav7Np9rVq3izhKcAwncAYuXEMd7qABTWCQwjO8wpuVWS/Wu/Uxb12xipkj+APr8weeFJFb</latexit> 1/3 … … … …
[Table]
Distributional Topology
Combined Data
[Table]
TRO
[Table]
Data Label
EVENT
Data Label
EVENT 2 Data Label
EVENT
Data ?
EVENT N <latexit sha1_base64="FQW66Tvdz8Opul+NaYFhJeCIpCY=">AAACFHicbZDLSsNAFIYn9VbrLerSzWARKkpJpKjLghuXFexFmhAmk0k7dCZJZyZCCX0IN76KGxeKuHXhzrdxknah1R8GPv5zDnPO7yeMSmVZX0ZpaXllda28XtnY3NreMXf3OjJOBSZtHLNY9HwkCaMRaSuqGOklgiDuM9L1R1d5vXtPhKRxdKsmCXE5GkQ0pBgpbXnmiZM543GKAgcHsZJOgVNHUu5wpIa+n7WmHq/1TuHdsWdWrbpVCP4Few5VMFfLMz+dIMYpJ5HCDEnZt61EuRkSimJGphUnlSRBeIQGpK8xQpxINyuOmsIj7QQwjIV+kYKF+3MiQ1zKCfd1Z76pXKzl5n+1fqrCSzejUZIqEuHZR2HKoIphnhAMqCBYsYkGhAXVu0I8RAJhpXOs6BDsxZP/Quesbp/XGzeNarMxj6MMDsAhqAEbXIAmuAYt0AYYPIAn8AJejUfj2Xgz3metJWM+sw9+yfj4Bic7ns0=</
[Table]
> { ··· }⇠ P m ( X,Y ) <latexit sha1_base64="vFganPQ5gh8CW2+aKxSpUZMfrXM=">AAACFHicbZDLSsNAFIYn9VbrLerSzWARKkpJSlGXBTcuK9iLNCFMJpN26OTSmYlQQh7Cja
[Table]
+srlXXaxubW9s7+u5eT0QJx6SLIxbxgYsEYTQkXUklI4OYExS4jPTdyVVe798TLmgU3spZTOwAjULqU4ykshz9xEqt6TRBnoW9SAqrwMwSNLACJMeum3Yyx2wMTuHdsaPXjaZRCP4Fs4Q6KNVx9E/Li3ASkFBihoQYmkYs7RRxSTEjWc1KBIkRnqARGSoMUUCEnRZHZfBIOR70I65eKGHh/pxIUSDELHBVZ76pWKzl5n+1YSL9SzulYZxIEuL5R37CoIxgnhD0KCdYspkChDlVu0I8RhxhqXKsqRDMxZP/Qu+saZ43WzetertVxlEFB+AQNIAJLkAbXIMO6AIMHsATeAGv2qP2rL1p7/PWilbO7INf0j6+AcrUnpE=</latexit> { ··· }⇠ P 1 ( X,Y ) …
Diffusion EMD
EVENT 1 EVENT
EVENT N-
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 connectives and interactions among 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 captures 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 enjoys strong explainability while being consistent with human knowledge and scientific plausibility.
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 the future work.
3. L EARNING ON TOPOLOGY : E XPLOIT TOPOLOGY FOR ROBUST OPTIMIZATION
Algorithm
Input: Data of E train , Step sizes η θ and η q Output: Learned model f
Topology Learning : if G physic exists then
G ← G physic else
Obtain the affinity matrix K from data
Q ← Diag (cid:16)(cid:80)
j K ij (cid:17)
M ← Q − 1 K − 1 Q −
D ← Diag (cid:16)(cid:80) j M ij (cid:17) P ← D − 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 | v ) is the number of ( s, t ) -paths that go through group e . Intuitively, c physic e 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 groups 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 rate 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 function in Theorem 1 and the convergence rate for non-convex loss function 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 are 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 ∥∇ f ( x ) ∥
≤ ϵ .
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: (cid:13)(cid:13)(cid:13) ∇
[Equation]
R (˜ θ, ˜ q ) (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 lower upper bound on the generalization risks over unseen distributions with the topological constraint.
Let H denote the family of losses associated to 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.
Following Mohri et al.
(2019), we define weighted Rademacher complexity for any hypothesis set F as: [Equation] n (
[Equation] = [Equation] i
[Equation]
(cid:35) , 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
In this section, we evaluate the proposed TRO in a wide range of tasks including classification, regression, and semantic segmentation.
We compare the 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 and additional results in the Appendix.
We include the source codes in the supplementary.
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 achieves the highest performance.
For DG-15 , we show the detailed results of all groups in Fig. 8.
We visualize the decision boundary of DG-15 and DG-60 in Appendix 8.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.
As seen, in DRO, group “1” is assigned with the highest weight while “1” is the most distant group to the test groups.
Instead, our method ignores group “1” and assigns the highest weight to group “6”, improving the model’s robustness against 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, the 24 southern states as test groups.
Test groups one hop away from the closest training group are defined as Hop-1 test 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) consistently outperforms TRO (physical) in both tasks, indicating the data topology may capture the distributional distances more accurately.
[Table] [Figure] Results . We show the results of TPT-48 in Tab. 2.
As observed, TRO achieves 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 achieves 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.
TRO achieves the best performance on Hopgroups, indicating its generalization capability under large distributional drifts.
Ablations study .
(1) Data-driven topology yields better performance than physical-based topology.
We show group centrality of physical and data topology on the task of North → South.
Results are shown in Fig.
As observed, “RI” and “NE” are identified as influential groups in physical topology while “NY” and “MA” are influential groups in data topology.
Ignoring either the worst-case (IW-ERM) or influential (DRO) groups would yield compromised performance.
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[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 as the weights for 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.
We can see that the results of IWERM is even worse than ERM and DRO.
We further show the group importance of DRO and TRO in Tab. 5.
As shown, TRO significantly reduces the generalization risks by striking a good balance between the worst-case and influential groups.
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
IND ESP
GHA
LKA
PAK
SOM
KHM NGA
Centrality
5. S EMANTIC S EGMENTATION
Datasets .
Sen1Floods11 (Bonafilia et al., 2020) is a public dataset for flood mapping at global scale.
The dataset provides global coverage of 4,831 chips of 512 x 512 10m satellite images acrossdistinct 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 test, and data of other events are split into training and validation sets with a random 80-20 split. [Figure]
[Table] TRO achieves the highest IoU on test set while achieves the highest validation [Table] I o U o n t h e V a li d a t i o n S e t
Results .
We show the results of Sen1Floods11 in Tab. 4.
As seen, TRO achieves the highest IoU on test set while ERM achieves the highest IoU on validation set.
Results demonstrate the strong OOD resilience of TRO in high-stake applications.
Ablations study .
(1) Data-driven distributional topology is consistent with human knowledge.
We visualize the distributional topology as well as group centrality in Fig.
6 (right).
We observe that “IND” and “NGA” are identified as the most influential events.
Both “IND” and “NGA” are aroused by heavy rainfall, the most prevalent disasters that cause floods.
We empirically found that the least influential events are usually caused by edge cases.
For example, “GHA” was caused by heavy rains and waters spilling from a dam in neighboring Burkina Faso, and “KHM” was caused by a dam that collapsed in southern Laos.
(2) Ablation study on the Lagrange multiplier λ .
We report IoU on the validation set of Sen1Floods11 .
Results are shown in Fig.
As seen, IoU remains stable during a wide range of λ , and λ = 0 .
01 achieves the highest IoU on the validation set.
C ONCLUSION
In this paper, we proposed 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.
Moreover, the distributional topology is highly explainable while consistent with human knowledge.  REPRODUCIBILITY STATEMENT We provide the source codes for reproduction in the supplementary material and describe the implementation details in Appendix 8.2.  A PPENDIX 8. 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.
8. 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}.
8. A DDITIONAL RESULTS
DG-15 and DG-60 .
We provide detailed classification results for each groups.
The results are shown in Fig.
We can see that, compared to other baselines, TRO significantly improves the generalization performance on groups which 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]
8. P ROOF OF T HEOREM
Proof .
By using the property of convex loss functions, we can obtain:
[Equation]
where ∀ t ≥ 1 :
[Equation]
By rearrange the terms in the above equation [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*.
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 / (cid:16) ˆ G q √ T (cid:17) .
8. 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 [Equation] we can obtain:
[Equation]
Next [Equation] we investigate the convergence of q :
[Equation]
By aggregating the difference at all time steps
[Equation]
we obtain:
[Equation]
Since the above equation holds for all q ∈ ∆ m
[Equation]
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.