Domain Generalization without Excess Empirical Risk
Anonymous Author(s)
AffiliationAddress
Abstract
Given data from diverse sets of distinct distributions, domain generalization aimsto learn models that generalize to unseen distributions.
A common approach isdesigning a data-driven surrogate penalty to capture generalization and minimizethe empirical risk jointly with the penalty.
We argue that a significant failure modeof this recipe is an excess risk due to an erroneous penalty or hardness in jointoptimization.
We present an approach that eliminates this problem.
Instead ofjointly minimizing empirical risk with the penalty, we minimize the penalty underthe constraint of optimality of the empirical risk.
This change guarantees that thedomain generalization penalty cannot impair optimization of the empirical risk,i.e., in-distribution performance.
To solve the proposed optimization problem, wedemonstrate an exciting connection to rate-distortion theory and utilize its toolsto design an efficient method.
Our approach can be applied to any penalty-baseddomain generalization method, and we demonstrate its effectiveness by applying itto three examplar methods from the literature, showing significant improvements.
Introduction
Empirical risk minimization (ERM) is the workhorse of machine learning.
In this setting, a learnercollects a set of data, assumed to be independent and identically sampled from a particular distributionof interest.
An empirical risk over these samples is later minimized, hoping that the resulting modelwill generalize to future unseen data from the same distribution.
Although ERM is ubiquitous, it isbrittle to minor shifts in the distribution and the resulting models suffer a significant performancedrop when evaluated beyond their training distributions.
Examples of such distribution shifts includedifferent hospitals for medical imaging [49], camera types for perception models [8] and time periodsfor natural language models [20].
The problem of learning models which generalize under suchdistribution 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 thatgeneralize to unseen domains.
Although the domain generalization literature is rich, many model-freemethods follow a common recipe.
A surrogate penalty that encapsulates generalization is designedand 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 anoptimal 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 unseendistribution can be bounded as the combination of an error on the training distributions (empiricalrisk) 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 onboth 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 domaingeneralization 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 generalizationmethods have a smaller gap than ERM on most problems.
This supports our hypothesis that existingalgorithms might be creating too much excess empirical risk.
Thus the motivating question of ourwork: Can we perform joint optimization without causing any excess empirical risk?
Solving joint optimization without sacrificing empirical risk is not generally possible unless surrogatepenalty term and empirical risk align perfectly.
Instead, we formulate an optimization problem withthe explicit constraint of no excess empirical risk.
In other words, we propose to minimize thesurrogate penalty only if it is not hurting the empirical risk.
Formally, instead of joint minimizationof empirical risk and surrogate penalty, we minimize the surrogate penalty under the constraint ofoptimality of the empirical risk.
To design an algorithm for the proposed optimization problem, we start with searching for a formaldefinition of the optimality of the empirical risk.
Since the applications of interest are non-convex, thiscan only be local optimality.
Hence, the constraint we consider is convergence to a stationary pointof the empirical risk.
We show that biased gradient descent converges to a stationary point as longas the bias is bounded and decreases through optimization.
Using this result, we design an iterativealgorithm that seeks the update minimizing the surrogate penalty under the constraint of boundeddeviation from the empirical gradient.
This construction is closely related to the rate-distortionproblem 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 beapplied to any penalty-based model with minor modifications.
We perform our primary empiricalevaluation on the CORAL model due to its simplicity [42].
We evaluate on the large-scale domaingeneralization problems in the WILDS benchmark of real distribution shifts [23].
Our evaluationusing the image, formal language, natural language, and graph modalities suggests that our methodsignificantly improves performance.
To further validate the generality of our method, we apply it tothe 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 baseddomain generalization methods.
We are interested in a learning problem in the space of X × Y , where X is the space of inputs and Yis 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 π jfor domain j with N j samples.
Consider the union of these datasets as the training set of domaingeneralization in the form of tuples { x, e, y } such that x is the input, e is the domain id, and y isthe label.
We further consider a parametric function family h ( · ; θ ): X → Y and a loss functionl ( · , · ): 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 riskand domain generalization penalty are jointly minimized.
Each model defines a penalty following adifferent 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 weare interested in the optimization of these models, not their construction, we treat Penalty as ageneric entity and consider the optimization problem defined by these methods.
These methods aimto solve the following joint optimization problem:
[Equation]
where L ( θ ; { x i , e i , y i } ) is the empirical risk defined as L ( θ ) ≜ 1 M (cid:80) Mj
=1 1 N j (cid:80) N j i =1 l ( h ( x ij ; θ ) , y ij ) .
2.2 Optimizing Domain Generalization Objectives
The optimization problem defined in (1) is a sound proxy in order to optimize the risk for the unseendomains.
Consider an unseen distribution π ⋆ , its corresponding population risk can be bounded usingthe 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 94 (cid:0)
Penalty ( θ ; { x i , e i , y i } ) (cid:1) approximates the generalization error.
However, this formulation has afailure point.
Consider a well-designed penalty term which successfully approximates the generalization gap.
Clearly the joint optimization will decrease the generalization error.
However, it caneffectively 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, thefinal performance on unseen domains can be worse than not doing any domain generalization.
Wehypothesize 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 enclosedin Appendix D. Take CORAL [42] as a representative method 1 .
When we look at out-of-distributionperformance, ERM outperforms CORAL on 7 out of 9 problems.
However, when we look at thegeneralization 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 CORALhave lower generalization error, the loss of in-distribution performance is so significant that theresulting 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 jointlyminimizing the empirical risk and the penalty, we minimize the penalty under the constraint of theoptimality 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 guaranteesno 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 ifperformance on seen domains and unseen domains are inherently contradictory.
For example, if thereis strong nuisance information in the seen domains, any model not using this nuisance informationwill 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 multipledatasets, this is not the case as a universal model solving all datasets (seen and unseen) exists.
Inother words, we address only a subset of domain generalization problems where seen and unseendomains do not compete.
iii)
Since the models we use are significantly overparameterized, the spaceof solutions to empirical risk minimization is quite large, likely including many domain-invariant anddomain-sensitive solutions.
Hence, this problem is still non-trivial from the numerical optimizationperspective.
iv)
Finally, one can mistakenly think that our reformulation and the original jointoptimization would find the same solution due to the Lagrangian function.
However, this is incorrectas the Lagrangian function requires optimization of the Lagrangian multiplier, which is intractable forlarge-scale deep learning and typically replaced with fixed hyper-parameters.
Moreover, we proposeto 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 apractical form for first-order optimization.
Then in Section 2.4 we discuss how (3) is closely relatedto rate-distortion theory, demonstrating an unexpected connection.
Finally in Section 2.5 we utilizetools 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 lossfunctions we are interested in is non-convex, the best we can hope is convergence to a stationarypoint.
Hence, we convert (3) into
[Equation]
While solving the problem in (4), we constrain ourselves to iterative methods to be compatible withstochastic gradient descent (SGD), which is commonly used to train deep neural networks.
Hence, wewould like to have an iterative procedure θ t +
=
θ t − ηG t which solves (4) in convergence.
Beforewe proceed, we note that when G t is the stochastic gradient (i.e., E [ G t ] = ∇ θ ˆ L ( θ ) ), convergence toa stationary point of the empirical risk can be shown under mild conditions.
However, we are alsointerested in minimizing the penalty.
Hence, we need to allow updates that are different than thegradients while still supporting eventual convergence.
Convergence of SGD with biased gradients has previously been analyzed by Sener & Koltun [37], Huet al.
[22], and Ajalloeian & Stich [1].
The common theme is the necessity for the bias to decreaseover time.
This requirement is explicit in [37] and implicit in [22] and [1].
In the implicit case, biasis 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 ofthe bias is bounded and this bound decreases during training faster than the rate 1 / √ t , it convergencesto a stationary point of a smooth and non-convex risk function.
We state the convergence result in thefollowing 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, 153 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 validfor common deep learning architectural choices [19].
Finally, the third assumption is on the biasdecreasing 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 longas we can bound the difference between the updates and the gradient of the empirical risk, we canguarantee convergence to a stationary point.
Hence, our iterative approach to domain generalizationstarts with an initialization θ 0 and applies the iteration θ t +
=
θ t − ηG t , with G t being the solutionof 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 theempirical risk.
In the meantime, we seek the updates minimizing the penalty among these satisfactoryones.
Since our approach utilizes satisfactory updates for empirical risk instead of the first-orderoptimal 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 theupdates instead of the update itself.
Although this is unconventional, it is intentional and crucial formaking the problem compatible with rate-distortion theory, which we utilize heavily in our solver.
2.4 Understanding Satisficing Domain Generalization
In this section, we provide observations on the proposed formulation (5) by utilizing tools fromrate-distortion theory.
Consider the function R ( D ) as the solution of (5) for a given constantfactor 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, wherethe (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 wouldhave been exactly equivalent.
Hence, we can further utilize tools from rate-distortion theory
[Figure]
to better understand (5).
We now state a fewsimple facts derived from rate-distortion theory.
( Proofs follow their rate-distortion counterpartsdirectly 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 distor- is unbounded ( D t = ∞ ), updates directlyminimize the penalty.
Moreover, our
formulation smoothly goes from updates mini- penalty to ERM updates during training.
Improvement guarantee: R ( D ) is a nonincreasing and convex function of D .
Hence,
R (0) ≥ R ( D t ) .
In other words, SDG decreasesthe penalty when compared with ERM.
Geometric picture: The satisficing requirementimplies a certain structure over the plane wherethe x -axis is the distortion and the y -axis is thepenalty.
We visualize this geometry and theinterpretation 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 designa tractable method.
We propose to regularize the objective with the mutual information between gradients ( G t ) and thedomain-ids ( E ).
This regularization is sensible as we would like gradients to carry little informationabout 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 technicalderivation 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 firstdefine auxiliary variables p G ti | e = p ( G t =
G ti | E = e ) , which are conditioned on domain-ids.
Bythis decomposition, we can solve for each domain conditional probability p G ti | e separately and latermarginalize 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 witha stronger constraint of uniform bounds over domains as E G t | E = e [ ∥ G t − ∇ θ ˆ L e ( θ ) ∥ 2 ] ≤ D / √ t andapply 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 tk | 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 whereall 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 solvefor each parameter independently, and ii) we only estimate the sign and copy the magnitude fromthe 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.2with 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 utilizethe 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 betweenworker nodes in multi-node optimization. Sign of the gradient suffices for convergence in non-convexproblems both theoretically and empirically.  Experimental Results
Implementation details.
We implement the proposed algorithm in PyTorch.
Following the analysisin Proposition 2.1, D t = D (cid:112) 1 / t .
Moreover, β is inversely related to D .
We choose β t = ( β ◦ / √ T ) √ tfor 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 relevantterm in the iterative update is exp (cid:104) − 1 γ
( InvP ( G tk ) + βd ( G tk , 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 Ddomains, then sample B examples per domain (named as group sampler in WILDS[23]).
We searchall hyperparameters with random sampling with the constraint that all methods have the same budgetfor the hyperparameter search.
Specifically, we use 20 random hyperparameter choices.
We alsoperform 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 ourmethod is slower, we perform fewer epochs than other methods.
We list all chosen hyperparametersin
Appendix G. We plan to share the source code, and evaluation runs for all experiments.
3.1 Evaluation on WILDS
Datasets Our main evaluation is using the WILDS benchmark [23].
WILDS is designed to evaluatedomain generalization and subpopulation shift in realistic problems.
Among these problems, we useseven problems that are either pure domain generalization problems or a combination of domaingeneralization 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 numberof domains from 5 to 120 K .
The summary of the benchmarks is in Appendix F.
Baselines We apply our proposed optimizer (SDG) to CORAL [42] and denote CORAL+SDG.
Wecompare CORAL+SDG with i)
CORAL: the standard implementation of CORAL which is jointly [Table] 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 CORALminimization.
It is important to note that AND-MASK[30] and SAND-MASK[38] have not beenapplied in conjunction with CORAL before.
Still, we compare with them since they also directlymodify the optimization procedure.
We compute the gradient of the empirical risk plus coral penaltyfor each domain separately and utilize their and-mask and sand-mask as corresponding updates.
3.1.1 Main Results
We summarize the out-of-distribution results in Table 1 and in-distribution results in Appendix Ein Table 8.
We also compute the generalization gap (difference between in-distribution and outof-distribution performance) and tabulate it in Table 2.
When evaluated on out-of-distributionperformance, 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, althoughthe CORAL outperforms ERM with a standard optimizer only for 1 out 7 benchmarks, combinedwith our optimizer, CORAL outperforms ERM for 6 out of 7 benchmarks.
Our method closes thegap for Camelyon17, FMoW, Amazon, Py150, and PovertyMap.
The only setting CORAL fails tooutperform ERM when combined with our optimizer is OGBMolPCBA.
We believe one reason forthis is the additional computational complexity of our method.
For a fair comparison, our methodperforms significantly fewer epochs than ERM.
When compared with AND-mask and SAND-mask,our method outperforms them in all benchmarks.
When evaluated for in-distribution performance in Table 8, our method improves upon CORAL in allbenchmarks.
This is rather expected as our approach explicitly ensures the optimality of empirical risk.
Gap (difference between indistribution and out-distribution performance) in the
WILDS benchmarks [23].
We denote the results which generalize better than ERM with Green color and worse than ERM with Red color.
[Table]
When the generalization gap (difference between in-distribution and out-ofdistribution) is evaluated in Table 2, ourmethod has a similar generalization gapwith CORAL (our method has a smallergap for 3 benchmarks whereas CORALhas a smaller gap for other 4).
Hence, ourproposed optimizer improves empiricalrisk without hurting the generalizationability as desired.
This confirms that different datasets in WILDS are not contradictary as improving both in-distributionand out-of-distribution is feasible.
3.2 Additional Empirical Analysis
Applicability beyond CORAL We test generality of our method beyond CORAL.
Specifically, weapply our method to FISH [39] model in WILDS benchmark [23] and report in Table 1.
Applicationto 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%.
[Table]
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] andfollow the official process.
We tabulate the results in Table 3.
The results suggest that our methodconsistently improves the performance of all tested methods.
Consistent improvement for 3 different domain generalization models on 14 different benchmarksover 2 experimental suites conclusively demonstrates the generality of the proposed method. Impact on training time.
needs domain needs number a limitation of method
[Table]
Although our method effectively learns generalizablemodels, it has a rather significant impact on computational complexity.
Consider an input batch of B × B D datapoints over B D domains.
Direct optimizationonly needs the gradient for the average risk, while ourmethod requires a gradient for each domain.
Hence,our method needs B D backward passes, whereas directoptimization needs only a single pass.
Effectively, ourtraining time is up to B D times slower.
It is important to note that our method has very limited impacton training time when the original algorithm computesgradient per domain.
For example, utilizing our methodwith IRM would result in small change in training time.
To empirically evaluate the computational impact, wecompute the average time per epoch per dataset and report it for our method and ERM in Table 4.
Our methodincreases the training time significantly.
Although thislimitation 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 numberof 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 datacome from different distributions, violating the common independent and identically distributed (iid.)assumption.
The training data comes from multiple distinct data distributions (called domains).
Thetypical 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 isdistributionally robust optimization .
Here a robust form of risk is optimized instead of the averagerisk [35, 45, 16].
The robustness is typically against a predefined amount of distributional shift aroundthe training data, hoping that it covers the test distribution.
Similarly, Krueger et al.
[25] enforceuniformity of risk functions to guarantee robustness following a causality-inspired approach.
Another common approach is adversarial feature learning .
Here a representation learning problem is posedsuch that an adversary cannot detect the domain information from the learned features.
Adversarialfeature 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 inspiredmethods is another thrust where the problem is posed in terms of finding features following the causaldirection.
The out-of-distribution generalization problem considered in causality-based methodsaddresses a larger class of problems beyond domain generalization since they are applicable beyondcovariate shift.
An early application of causality to out-of-distribution generalization was for linearmodels [31], later extended to the non-linear case [5].
Final category which we discuss is gradientalignment .
Here gradient-based learning is directly addressed and algorithms to align gradients ofrisk functions from different domains are developed.
This class of methods is the most similar toour case.
Koyama & Yamaguchi [24] enforce gradient similarity by minimizing the variance of thegradients and Shi et al.
[39] minimize the cosine distance between gradients.
Similarly, [50] findsdomain transferrable solutions using meta-learning.
Parascandolo et al.
[30] combine these ideaswith causality-based motivation and propose a method that directly manipulates gradients.
Amongthese methods, ANDMask [30] and SANDMask [38] are the most relevant to us since they are theonly ones that explicitly modify gradients and use sign information.
Rate-distortion approaches.
Rate-distortion theory addresses the problem of lossy compressionwithin information theory [13].
Among approaches that are numerical and directly applicable toany distribution, the Blahut-Arimoto algorithm is a common choice [11, 4].
Although there areimproved variants [46, 48, 27], we find the original algorithm simple and effective.
To the best ofour knowledge, our work is the first application of the Blahut-Arimoto algorithm to the problem ofdomain generalization.
The only application of Blahut-Arimoto within machine learning that we areaware of is the BLAST algorithm
[6], which uses the information-theoretic analysis of Thompsonsampling [34] for better sample efficiency.
Ideas from rate-distortion theory have recently beenapplied to the problem of learning invariant representations by Zhao et al.
[51] to understand theunderlying tradeoffs of transfer learning theoretically.
Different from [51], we analyze convergenceproperties in addition to the information-theoretic motivation.
We also propose a practical methodfor 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 keepingthe mutual information between representation and label maximal.
Although these methods havepartially been applied to domain adaptation and generalization setting [14, 41, 15, 52], they typicallyrequire stochastic neural networks and additional modeling assumptions; hence, they are not directlyapplicable to our case.
To the best of our knowledge, these methods have not yet been applied tolarge-scale problems on the scale that we are interested in.
Our approach is significantly differentfrom these approaches as our method stems from an analysis of gradient descent and an explicitformulation of satisficing gradient descent.
Conclusion
We considered penalty-based domain generalization and reformulated it to find updates minimizingthe surrogate penalty under the constraint that applying these updates eventually converges to thestationary point of the empirical risk.
We developed an optimizer that is applicable to existing domaingeneralization models.
Our extensive empirical analysis on the WILDS and DomainBed benchmarkvalidates 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 perdomain.
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 tobetter understand and improve domain generalization.
Checklist
For all authors...(
a)
Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes](
b)
Did you describe the limitations of your work?
[Yes] See Section 2.
2 (remarks on the
reformulation) and Section 3.2.1.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Appendix H.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes] See thestatement of propositions.
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi-
mental results (either in the supplemental material or as a URL)?
[Yes] As part of the
supplemental material.
(b) Did you specify all the training details (e.g., data splits, hyperparameters,
how they
were chosen)?
[Yes] All hyperparameters which are different than the default ones in
WILDS[23] and DomainBed[18] are given in Appendix G.
(c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] In Appendix G
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes] We are using datasets
from WILDS[23] and DomainBed[18]
(b) Did you mention the license of the assets?
[N/A] See WILDS[23] and DomainBed[18]
for discussion on licenses.
(c) Did you include any new assets either in the supplemental material or as a URL?
[N
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A] We are not curating any new asset.
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A] See WILDS[23] and DomainBed[18] for
discussion on personally identifiable information or offensive content.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
[Table]
applicable?
[N/A]
We did NOT use crowdsourcing or conducted research with human
[Table]
subjects,
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
We did NOT use crowdsourcing or
conducted research with human subjects,
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]
We did NOT use crowdsourcing or conducted
research with human subjects,
A Proof of Proposition 2.1
{+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 η , we obtain
T
[Equation]
=
[Equation]
Using the fact that the distortions are bounded as D t =
∥
E [ G t ] − ∇R ( θ t ) ∥ 2 ≤ D / √ t , we can boundthe sum (cid:80) Tt =1 D t with 2 D √ T .
By solving for the optimal η as;
[Equation]
The final bound becomes,
T (cid:88) 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.4
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]
[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 thedistortion is unbounded ( D t = D = ∞ ), updates follow a distribution with the minimum penalty inexpectation.
Proof for i. When D t = D = 0 , the constraint in (14) becomes,
[Equation]
[Equation]
The norm is nonnegative, and a nonnegative random variable has zero expectation if and only if it iszero almost surely.
Hence,
[Equation]
Proof for ii.
When D t = D = ∞ , the constraint in (14) is trivially satisfied for all distribution as weassume 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]
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 informationto 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 ) asp
[Equation]
G
[Equation]
) =
[Equation]
e p ( G t | E = e ) p ( E = e ) .
Denote d ( G t , e ) =
∥
G t − ∇L e ( θ ) ∥ 2 where L e is the lossfor domain e and | E | as the number of domains.
Moreover, we assume all domains are equal inimportance 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
[Equation]
[Equation]
here we replaced E G t (cid:104)
∥
G t − ∇ θ ˆ L ( θ ) ∥ 2 (cid:105) ≤ D √ t with a stronger constraint of uniform boundsover domains as E G t | E =
e
[Equation]
∥ G t − ∇ θ ˆ L e
[Equation]
θ ) ∥ 2 (cid:105) ≤ D √ t .
We further compute the Lagrangian of thisoptimization 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.
Startwith initialization p ◦ ( G k | E = e ) , which is typically uniform unless a prior information exists.
Then,the following iterations sove (22),
[Equation]
[Equation]
[Equation]
+
[Equation]
| E
[Equation]
e )
C.2 Applying Blahut-Arimoto Style Method to Penalty Based Deep Domain 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 approximationof 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 substitutethis approximation in the Blahut-Arimoto iteration eventually obtaining;
[Equation]
[Equation]
p l +1 ( G k )
= 1 / | E |
[Equation]
E = e p l +1 ( G k
[Equation]
E = e )
We further apply two approximations, i) solving each parameter independently, and ii) solving onlyfor the direction of the gradient (positive or negative).
When these approximations are applied, theresulting 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
E |}
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 fewdomains to obtain ˜ θ t , ˜ θ t − θ t approximates the gradient of the penalty.
Hence, we utilize this fact toperform the iterations as
{+• ˆ p l +
+}
( G k | E = e ) =
p l ( G k ) exp (cid:104) − 1 γ (cid:16) βd
( G k
[Equation]
e ) + (˜ θ t − θ ) ⊺
G k (cid:17)(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 )
D Generalization Gap in WILDS [23]
We argued that when the empirical estimation error of the invariance penalty is high, the resultingoptimization problem might hurt empirical risk (in-distribution performance) and fails to performout-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-distributionresults of various benchmarked algorithms from [23] in Table 5&6.
Moreover, we compute thegeneralization gap as the difference between in-distribution and out-of-distribution performance andtabulate them in Table 7.
The results confirm the failure mode we hypothesized.
Specifically, ERM outperforms existingdomain generalization algorithms in almost all benchmarks when the metric is out-of-distributionperformance.
However, existing domain generalization methods generalize better (generalization gapis smaller than ERM).
Hence, the failure of these methods is not due to the lack of generalization butdue 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 Additional Empirical Results
We provide additional empirical results which we did not include in the main text due to the spacelimitations.
[Table]
F Details on Used Datasets
iWildCam [8]: The problem is predicting animal species from camera trap images over differentlocations.
The dataset includes 323 different camera traps, considered as domains.
Among these, 243are 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 isthe 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,084scaffolds, considered as domains.
The largest 44,390 scaffolds are used for training, the next largest31,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 asatellite 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 accuracyover the worst region to evaluate the sub-population shift.
Amazon Reviews [29]: The problem is predicting sentiment from product reviews over 3920 differentusers considered as domains.
Among those, 1252 users are used for training, 1334 for validation, and1334 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 isthe accuracy of predicting class/method name tokens.
PovertyMap [47]: The problem is predicting asset wealth from satellite images over differentcountries.
Evaluation is over 5 different folds where 13-14 countries are used for training, 4-5 forvalidation, and 4-5 for the test.
The metric is the Pearson correlation averaged over the rural regions.
G Hyper Parameters and the Source Code
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 withoutany modifications.
We perform a hyperparameter search for the WILDS experiments with the same budget reported in[23].
The chosen hyperparameters are given in Table 9.
All the experiments are conducted on aninternal cluster of Nvidia RTX 3090 GPUs and Intel Xeon CPUs.
The total amount of compute forall presented results corresponds to roughly 250 days of GPU time.
[Table]
H Societal Impact
[Table]
Our work studies the optimization of domain generalization algorithms by re-formulating the jointoptimization problem of domain generalization as constrained optimization of penalty under theconstraint of optimality of the empirical risk.
Our approach is purely algorithmic and applicable toany penalty-based method.
Hence, it does not change the societal impact of using the original methodand dataset.
It is still beneficial to re-iterate some aspects of the societal impact of the state of thedomain generalization algorithms.
Our work, among others, shows that immediate deployment ofmachine learning models without investigating the fairness and safety when applied to the targetdistribution might not be feasible.
Performance and behavior on unseen distributions are still farfrom being predictable.
Hence, particular care should be spent studying and monitoring the targetdistribution 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 thiscomprehensive analysis.