On-Demand Sampling: Learning Optimally from Multiple Distributions ∗
Abstract
Societal and real-world considerations such as robustness, fairness, social welfare and multi-agent tradeoffs have given rise to multi-distribution learning paradigms, such as collaborative [6], group distributionally robust [38], and fair federated learning [29].
In each of these settings, a learner seeks to minimize its worstcase loss over a set of n predefined distributions, while using as few samples as possible.
In this paper, we establish the optimal sample complexity of these learning paradigms and give algorithms that meet this sample complexity.
Importantly, our sample complexity bounds exceed that of the sample complexity of learning a single distribution only by an additive factor of n log( n ) ε.
These improve upon the best known sample complexity of agnostic federated learning by Mohri et al.
[29] by a multiplicative factor of n , the sample complexity of collaborative learning by
Nguyen and Zakynthinou [31] by a multiplicative factor log n ε 3 , and give the first sample complexity bounds for the group DRO objective of Sagawa et al.
To achieve optimal sample complexity, our algorithms learn to sample and learn from distributions on demand.
Our algorithm design and analysis extends stochastic optimization techniques to solve zero-sum games in a new stochastic setting.
Introduction
Pervasive needs for robustness, fairness and multi-agent collaboration in learning have given rise to multi-distribution learning paradigms (e.g., [6, 38, 29, 13]).
In these settings, we seek to learn a model that performs well on any distribution in a pre-defined set of interest.
For fairness considerations, these distributions may represent heterogeneous populations of different protected or socio-economic attributes; in robustness applications, they may capture a learner’s uncertainty regarding the true underlying task; and in muti-agent collaborative or federated applications, they may represent agentspecific learning tasks.
In these applications, the performance and optimality of a model is measured by its worst test-time performance on a distribution in the set.
We are concerned with this fundamental problem of designing sample-efficient multi-distribution learning algorithms.
The sample complexity of multi-distribution learning differs from that of learning a single distribution in several ways.
On one hand, learning tasks of varying difficulty require different numbers of samples.
On the other hand, similarity or overlap among learning tasks may obviate the need to sample from some distributions.
This makes the use of a fixed per-distribution sample budget highly inefficient and suggests that optimal multi-distribution learning algorithms should sample on demand .
That is, algorithms should take additional samples whenever they need them and from whichever distribution they want them.
On-demand sampling is especially appropriate when some population data is scarce (as in fairness mechanisms in which samples are amended [34]); when the designer can actively
[Table]
perturb datasets towards rare or atypical instances (such as in robustness applications [22, 46]); or when sample sets represent agents’ contributions to an interactive multi-agent system [29, 7].
Blum et al.
[6] demonstrated the benefit of on-demand sampling in the collaborative learning setting, where all data distributions are realizable with respect to the same target classifier.
This line of work established that learning n distributions on-demand takes (cid:101) O (log( n )) times the sample complexity of learning a single realizable distribution [6, 9, 31], whereas relying on batched uniform convergence takes (cid:101) Ω ( n ) times that of learning a single distribution [6].
However, beyond the realizable setting, the best known multi-distribution learning results fall short of this promise: existing on-demand sample complexity bounds for agnostic collaborative learning have highly suboptimal dependence on ε , requiring (cid:101) O (cid:0) log( n ) /ε 3 (cid:1) times the sample complexity of agnostically learning a single distribution [31].
On the other hand, agnostic federated learning bounds [29] have been studied only for algorithms that sample in one large batch and thus require (cid:101) Ω ( n ) times the sample complexity of a single learning task.
Moreover, the test-time performance of some key multi-distribution methods, such as group distributionally robust optimization [38], have not been studied from a provable or mathematical perspective before.
In this paper, we give a general framework for obtaining optimal and on-demand sample complexity for three multi-distribution learning settings.
Table 1 summarizes our results.
All three settings consider a set D of n distributions and a model class H .
They evaluate the performance of a model h (or a distribution over models) by its worst-case performance, max D ∈D Risk D ( h ) .
As a benchmark, they consider the worst-case loss of the best model, i.e., OPT = min h ∗ ∈H max D ∈D Risk D ( h ∗ ) .
Importantly, all of our sample complexity upper bounds demonstrate only an additive increase of ε − 2 n log( n/δ ) over the sample complexity of a single learning task , compared to the multiplicative factor increase required by existing works.
- Collaborative learning of Blum et al.
[6]: For agnostic collaborative learning, our Theorem 4.1 gives a randomized and a deterministic model that achieve performance guarantees of OPT + ε and 2 OPT + ε , respectively.
Our algorithms have an optimal sample complexity of
O (cid:0)
ε 2 (log( | H | ) + n log( nδ )) (cid:1) .
This improves upon the work of Nguyen and Zakynthinou [31] in two ways.
First, it provides error bounds of OPT + ε for randomized classifiers, where only 2 OPT + ε was previously established.
Second, it improves the upper bound of Nguyen and
Zakynthinou [31] by a multiplicative factor of log( n ) /ε 3 .
In Theorem 4.2, we give a matching lower bound on this sample complexity, thereby establishing the optimality of our algorithms.
- Group distributionally robust learning (group DRO) of Sagawa et al.
[38]: For group DRO, we consider a convex and compact model space H .
Our Theorem 5.1 studies a model that achieves an OPT + ε guarantee on the worst-case test-time performance of the model with an on-demand sample complexity of O (cid:0)
ε 2 ( D H + n log( nδ ) (cid:1) .
Our results also imply a high-probability bound for the convergence of group DRO training error that improves upon the (expected) convergence guarantees of Sagawa et al.
[38] by a factor of n .
- Agnostic federated learning of [29]: For agnostic federated learning, we consider a finite class of hypotheses.
Our Theorems 4.1 and 5.1 show that on-demand sampling can accelerate the generalization of agnostic federated learning by a factor of n compared to batch results established by Mohri et al.
Our results also imply matching high-probability bounds to Mohri et al.
[29] on the convergence of the training error in the batched setting.
To achieve these results, we contribute new insights and techniques for solving stochastic zero-sum games with sources of randomization that differ in both cost and quality.
We frame the multidistribution learning problems as a stochastic zero-sum game with uncertain payoffs and utilize stochastic mirror descent and a variational perspective to solve the game.
In this case, the maximizing player can be interpreted as a weight vector for distributions D , specifying from which distributions future on-demand samples should be taken.
These on-demand samples form a stochastic gradient for the players.
However, the quality of these estimators, the number of samples needed for them, and whether they can be reused later on, differs between the two players.
We extend the Stochastic Mirror
Descent framework to optimally trade off these asymmetric needs for samples.
In Section 3 we give an overview of this approach and its technical challenges and contributions.
1.1 Related Work
Learning models.
There are many lines of work that study multi-distribution learning but which have evolved independently in separate communities.
The field of collaborative learning concerns the learning of a shared machine learning model by multiple stakeholders that each desire a model with low error on their own data distribution.
The line of work initiated by Blum et al.
[6] studies on-demand sample complexity bounds for realizable collaborative learning and was later extended to several related settings (e.g., [31, 9, 8, 32]).
The agnostic federated learning framework of Mohri et al.
[29] poses an equivalent of the multi-distribution learning objective as a fair and intuitive target for federated learning algorithms, and studies it in the offline setting with data-dependent analysis.
Multi-distribution learning also arises in distributionally robust optimization [5] as the Group DRO problem [18], which is is motivated by deep learning applications with multiple deployment domains or protected demographics.
These works focus on an empirical perspective, but have discussed training error convergence in offline settings [18, 38, 39].
Multi-distribution learning is also related to a line of work on multi-group fairness notions (e.g., [37, 40, 14]).
We describe these parallel threads in more detail in Section A.
Stochastic game equilibria.
Our approach relates to a line of research on using online algorithms to find min-max equilibria by playing no-regret algorithms against one another [36, 16, 33, 10, 11].
Online mirror descent (OMD) is one well-studied family of methods that can find approximate minima of convex functions, and also approximate min-max equilibria of convex-concave games, with high probability using noisy first-order information [35, 30, 17, 3].
We bring these online learning tools to bear on the problem of finding saddle points in robust optimization formulations.
The primary technical difference between multi-distribution learning and traditional saddle-point optimization problems is that we have sample access to distributions instead of noisy local gradients.
Preliminaries
Let X be an instance space, Y a label space, and Z = X ×Y a space of datapoints.
A data distribution
D is a joint probability distribution over Z .
We consider a hypothesis class H of a subset of functions mapping X to Y .
We work with loss functions ℓ : H ×Z → [0 , 1] that measure the loss of hypothesis h on data point z ∈ Z .
When Y = { 0 , 1 } , ℓ is the misclassification error.
We denote the expected loss, i.e. risk, of a hypothesis h ∈ H under a data distribution D ∈ D by:
[Equation]
For a distribution over the hypothesis class, p ∈ ∆ H , and a distribution over data distributions, q ∈ ∆ D , we refer to their expected loss by Risk q ( p ) :
= E D
Collaborative Learning.
We will use the collaborative PAC learning model of Blum et al.
[6] and its agnostic extensions by Nguyen and Zakynthinou [31].
The overall goal of this setting is to guarantee small risk for every distribution in a collection of distributions.
Formally, we consider a set of data distributions D :=
{ D 1 , . . .
, D n } .
The goal of the learner is to learn a hypothesis h such that, with probability 1 − δ ,
[Equation]
, where
[Equation]
Group Distribution Robustness.
We will also study the closely related setting of group distributionally robust optimization (Group DRO) of Sagawa et al.
Formally, the group DRO setting considers a model set Θ that is a convex compact subset of the Euclidean space and a convex loss function ℓ : Θ × Z → [0 , 1] that is assumed to be differentiable over Θ .
Given a set of data distributions D := { D 1 , . . .
, D n } , the learner seeks a model θ ∈ Θ , such that, with probability 1 − δ , max D ∈D
[Equation]
[Equation] [Equation]
, where R
[Equation]
OPT
[Equation]
=
[Equation]
There is a close relationship between the Group DRO setting and collaborative learning.
In particular, when Θ = ∆( H ) and H is finite, the two goals are analogous but with two exceptions: first, the Group
DRO could return a distribution over functions while collaborative learning requires the solution to be a deterministic function, and second, R-OPT is potentially more competitive than OPT since it allows randomization.
We note that the group DRO setting is equivalent to the agnostic federated learning framework of [29], thus our results for DRO extend to that setting as well.
Sample complexity.
We are interested in the design of algorithms that achieve the above goals while using a small number of samples from distributions D 1 , . . .
, D n .
We formalize the sample complexity by the total number of calls made to example oracles EX ( D i ) .
Each call EX ( D ) produces an i.i.d.
sample from D .
We note that these example oracles also allow us to sample from any mixture distribution q ∈ ∆ D , e.g., by first selecting a D i according to the mixture and then calling EX ( D i ) .
2.1 Technical Background
We will use tools and definitions from the literature on zero-sum games and no-regret learning throughout the paper.
This section provides a brief overview of these concepts.
Zero-Sum Games.
A finite two-player zero-sum game is described by the tuple ( A - , A + , ϕ ) where
A - = { 1 , . . . , n } and A
+ = { 1 , . . .
, m } are finite sets of actions and where ϕ : A - × A + → [0 , C ] .
In this game, the players choose mixed strategies over actions sets.
These are distributions that are denoted by a vector of probabilities p
∈ ∆ A - and q ∈ ∆ A + .
The expected payoff of mixed strategies is denoted by ϕ ( p, q ) =
E i ∼ p,j ∼ q [ ϕ ( i, j )] .
The goal of the minimizing player is to minimize this expected payoff and the maximizer seeks to maximize the expected payoff; that is, to solve min
p ∈ ∆ A - max q ∈ ∆
A + ϕ ( p, q ) .
A pair ( p, q ) that solves this optimization problem is called a min-max equilibrium .
Similarly, a solution is called an ε -min-max equilibrium if neither player can unilaterally improve their objective by more than ε .
Formally, ( p, q ) is an ε -min-max equilibrium if both players’ regrets are at most ε , i.e., Reg-Min ( p, q ) :
= ϕ ( p, q ) − min i ∗ ∈
A - ϕ ( i ∗ , q ) ≤ ε and Reg-Max ( p, q ) :
= max j ∗ ∈
A + ϕ ( p, j ∗ ) − ϕ
( p, q ) ≤ ε.
We will next describe methods that find ε -min-max equilibria by finding solutions ( p, q ) for which Reg-Min ( p, q ) + Reg-Max ( p, q ) is at most ε .
We describe a more general formulation for convex-concave zero-sum games in Appendix B.1 which we will use for the Group DRO problem.
No-Regret Learning.
We consider an online setting where an arbitrary set of operators , g (1) , . . .
g ( T ) ∈
E ∗ , is revealed sequentially to a learner who must choose a matching sequence of actions, w (1) , . . .
w ( T ) , from a convex compact set Z ⊆ E .
Here, E and E ∗ respectively refer to an arbitrary Euclidean space and its dual.
We focus on a setting where an online learner commits to action w ( t ) ∈
Z before seeing g ( t ) , g ( t +1) , . . .
and aims to achieve vanishing variational error
Errv ( w (1: T ) ) defined by
[Equation]
We will denote no-regret algorithms by their update rule Q : { Z × E ∗ } → Z , where { Z × E ∗ } denotes the space of arbitrary length sequences of action-operator pairs.
Given a history sequence w (1) , . . .
, w ( t )
∈
Z and operator sequence g (1) , . . .
, g ( t )
∈
E ∗ , the algorithm returns w ( t +1) =
Q (cid:0)(cid:8) w (1) , g (1) (cid:9) , . .
, (cid:8) w ( t ) , g ( t ) (cid:9)(cid:1) .
When the history is clear from context, we write w ( t +1) =
Q (cid:0) w ( t ) , g ( t ) (cid:1) as shorthand.
For the particular case where Z = ∆ n is a probability simplex, one such algorithm is Exponential Gradient Descent (also known as Hedge):
[Equation]
where η is a user-defined step size, and w 1 is a user-defined initial iterate.
By default, we take w 1 = (cid:2) 1 n (cid:3) n .
The following lemma is a classical result on the variational error of exponential gradient descent.
Lemma 2.1 ([42]) .
Let g (1) , . .
, g ( T ) ∈
R n and Z
= ∆ n .
Further assume (cid:13)(cid:13) g
( t ) (cid:13)(cid:13) ∞ ≤ C for all timesteps t = 1 , . . .
, T .
Choosing η = (cid:112) log n/T , after T iterations of exponential gradient descent, the outputs w (1) , . .
. , w ( T ) satisfies,
[Equation]
Technical Overview of Our Approach
In this section, we provide an overview of our technical approach for addressing the sample complexity of collaborative learning and group DRO problems.
In later sections, we will refer to the approach outlined in this section to sketch proofs and design algorithms.
We will focus our exposition on collaborative learning and briefly indicate how the same approach applies to the group DRO setting.
At a high level, we first frame collaborative learning as a zero-sum game with uncertain payoffs and aim to use a variational perspective to learn its minmax equilibrium.
We specifically choose the variational perspective (instead of an arbitrary online learning approach), since it allows us to linearize the effect of uncertain payoffs on the resulting error.
We then use stochastic gradients to solve the variational problem.
Our stochastic gradients will rely on i.i.d.
samples from the distributions to estimate gradients both with respect to distributions over H and mixtures over D but with an asymmetric bound on the bias and variance of the estimates.
Along the way, we develop tools and formalisms that handle the asymmetric cost of stochastic gradients and obtain optimal sample complexity results.
We now address these steps in more detail.
Collaborative Learning as Zero-Sum Games.
When the hypothesis class H is finite, the collaborative learning problem with distribution set D corresponds to a zero-sum game ( A - , A + , ϕ ) with
A - = H , A + = D , and ϕ ( i, j ) =
Risk j
( i ) , where i ∈
A - and j ∈ A + .
Observe that the value of the min-max solution is equivalent to R-OPT.
It is not hard to see that any ε -min-max equilibrium ( p, q ) of this game corresponds to a 2 ε collaborative learning solution, i.e.,
[Equation]
This enables us to use tools that have been developed for solving zero-sum games in order to address collaborative learning and group DRO settings.
We will use a similar construction when hypothesis class H has finite VC dimension, where A - will instead refer to an appropriate ε -cover of H .
Using the Variational Error to deal with Payoff Uncertainty.
A sufficient condition for minimizing regret, and thus finding ε -min-max equilibrium, is minimizing the variational error (Equation 3).
In particular, for any finite zero-sum game ( A - , A + , ϕ ) , defining Z = [∆ A - , ∆ A + ] and operators
[Equation]
ensures that variational error provides an upper bound on regret: Errv ( w (1: T ) ) ≥ Reg-Min ( p, q ) +
Reg-Max ( p, q ) , where w = ( p, q ) (see Fact C.1).
In collaborative learning, when p ( t ) is the minplayer’s distribution over hypotheses and q ( t ) is max-player’s distribution over the mixtures, the gradient vectors refer to the risks of each hypothesis or distribution under q ( t ) or p ( t ) respectively:
[Equation]
In the collaborative learning setting, we can only create noisy estimates (cid:98) g for these gradients from samples.
No-regret algorithms are advantageous in this setting as they choose their t th iterate w ( t ) before seeing the t th gradient g ( t ) .
This means that w ( t ) is independent of gradient noise, ε ( t ) :
= g ( t ) − (cid:98) g ( t ) .
We can thus linearize the noise and decompose variational error into the training and generalization errors as follows
[Equation]
In contrast, generic no-regret algorithms that do not solve the variational inequality (e.g., when one player plays Hedge and another plays clairvoyant best-response as used in existing work in collaborative learning due to Blum et al.
[6], Nguyen and Zakynthinou [31], Chen et al.
[9]) nest the generalization and training errors which leads to a multiplicative increase in sample complexity.
Leveraging Noisy Stochastic Gradients.
We will work with stochastic estimators of g .
These are functions (cid:98) g : ξ × ∆ A - × ∆ A + of some external source of randomness, ξ ∈ ξ , and a strategy profile of interest.
For collaborative learning, the randomness source ξ is an i.i.d.-sampled data point from an appropriate mixture of distributions and the estimator (cid:98) g is then the empirical loss on this sample, which is an unbiased and bounded estimator in the range of the loss function, i.e., [0 , 1] .
Interestingly, estimators of these stochastic gradients have an asymmetric need for data.
As seen in
Equation 7, the min-player’s gradient g - ( p, q ) includes the risk of every hypothesis h ∈ H for the same data distribution q .
Therefore, an unbiased estimator (cid:98) g - ( p, q ) can be constructed from a single call to an example oracle EX ( q ) .
We call this source of randomness ξ q and say that its cost is r − = 1 .
While ξ q costs 1 unit, the randomness it provides is specialized to the point of inquiry, that is, it cannot be used for estimating other (cid:98) g - ( p, q ′ ) .
We call this source of randomness and its associated unbiased estimation a locally unbiased estimator.
On the other hand, the max-player’s gradient g + ( p, q ) includes the risk of the same hypothesis p on every distribution D ∈ D .
Therefore, an unbiased estimator (cid:98) g
+ ( p, q ) requires n samples, i.e., a call to every example
oracle EX
( D i ) .
We call this source of randomness producing n samples ξ p and say that its cost is r + = n .
Importantly, while ξ p costs n unit, the randomness it provides can be reused for estimating other gradients, that is, it can provide an unbiased and bounded estimators for all (cid:98) g
+ ( p ′ , q ′ ) .
We call this source of randomness and its associated unbiased estimator a globally unbiased estimator.
To emphasize the fact that this source of randomness is agnostic to ( p, q ) we refer to it by ξ ⊥ hereafter.
We refer the reader to Appendix B.2 for a more formal definition and description of these asymmetries.
Minimizing Regret with Asymmetric Cost.
With the goal of minimizing sample complexity in mind, it is essential that we reuse randomness ξ ⊥ across n time steps of variational algorithms.
To do this, we introduce a stochastic variational approach in Algorithm 1 that accommodates different sampling frequencies for the minimizing and maximizing players.
This will decouple the sample complexity of the minimizing agent (who requires a time horizon of at least log( A - ) ≈ log( H ) ) and the maximizing agent.
Lemma 3.1 proves this decoupling allows us to find an ε -min-max equilibrium with an additive n + log( H ) sample complexity instead of a multiplicative n log( H ) .
Algorithm 1 uses the same randomness ξ ⊥ ( a ) of cost r for estimating g + ( p t , q t ) for all t ∈ [ ar + 1 , . .
. , a ( r + 1)] .
On the other hand, the algorithm uses fresh randomness ξ ( t ) of cost 1 to estimate g - ( p t , q t ) for every time step t .
The total randomness cost of this algorithm is thus 2 T because iteration of the outer loop incurs 2 r cost.
Lemma 3.1.
Let ( A - , A + , ϕ ) be a finite zero-sum game.
Assume there exists ξ q ( t ) of cost 1 providing locally unbiased estimates (cid:98) g - ( · ) and there exists ξ ⊥ ( a ) of cost r providing globally unbiased estimates (cid:98) g + ( · ) .
With probability 1 − δ , Algorithm 1 returns an ε -min-max equilibrium of the game, so long as
T
[Equation] [Equation]
Moreover, the total cost of randomness incurred by the algorithm is at most 2 T .
Output: Mixed strategy profile ( p, q )
∈
∆ A - × ∆ A + ;
Input do
Realize ξ at cost r ; //
Sample datapoints from for t ar
1 − ar do
Realize ξ q ( t ) at cost 1 ; //
Sample from adversary-selected distribution.
Estimate gradients: (cid:98) g ( t ) + = (cid:98) g
+ (cid:16) ξ ⊥
( a ) , p ( t ) , q ( t ) (cid:17) , (cid:98) g ( t ) - = (cid:98) g - (cid:16) ξ q ( t ) , p ( t ) , q ( t ) (cid:17) ;
Run Hedge updates: p ( t +1) =
Q hedge (cid:16) p ( t ) , (cid:98) g ( t ) + (cid:17) , q ( t +1) =
Q hedge (cid:16) q ( t ) , (cid:98) g ( t ) + (cid:17) ; end for end for
Return the uniformly mixed strategies p = 1 T (cid:80) Tt =1 p ( t ) and q = 1 T (cid:80)
Tt =
1 q ( t ) ;
Proof sketch.
Our approach uses Equation 8 to decompose the variational error into training error and generalization error.
Since exponential gradient descent is known to bound the training error (as shown in Lemma C.4), it only remains to bound the generalization error (the second term in
Equation 3).
We note that in expectation each summand (cid:10) ε ( t ) , w ( t ) − w ∗ (cid:11) is zero.
This is because ε ( t ) =
g ( t ) − (cid:98) g ( t ) and (cid:98) g ( t ) are unbiased estimators.
Therefore, the sum of these terms has an intuitive martingale interpretation and could be bounded by the Azuma-Hoeffding inequality.
There is a subtlety here, however.
When we reuse the maximizing player’s randomness over r rounds, we create correlations between these terms in the generalization error that cannot be directly accommodated by a martingale.
The trick here is to note that these correlations are entirely contained in r -length periods.
So, we can partition our sequence to r martingales and bound each one.
This completes the proof.
See Appendix C.1 for detailed proof of this lemma.
Derandomization.
The ε -min-max equilibria ( p, q ) returned by Exponentiated Gradient Descent gives a probability distribution p over the hypothesis class H that achieves the collaborative learning bound.
To obtain a deterministic hypothesis, we can instead work with h Majp whose predictions are p -weighted majority votes over the hypotheses in H .
As stated below, the error of this deterministic classifier is approximately bounded by the expected error of p .
Lemma 3.2.
For any p ∈ ∆ H , max D ∈D Risk D ( h Majp ) ≤ 2 max D ∈D
Risk D ( p ) .
This lemma in particular implies that for any ε -min-max equilibria ( p, q ) , we have
[Equation]
Collaborative Learning Bounds
In this section, we characterize the sample complexity of collaborative learning by providing tight upper and lower bounds for this problem.
We describe Algorithm 2, which attains near-optimal sample complexity by on-demand sampling: iteratively selecting distributions to sample from.
4.1 Sample Complexity Upper Bounds
We are now prepared to describe our collaborative learning algorithm and guarantees, using the tools we developed in Section 3.
Algorithm 2 is a direct application of Algorithm 1 to a zero-sum game with action sets A - = H , A + = D and payoff ϕ ( h, D ) =
Risk D ( h ) .
Here, ξ q ( t ) makes one call to
EX ( q ( t ) ) and ξ ⊥ ( a ) makes one call to EX ( D ) for each D ∈ D .
In other words, Algorithm 2 constructs distributions p ( t ) ∈
∆ H and q ( t ) ∈ ∆ D by running the Hedge update.
The gradient estimators used by Hedge are the empirical losses on a set of independent random variables.
In particular, the minimizing player uses gradients ℓ ( h, z ( t ) ) for all h ∈ H for a single sample z ( t ) ∼ EX ( q ( t ) ) and the maximizing player uses gradients ℓ ( p ( t ) , z aD ) for all distributions D ∈ D where a single sample z aD ∼ EX ( D ) is drawn per distribution and is reused for all time steps t ∈
[( a − 1) n + 1 , . .
. , an ] .
Our main result in this section bounds the sample complexity of Algorithm 2.
Theorem 4.1.
For any finite hypothesis class H and unknown set of distributions D , with probability 1 − δ , Algorithm 2 returns a distribution
p ∈
∆ H such that
[Equation]
using a number of samples that is O (cid:16)
[Equation]
n
[Equation]
Proof sketch.
By construction, Lemma 3.1 guarantees that with probability at least 1 − δ , the pair ( p, q ) is an ε/ 2 -min-max equilibrium for the corresponding zero-sum game.
As shown by Equation 5, p is a randomized classifier that meets the collaborative learning objective, i.e., its expected worst-case error is OPT + ε .
By Lemma 3.2, the corresponding deterministic classifier h Majp has worst-case error of 2 OPT + ε .
This bounds the error of the resulting classifier.
To bound the sample complexity, Lemma 3.1 shows that the randomness cost of Algorithm 1 is at most 2 t .
Since the cost of randomness is exactly the total number of samples we take from our example oracles, the total sample complexity of Algorithm 2 is 2 t ∈ O (cid:0) ε − 2 (log |H| + n log( n/δ )) (cid:1) .
An analogue of Theorem 4.1 (Theorem C.3) holds for the case of infinite hypothesis classes of bounded Littlestone dimension with a sample complexity of O (cid:0) ε − 2 ( Little ( H ) + n log( n/δ )) (cid:1) .
When the VC dimension (which is smaller than the Littlestone dimension) is finite, a similar result also holds under some additional mild assumptions.
In this case, one can instead run Algorithmwith a hypothesis class H
′ that is an ε -net with respect to every distribution in D .
We note that such ε -nets of size nε − 2VCD( H ) necessarily exist (see, e.g., [2]); for example, the union of ε -nets with respect to each distribution D ∈ D .
When such H ′ is known in advance, we may run
Algorithm 2 with H ′ and incur a sample complexity that now replaces log( |H ′ | )
=
O ( d log(1 /ε )) in the sample complexity of Theorem 4.1.
We remark that it is not strictly necessary to know an ε -net in advance.
Instead, one can compute a net from samples or from other information about distributions in D .
In Appendix C.5, we explore a range of assumptions that allow us to compute such an ε -net from samples, without incurring a significant increase in the sample complexity of Theorem 4.1.
We end this subsection with two remarks about our sample complexity upper bound.
Remark 4.1.
Theorem 4.1 improves over the best-known sample complexity for agnostic collaborative learning by Nguyen and Zakynthinou [31] in two ways, giving an OPT + ε bound for randomized classifiers instead of their 2 OPT + ε bound, and improving their sample complexity of
O (cid:0) 1 ε 5 (cid:0) log( n ) log( |H| ) log (cid:0) 1 ε (cid:1) + n log (cid:0) n δ (cid:1)(cid:1)(cid:1) by a multiplicative factor of 1 ε 3 log ( n ) log (cid:0) 1 ε (cid:1) .
Remark 4.2.
For constants ε and δ , our sample complexity of O (log( |H| ) + n log n ) appears to violate the lower bound of Ω (log( |H| ) log n + n log log |H| )
due to Chen, Zhang, and Zhou [9].
This discrepancy is due to a small error in the proof of that lower bound, which we have verified in private communications with the authors.
In the next subsection, we give lower bounds on the sample complexity of collaborative learning that match our upper bounds.
4.2 Sample Complexity Lower Bound
We now provide matching lower bounds for agnostic collaborative learning.
Our lower bounds hold for collaborative learning algorithms obtaining error of R-OPT + ε , using a randomized or deterministic hypothesis.
We call an algorithm an ( ε, δ ) -collaborative learning algorithm if for any collaborative instances it attains an error of R-OPT + ε with probability at least 1 − δ .
Theorem 4.2.
Take any n, d ∈ Z
+ , ε, δ ∈ (0 , 1 / 8) , and ( ε, δ ) -collaborative learning algorithm A .
There exists a collaborative learning problem ( H , D ) with |D| = n and |H| =
2 d , on which A takes at least Ω (cid:0) 1 ε 2 (log |H| + |D| log(min {|D| , log |H|} /δ ))
(cid:1) samples.
Proof sketch.
We defer the formal proof of this theorem to Appendix C.3 and sketch the main ideas here.
Let X = { 1 , . . .
, d } , Y = { + , −} , and H be the set of all functions X
→ Y .
Our construction combines two sets of hard distributions.
Consider the case when n = d · η for some η ∈ Z .
First, we can reduce to a d -armed multi-arm bandit exploration problem giving us an Ω (cid:0) d log(1 /δ )
/ε 2 (cid:1) .
Second, we construct η hard instances on η corresponding points.
Since the learning algorithms has to solve each problem it has to incur a loss of η · d log( d/δ ) /ε 2 .
Group DRO and Agnostic Federated Learning
The results we describe in the collaborative learning setting can be generalized to the group DRO setting, and equivalently, agnostic federated learning.
Theorem 5.1.
Consider a group distributionally robust problem (Θ , D ) with convex
compact unitdiameter parameter space Θ of Bregman radius D Θ (Definition B.11), and convex loss ℓ : Θ × Z → [0 , C ] .
A variant of Algorithm 2 (in particular Algorithm 4 in Appendix 4.1), returns θ ∈ Θ such that
[Equation]
z )] ≤ R-OPT + ε, using a number of samples that is O (cid:16) D Θ C 2 + nC 2 log( n/δ ) ε(cid:17) .
The proof of this lemma is deferred to Appendix 4.1 and is similar to the proof of Theorem 4.
except that it uses a generalization of Lemma 3.1 for general convex-concave games.
This theorem establishes a generalization bound for the problem of group distributionally robust optimization [38] and improves, by a factor of n , existing sample complexity bounds for agnostic federated learning [29].
This improvement is attained by sampling data on-demand, whereas [29] only chooses a fixed distribution over groups/clients to sample from; this highlights the importance of adapting one’s sampling strategy on-the-fly when learning robust models.
Another important question is how fast the training error of stochastic gradient descent converges for the group DRO/AFL settings and was considered by Sagawa et al.
We can transfer our generalization guarantees for on-demand settings into batch settings and achieve the following corollary, which improves on the convergence guarantees of Sagawa et al.
[38] by a factor of n .
Corollary 5.2.
Under the same assumptions of Theorem 5.1, we give a procedure (see Appendix 4.1) that minimizes GDRO/AFL training error within ε of R-OPT with probability at least 1 − δ in fewer samples than O (cid:16) D Θ C 2 + nC 2 log( n/δ )
ε
[Equation]
Empirical Analysis of On-Demand Sampling for Group DRO
This section describes experiments where we adapt our on-demand sampling-based multi-distribution learning algorithm for deep learning applications.
In particular, we compare our algorithm against the de-facto standard multi-distribution learning algorithm for deep learning, Group DRO (GDRO)
As GDRO is designed for use with offline-collected datasets, to provide an accurate comparison, we modify our algorithm to work on offline datasets (i.e., with no on-demand sample access).
Resampling Multi-Distribution Learning (R-MDL).
To adapt our multi-distribution learning algorithm, Algorithm 2, for deep learning applications, we replace its hypothesis-selecting no-regret learning algorithm with a minibatch gradient descent algorithm.
We can further adapt our algorithm to offline datasets by simulating on-demand sampling on the empirical distributions of datasets.
This modified algorithm, R-MDL, is described in full in Algorithm 5.
In contrast, the GDRO algorithm is also minibatch gradient descent but samples minibatches uniformly from all distributions.
Datapoints in each minibatch are importance weighted according to their distribution of origin, where a no-regret algorithm adversarially weights each distribution.
Though effective, this GDRO method is brittle and requires tricks like unconventionally strong regularization [38].
Our theory of on-demand sampling suggests that R-MDL should mollify this brittleness.
Experiment Setting In Table
[Table]
we replicate the Group DRO experiments of Sagawa et al.
[38] and compare the standard GDRO algorithm with our R-MDL algorithm (Algorithm 5).
We finetune Resnet-50 models (convolutional neural networks) [19] and BERT models (transformer-based network) [12] on the image classification datasets Waterbirds [38, 43] and CelebA [25] and the natural language dataset MultiNLI [44] respectively.
We train these models in 3 settings: with standard hyperpameters, under strong weight decay ( ℓ -2) regularization, or under early stopping
[Table]
R-MDL consistently outperforms GDRO and ERM.
In every dataset and in almost every setting,
R-MDL significantly outperforms GDRO and ERM in worst-group accuracy.
In addition, whereas
GDRO and ERM have large gaps between worst-group accuracy and average accuracy, R-MDL has almost matching worst-group and average accuracies.
This indicates that R-MDL is more effective at prioritizing learning on difficult groups.
R-MDL is robust to regularization strength.
R-MDL retains high worst-group accuracy even without strong regularization.
These results challenge the observation of Sagawa et al.
[38] that strong regularization is critical for the performance of Group DRO methods.
This suggests that the brittleness of GDRO is due to reweighting rendering the adversary too weak.
In contrast, R-MDL provides a robust multi-distribution learning method with significantly less hyperparameter sensitivity.
Acknowledgments
This work was supported in part by the National Science Foundation under grant CCF-2145898, a
C3.AI Digital Transformation Institute grant, and the Mathematical Data Science program of the
Office of Naval Research.
This work was partially done while Haghtalab and Zhao were visitors at the Simons Institute for the Theory of Computing.
C Omitted Proofs
C.1 Proof of Lemma C.1 (Generalization of Lemma 3.1) . . . . . . . . . . . . . . . .
C.2 Proof of Theorem C.2 (Generalization of Theorems 4.1 and 5.1) . . . . . . . . . .
C.3 Proof of Theorem 4.2 . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.4 Proof of Lemmas C.4 and C.10 . . . . . . . . . . . . . . . . . . . . . . . . . . . .
C.5 Proof of Theorem C.
D Experiment Details
A Extended Related Work
There are many lines of work that study multi-distribution learning but which have evolved independently in separate communities.
Collaborative and Federated Learning Blum, Haghtalab, Procaccia, and Qiao [6] posed the first fully general description of multi-distribution learning, motivated by the application of collaborative
PAC learning .
The field of collaborative learning concerns the learning of a shared machine learning model by multiple stakeholders that each desire a model with low error on their own data distribution.
The line of work studies on-demand sample complexity bounds for the setting where stakeholders collect data so as minimize the error of the worst-off stakeholder [6, 31, 9, 8].
This setting, stated in its full generality, yields the multi-distribution learning problem as presented in this paper.
Blum et al.
[6] established a log( n ) factor blowup for the realizable case and the best known sample complexity guarantees for the general agnostic setting experiences a factor log( k ) /ε
3 blowup and is due to
Nguyen and Zakynthinou [31].
In comparison, our work establishes a tight additive increase in the sample complexity (which is comparable to log( k ) multiplicative factor blowup with no dependence on ε ).
A related line of work concerns the strategic considerations of collaborative learning and seeks incentive-aware mechanisms for collecting data used in collaborative learning [7].
The field of federated learning concerns a closely related motivating application where the goal is to learn a model from data dispersed across multiple devices but querying data from each device is expensive [28].
The agnostic federated learning framework of Mohri, Sivek, and Suresh [29] poses (an equivalent of) the multi-distribution learning objective as a fair and intuitive target for federated learning algorithms, and studies it in the offline setting with data-dependent analysis.
Group Distributionally Robust Optimization (Group DRO) Multi-distribution learning also arises in distributionally robust optimization [5] under the name of Group DRO, a class of DRO problems where the distributional uncertainty set is finite [18].
Group DRO literature is motivated by applications where these distributions correspond to deployment domains or protected demographics that a machine learning model should avoid spuriously linking to labels [18, 38, 39].
Although Group
DRO—like collaborative learning—is mathematically an instance of multi-distribution learning, prior work on group DRO focus on training error convergence in offline settings as they focus on deep learning applications.
As we later discuss, theoretical aspects of online multi-distribution learning can translate into actionable insights for Group DRO applications.
Multi-Group Fairness and Learning Notions Multi-distribution learning is also related to the fields of multi-group learning [37, 40] and multi-group fairness [14, 20].
Most of these work are formulated for the supervised learning setting on an existing set of samples, exceptions include the work of Rothblum and Yona [37] that implicitly conducts on-demand sampling for a slightly different multi-distribution problem involving distribution-wise optimality and experience a sample complexity blowup of log( n ) /ε 8 .
Multi-Source Domain Adaptation and Learning Multi-source domain adaptation, or multi-task, learning is another related line of work that concerns using data from multiple different training distributions to learn some target distribution, under the assumption that the training and target distributions share some task relatedness [4, 26].
Multi-distribution learning can be framed similarly as using a finite set of training distributions to simultaneously learn the convex hull of the training distributions.
Interestingly, multi-distribution learning’s requirement of learning the entire convex hull implicitly obviates the task relatedness assumptions of multi-source learning.
Stochastic game equilibria.
Our technical approach to multi-distribution learning relates to a line of research on using online algorithms to find min-max equilibria using stochastic feedback [36, 16, 33, 10, 11].
Online mirror descent (OMD) is one well-studied family of methods that can find approximate minima of convex functions, and also approximate min-max equilibria of convexconcave games, using noisy first-order information [35, 30, 17, 3].
However, applying OMD analysis to the saddle-points in multi-distribution learning is a non-trivial affair.
While optimization theory typically thinks about noisy zeroth/first-order oracles (e.g., Juditsky et al.
[21]), our learning-theoretic analysis uses oracles that provide noisy semi-local information (i.e., datapoints).
B Full Formulation
In this section, we formally describe our formulations of stochastic convex-concave games and multi-distribution learning problems.
B.1 Convex-Concave Zero-Sum Game
In this subsection, we give a formal definition of a convex-concave zero-sum game and its min-max equilibria.
We also introduce assumptions on these games for efficiently finding saddle-points.
Definition B.1.
A convex-concave two-player zero-sum game is described by the tuple ( A - , A + , ϕ ) , where A - ⊂ E - is a subset of Euclidian space E - , A + ⊂ E + is a subset of Euclidian space E + , and ϕ : A - × A + → R is a Lipschitz continuous convex-concave function.
On a convex-concave two-player zero-sum game ( A - , A + , ϕ ) , we can define both exact and approximate notions of min-max equilibria in terms of player regrets.
Definition B.2.
The minimizing and maximizing player’s regrets at a strategy profile ( p, q )
∈
A - × A + are denoted Reg-Min , Reg-Max respectively, and defined as,
Reg-Min ( p, q )
[Equation]
Definition B.3.
A strategy profile ( p, q ) ∈
A - × A
+ is a min-max equilibrium if both players have zero regret: Reg-Min ( p, q ) =
and Reg-Max ( p, q ) = 0 .
More weakly, ( p, q ) ∈
A - × A
+ is an ε -min-max equilibrium if both players have at most ε regret: Reg-Min ( p, q ) ≤ ε and Reg-Max ( p, q ) ≤ ε .
In this paper, we may also impose the following assumptions on a convex-concave zero-sum game.
Assumption 1.
The action sets A - , A + are compact, convex, and have diameters R - , R + respectively:
[Equation]
Assumption 2.
At any p, q ∈ A - × A + , the partial subdifferential of the payoff function ϕ is non-empty.
Furthermore, every partial subgradient vector has a bounded norm:
[Equation]
B.2 Stochastic Settings
In this subsection, we give a formal definition of an asymmetric stochastic setting for a zero-sum game.
Our formulation of stochastic first-order oracles observes the convention of representing all randomness in stochastic oracles—and by extension, in any stochastic optimization process—in terms of an i.i.d.
sequence of random variables.
One nuance our formulation addresses is how randomness can be re-used by stochastic first-order oracles.
We do this by formalizing our stochastic setting in terms of multiple i.i.d.
sequences of random variables, where the sequence to which a random variable belongs specifies how randomness corresponding to the random variable can be used.
We begin by introducing the notion of a coupled random variable.
In the context of a two-player game, a random variable may be coupled to a minimizing player’s strategy profile, a maximizing player’s strategy profile, neither or both.
Our definition formalizes the notion that a random variable can only be interpreted in the context of the mixed strategy to which it is coupled.
Definition B.4.
For any p ∈ A - , we define a random variable η to be p -coupled if its range is a measurable space E p defined by p .
Similarly, for any q ∈ A + , we define a random variable η to be q -coupled if it’s range is a measurable space defined by q .
A random variable η is ( p, q ) -coupled
if it’s range is a measurable space defined by ( p, q ) .
For convenience, we will denote p -coupled random variables with superscript η p and, similarly, q -coupled random variables with superscript η q .
Random variables that are not coupled will be denoted by η ⊥ when such clarification is necessary.
We will now define stochastic first-order oracles that express their randomness in terms of sequences of i.i.d.
coupled random variables.
Definition B.5.
In a zero-sum two-player game, the minimizing player’s randomness source is defined as a set ξ - ⊆ { ξ q - |
q ∈ A + (cid:83) {⊥}} , where ξ q - := (cid:8) ξ q - ; i (cid:9) i ∈
Z is a sequence of i.i.d.
random variables all coupled with q ∈ A + .
In addition, all random variables in all sequences in ξ - are independent.
Definition B.6.
In a zero-sum two-player game, the maximizing player’s randomness source is defined as a set ξ + ⊆ { ξ p + | p ∈
A - (cid:83) {⊥}} , where ξ p + := (cid:8) ξ p + ; i (cid:9) i ∈
Z is a sequence of i.i.d.
random variables all coupled with p ∈
A + .
In addition, all random variables in all sequences in ξ + are independent.
Definition B.7.
For any q ∈ A + , consider the function (cid:98) g q - : E q × A - × A + → E ∗ - .
The minimizing player has a locally unbiased first-order oracle if there exists, for all q ∈ A
+ , a (cid:98) g q - such that for all p ∈ A - and i ∈
Z :
[Equation]
We analogously define locally unbiased oracles for the maximizing player.
When q is clear from context, we write (cid:98) g q - as (cid:98) g - .
We can also define a globally unbiased oracle.
Definition B.8.
For any q ∈ A + , consider the function (cid:98)
g ⊥ - : A - × A + → E ∗ - .
The minimizing player has a globally unbiased first-order oracle if there exists (cid:98) g ⊥ - where for all q ∈ A + and p ∈
A - and i ∈
Z :
[Equation]
We analogously define globally unbiased first-order oracles for the maximizing player.
Finally, we may impose the following norm-bound assumption on the first-order oracles we discuss.
Assumption 3.
Every globally unbiased first-order oracle has a range with bounded norm: (cid:13)(cid:13)(cid:98)
g ⊥ - ( · ) (cid:13)(cid:13) E ∗ - ≤ C - , (cid:13)(cid:13)(cid:98) g ⊥ + ( · ) (cid:13)(cid:13) E ∗ + ≤ C + .
Furthermore, every locally unbiased first-order oracle also has a range with bounded norm: for all p, q ∈ A - , A + , ∥ (cid:98) g q
[Equation] [Equation]
∥ E ∗
[Equation]
≤
[Equation]
∥
[Equation]
g
[Equation]
+
[Equation]
∥ E ∗ + ≤ C + ,
B.3 Multi-Distribution Learning
In this subsection, we give a formal definition of multi-distribution learning that unifies the problem formulations of collaborative learning [6], agnostic federated learning [29], and group DRO [38].
We further introduce assumptions that characterize two special cases of multi-distribution learning: convex multi-distribution learning and binary classifier multi-distribution learning .
We begin by reviewing some common definitions from convex optimization.
Definition B.9.
Let Z be a convex compact subset of a Euclidian space E with norm ∥·∥ .
A distance generating function on Z is a function ω : Z → R , where:
1. ω is continuous and strongly convex, modulus 1, w.r.t to ∥·∥ on Z . 2.
There exists a non-empty subset Z o ⊂ Z where the subdifferential ∂ω is non-empty and ∂ω admits a continuous selection on Z o .
Furthermore, the center of Z w.r.t.
ω is defined as z c := arg min z ∈ Z o ω ( z ) .
Definition B.10.
The prox function V : Z o × Z
→ R + associated with a distance generating function ω : Z → R is defined as:
[Equation]
The prox function is also known as the Bregman divergence.
Definition B.11.
Given a convex set Z with a distance generating function ω satisfying Definition
B.9, the Bregman radius is defined as max u ∈ Z V ( z c , u ) ≤ D Z , where z c is the center of Z as defined in Definition B.9.
We state our most general formulation of multi-distribution learning as follows.
Definition B.12.
Let Z =
X ×Y be a space of datapoints and D
=
{ D i } ni =1 be a finite set of n joint probability distributions over Z .
Let Θ denote a set of parameters and ℓ : Θ × Z → [0 , L ] be a loss function.
Then the tuple (Θ , D , ℓ ) describes a multi-distribution learning problem, w.r.t.
X , Y , Z .
One case of multi-distribution learning we study in this paper is convex multi-distribution learning , which includes as special cases the problem formulations of Sagawa et al.
[38] and Mohri et al.
Convex multi-distribution learning also encompasses the problem formulation of Blum et al.
[6] for finite hypothesis spaces, i.e., when |H| < ∞ .
Definition B.13.
The tuple (Θ , D , ℓ ) describes a convex multi-distribution learning problem when Θ is convex compact, ℓ is convex in Θ , and there exists a distance generating function ω :
Θ → R on our parameter space Θ .
Definition B.14.
The diameter of the parameter space Θ is an R Θ > 0 satisfying:
[Equation]
Assumption 4.
Given a convex multi-distribution learning problem (Θ , D , ℓ ) , we assume that, for any datapoint z in the supports of the distributions in D and any θ ∈ Θ o , the partial subgradient of ℓ ( θ, z ) w.r.t.
θ has bounded norm:
[Equation]
Assumption 5.
Given a convex multi-distribution learning problem (Θ , D , ℓ ) , we assume there exists a distance generating function ω where Θ has bounded Bregman radius D Θ .
Remark B.1.
As ω is strongly convex modulus 1 by definition, any Θ satisfying Assumption 5 has a finite diameter R Θ ≤ 2 √ 2 D Θ
Another important case of multi-distribution learning is binary classifier multi-distribution learning , which includes the problem formulations of Blum et al.
[6] both for finite hypothesis spaces ( |H| < ∞ ) and finite VC dimension hypothesis spaces ( VCD( H ) < ∞ ).
Definition B.15.
The tuple (Θ , D ) describes a binary classifier multi-distribution learning problem when Z = X × { 0 , 1 } , Θ is the set of probability distributions over a set of binary classification rules H :
X → { 0 , 1 } and ℓ ( θ, ( x, y )) :
= E h ∼ θ [1[ h ( x ) =
y ]] .
Remark B.2.
A binary classifier multi-distribution learning problem (Θ , D , ℓ ) is equivalent to a convex multi-distribution learning problem (Θ , D , ℓ ) when the support of Θ is finite, i.e., Θ is a probability distribution over a finite number of binary classification rules.
Finally, we can define a multi-distribution analogue to probably-approximately-correct learning [41].
Definition B.16.
An example oracle EX ( D ) is an infinite set of i.i.d.
samples from a probability distribution D over datapoints.
Colloquially, a “new call” to example oracle EX ( D ) refers to realizing a previously unrealized sample in EX ( D ) .
Definition B.17.
A learning algorithm
A is an ( ε, δ ) multi-distribution learning algorithm for a set of multi-distribution learning problems V :=
{ (Θ i , D i , ℓ i ) } i if, given any problem (Θ i , D i , ℓ i ) ∈ V , accessing only the tuple (Θ i , ℓ i , { EX ( D ) | D ∈ ∆ D} ) , A outputs a parameter θ ∈ Θ i that satisfies, with probability at least 1 − δ :
[Equation]
We use ( ε, δ ) -algorithm as a shorthand for ( ε, δ ) multi-distribution learning algorithm.
Definition B.18.
A multi-distribution learning algorithm
A has a sample complexity of N (or “takes
N samples”) on a set of multi-distribution learning problems V :=
{ (Θ i , D i , ℓ i ) } i if N is the smallest integer such that, given any problem V ∈ V , the event that A takes more than N samples is measure-zero.
If no such N exists, we say A has infinite sample complexity.
C Omitted Proofs C.1 Proof of Lemma C.1 (Generalization of Lemma 3.1)
Input
Action sets no regret learning algorithms Q
: { A
E
} →
A
, Q + : { A + × E
} → A + .
for t = 1 , 2 , . . .
, T do
Realize randomness ξ q ( t ) - ; t and ξ ⊥ + ; ⌈ tr ⌉ .
Take estimates (cid:98) g ( t ) - := (cid:98) g - (cid:16) ξ q ( t ) - ; t , p ( t ) , q ( t ) (cid:17) and (cid:98) g ( t ) + := (cid:98) g
+ (cid:16) ξ ⊥ + ; ⌈ tr ⌉ , p ( t ) , q ( t ) (cid:17) .
Run the no-regret updates: p ( t +1)
=
Q - (cid:16)(cid:110) p (1) , (cid:98) g (1) - (cid:111) , . . .
(cid:110) p ( t ) , (cid:98) g
t ) - (cid:111)(cid:17) q
( t +1)
=
Q + (cid:16)(cid:110) q (1) , (cid:98) g (1) + (cid:111) , . . .
(cid:110) p ( t ) , (cid:98) g ( t ) + (cid:111)(cid:17) end for
Return the uniformly mixed strategies p = 1 T (cid:80) Tt =1 p ( t ) and q = 1 T (cid:80)
Tt =
1 q ( t ) .
We first recall the following lemma from Section 3.
Lemma 3.1.
Let ( A - , A + , ϕ ) be a finite zero-sum game.
Assume there exists ξ q ( t ) of cost 1 providing locally unbiased estimates (cid:98) g - ( · ) and there exists ξ ⊥ ( a ) of cost r providing globally unbiased estimates (cid:98) g + ( · ) .
With probability 1 − δ , Algorithm 1 returns an ε -min-max equilibrium of the game, so long as
T
[Equation]
Moreover, the total cost of randomness incurred by the algorithm is at most 2 T .
We will prove a more general result, Lemma C.1, that implies Lemma 3.1 as a special case.
Lemma C. provides sample complexity upper bounds for Algorithm 3, an algorithm for approximating the saddle-point of a convex-concave game with high-probability.
Algorithm 3 is also a generalization of
Algorithm 1.
Lemma C.1 (Generalization of Lemma 3.1) .
Let ( A - , A + , ϕ ) be a convex-concave game satisfying
Definition B.1 and Assumptions 1 and 2.
Suppose the minimizing player has a locally unbiased first-order oracle (cid:98) g - and the maximizing player has a globally unbiased first-order oracle (cid:98) g
+ , with both oracles satisfying Assumptions 3.
Take Q - to be any no-regret algorithms with the guarantee that for, any sequence g (1) , . .
, g ( T )
∈
E ∗ - , if (cid:13)(cid:13) g
( i )
(cid:13)(cid:13)
E ∗ - ≤ C for all i ∈
[ T ] , the Q - -learned sequence w (1) , . . .
, w ( T ) satisfies:
[Equation]
T
[Equation]
Take Q + to be any no-regret algorithms with the guarantee that for, any sequence g (1) , . .
, g ( T )
∈
E ∗ + , if (cid:13)(cid:13) g ( i ) (cid:13)(cid:13)
E ∗ + ≤ C for all i ∈
[ T ] , the Q + -learned sequence w (1) , . . .
, w ( T ) satisfies:
[Equation]
Then, the mixed strategy profile ( p, q ) outputted by Algorithm 3 is an ε -min-max equilibrium with probability at least 1 − δ so long as:
[Equation]
Moreover, exactly T elements of ξ - and ⌈ T/r ⌉ elements of ξ + (defined in Definitions B.7 and B.8) will be realized.
This means that if sampling from ξ - incurs a unit cost and sampling from ξ + incurs at most r unit cost, total cost will be at most 2 r ⌈ T/r ⌉ .
Before proving Lemma C.1, we review the following technical results.
First, we note an immediate consequence of working with a convex payoff function.
Fact C.1.
Let ϕ : Z → R be a convex function on a convex compact domain Z and g ( t ) =
∂ϕ ( w ( t ) ) be a partial subgradient of ϕ at w ( t ) .
Then, for any [ w ( t ) ]
Tt =1 ∈ Z :
[Equation]
Proof.
Fix any w ∗ ∈ Z .
By our choice of g , we know that
[Equation]
By the convexity of ϕ , it follows that
[Equation]
with equality when ϕ is bilinear.
Fact C.2.
Let ϕ : A - × A + → R be a convex-concave function on convex compact domains A - , A + and define the operators g ( t ) - :=
∂ p ( t ) ϕ (cid:0) p ( t ) , q ( t ) (cid:1) and g ( t ) + := ∂ q ( t ) ϕ (cid:0) p ( t ) , q ( t ) (cid:1) .
Given sequences p (1) , . . .
, p ( T ) ∈
A - and q (1) , . . .
, q ( T )
∈
A + , their ergodic averages p (1: T ) :
= 1 T (cid:80)
Tt =1 p ( t ) and q (1: T ) :
= 1 T (cid:80)
Tt =1 q ( t ) constitute an ε -equilibrium if Errv ( p (1: T ) ) ≤ ε and Errv ( q (1: T ) ) ≤ ε .
Proof.
By Fact C.1, when variational errors are bounded as Errv ( p (1: T ) ) ≤ ε and Errv ( q (1: T ) ) ≤ ε , we know player regrets are bounded: Reg-Min ( p (1: T ) , q (1: T ) ) ≤ ε and Reg-Max ( p (1: T ) , q (1: T ) )
≤ ε .
This satisfies our Definition B.3 for an ε -min-max equilibria.
We now claim concentration results for locally unbiased and globally unbiased first-order oracles.
Fact C.3.
Let ( A - , A + , ϕ ) be a convex-concave game satisfying Definition B.1 and Assumptionsand 2.
Without loss of generality, let our player of interest be the minimizing player.
Consider a play sequence (cid:8) p ( t ) , q ( t ) (cid:9) T t =1 with some complementary sequence (cid:8) y ( t ) (cid:9) T t =1 ∈ A - .
Suppose, at each timestep, the minimizing player uses a random variable (cid:98) g ( t ) - to estimate g ( t ) - :=
∂ p ( t ) ϕ (cid:0) p ( t ) , q ( t ) (cid:1) .
If the following assumptions hold:
For every t ∈ [ T ] , the subsequences (cid:8) p ( τ ) , q ( τ ) , y ( τ ) (cid:9) t τ =1 is independent of (cid:98) g ( t ) - , . . .
(cid:98) g ( T ) - . 2.
All estimates (cid:98) g (1) - , . . .
, (cid:98) g ( T ) - are independent.
(cid:98) g ( t ) - is an unbiased estimate of g ( t ) - and additionally satisfies Assumption 3.
We can then bound the error of the stochastic oracle, ε ( t ) -
[Equation]
−
[Equation]
g ( t ) - , with respect to our play sequence as follows.
With probability at least 1 − δ ,
[Equation]
T
[Equation]
Proof.
Define the filtration (cid:8) F ( t ) (cid:9) T t =0 as the sigma algebra generated by (cid:8)(cid:98) g ( t ) - (cid:9) T t =1 , with F (0) being a singleton containing only the superset of our sigma algebra.
Observe that ε ( t ) - is independent of (cid:8) p ( τ ) (cid:9) t τ =1 and (cid:8) y ( τ ) (cid:9) t τ =1 by assumption.
As (cid:98) g - ( · ) is unbiased, for any t ′ = 0 , . .
, t − 1 :
E (cid:104)(cid:68) ε ( t ) - , p ( t ) − y ( t ) (cid:69) | F ( t ′ ) (cid:105)
= 0 .
We can thus construct the Doob martingale:
[Equation]
and bound the difference sequence accordingly.
For any t ′ ∈ [ T ] , we have the deterministic bound:
[Equation]
where the final inequality is Holder’s.
Since, by Assumption 1, the diameter of our action sets are bounded by R - , we know (cid:13)(cid:13)(cid:13) p ( t ′ )
− y
[Equation]
t ′ )
(cid:13)(cid:13)(cid:13) ≤ R - .
Invoking Assumptions 2 and 3, we know (cid:13)(cid:13)(cid:13)
[Equation] [Equation]
≤ 2 C - .
By the Azuma-Hoefdding inequality, we can thus bound, for any ε > 0 ,
[Equation]
Fact C.4.
Let ( A - , A + , ϕ ) be a convex-concave game satisfying Definition B.1 and Assumptionsand 2.
Without loss of generality, let our player of interest be the minimizing player.
Consider a play sequence (cid:8) p ( t ) , q ( t ) (cid:9) T t =1 .
If the following assumptions hold:
(cid:98) g ( t ) - is an estimate of g ( t ) -
[Equation]
p
[Equation]
) , q ( t ) (cid:1) that satisfies Assumption 3. 2.
There exists the no-regret algorithm Q - satisfying the assumptions of Lemma C.1.
We can then bound the error of the stochastic oracle, ε ( t ) - := g ( t ) - − (cid:98) g ( t ) - , with respect to our play sequence as follows.
Define y ( t +1) :
= Q - (cid:16)(cid:8) y
( τ ) , ε ( τ ) - (cid:9) t τ =1 (cid:17) .
With probability at least 1 − δ ,
[Equation]
T
[Equation]
Proof.
We can rewrite Equation 11 with respect to a sequence max p ∗ ∈
A -
[Equation]
We will first bound the summand:
[Equation]
By definition, our sequence y (1) , . . .
y ( T ) is a Q - -learned sequence for the operator errors ε (1) - , . . .
, ε ( T ) - .
By Assumptions 2 and 3, we enforce that all operators and operator estimates have bounded norm, i.e.,
∥ g - ( · ) ∥
E ∗ ≤ C - and ∥ (cid:98) g - ( · )
∥
E ∗ ≤ C - .
By triangle inequality, we can bound (cid:13) (cid:13)(cid:13) ε ( t ′ ) - (cid:13) (cid:13)(cid:13) ∗ ≤ 2 C - .
Hence, the guarantees of Q - imply:
[Equation]
We now prove our general claim.
Proof of Lemma C.1.
By Fact C.2, it suffices to prove the variational error bounds for each player:
[Equation]
with respect to the operators,
[Equation]
In Algorithm 3, we estimate the true operators (cid:8) g ( t )
[Equation]
T t =1 ,
[Equation]
g
[Equation]
t )
+
[Equation]
T t =1 with the stochastic estimates:
[Equation]
Let ε ( t ) - := g ( t ) - − (cid:98) g ( t ) - and let ε ( t ) + := g ( t ) + − (cid:98) g ( t )
+ denote the difference between our true and estimated operators at each timestep.
We can thus divide each variational error into a training error and generalization error component:
[Equation]
We handle the training error first.
Recall that p (1) , . . .
p ( T ) is a Q - -learned sequence for the operator sequence (cid:98) g (1) - , . . .
, (cid:98) g ( T ) - .
By Assumption 3, we enforce that all operator estimates have bounded norm, i.e., ∥ (cid:98) g - ( · )
∥
E ∗ ≤ C - .
Hence, the guarantees of Q - imply:
[Equation]
Similarly, q (1) , . . .
q ( T ) is a Q + -learned algorithms enjoying Q + ’s guarantee:
[Equation]
We now handle the generalization error.
We first consider the minimization player.
Observe that, for every t ∈ [ T ] , the play sequence (cid:8) p ( τ ) , q ( τ ) (cid:9) t τ =1 is measurable by (cid:110) ξ q ( τ ) - ; τ , ξ ⊥ + ; ⌈ τ/r ⌉ (cid:111) t − 1 τ =1 , which ξ q ( t ) - ; t is independent of by construction.
We can thus invoke Facts C.3 and C.4 to bound, with probability at least
[Equation]
We now consider the maximization player.
First, we invoke Fact C.4 to separate:
[Equation]
For notional convenience, let i ( j ) =
( j − 1) r
+ i denote the i th timestep of the j th period.
Also let m i :=
(cid:12)(cid:12)(cid:12) { i ( j ) } ∞
j =1 (cid:83) [ T ] (cid:12)(cid:12)(cid:12) denote the number of valid timesteps that can be written as i ( j ) .
Observe that m i ≤ ⌈ T/r ⌉ .
Fix a choice of i ∈ [ r ] .
Observe that, for every k ∈ [ m i ] , the play sequence (cid:8) p ( i ( j )) , q ( i ( j )) (cid:9) k
j =1 is measurable by (cid:110)
ξ q ( i ( j )) - ; i ( j ) , ξ ⊥ + ; j − 1 (cid:111)
k − 1 j =1 , which ξ ⊥ + ; j is independent of by construction.
We can thus again invoke Fact C.3 to bound, with probability at least 1 − δ :
[Equation]
j = (cid:68) ε ( i ( j ))
+ , q ( i ( j ))
− y
[Equation]
i ( j )) (cid:69) ≤ (cid:115) 8 R
[Equation]
C
[Equation]
+ log (cid:0) rδ (cid:1) m i .
Taking a union bound over said Azuma inequality for all i ∈
[ r ] , we have that with probability at least 1 − δ ,
[Equation]
Gluing together our bounds on training error (Equation 12, Equation 13) and generalization error (Equation 14, Equation 15, Equation 16) with triangle inequalities and union bounds, we have with probability at least 1 − δ ,
[Equation]
C.2 Proof of Theorem C.2 (Generalization of Theorems 4.1 and 5.1)
Fact C.5.
Let (Θ , D , ℓ ) be a multi-distribution learning problem satisfying Definition B.12.
Define a corresponding convex-concave game ( A - , A + , ϕ ) where:
[Equation]
The minimizing player’s mixed strategy p in any ε -min-max equilibria (Definition B.3) constitutes an 2 ε -error solution to (Θ , D , ℓ ) .
Proof.
If ( p, q ) is an ε -min-max equilibria, the following holds by definition
[Equation]
Equivalently, by the min-max theorem,
[Equation]
Input: Parameter space Θ with distance generating function ω , distribution set D with n := |D| , and loss function ℓ : Θ × Z → [0 , L ] , all satisfying Definition B.13 and Assumptions 5 and 4;
Initialize: minimizing iterate p (1) =
Θ o where θ o is as defined in Definition B.9, maximizing iterate q (1) =
[1 /n ] n , and iteration cap:
[Equation]
; for a = 1 , 2 , . . . , ⌈ T/n ⌉ do
For all D ∈ D , sample datapoint z aD from EX ( D ) ; for t = an + 1 − n, . . .
, an do
Sample datapoint z ( t ) from EX
; Define
end end for
Theorem C.2 (Generalization of Theorems 4.1 and 5.1) .
Algorithm 4 is an ( ε, δ ) multi-distribution learning algorithm for any convex multi-distribution learning problem (Θ , D , ℓ ) satisfying Definitions
B.12 and B.13 and Assumptions 5 and 4.
In other words, Algorithm 4 returns an θ ∈ Θ such that:
[Equation]
Furthermore, the sample complexity of Algorithm 4 is in O (cid:16) D Θ C
+( R 2 − C 2 + nL 2 )log( n/δ ) ε(cid:17) .
Proof.
The sample complexity of Algorithm 4 is immediate from its construction.
Every period a , Algorithm 4 samples n datapoints.
Every iteration t , Algorithm 4 samples 1 datapoint.
Thus,
Algorithm 4 samples 2 n ⌈ T/n ⌉ datapoints exactly.
We now prove that Algorithm 4 is an ( ε, δ ) -learning algorithm for any convex multi-distribution learning problem.
We begin by constructing the following convex-concave game ( A - , A + , ϕ ) where:
[Equation]
We observe that this game satisfies Definition B.1 and Assumptions 1 and 2:
Definition B.13 defines Risk q ( p ) —and by extension ϕ ( p, q ) —to be convex in p .
As Risk q ( p ) :
= (cid:80) D ∈D q D Risk D ( p ) by definition, ϕ ( p, q ) is linear and thus also concave in q . 3.
In the l-1 norm, ∆ D satisfies Assumption 1 with diameter R + = 2 . 4.
Since Θ has finite Bregman radius of D Θ by Assumption 5 and ω is strongly convex modulus 1 by definition, Θ satisfies Assumption 1 with a finite R Θ ≤ 2 √ 2 D Θ . 5.
Since ∂ q ϕ ( p, q ) =
[Risk D ( p )]
D ∈D and the range of ℓ is [0 , L ] , Assumption 2 is satisfied for ∂ q ϕ ( p, q ) by C + ≤ L in the l-infinity norm.
∂ p ϕ ( p, q ) satisfies Assumption 2 for some finite C - = C directly by Assumption 4.
We now define a stochastic setting for our game.
Let the minimizing player’s randomness source be given by the sequences ξ q - = { EX ( q ) i } ∞ i =1 ; recall that EX ( D ) k refers to the k th call to an example oracle for a D ∈ ∆ D .
Let the maximizing player’s randomness source be given by the sequence ξ ⊥ +
= { [ EX ( D ) i ]
D ∈D } ∞ i
=1 .
Next, define the first-order oracle estimators:
[Equation]
We can observe that (cid:98) g - is a locally unbiased first-order oracle (satisfying Definition B.7) and (cid:98) g + is a globally unbiased first-order oracle (satisfying Definition B.8), with both (cid:98) g - and (cid:98) g + satisfying
Assumptions 3.
By the unbiasedness of empirical risk estimates, (cid:98) g
+ is globally unbiased as returns an empirical risk sample for each D ∈ D .
Similarly, by the unbiasedness of empirical risk estimates and linearity of derivatives, (cid:98) g - is locally unbiased.
As the range of loss function ℓ is in [0 , L ] , empirical loss is also bounded in [0 , L ] , (cid:98) g
+ satisfies Assumption 3 with C + ≤ L in the l-infinity norm.
By Assumption 4, empirical partial subgradients are norm-bounded by some finite C , so (cid:98) gsatisfies Assumption 3 with some finite C - = C .
Finally, we observe that Algorithm 4 is equivalent to instantiating Algorithm 3 on our constructed game ( A - , A + , ϕ ) for our constructed stochastic setting.
We will now rewrite the iteration complexity requirement of Lemma C.1 given by Equation 10 (copied below):
[Equation]
In particular, we aim to show that the default iteration setting of Algorithm 4 satisfies it for ε ′ =
ε/ 2 .
By Lemmas C.4 and C.10, we can bound the efficacy of our no-regret algorithms Q omd ,, Q hedge by:
[Equation]
where γ - ( T, A - , C - ) and γ + ( T, A + , C + ) are as defined in Lemma C.1.
Accounting for our previous derivations of C - , C + , ε ′ , R + , to satisfy Equation 10, it suffices to set:
[Equation]
or simplified further:
[Equation]
Thus, by Lemma C.1, q := 1 T (cid:80) Tt =1 q ( t ) and θ , the output of Algorithm 4, form an ε 2 -min-max equilibria of our game ( A - , A + , ϕ ) with probability at least 1 − δ .
The Theorem then follows by Fact
C.5.
The following theorems, which are restated from the main text, are immediate corollaries of Theorem
C.2.
Theorem 4.1.
For any finite hypothesis class H and unknown set of distributions D , with probability 1 − δ , Algorithm 2 returns a distribution
p ∈
∆ H such that
[Equation]
using a number of samples that is O (cid:16) log |H| + n log
[Equation]
n
[Equation]
δ
[Equation]
Proof.
Observe that this finite multi-distribution learning problem can be re-written as the convex multi-distribution learning problem (∆ H , D , ℓ ) .
Since ∆ H is a probability simplex of dimension |H| , we know it is compact, convex, with C = 1 and L = 1 (as the range of ℓ is in [0 , 1] ), and with
R Θ ≤ 2 .
We can then directly apply Theorem C.2, observing that Algorithm 2 is equivalent to
Algorithm 4 in this setting.
Theorem 5.1.
Consider a group distributionally robust problem (Θ , D ) with convex
compact unitdiameter parameter space Θ of Bregman radius D Θ (Definition B.11), and convex loss ℓ : Θ × Z → [0 , C ] .
A variant of Algorithm 2 (in particular Algorithm 4 in Appendix 4.1), returns θ ∈ Θ such that
[Equation]
, using a number of samples that is O (cid:16)
[Equation]
Proof.
Similarly to Theorem 4.1, this claim follows immediately from Theorem C.2 for unit diameter
R Θ = 1 and loss bound L = C .
Corollary 5.2 follows in a similar fashion, running Algorithm 4 on empirical data distributions.
The following proposition re-states this formally.
Proposition C.1 (Generalization of Corollary 5.2) .
Let (Θ , D , ℓ ) be a convex multi-distribution learning problem satisfying Definitions B.12 and B.13 and Assumptions 5 and 4.
For every D ∈ D , let B D ∼ D be a non-empty batch of i.i.d.
datapoint samples.
Define D ′ =
{ D ′ } D ∈D , where
D ′ is the empirical distribution of B D .
It follows that (Θ , D ′ , ℓ ) also satisfies Definitions B. and B.13 and Assumptions 5 and 4 with identical parameters.
Thus, Algorithm 4, when applied to (Θ , D ′ , ℓ ) , with probability at least 1 − δ returns an θ with a multi-distribution training error of at most ε .
Furthermore, the number of iterations—and accordingly partial derivative operations—is in
O
[Equation]
A similar proof as Theorem C.2 gives an analagous result for binary classification problems with a hypothesis class of finite Littlestone dimension.
Theorem C.3 (Littlestone Variant of Theorem 4.1) .
Let Q Little denote the no-regret algorithm that achieves a regret of Θ( (cid:112) Little ( H ) T ) in any online learning setting with Littlestone dimension
Little ( H ) ; such an algorithm exists by Theorem 2.4 in Alon et al.
Replacing Q omd ,ω with Q Little in Algorithm 4 and updating the iteration cap to,
[Equation]
yields an ( ε, δ ) multi-distribution learning algorithm for any binary classification multi-distribution learning problem ( H , D , ℓ ) satisfying Definitions B.12 and B.15.
Further assume that the hypothesis set has finite Littlestone dimension Little ( H ) < ∞ .
The sample complexity of Algorithm 4 is in
O (cid:16) Little ( H )+ n log
[Equation]
n
[Equation]
Proof.
The sample complexity of our modified Algorithm 4 remains immediate from its construction.
Algorithm 4 samples 2 n ⌈ T/n ⌉ datapoints exactly.
Theorem C.2’s proof that Algorithm 4 is an ( ε, δ ) -learning algorithm for multi-distribution learning also continues to hold.
However, we must update the iteration complexity requirement of Lemma C.1 given by Equation 10 (copied below):
[Equation]
In particular, we aim to show that the default iteration setting of Algorithm 4 satisfies it for ε ′ =
ε/ 2 .
By Lemmas C.4 and by assumption on Q Little , we can bound the efficacy of our no-regret algorithms
Q Little , Q hedge by:
[Equation]
where γ - ( T, A - , C - ) and γ + ( T, A + , C + ) are as defined in Lemma C.1, and C ′ is some universal constant.
Accounting for our previous derivations of C - , C + , ε ′ , R + for Theorem 4.1, it suffices to set:
[Equation]
Thus, by Lemma C.1, q := 1 T (cid:80) Tt =1 q ( t ) and θ , the output of Algorithm 4, form an ε 2 -min-max equilibria of our game ( A - , A + , ϕ ) with probability at least 1 − δ .
The Theorem then follows by Fact
C.5.
We remark that a similar sample complexity bound can be achieved using the Standard Optimal Algorithm (SOA) of Littlestone [24] instead of the implicit algorithm of Alon et al.
[1], as SOA achieves a regret of O ( (cid:112) Little ( H ) T log( T )) .
C.3 Proof of Theorem 4.
We now provide matching lower bounds for collaborative PAC learning.
We first define a notion of expected sample complexity.
Take any multi-distribution learning problem
V =
(Θ , D , ℓ ) .
Recall that, on this problem, the input to any multi-distribution learning algorithm is a random variable of form (cid:98) V
=
(Θ i , ℓ i , { EX ( D )
} D ∈ ∆ D ) .
Also recall that each example oracle
EX ( D ) is an infinite sequence of i.i.d.
samples from D .
We will let X V denote the probability distribution of the random variable tuple (Θ i , ℓ i , { EX ( D ) } D ∈ ∆ D ) .
Further let N A ( (cid:98) V ) denote the expected sample complexity of A given inputs (cid:98) V , where expectation is taken over any randomness from the algorithm A itself.
We can now define a general notion of expected sample complexity.
Definition C.1.
Let A be a multi-distribution learning algorithm and P a probability distribution over a set of multi-distribution learning problems V :=
{ (Θ i , D i , ℓ i ) } i .
We define the expected sample complexity N A ( P ) as:
[Equation]
The outer expectation is taken over the randomness of the problem selection, the inner expectation is taken over the randomness of datapoints, and N A ( (cid:98) V ) takes an expectation over the internal randomness of the algorithm A .
Unless otherwise specified, we will use the shorthand: E V ∼ P (cid:104) E (cid:98) V ∼ X V [ · ] (cid:105)
=
E (cid:98) V [ · ] .
We recall the following theorem from Section 4.2.
Theorem 4.2.
Take any n, d ∈ Z
+ , ε, δ ∈ (0 , 1 / 8) , and ( ε, δ ) -collaborative learning algorithm A .
There exists a collaborative learning problem ( H , D ) with |D| = n and |H| =
2 d , on which A takes at least Ω (cid:0) 1 ε 2 (log |H| + |D| log(min {|D| , log |H|} /δ ))
(cid:1) samples.
We now prove two lemmas, Lemma C.4 and Lemma C.5, that directly imply Theorem 4.2.
These lower-bound constructions are fairly generous and allow all distributions to share the exact same feature distribution and all but one distribution to share the exact same label distribution.
Lemma C.4.
Take any n, d ∈ Z
+ , ε, δ ∈ (0 , 1 / 8) , and ( ε, δ ) -collaborative learning algorithm A .
There exists a set of collaborative learning problems V on which A takes at least Ω (cid:16) log |H| ε 2 (cid:17) samples and where, for every ( H , D ) ∈ V , |D| = n and |H| = 2 d .
Proof.
This claim follows directly from the lower bound on sample complexity of agnostic probablyapproximately-correct (PAC) learning [41].
Let ( H , D ) be an agnostic PAC learning problem.
Accordingly define the collaborative learning problem ( H , D ) , where D = { D } ni =1 .
Observe that min D ′ ∈D
Risk D ′ ( h ) =
Risk D ( h ) for all well-defined choices of h .
Thus, given an algorithm A that ( ε, δ ) solves ( H , D ′ ) with at most m samples, we can design an algorithm B that ( ε, δ ) solves ( H , D ) : run algorithm A by simulating samples from any D ′ ∈ D with samples from D and return the output of A .
We can thus invoke the well-known lower bound of agnostic PAC learning to observe that there exists an agnostic PAC learning problem V ′
such that any ( ε, δ ) -learning algorithm has a sample complexity of Ω (cid:16) log |H| ε 2 (cid:17) [15]; we defer interested readers to Zhang [47] for a constructive proof of the existence of V ′ .
Thus, there exists a V satisfying the assumptions of Lemma C.4 where any ( ε, δ ) collaborative learning algorithm has a sample complexity of Ω (cid:16) log |H| ε 2 (cid:17) .
Lemma C.5.
Take any n, d ∈ Z
+ , ε, δ ∈ (0 , 1 / 8) , and ( ε, δ ) -collaborative learning algorithm A .
There exists a set of collaborative learning problems V on which A takes at least Ω (cid:0)
ε 2 ( |D| log( k/δ )) (cid:1) samples and where, for every ( H , D ) ∈ V , |D| = n and |H| =
2 d with k :=
min { n, d } .
Proof.
We prove this constructively.
We begin by defining collaborative learning problem sets V u,w for all w, η ∈ N with u = w · η .
Problems in V u,w share a feature space X
= { 1 , . .
. , w } , label space
Y = { + , −} , and hypothesis class H = { f : X → Y} (the set of all deterministic binary labeling functions).
For every i ∈ [ w ] , define distributions D − i and D
+ i as:
[Equation]
Let P ′
i be a probability distribution over possible choices of η collaborative learning problems.
With 1 / 2 probability, P ′ i returns a set of η copies of D − r .
With 1 / 2 probability, P i returns a uniformly sampled shuffling of the set of η − 1 copies of D − r and one copy of D + r .
We then define P u,w = P ′ 1 × · · · × P ′ w with V u,w being the support of P u,w .
Observe that P u,w is a distribution over collaborative learning problems each with |H| = 2 w and u distributions.
The following claims characterize sample complexity lower bounds on P u,w .
Claim C.1.
Choose ε ∈ (0 , 1 / 2) and δ ∈ (0 , 1) .
Take any ( ε, δ ) -learning algorithm A on V η, 1 under P 1 , 1 , or equivalently, on P ′ 1 .
Then, the expected sample complexity (see Definition C.1) of A is at least
η 256 ε 2 log(1 / 2 δ )
Claim C.2.
Choose ε ∈ (0 , 1) and δ ∈ (0 , 1 / 8) Suppose there exists an ( ε, δ ) learning algorithm A for V u,w with an expected sample complexity of m under P u,w .
Then there exists an ( ε, 8 δ 7 w ) learning algorithm A ′ for V 1 , 1 under P 1 , 1 with an expected sample complexity on P 1 , 1 of 8 7 w m .
Since our desired lower bound are trivially weakly monotonic in n, d , fix a choice of η, d ∈ Z + and ε, δ ∈ (0 , 1 / 8) with n = η · d .
Combining claims C.1 and C.2, we see that any ( ε, δ ) collaborative learning algorithm A for V n,d has an expected sample complexity on P n,d of at least m ≥ 7 n 2048 ε 2 log (cid:0) 7 d 16 δ (cid:1) .
By the probabilistic method, for at least some collaborative learning problem
V ∈ V n,d , A must have a sample complexity in Ω (cid:0) 7 n 2048 ε 2 log (cid:0) 7 d 16 δ (cid:1)(cid:1) .
Proof of Claim C.1.
Consider the following problem: take η identical-looking coins.
Under the
H 0 hypothesis, all coins are biased tails with probability 1 / 2 + 2 ε .
Under the H i hypothesis, the i th coin is instead biased heads with probability 1 / 2 + 4 ε .
Let Pr be a probability distribution on
H ∈ { H i } ηi =0 with Pr( H 0 )
=
1 / 2 and Pr( H 1 ) =
· · · = Pr( H η ) = 12 η .
Suppose that A is an ( ε, δ ) -algorithm for V η, 1 with an expected sample complexity of m under P ′ 1 .
Then, we can construct a coin algorithm A ′ that, conditioned on any H ∈ { H i } ηi =0 , with probability at least 1 − δ , can identify whether H 0 is false.
Furthermore, under hypothesis distribution Pr , A ′ will also have an expected sample complexity of exactly m .
To see this, have A ′ run A by simulating draws from the i th distribution by flipping the i th coin.
If all coins are biased tails with probability 1 / 2 + 2 ε , any ε -error hypothesis h must satisfy Pr( h (1) =
+) > 1 / 2 .
Conversely, if one coin is biased heads, any ε -error hypothesis h must satisfy Pr( h (1) =
+) < 1 / 2 .
It thus suffices to lower bound the sample complexity of A ′ , which is reminiscent of the pure-exploration multi-armed bandit problem.
Suppose A ′ , conditioned on H 0 , correctly predicts H 0 with probability at least 1 − δ .
Then, suppose
A ′ , under H 0 , takes no more than T i flips from the i th coin.
Let p i,j 1 : j 2 be a probability distribution over { 0 , 1 } corresponding to the outcomes of the j 1 st to j 2 nd coin toss by A ′ under H i .
Let p ∗ j be a uniform distribution over { 0 , 1 } j .
We observe that p i,j : j and p ∗ j are Bernoulli distributions with a parameter within 4 ε of 1 / 2 .
A standard information-theoretic result is that, for ε < 1 / 2 ,
KL ( p i,j : j , p ∗ 1 ) <
128 ε 2 ; e.g. see Zhang [47] and Karp and Kleinberg [23].
By the tensorization of KL Divergence, KL ( p i, 1: j , p ∗ j ) < 128 jε 2 .
By Pinsker’s inequality, we can bound total variation distance by TV ( p i, 1: j , p ∗ j ) ≤ 8 ε √ j .
Let E be the set of outcomes of T i flips under which A ′ predicts
H 0 ; by correctness under H 0 , we have that Pr H
( E ) ≥ 1 − δ .
Thus, total variation distance implies 1 − δ − 8 ε √ j <
Pr H
i ( E ) .
Since Pr H i ( E ) < δ , we have that
[Equation]
Thus, if A ′ is δ accurate under all hypotheses, under H 0 , A ′ must take at least
[Equation]
samples from each distribution.
Thus, the average sample complexity of A ′ under Pr —and similarly the sample complexity of A under P ′ 1 , must be at least
[Equation]
Proof of Claim C.2.
This claim is similar to the lower bounds of Blum et al.
[6] and Karp and
Kleinberg [23].
We construct A ′ as follows.
Let I j
=
[( j − 1) η + 1 , jη ] .
Suppose a problem
V ′ ∈ V η, 1 is drawn.
A ′ draws a problem ( H , D )
∈ V u,w and chooses an index i ∈ [ w ] uniformly at random.
A ′ simulates the algorithm A on ( H , D ) ; when A tries to sample a datapoint from distribution D r with r ∈
I i , return a sampled datapoint from the ( r − ( i − 1) η ) st data distribution of V ′ .
When A terminates and returns a classifier h , A ′ checks whether, for every j ̸
= i : max r ∈ I
j Risk D r ( h ) < 12 .
If this condition is satisfied, A ′ returns h ′ (1) =
h ( ℓ ) .
If not, we repeat from Step 1.
We denote the number of total iterations with T .
Consider the probability p i
that, in the third step, for every j ̸ = i we have max r ∈
I j Risk D r ( h ) <
but max r ∈
I i Risk D r ( h ) ≥ 12 .
Let E t denote the event that A ′ returns an at least ε -error hypothesis after t iterations of our procedure.
Noting that E t can only occur if A failed all t − 1 iterations before and at the t th iteration, Step 3 fails to catch the bad hypothesis for D i .
By assumption, δ ≥ (cid:80) wi =1 p i .
By symmetry of our construction V and recalling δ < 1 / 8 :
[Equation]
Thus, A ′ is an ( ε, 8 δ 7 w ) -algorithm for P η, 1 .
We now bound the sample complexity of A ′ .
Let N A ′ ( t ) denote the number of samples that A ′ takes from V ′ on the t th iteration.
Note that N A ′ (1) , N A
are i.i.d.
In addition, by the symmetry of V and linearity of expectation, E V ′ ∈
P η, 1 [ N A ′ ( t )]
= m/w .
Thus, we can write:
[Equation]
We can upper bound T by observing that our procedure only repeats if A fails.
Thus,
[Equation]
Thus, A ′ has an expected sample complexity of at most 8 m 7 w .
The following is a more general restatement of Theorem 4.2 in terms of the terminology of Section
B. It follows by observing the difficult cases described in Theorem 4.2 constitute challenging cases for both convex multi-distribution learning (Definition B.13) and binary classifier multi-distribution learning (Definition B.15).
Corollary C.6.
Take any n, m ∈ N and ε, δ ∈ (0 , 1 / 8) .
There exists a finite set V of multi-distribution learning problems where:
Every (Θ , D , ℓ ) ∈
V satisfies Definitions B.13 and B.15, with |D| = n and D Θ = log( m ) .
Every ( ε, δ ) -algorithm
A has a sample complexity in (cid:16) D Θ + n log(min { n,D Θ } /δ )
ε
[Equation]
C.4 Proof of Lemmas C.4 and C.
For completeness, this section includes standard results on exponentiated gradient descent and mirror descent more generally, proofs of which can be found in [42].
Lemma C.7.
Pinsker’s inequality For any two vectors from the same probability simplex w,
w ′ ∈ ∆ n , we can bound their generalized Kullback-Leibler divergence as,
[Equation]
Lemma C.8.
Law of Cosines Define x, y, z ∈ Z where Z is a convex set, and let H :
Z → R be a strictly convex differentiable distance generating function and D H the Bregman divergence of H .
Then,
[Equation]
Lemma C.9.
Pythagorean theorem for Bregman Divergence Define x, y ∈
Z 0 where Z 0 is a closed subset of a convex set Z , z ∈ Z .
Let H : Z → R be a strictly convex differentiable distance generating function, and let D H be the Bregman divergence of H .
If y = arg
min u ∈ Z 0 D H ( u, z ) , then D H ( x, y )
+ D H
( y, z ) ≤ D H ( x, z ) .
We now turn to proving Lemma (restated below), which concerns exponentiated gradient descent with bounded gradients.
Lemma 2.1 ([42]) .
Let g (1) , . .
, g ( T ) ∈
R n and Z
= ∆ n .
Further assume (cid:13)(cid:13) g
( t ) (cid:13)(cid:13) ∞ ≤ C for all timesteps t = 1 , . . .
, T .
Choosing η = (cid:112) log n/T , after T iterations of exponential gradient descent, the outputs w (1) , . .
. , w ( T ) satisfies,
[Equation]
Proof.
This proof closely follows that of Theorem 7.5 in [42].
Fix t ∈ 1 , . .
, T .
First, we provide an expression for g ( t ) in terms of (cid:101) w ( t +1) and w ( t ) , where (cid:101) w is as defined in Equation 4.
For all i ∈ 1 , . . . , n , we have by Equation 4:
[Equation]
Equivalently,
[Equation]
Letting H ( x ) =
(cid:80) n i =
1 x i log x
i − x i denote our distance generating function, generalized negative entropy, we can also write,
[Equation]
where logs are applied coordinate-wise.
Defining KL ( ·||· ) as generalized Kullback–Leibler divergence: the Bregman divergence with respect to our choice of H as a distance generating function.
Since Z is already closed, by Lemma C.8, for any w ∗ ∈ Z ,
[Equation]
Generalized Pythagorean Theorem, e.g. Theorem 7.7 in [42] gives,
[Equation]
Then we can bound,
[Equation] [Equation]
To bound the second term, we again apply the law of cosines, this time in reverse order, recovering,
[Equation]
As w ( t +1) , w ( t ) ∈
Z , by Pinsker’s inequality (Lemma C.7),
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
where the final inequality follows from maximizing the quadratic ηCz − z 2 2 , attained at z = (cid:13)(cid:13) w ( t ) − w ( t +1)
(cid:13)(cid:13)
= Cη .
The claim follows by plugging Equation 18 into Equation 17.
Exponentiated gradient descent is a special case of mirror descent in the Euclidian space E =
R n equipped with an L1-norm ∥·∥ 1 , over the probability simplex Z = ∆ n , and using entropy as a distance generating function.
The following lemma generalizes Lemma C.4 to more general Euclidian spaces, choices of convex compact subsets, and strongly-convex distance generating functions.
As the proof closely mirrors that of Lemma C.4, we defer interested readers to Beck and Teboulle [3].
Lemma C.10 (Generalization of Lemma C.4 [3]) .
Let Z be a convex compact subset of a Euclidean space E with distance generating function ω satisfying Definition B.9.
Let g (1) , . .
, g ( T )
∈
E ∗ .
Further assume (cid:13)(cid:13) g
( t )
(cid:13)(cid:13)
E ∗ ≤ C for all timesteps t = 1 , . . .
, T .
Choose step size η = (cid:113) D Z T where
D Z is the Bregman radius of Z .
After T iterations of online mirror descent [3], the output { w } Tt = satisfies,
[Equation]
C.5 Proof of Theorem C.
This section discusses the implications of our results for finite VC dimension problems.
First, we will use D X to denote the marginal distribution of a data distribution D , i.e. the distribution of D over its feature space.
We also introduce the following definitions.
Definition C.2.
The Renyi divergence D α ( P || Q ) between discrete distributions P, Q is defined by:
[Equation]
and between continuous distributions P, Q as:
[Equation]
We will write d α ( P || Q ) :
= 2 D α
( P || Q ) .
Remark C.1.
Denoting the support of Q as X Q , we can write sup x ∈X Q P ( x ) Q ( x ) =
d ∞ ( P || Q ) .
Recall that in Theorem C.2 we describe a multi-distribution learning algorithm
(Algorithm 4) with provably tight sample complexity upper bounds for convex multi-distribution learning problems (Definition B.13).
We note that there is one class of multi-distribution learning problems, non-convex finite VC multi-distribution learning , that has been previously studied by [6, 31, 9] but does not satisfy the assumptions of convex multi-distribution learning.
A non-convex finite VC multi-distribution learning problem is a binary-classification multi-distribution learning problem (Definition B.15) that satisfies three criteria: the hypothesis space H is non-convex, of infinite size, and of finite
VC dimension VCD( H ) < ∞ .
[6, 31, 9] provide upper bounds for non-convex finite VC multidistribution learning that are identical to their upper bounds in Table 1 but replacing log |H| with
VCD( H ) .
In contrast, our Theorem C.2 upper bounds do not directly apply to non-convex finite VC multidistribution learning.
However, a similar result can be obtained by running our Algorithm 4 on a probability simplex ∆ H ′ over some ε -covering H ′ of H .
Such ε -nets of size nε − 2VCD( H ) necessarily exist (see, e.g., [2]).
For example, given an ε -net for each distribution D ∈ D , we may take their union as the covering H ′ and run Algorithm 2.
This directly inherits a favorable upper bound from
Theorem C.2.
Corollary C.11 (VC Dimension Corollary of Theorem C.2) .
Consider any binary classification multi-distribution learning problem ( H , D ) where the hypothesis set H is of finite VC dimension d and the unknown distribution set is of size |D| = n .
There is an algorithm that, given an ε -net of size poly (cid:0) ε d , ε, d, n (cid:1) for each distribution, with probability 1 − δ , returns a distribution
p ∈
∆ H with,
[Equation]
D
[Equation]
using a number of samples that is O (cid:16) d log(1 /ε )
+ n log( n/δ )
ε
[Equation]
It is also not strictly necessary to know an ε -net in advance.
Instead, one can compute a net from samples or from other information about distributions in D .
Theorem C.12 formalizes this claim.
Theorem C.12.
For any H of VC dimension d and unknown set of distributions D for which
Assumption 1 or 2 is met, there is an algorithm that, with probability 1 − δ , returns a distribution
p ∈
∆ H with,
[Equation]
a number of samples that is O
[Equation]
The following lemmas directly imply Theorem C.12.
Lemma C.13 (Assumption 1) . Consider any binary classification multi-distribution learning problem ( H , D ) where the hypothesis set H is of finite VC dimension d and the unknown distribution set is of size |D| = n .
There is an algorithm that, given access to the marginal distribution D X of every
D ∈ D , with probability 1 − δ , returns a distribution
p ∈
∆ H with,
[Equation]
, using a number of samples that is O (cid:16) d log(1 /ε )
+ n log( n/δ ) ε(cid:17) .
Proof.
By uniform convergence, sampling Θ (cid:16) d log( d/ε )+log(1 /δ )
ε(cid:17) i.i.d.
samples from a distribution D X , with probability at least 1 − δ , yields an ε -covering on D .
By Sauer-Shelah’s lemma, the resulting covering H
′ D is of size O (cid:0) (log( d/ε ) + 1 d log(1 /δ )
/ε 2 ) d (cid:1) .
Repeating this procedure for each D ∈ D , with probability at least 1 − nδ , we have an ε -covering of D of size
O (cid:0) n (log( d/ε ) + 1 d log(1 /δ )
/ε
) d (cid:1) ∈ poly (cid:0) ε d , ε, d, n (cid:1) and can appeal to Corollary C.11.
Lemma C.14 (Assumption 2) .
Consider any binary classification multi-distribution learning problem ( H , D ) where the hypothesis set H is of finite VC dimension d and the unknown distribution set is of size |D| = n .
We say an algorithm has weak unlabeled access if the algorithm can access, for each
D ∈ D , a marginal distribution D ′X such that d ∞ ( D ′X || D X ) ∈
poly(1 /ε, d, n ) , with probability 1 − δ .
There is an algorithm that, given weak unlabeled access, with probability 1 − δ , returns a distribution
p ∈
∆ H with,
[Equation]
using a number of samples that is O (cid:16) d log( dn/
[Equation]
+ n log
[Equation]
n/ [Equation]
Proof.
Observe that when d ∞ ( D ′X || D X ) < γ , D ′X can be written as a mixture over D X with probability at least 1 γ and some other distribution (cid:101) D X with probability at most 1 − 1 γ .
Once again invoking uniform convergence, we observe that sampling Θ (cid:16)
[Equation]
samples from distribution D ′X , with probability at least 1 − δ , yields an ε -covering on D .
By
Sauer-Shelah’s lemma, the resulting covering H
′ D is of size O (cid:0) (poly(1 /ε, d, n ))
d (cid:1) .
Repeating this procedure for each D ∈ D , with probability at least 1 − nδ , we have an ε
-covering H ′ of D of size |H ′
| ∈ O (cid:0) n (poly(1 /ε, d, n ))
d (cid:1) .
We can then appeal directly to Theorem 4.1 for a sample complexity bound on learning ( H ′ , D ) .
Input: Parameter
, initial parameter
, training set validation set
D
D where n
= |D| , iterations T , minibatch size B , adversary minibatch size B
, blackbox minibatch learning algorithm A ;
Initialize: q (0) =
[1 /n ] n ; for iterations t
= 1 ,
, T do
For each D ∈ D , sample B ′ /n points X ′ ( t ) D from the empirical distribution of X ′ D .
In other words, sample uniformly with replacement from X ′ D .
Run Hedge update on adversary using datasets [ ℓ ( θ, X ′ ( t ) D )]
D ∈D .
In other words, letting exp be component-wise exponentiation, · be component-wise multiplication, and η some learning rate:
[Equation]
[Equation]
Sample B datapoints from the mixture over the empirical distributions of {X D } D ∈D weighted by q ( t ) .
In other words, sample with replacement from X D with probability q ( t ) D .
Run minibatch update on learner, with θ ( t ) =
A ( θ ( t − 1) , X ( t ) ) .
end for
Return: θ ( T ) .
D Experiment Details
R-MDL Algorithm
The R-MDL algorithm is defined in full in Algorithm 5.
This algorithm is implemented in the Github repository ericzhao28/multidistributionlearning.
Additional Observation: R-MDL converges faster than ERM or GDRO.
The R-MDL methods in Table 2 used a fraction of the training epochs that their GDRO counterparts used.
The ratio of
R-MDL to GDRO training epochs is 1:3, 2:5, 1:2 on the Waterbirds, CelebA, and MultiNLI datasets respectively.
This fast convergence rate is predicted by our theory, particularly Corollary 5.2.
In our Figure 1, we also replicate the Figure 2 of Sagawa et al.
[38], appending our additional results on R-MDL.
We again see a trend of faster test error convergence (solid lines) and more uniform per-group risks by the R-MDL algorithm.
Datasets Our experiments were performed on three datasets: Multi-NLI, CelebA, and Waterbirds [38].
We use identical preprocessing settings and dataset splits as Sagawa et al.
Our experiments, unless otherwise specified, replicate the exact hyperparameter settings adopted by Sagawa et al.
[38] for their Table 2 experiments.
This includes the choice of random seeds, batch sizes, learning rates, learning schedules, and regularization.
We defer readers to Sagawa et al.
[38] or to our public source code for replication details.
The Multi-NLI dataset [44] concerns the following natural language inference task: determine if one statement is entailed by, neutral with, or contradicts a given statement.
This dataset is challenging because traditional ERM models are prone to spuriously correlating “contradiction” labels with the existence of negation words.
The dataset is divided into 6 distributions: the Cartesian product of the label space (entailment, neutral, contradiction) and an indicator of whether the sentence contains a negation word.
The label space annotations were annotated by [44] while negation labels were annotated by Sagawa et al.
There are 206,175 datapoints available in the Multi-NLI dataset; the smallest distribution (entailment + negation) is represented by only 1,521 datapoints.
We use a randomly shuffled 50-20-30 training-validation-testing split.
The CelebA dataset is a dataset of celebrity face images and a label space of potential physical attributes [25].
This dataset is challenging because traditional ERM models are prone to spuriously correlating attribute labels with demographic information such as race and gender.
Following Sagawa et al.
[38], we divide the dataset into 4 distributions: the Cartesian product of the blond vs dark hair attribute label (“Blond_Hair”) with the “gender” attribute label (“Male”).
Note that the authors of
Liu et al.
[25] limited the “gender” attribute label to binary options of male and not male.
There are addition, accuracy all due because most so eventually
162,770 datapoints available in the CelebA dataset; the smallest distribution (blond-hair + male) is represented by only 1,387 datapoints.
We use the official training-testing-validation dataset split.
The Waterbirds dataset is a dataset by Sagawa et al.
[38] curated from a larger Caltech-UCSD
Birds-200-2011 (CUB) dataset [43].
It concerns the task of predicting whether a bird is of some waterbird (sub)species from an image of said bird.
This dataset is challenging because traditional
ERM models are prone to spuriously correlating backgrounds with foreground subjects; for instance, a model may often predict that a bird is a waterbird only because the image of the bird was taken at a beach.
The dataset has 4 distributions: the Cartesian product of the waterbird vs not waterbird label with whether the background of the picture is over water.
There are 4,795 datapoints available in the
Waterbirds dataset; the smallest distribution (waterbirds on land) is represented by only 56 examples.
Models We use two classes of models in our experiments: Resnet-50 [19] and BERT [12].
We use the torchvision [27] implementation of the convolutional neural network Resnet-50, with a default choice of a stochastic gradient descent optimizer with momentum 0.9 and batch size 128.
Batch normalization is used; data augmentation and dropout are not used.
We use the HuggingFace [45] implementation of the language model BERT, with a default choice of an Adam optimizer with dropout and batch size 32.
Hyperparameters In the Standard Regularization experiments, we use a Resnet-50 model with an ℓ -2 regularization parameter of λ = 0 . 0001 and a fixed learning rate of α = 0 .
001 for both
Waterbirds and CelebA datasets.
The ERM and Group DRO baselines are trained on CelebA forepochs and Waterbirds for 300 epochs.
Our multi-distribution learning method is trained on CelebA for only 20 epochs and Waterbirds for 100 epochs; this is due to the faster training error convergence of our method.
For the MultiNLI dataset, we use a BERT model with a linearly decaying learning rate starting at α 0 = 0 .
00002 and no ℓ -2 regularization.
The ERM and Group DRO baselines are trained on Multi-NLI for 20 epochs.
Our multi-distribution learning method is trained on Multi-NLI for only 10 epochs.
Our multi-distribution learning method uses adversary learning rates η + of 1, 1, 0.2 on Waterbirds, CelebA and MultiNLI respectively.
In the Strong Regularization experiments, we follow similar settings to the Standard Regularization experiments.
The only change is that an ℓ -2 regularization parameter of λ = 1 is used for Waterbirds and an ℓ -2 regularization parameter of λ = 0 .
1 is used for CelebA.
Our multi-distribution learning method uses adversary learning rates η + of 1 and 0.2 on Waterbirds and CelebA respectively.
In the Early Stopping experiments, we follow similar settings to the Standard Regularization experiments.
The only change is that all CelebA and Waterbird experiments are run for a single epoch.
MultiNLI experiments are run for 3 epochs.
Our multi-distribution learning method uses adversary learning rates η + of 1, 1, 1 on Waterbirds, CelebA and MultiNLI respectively.
The only hyperparameters we use that differ from prior literature are the number of training epochs and the adversary learning rates of our method (R-MDL).
The choice of epoch was not fine-tuned, and was selected due to our observation of early training error convergence.
We selected our adversary learning rate η - by training our method, on each dataset, for both η - = 1 and η - = 0 . 2 and selecting the η - yielding the highest validation-split worst-group accuracy.
Compute The total amount of compute run for the experiments in this section is approximately
GPU hours.
A “n1-standard-8” machine was leased from the Cloud computing service Google Cloud; the “n1-standard-8” machine was equipped with 8 Intel Broadwell chips and 1 NVIDIA Tesla V
GPU.
The cost of these computing resources totaled approximately USD $2 per hour, with a total cost of approximately USD $100.
All results described in this section, with the exception of existing results cited from other works, were obtained with experiments on said machine.
All experiments were implemented in Python and PyTorch.