Domain Generalization without Excess Empirical Risk
Abstract
Given data from diverse sets of distinct distributions, domain generalization aims to learn models that generalize to unseen distributions.
A common approach is designing a data-driven surrogate penalty to capture generalization and minimize the empirical risk jointly with the penalty.
We argue that a significant failure mode of this recipe is an excess risk due to an erroneous penalty or hardness in joint optimization.
We present an approach that eliminates this problem.
Instead of jointly minimizing empirical risk with the penalty, we minimize the penalty under the constraint of optimality of the empirical risk.
This change guarantees that the domain generalization penalty cannot impair optimization of the empirical risk, i.e., in-distribution performance.
To solve the proposed optimization problem, we demonstrate an exciting connection to rate-distortion theory and utilize its tools to design an efficient method.
Our approach can be applied to any penalty-based domain generalization method, and we demonstrate its effectiveness by applying it to three examplar methods from the literature, showing significant improvements.
Introduction
Empirical risk minimization (ERM) is the workhorse of machine learning.
In this setting, a learner collects a set of data, assumed to be independent and identically sampled from a particular distribution of interest.
An empirical risk over these samples is later minimized, hoping that the resulting model will generalize to future unseen data from the same distribution.
Although ERM is ubiquitous, it is brittle to minor shifts in the distribution and the resulting models suffer a significant performance drop when evaluated beyond their training distributions.
Examples of such distribution shifts include different hospitals for medical imaging [49], camera types for perception models [8] and time periods for natural language models [20].
The problem of learning models which generalize under such distribution shifts is studied in the literature under the rubric of domain generalization .
Given a collection of domains (data distributions) domain generalization aims to find models that generalize to unseen domains.
Although the domain generalization literature is rich, many model-free methods follow a common recipe.
A surrogate penalty that encapsulates generalization is designed and a joint optimization problem of minimizing the empirical risk with the surrogate penalty is solved.
Typical such penalties are invariance to an adversarial domain classifier [17, 42], invariance of an optimal predictor [5], and uniformity of risk [35].
Although these approaches are technically sound, their effectiveness has been challenged by recent empirical benchmarks [23, 18].
Joint optimization of the empirical risk and penalty is a sound method since the error on the unseen distribution can be bounded as the combination of an error on the training distributions (empirical risk) and the gap between performance on training and unseen distributions (generalization gap).
Although domain generalization literature focuses mainly on controlling the generalization gap, we hypothesize that this joint optimization can increase empirical risk, effectively hurting performance on both training and unseen distributions.
In other words, the lacking empirical performance of existing algorithms might be because they fail to optimize the empirical risk (i.e., in-distribution performance) even though they learn to control generalize gap.
Consider the extensive evaluation conducted on the
WILDS benchmark [23].
ERM yields higher out-of-distribution performance than dedicated domain generalization methods on most problems.
In contrast, when we look at the generalization gap (i.e., the difference between in-distribution and out-of-distribution performance), domain generalization methods have a smaller gap than ERM on most problems.
This supports our hypothesis that existing algorithms might be creating too much excess empirical risk.
Thus the motivating question of our work: Can we perform joint optimization without causing any excess empirical risk?
Solving joint optimization without sacrificing empirical risk is not generally possible unless surrogate penalty term and empirical risk align perfectly.
Instead, we formulate an optimization problem with the explicit constraint of no excess empirical risk.
In other words, we propose to minimize the surrogate penalty only if it is not hurting the empirical risk.
Formally, instead of joint minimization of empirical risk and surrogate penalty, we minimize the surrogate penalty under the constraint of optimality of the empirical risk.
To design an algorithm for the proposed optimization problem, we start with searching for a formal definition of the optimality of the empirical risk.
Since the applications of interest are non-convex, this can only be local optimality.
Hence, the constraint we consider is convergence to a stationary point of the empirical risk.
We show that biased gradient descent converges to a stationary point as long as the bias is bounded and decreases through optimization.
Using this result, we design an iterative algorithm that seeks the update minimizing the surrogate penalty under the constraint of bounded deviation from the empirical gradient.
This construction is closely related to the rate-distortion problem from information theory [13], and we draw on this vast literature to design an algorithm.
Specifically, we propose a Blahut-Arimoto-style method [11, 4] to optimize the proposed objective.
Our method only considers the optimization aspect of the domain generalization problem and can be applied to any penalty-based model with minor modifications.
We perform our primary empirical evaluation on the CORAL model due to its simplicity [42].
We evaluate on the large-scale domain generalization problems in the WILDS benchmark of real distribution shifts [23].
Our evaluation using the image, formal language, natural language, and graph modalities suggests that our method significantly improves performance.
To further validate the generality of our method, we apply it to the FISH model [39] and VRex [25] showing similar gains on WILDS [23] and DomainBed [18].
Method 2.1 Formulation and the Preliminaries
In this section, we introduce our notation and present a unified treatment of existing penalty based domain generalization methods.
We are interested in a learning problem in the space of X × Y , where X is the space of inputs and Y is the space of labels.
In domain generalization, a collection of different domains ˆ π = { π 1 , . . .
, π M } as well as an independent and identical (iid.) dataset for each domain are given as { x ij , y ij } ∼ P π j for domain j with N j samples.
Consider the union of these datasets as the training set of domain generalization in the form of tuples { x, e, y } such that x is the input, e is the domain id, and y is the label.
We further consider a parametric function family h ( · ; θ ): X → Y and a loss function l ( · , · ): Y × Y → R + .
Throughout our discussion, we denote random variables with capital letters (e.g. X, E, Y ) and their realizations with lower-case letters (e.g. x, e, y ).
We are specifically interested in penalty based domain generalization methods where empirical risk and domain generalization penalty are jointly minimized.
Each model defines a penalty following a different set of assumptions.
For example, penalty is defined in terms of second-order statistics for
CORAL [42] and in terms of sub-optimality of the domain-specific classifier for IRM [5].
Since we are interested in the optimization of these models, not their construction, we treat Penalty as a generic entity and consider the optimization problem defined by these methods.
These methods aim to solve the following joint optimization problem:
[Equation]
where L ( θ ; { x i , e i , y i } ) is the empirical risk defined as L ( θ ) ≜
[Equation] Although, the penalty term is typically weighted, weighting can be considered as part of its definition.
2.2 Optimizing Domain Generalization Objectives
The optimization problem defined in (1) is a sound proxy in order to optimize the risk for the unseen domains.
Consider an unseen distribution π ⋆ , its corresponding population risk can be bounded using the following straightforward decomposition;
[Equation]
Following this decomposition, the optimization problem in (1) is sensible since the empirical risk (cid:0) L ( θ ; { x i , e i , y i } ) (cid:1) approximates the risk on seen domains and a the penalty term (cid:0)
Penalty ( θ ; { x i , e i , y i } ) (cid:1) approximates the generalization error.
However, this formulation has a failure point.
Consider a well-designed penalty term which successfully approximates the generalization gap.
Clearly the joint optimization will decrease the generalization error.
However, it can effectively increase the risk on seen domains due to the difficulties in optimization process.
Moreover, if the reduction in generalization error is lower than the unintended excess risk on seen domains, the final performance on unseen domains can be worse than not doing any domain generalization.
We hypothesize that this behavior might explain the recently demonstrated behavior from [23, 18].
As a motivating empirical finding, consider the benchmarking results from [23], which are enclosed in Appendix D. Take CORAL [42] as a representative method 1 .
When we look at out-of-distribution performance, ERM outperforms CORAL on 7 out of 9 problems.
However, when we look at the generalization gap (i.e., the difference between in-distribution and out-of-distribution performance),
CORAL has lower gap than ERM on 7 out of 9 problems.
Although the models learned with CORAL have lower generalization error, the loss of in-distribution performance is so significant that the resulting models do not perform well when tested out-of-distribution.
In order to partially address this issue, we propose a constraint of no excess risk.
Instead of jointly minimizing the empirical risk and the penalty, we minimize the penalty under the constraint of the optimality of the empirical risk.
More formally, we consider the following optimization problem:
[Equation]
Remarks on the reformulation: i)
This reformulation eases the aforementioned issue as it guarantees no excess risk over the seen domains.
Since the optimality of the empirical risk is directly a constraint, this formulation would not harm in-distribution performance.
ii)
This reformulation is not sensible if performance on seen domains and unseen domains are inherently contradictory.
For example, if there is strong nuisance information in the seen domains, any model not using this nuisance information will fail to optimize the empirical risk, effectively making excess risk a requirement for generalization.
However, for the benchmarks we are interested in, which involve supervised learning from multiple datasets, this is not the case as a universal model solving all datasets (seen and unseen) exists.
In other words, we address only a subset of domain generalization problems where seen and unseen domains do not compete.
iii)
Since the models we use are significantly overparameterized, the space of solutions to empirical risk minimization is quite large, likely including many domain-invariant and domain-sensitive solutions.
Hence, this problem is still non-trivial from the numerical optimization perspective.
iv)
Finally, one can mistakenly think that our reformulation and the original joint optimization would find the same solution due to the Lagrangian function.
However, this is incorrect as the Lagrangian function requires optimization of the Lagrangian multiplier, which is intractable for large-scale deep learning and typically replaced with fixed hyper-parameters.
Moreover, we propose to directly solve (3) without resorting to Lagrangian function.
In the rest of this section, we first discuss in Section 2.3 how to convert the constraint in (3) into a practical form for first-order optimization.
Then in Section 2.4 we discuss how (3) is closely related to rate-distortion theory, demonstrating an unexpected connection.
Finally in Section 2.5 we utilize tools from rate-distortion theory to numerically solve (3).
2.3 Stationarity of the ERM as a Constraint
The overarching objective of θ ∈ arg min θ L ( θ ) is rather too ambitious.
Since the family of loss functions we are interested in is non-convex, the best we can hope is convergence to a stationary point.
Hence, we convert (3) into
[Equation]
While solving the problem in (4), we constrain ourselves to iterative methods to be compatible with stochastic gradient descent (SGD), which is commonly used to train deep neural networks.
Hence, we would like to have an iterative procedure θ t +
=
θ t − ηG t which solves (4) in convergence.
Before we proceed, we note that when G t is the stochastic gradient (i.e., E [ G t ] = ∇ θ ˆ L ( θ ) ), convergence to a stationary point of the empirical risk can be shown under mild conditions.
However, we are also interested in minimizing the penalty.
Hence, we need to allow updates that are different than the gradients while still supporting eventual convergence.
Convergence of SGD with biased gradients has previously been analyzed by Sener & Koltun [37], Hu et al.
[22], and Ajalloeian & Stich [1].
The common theme is the necessity for the bias to decrease over time.
This requirement is explicit in [37] and implicit in [22] and [1].
In the implicit case, bias is assumed to be linearly bounded by the gradient norm, later shown to decrease.
We adapt the analysis of SGD with bias from Sener & Koltun [37] and show that as long as the norm of the bias is bounded and this bound decreases during training faster than the rate 1 / √ t , it convergences to a stationary point of a smooth and non-convex risk function.
We state the convergence result in the following proposition and defer the proof to Appendix A.
Proposition 2.1.
Apply stochastic updates θ t +
=
θ t − ηG t for T steps.
Assume i) R ( θ ) is a nonconvex, L -Lipschitz, and µ -smooth function bounded by ∆ ; ii) Updates are bounded in expectation, E [ ∥ G t ∥ 22 ] ≤ V ; iii) Bias of the updates is decreasing, ∥ E [ G t ] − ∇R ( θ t ) ∥ 2 ≤ D / √ t .
Then,
[Equation]
In this proposition, we assume that the risk function is non-convex, L -Lipschitz, and µ -smooth 2 , and that the stochastic updates are bounded.
These are standard assumptions and shown to be valid for common deep learning architectural choices [19].
Finally, the third assumption is on the bias decreasing with the rate of 1 / √ t , and we explicitly handle it in our solver, which we discuss below.
With Proposition 2.1, we operationalize the definition of a satisfactory iterative update.
As long as we can bound the difference between the updates and the gradient of the empirical risk, we can guarantee convergence to a stationary point.
Hence, our iterative approach to domain generalization starts with an initialization θ 0 and applies the iteration θ t +
=
θ t − ηG t , with G t being the solution of the following optimization problem:
[Equation]
The constraint of our final formulation in (5) guarantees the assumption (iii) from Proposition 2.1.
Hence, applying the updates G t , which are the solution of (5), would satisfactorily minimize the empirical risk.
In the meantime, we seek the updates minimizing the penalty among these satisfactory ones.
Since our approach utilizes satisfactory updates for empirical risk instead of the first-order optimal ones, following Herbert Simon’s characterization of satisficing [40], we call this approach
Satisficing Domain Generalization (SDG).
Finally, the minimization is over the distribution of the updates instead of the update itself.
Although this is unconventional, it is intentional and crucial for making the problem compatible with rate-distortion theory, which we utilize heavily in our solver [Figure] [Figure]
2.4 Understanding Satisficing Domain Generalization
In this section, we provide observations on the proposed formulation (5) by utilizing tools from rate-distortion theory.
Consider the function R ( D ) as the solution of (5) for a given constant factor D .
The problem (5) is simply minimizing the penalty with the constraint of the limited distortion in updates.
This form is strikingly similar to the rate-distortion function, where the (compression) rate is optimized with the constraint of acceptable degradation in decoding.
If the penalty were the mutual information between the true gradient and G t , it would have been exactly equivalent.
Hence, we can further utilize tools from rate-distortion theory to better understand (5).
We now state a few simple facts derived from rate-distortion theory.
( Proofs follow their rate-distortion counterparts directly and are provided in Appendix B. )
Extreme cases: When D t = D / √ t = 0 , E [ G t ] =
∇L ( θ t ) .
Hence, at the end of training, we perform ERM updates.
When the distortion is unbounded ( D t = ∞ ), updates directly minimize the penalty.
Moreover, our proposed formulation smoothly goes from updates minimizing penalty to ERM updates during training.
Improvement guarantee: R ( D ) is a non-increasing and convex function of D .
Hence, R (0) ≥ R ( D t ) .
In other words, SDG decreases the penalty when compared with ERM.
Geometric picture: The satisficing requirement implies a certain structure over the plane where the x -axis is the distortion and the y -axis is the penalty.
We visualize this geometry and the interpretation of our method in Figure 1.
2.5 Solving Satisficing Domain Generalization
We present a numerical algorithm to solve (5).
We first give an iterative method similar to BlahutArimoto (BA) [11, 4] with intractable complexity.
We later make simplifying assumptions to design a tractable method.
We propose to regularize the objective with the mutual information between gradients ( G t ) and the domain-ids ( E ).
This regularization is sensible as we would like gradients to carry little information about the domain-ids.
The final numerical optimization problem we tackle is
[Equation]
The information regularization is not only desired but also crucial as it enables us to apply technical derivation similar to Blahut-Arimoto [11, 4].
To solve (6), We only consider the discrete and finite probability mass functions where G t ∈ { G t 1 , . .
, G tK } ; hence, we only need to solve for the point masses p G ti = p ( G t = G ti ) .
We first define auxiliary variables p G ti | e = p ( G t =
G ti | E = e ) , which are conditioned on domain-ids.
By this decomposition, we can solve for each domain conditional probability p G ti | e separately and later marginalize them via p G ti = 1 / M (cid:80) e p G ti | e .
We further replace E G t [ ∥ G t − ∇ θ ˆ L ( θ ) ∥ 2 ] ≤ D / √ t with a stronger constraint of uniform bounds over domains as E G t | E = e
[ ∥ G t − ∇ θ ˆ L e ( θ ) ∥ 2 ] ≤ D / √ t and apply a similar derivation as Blahut [11] and Arimoto [4].
We defer the complete derivation to the
Appendix C.1 and state the update steps directly.
Consider the notation d ( G, e )
=
∥ G − ∇L e ( θ ) ∥ 2 , and Pen ( G ) =
Penalty ( θ t + G ; { x i , e i , y i } ) .
We initialize the conditional probabilities uniformly as p G ti | e = 1 / K and apply the iterations:
• Domain Specific Gradients:
ˆ p l +1 G tk | e = p lG tk exp (cid:104) − 1 γ ( Pen ( G tk ) + βd ( G tk , e ))
(cid:105) • Normalization: p l +1 G t k | e = ˆ p l +1 Gtk | e (cid:80) ˆ k ˆ p l +
Gt
ˆ k | e • Marginalization: p l +1 G tk = 1 / M (cid:80) e p l +1 G tk | e
Here β t is a parameter that inversely depends on D t .
These iterations converge to a solution where all the conditional distributions are equal.
Applying these updates to the space of gradients is intractable since the space of G t is highdimensional.
We propose two simplifications to make the application of BA tractable: i) we solve for each parameter independently, and ii) we only estimate the sign and copy the magnitude from the ERM.
Specifically, we solve the BA problem using the discrete set ( G t ) p ∈ { ( G t 1 ) p
= ( ∇L ) p , ( G t 2 ) p = − ( ∇L ) p } for ( G t ) p , which is the update for the p th parameter.
We use uniform initialization: p 0 (( G t 1 ) p ) =
p 0 (( G t 2 ) p ) = 0 .
We present a more detailed discussion in Appendix C.1&C. with complete derivation and give pseudocode in Algorithm 1.
After the BA iterations are completed, we sample the update G t ∼ p ( G t ) and apply it to the model.
To understand the proposed simplifications, we consider their implications for the rest of the pipeline.
We feed the resulting estimated gradients to first-order numerical optimizers.
Moreover, we utilize the first-order approximation of the penalty function in our algorithms in Appendix C.1&C.2.
Hence, the rest of the pipeline is first order, not utilizing higher-order relationships between coordinates, justifying the first simplification.
The second simplification is similar to SignSGD [10] like methods, where only the sign of the gradient is used without its magnitude for efficient communication between worker nodes in multi-node optimization.
Sign of the gradient suffices for convergence in non-convex problems both theoretically and empirically.
Experimental Results 3.1 Implementation details.
We implement the proposed algorithm in PyTorch.
Following the analysis in Proposition 2.1, D t = D (cid:112) 1 / t .
Moreover, β is inversely related to D .
We choose β t = ( β ◦ / √ T ) √ t for epoch t from t = 1 to t = T .
We treat β ◦ as a hyperparameter and search in β ◦ ∈ { 0 . 01 , 0 . 1 , 1 . 0 } together with all other hyperparameters.
Instead of treating γ (weight of the information regularization) as a hyper-parameter, we present a simple heuristic to set it.
Specifically, since the relevant term in the iterative update is exp (cid:104) − 1 γ ( InvP ( G t k )
+ βd ( G t k , e )) (cid:105) , we keep track of the average of ( InvP ( G tk ) + βd ( G tk , e )) throughout training and set it as γ at each instant.
In all of our experiments, we used 25 BA iterations.
Since using all domains at each batch is not feasible, we first sample B D domains, then sample B examples per domain (named as group sampler in WILDS[23]).
We search all hyperparameters with random sampling with the constraint that all methods have the same budget for the hyperparameter search.
Specifically, we use 20 random hyperparameter choices.
We also [Table] [Table] perform early stopping and choose the best epoch using validation domain performance.
Moreover, to ensure a fair comparison, all methods are run for the same amount of wall-clock time.
Since our method is slower, we perform fewer epochs than other methods.
We list all chosen hyperparameters in Appendix F.
3.2 Evaluation on WILDS
Datasets Our main evaluation is using the WILDS benchmark [23].
WILDS is designed to evaluate domain generalization and subpopulation shift in realistic problems.
Among these problems, we use seven problems that are either pure domain generalization problems or a combination of domain generalization and subpopulation shift.
These problems cover a wide range, including outdoor images (iWildCam), medical images (Camelyon), satellite images (FMoW, PovertyMap), natural language (Amazon), formal language (py150), and graph-structured data (OGB-MolPCBA) with the number of domains from 5 to 120 K .
The summary of the benchmarks is in Appendix E.
Baselines We apply our proposed optimizer (SDG) to CORAL [42] and denote CORAL+SDG.
We compare CORAL+SDG with i) CORAL: the standard implementation of CORAL which is jointly minimizing the loss and the penalty, ii) CORAL+AND-MASK:
Applying the AND-MASK[30] to
CORAL minimization, iii) CORAL+SAND-MASK:
Applying the SAND-MASK[38] to CORAL minimization.
It is important to note that AND-MASK[30] and SAND-MASK[38] have not been applied in conjunction with CORAL before.
Still, we compare with them since they also directly modify the optimization procedure.
We compute the gradient of the empirical risk plus coral penalty for each domain separately and utilize their and-mask and sand-mask as corresponding updates.
3.2.1 Main Results
We summarize the out-of-distribution results in Table 1 and in-distribution results in Table 2.
We also compute the generalization gap (difference between in-distribution and out-of-distribution performance) and tabulate it in Table 3.
When evaluated on out-of-distribution performance, our method improves CORAL in all settings.
The improvement is significant for Camelyon17 with more than 15% and OGB-MolPCBA with more than 7%.
Moreover, although the CORAL outperforms ERM with a standard optimizer only for 1 out 7 benchmarks, combined with our optimizer, CORAL outperforms ERM for 6 out of 7 benchmarks.
Our method closes the gap for Camelyon17, FMoW, Amazon, Py150, and PovertyMap.
The only setting CORAL fails to outperform ERM when combined with our optimizer is OGBMolPCBA.
We believe one reason for this is the additional computational complexity of our method.
For a fair comparison, our method performs significantly fewer epochs than ERM.
When compared with AND-mask and SAND-mask, our method outperforms them in all benchmarks.
[Table] When evaluated for in-distribution performance in Table 2, our method improves upon CORAL in all benchmarks.
This is rather expected as our approach explicitly ensures the optimality of empirical risk.
When the generalization gap (difference between in-distribution and out-ofdistribution) is evaluated in Table 3, our method has a similar generalization gap with CORAL in average (our method has a smaller gap for 3 benchmarks, similar gap within the standard deviation for 2, and a larger gap for other 2).
Hence, our proposed optimizer improves empirical risk without hurting the generalization ability on 5 out of 7 datasets.
Even when our method hurts generalization, the improvement on the empirical risk is so large the resulting out-of-distribution performance improves.
3.3 Additional Empirical Analysis
In addition to our main study of applying our method to CORAL in WILDS benchmark, here we test the generality of our method for penalty function choice (i
[Table]
e. generality beyond CORAL) as well as application (i.e. generality beyond WILDS).
Applicability beyond CORAL We test generality of our method beyond CORAL.
Specifically, we apply our method to FISH [39] model in WILDS benchmark [23] and report in Table 1.
Application to FISH [39] requires modifications as it is not directly penalty based and we explain them in
Appendix C.3.
Results on FISH [39] suggest that our method improves FISH on 5 out of 7 benchmarks.
Moreover, the improvement for iWildCam and PovertyMap is significant, more than 9%.
Applicability beyond WILDS We test generality of our method beyond WILDS benchmark.
Specifically, we apply our method to FISH [39], VRex [25], and CORAL [42] model in DomainBed [18] benchmark.
We directly use the benchmark code shared by DomainBed [18] and follow the official process.
We tabulate the results in Table 4.
The results suggest that our method consistently improves the performance of all tested methods.
over 2 experimental suites conclusively demonstrates the generality of the proposed method. [Table]
3.3.1 Impact on training time.
Since our method needs gradients perdomain, it needs number of domains in the batch times backward pass during training.
This is a limitation of our method.
[Table]
Although our method effectively learns generalizable models, it has a rather significant impact on computational complexity.
Consider an input batch of B × B D datapoints over B D domains.
Direct optimization only needs the gradient for the average risk, while our method requires a gradient for each domain.
Hence, our method needs B D backward passes, whereas direct optimization needs only a single pass.
Effectively, our training time is up to B D times slower.
It is important to note that our method has very limited impact on training time when the original algorithm computes gradient per domain.
For example, utilizing our method with IRM would result in small change in training time.
To empirically evaluate the computational impact, we compute the average time per epoch per dataset and report it for our method and ERM in Table 5.
Our method increases the training time significantly.
Although this limitation can be easied using some approximation techniques [36], we leave the potential computational improvements as future work.
Moreover, it is important to clarify that we decrease the number of epochs for our method to ensure all numbers are fair in terms of used compute.
Related Work
Domain generalization.
Domain generalization addresses settings where the training and test data come from different distributions, violating the common independent and identically distributed (iid.) assumption.
The training data comes from multiple distinct data distributions (called domains).
The typical assumption to tackle domain generalization is covariate shift [9], where only the data distribution changes between domains and the conditional distribution of label given data stays the same.
We are specifically interested in penalty/invariance based domain domain generalization.
However, the approaches for domain generalization are more general and include domain conditioning [28], self-supervised pre-training [3], domain-specific decomposition [32], and data augmentation [44].
We treat these works as orthogonal to us.
One approach for invariance based domain generalization is distributionally robust optimization .
Here a robust form of risk is optimized instead of the average risk [35, 45, 16].
The robustness is typically against a predefined amount of distributional shift around the training data, hoping that it covers the test distribution.
Similarly, Krueger et al.
[25] enforce uniformity of risk functions to guarantee robustness following a causality-inspired approach.
Another common approach is adversarial feature learning .
Here a representation learning problem is posed such that an adversary cannot detect the domain information from the learned features.
Adversarial feature learning is achieved via game-theoretic min-max formulations [17, 26], distributional distances such as MMD [2], and divergence of second-order statistics of features [42].
Causality inspired methods is another thrust where the problem is posed in terms of finding features following the causal direction.
The out-of-distribution generalization problem considered in causality-based methods addresses a larger class of problems beyond domain generalization since they are applicable beyond covariate shift.
An early application of causality to out-of-distribution generalization was for linear models [31], later extended to the non-linear case [5].
Final category which we discuss is gradient alignment .
Here gradient-based learning is directly addressed and algorithms to align gradients of risk functions from different domains are developed.
This class of methods is the most similar to our case.
Koyama & Yamaguchi [24] enforce gradient similarity by minimizing the variance of the gradients and Shi et al.
[39] minimize the cosine distance between gradients.
Similarly, [50] finds domain transferrable solutions using meta-learning.
Parascandolo et al.
[30] combine these ideas with causality-based motivation and propose a method that directly manipulates gradients.
Among these methods, ANDMask [30] and SANDMask [38] are the most relevant to us since they are the only ones that explicitly modify gradients and use sign information.
Rate-distortion approaches.
Rate-distortion theory addresses the problem of lossy compression within information theory [13].
Among approaches that are numerical and directly applicable to any distribution, the Blahut-Arimoto algorithm is a common choice [11, 4].
Although there are improved variants [46, 48, 27], we find the original algorithm simple and effective.
To the best of our knowledge, our work is the first application of the Blahut-Arimoto algorithm to the problem of domain generalization.
The only application of Blahut-Arimoto within machine learning that we are aware of is the BLAST algorithm
[6], which uses the information-theoretic analysis of Thompson sampling [34] for better sample efficiency.
Ideas from rate-distortion theory have recently been applied to the problem of learning invariant representations by Zhao et al.
[51] to understand the underlying tradeoffs of transfer learning theoretically.
Different from [51], we analyze convergence properties in addition to the information-theoretic motivation.
We also propose a practical method for transfer learning.
A closely related problem to rate-distortion is the information bottleneck [43], where the mutual information between representation and input data is minimized while keeping the mutual information between representation and label maximal.
Although these methods have partially been applied to domain adaptation and generalization setting [14, 41, 15, 52], they typically require stochastic neural networks and additional modeling assumptions; hence, they are not directly applicable to our case.
To the best of our knowledge, these methods have not yet been applied to large-scale problems on the scale that we are interested in.
Our approach is significantly different from these approaches as our method stems from an analysis of gradient descent and an explicit formulation of satisficing gradient descent.
Conclusion
We considered penalty-based domain generalization and reformulated it to find updates minimizing the surrogate penalty under the constraint that applying these updates eventually converges to the stationary point of the empirical risk.
We developed an optimizer that is applicable to existing domain generalization models.
Our extensive empirical analysis on the WILDS and DomainBed benchmark validates the efficacy of our method for CORAL, FISH and VRex models.
A major drawback of our method is the additional computational complexity as it needs gradients per domain.
A possible future direction is applying approximation techniques to reduce this complexity.
Moreover, we provided an exciting connection between rate-distortion theory and domain generalization.
Another interesting future direction is utilizing ideas from the rate-distortion theory literature to better understand and improve domain generalization.
A Proof of Proposition 2.
Consider applying the stochastic updates θ t
+
=
θ t − ηG t to a function R ( θ ) .
Using the µsmoothness of the function R ,
[Equation]
Taking the expectation of the inequality and using the fact that E [ ∥ G t ∥ 22 ] ≤ V ,
[Equation]
Using the Lipschitz smoothness of the R and re-ordering the terms,
[Equation]
Summing up from t = 1 to T and dividing by η [Equation] we obtain
[Equation]
Using the fact that the distortions are bounded as D t =
∥
E [ G t ] − ∇R ( θ t ) ∥ 2 ≤ D / √ t , we can bound the sum (cid:80) Tt =1 D t with 2 D √ T .
By solving for the optimal η as;
[Equation]
The final bound becomes,
T
[Equation]
We plug the resulting step size in the bound and divide it with T , to obtain,
[Equation]
B Proofs of Statements from Section 2.
In this section, we prove the facts that we state in Section 2.4 without proof.
These proofs straightforwardly follow their rate-distortion version from [13].
We are interested in understanding the following function;
[Equation]
We state the formal versions of the informal statements in Section 2.4 and their proofs as follows;
Proposition B.1 (Extreme Cases) .
i)
When D t = D / √ t = 0 , E [ G t ] = ∇L ( θ t ) .
ii) When the distortion is unbounded ( D t = D = ∞ ), updates follow a distribution with the minimum penalty in expectation.
Proof for i. When D t = D = 0 , the constraint in (14) becomes,
[Equation]
The norm is nonnegative, and a nonnegative random variable has zero expectation if and only if it is zero almost surely.
Hence,
[Equation]
Proof for ii.
When D t = D = ∞ , the constraint in (14) is trivially satisfied for all distribution as we assume G t is finite.
Hence, (14) is equivalent to;
[Equation]
where the solution is a distribution with the minimum penalty in expectation.
Proposition B.2 (Monotonicity and Convexity) .
R ( D ) is a non-increasing and convex function of D .
Proof.
R ( D ) is minimum of the penalty over increasingly larger sets as D increases.
Hence, R ( D ) is non-increasing in D .
In order to prove convexity, consider two solutions p ( G t 1 ) and p ( G t 2 ) , minimizing R ( D ) for given
D t 1 , D t 2 .
Now, consider the distribution p ( G tλ )
=
λp ( G t 1 ) + (1 − λ ) p
G t 2 ) and D tλ =
λD t 1 + (1 − λ ) D t 2 .
Since the R ( D tλ ) is minimization over all distributions, it is smaller than penalty for p ( G tλ ) .
In other words,
[Equation]
Since expectation is a linear operator,
[Equation]
t
[Equation] i [Equation] Since p ( G t 1 ) and p ( G t 2 ) are the solutions for the corresponding minimization problems,
[Equation]
Combining these two statements proves the convexity as;
[Equation]
C Missing Derivations for Blahut-Arimoto and its Application to CORAL, FISH and VRex. C.1 Deriving Blahut-Arimoto Style Method to Solve (6)
We develop a Blahut-Arimoto [11, 4] style method to solve the following problem:
[Equation]
Before we apply the Blahut-Arimoto technique, we first utilize the definition of mutual information to convert (22) into
[Equation]
where D KL is the KL-Divergence.
The key technique used in the Blahut-Arimoto method is decomposing p ( G t ) as p ( G t ) =
(cid:80) e p ( G t | E = e ) p ( E = e ) .
Denote d ( G t , e ) =
∥
G t − ∇L e ( θ ) ∥ 2 where L e is the loss for domain e and | E | as the number of domains.
Moreover, we assume all domains are equal in importance as p ( E = e )
= 1 / | E | .
Using the fact that the problem is discrete, G t ∈ { G 1 , . .
, G K } , the problem in (22) can be transformed into min
[Equation] [Equation]
here we replaced E G t (cid:104)
∥
G t − ∇ θ ˆ L ( θ ) ∥ 2 (cid:105) ≤ D √ t with a stronger constraint of uniform bounds over domains as E G t | E =
e (cid:104) ∥
G
[Equation]
− ∇ θ ˆ L e ( θ ) ∥ 2 (cid:105) ≤ D √ t .
We further compute the Lagrangian of this optimization problem with the Lagrange multiplier β as;
[Equation]
We take the derivative with respect to p ( G k | E = e ) and equate it to 0 .
Ignoring the constant factor,
[Equation]
Since we know the probabilities sum to 1 , we can handle the normalization factor separately.
Start with initialization p ◦ ( G k | E = e ) , which is typically uniform unless a prior information exists.
Then, the following iterations sove (22),
• ˆ
[Equation]
[Equation]
[Equation]
Penalty
[Equation]
• p l +1 ( G k | E = e ) =
[Equation]
p
[Equation]
+1 ( G k | E [Equation] e
[Equation]
k ˆ p l +1 ( G ˆ k | E = e ) • p l
{++
 G k ) = 1 / | E | (cid:80) E = e p l +1 (+} G k | E = e )
Generalization
In this section, we develop the CORAL+SDG and VREX+SDG method considering an arbitrary penalty function.
To apply the derivation in Section C.1, we need to compute
Penalty ( θ t + G k ; { x i , e i , y i } ) for an arbitrary G k .
We consider the first order approximation of this penalty as,
[Equation]
By this approximation, we do not need to perform additional computations for Penalty for each
G k , instead we can compute the derivative once and perform dot-products for the rest.
We substitute this approximation in the Blahut-Arimoto iteration eventually obtaining;
{+• ˆ p l +
 ( G k | E = e ) = p l ( G k ) exp (cid:104) − 1 γ ( βd ( G k , e ) + ∇ θ Penalty+}
θ t ) ⊺ G k ) (cid:105) • p l
+1 ( G k | E = e )
=
ˆ p l +1 ( G k | E = e ) (cid:80) ˆ k ˆ p l +1 ( G ˆ k | E = e ) •
p l
+1 ( G k ) = 1 / | E | (cid:80) E = e p l +1 ( G k | E = e ) We further apply two approximations, i) solving each parameter independently, and ii) solving only for the direction of the gradient (positive or negative).
When these approximations are applied, the resulting iterations can be efficiently performed via vector operations and summarized in Algorithm 1.
Algorithm
G + = 0 , G − =
▷ Initialize positive grad, and negative grad with
Vector
G penalty = ∇ θ L penalty ( θ t )
▷ Gradient of the penalty for e ∈ { 1 , . .
. , | E |} do ▷ For each domain
G e = 1 | E | ∇ θ L e ( θ t )
G + = G + + 1 [ G e > 0] · G e ▷ Parameters with positive gradients
G − =
G − + 1 [ G e < 0] · G e ▷ Parameters with negative gradients end for
Prob + = B LAHUT
A RIMOTO S TYLE S OLVER ( G 1 , . .
, G | E | , G penalty )
P + ∼ Bern ( Prob + )
▷ Sample the gradient directions
G = P + · G + + (1 − P + ) · G − θ t +1 =
θ t − ηG procedure B LAHUT
A RIMOTO S TYLE S OLVER ( G 1 , . .
, G | E | , G penalty )
Initialize: Prob + = Prob − = 0 .▷
Initialize probabilities uniformly for iter = 1 , . .
. , IterCount do for e ∈ { 1 , . .
. , | E |} do ▷ All Operations are Elementwise during the Loop ˜ Prob + ,e = Prob + · exp ( − 1 / γ [ β ( G e − G + ) · ( G e − G + ) + G penalty · G + ]) ˜ Prob − ,e = Prob − · exp ( − 1 / γ [ β ( G e − G + ) · ( G e − G + ) + G penalty · G + ]) Prob + ,e = ˜
Prob
+ ,e / ( ˜ Prob + ,e + ˜ Prob − ,e ) Prob − ,e = 1 − Prob + ,e end for Prob + = 1 / | E | (cid:80) e Prob + ,e
Prob − = 1 − Prob + end for return Prob + end procedure
C.3 Applying Blahut-Arimoto Style Method to Fish [39]
Fish [39] is not using a direct penalty function; hence, applying SDG requires some modifications.
In order to apply the SDG to Fish, we simply use the approximation provided in the original work [39].
Specifically, Shi et al.
[39] shows that when an inner update of SGD is applied to θ t for a few domains to obtain ˜ θ t , ˜ θ t − θ t approximates the gradient of the penalty.
Hence, we utilize this fact to perform the iterations as
[Equation]
• p l +1 ( G k | E = e ) =
[Equation]
• p l +1 ( G
[Equation]
=
[Equation]
E
[Equation]
D Generalization Gap in WILDS [23]
We argued that when the empirical estimation error of the invariance penalty is high, the resulting optimization problem might hurt empirical risk (in-distribution performance) and fails to perform out-of-distribution even if it generalizes well.
To further quantify this motivation, we look at the benchmarking results of [23].
We specifically enclosed the out-of-distribution and in-distribution results of various benchmarked algorithms from [23] in Table 6&7.
Moreover, we compute the generalization gap as the difference between in-distribution and out-of-distribution performance and tabulate them in Table 8.
The results confirm the failure mode we hypothesized.
Specifically, ERM outperforms existing domain generalization algorithms in almost all benchmarks when the metric is out-of-distribution performance.
However, existing domain generalization methods generalize better (generalization gap is smaller than ERM).
Hence, the failure of these methods is not due to the lack of generalization but due to the lack of effective minimization of the empirical risk.
ERM overwhelmingly outperforms existing domain generalization algorithms
[Table] [Table]
in the WILDS benchmarks [23].
We denote the results which generalize better than ERM with Green color and worse than ERM with Red color.
Domain generalization algorithms typically generalize better than ERM.
Hence, although they are good generalizers, they fail to perform out-of-distribution due to poor in-distribution performance
[Table]
E Details on Used Datasets
iWildCam [8]: The problem is predicting animal species from camera trap images over different locations.
The dataset includes 323 different camera traps, considered as domains.
Among these,are used for training, 32 for validation, and 48 for testing.
The evaluation metric is the F1-scores.
Camelyon17 [7]: The problem is predicting tumors from tissue slides over 5 different hospitals, treated as domains.
Among these, 3 are used as training, 1 for validation, and 1 for test.
The metric is the accuracy over balanced patches (i.e. same number of positive and negative patches).
OGB-MolPCBA [21]: The problem is predicting bioassays from molecular graphs over 120, scaffolds, considered as domains.
The largest 44,390 scaffolds are used for training, the next largest 31,361 are for validation, and the smallest 43,793 are for testing.
The metric is average precision.
Functional Map of the World (FMoW) [12]: The problem is predicting the land usage from a satellite image from 5 regions over 16 years, a total of 5 × 16 = 80 domains.
From 16 years of data, 11 years are used for training, 3 for validation, and 2 for testing.
The evaluation metric is the accuracy over the worst region to evaluate the sub-population shift.
Amazon Reviews [29]: The problem is predicting sentiment from product reviews over 3920 different users considered as domains.
Among those, 1252 users are used for training, 1334 for validation, and 1334 for testing.
The evaluation metrics are accuracy averaged over users in the 10th percentile. py150 [33]: The problem is autocompletion of Python code over 8421 git repositories as domains.
Among them, 5477 are used for training, 261 for validation, and 2471 for testing.
The evaluation is the accuracy of predicting class/method name tokens.
PovertyMap [47]: The problem is predicting asset wealth from satellite images over different countries.
Evaluation is over 5 different folds where 13-14 countries are used for training, 4-5 for validation, and 4-5 for the test.
The metric is the Pearson correlation averaged over the rural regions.
F Hyper Parameters
We share the implementation of our method and experimental setup for all WILDS experiments.
In order to reproduce the experiments, the scripts can be run with the following hyperparameters.
Any hyperparameter not shared here, directly follows the default one shared by [23].
For DomainBed experiments, we directly used the recommended hyperparameter search pipeline without any modifications.
We perform a hyperparameter search for the WILDS experiments with the same budget reported in [23].
The chosen hyperparameters are given in Table
[Table]
[Table]
G Societal Impact
Our work studies the optimization of domain generalization algorithms by re-formulating the joint optimization problem of domain generalization as constrained optimization of penalty under the constraint of optimality of the empirical risk.
Our approach is purely algorithmic and applicable to any penalty-based method.
Hence, it does not change the societal impact of using the original method and dataset.
It is still beneficial to re-iterate some aspects of the societal impact of the state of the domain generalization algorithms.
Our work, among others, shows that immediate deployment of machine learning models without investigating the fairness and safety when applied to the target distribution might not be feasible.
Performance and behavior on unseen distributions are still far from being predictable.
Hence, particular care should be spent studying and monitoring the target distribution and its changes over time.
Moreover, a critical aspect of societal impact is the datasets.
Fortunately, WILDS benchmark [23] provides a comprehensive analysis for the datasets we use.
We refer the interested reader to this comprehensive analysis.