What You See is What You Get: Distributional Generalization for Algorithm Design in Deep Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
We investigate and leverage a connection between Differential Privacy (DP) andthe recently proposed notion of Distributional Generalization (DG).
Applyingthis connection, we introduce new conceptual tools for designing deep-learningmethods that bypass “pathologies” of standard stochastic gradient descent (SGD).
First, we prove that differentially private methods satisfy a “What You See Is What
You Get (WYSIWYG)” generalization guarantee: whatever a model does on itstrain data is almost exactly what it will do at test time.
This guarantee is formallycaptured by distributional generalization.
WYSIWYG enables principled algorithmdesign in deep learning by reducing generalization concerns to optimization ones:in order to mitigate unwanted behavior at test time, it is provably sufficient tomitigate this behavior on the train data.
This is notably false for standard (non-
DP) methods, hence this observation has applications even when privacy is notrequired.
For example, importance sampling is known to fail for standard SGD,but we show that it has exactly the intended effect for DP-trained models.
We usethese insights to construct simple algorithms which match or outperform SOTAin several distributional robustness applications, and to significantly improve theprivacy vs. disparate impact tradeoff of DP-SGD.
Finally, we also improve onknown theoretical bounds relating DP, stability, and distributional generalization.
Introduction
Much of machine learning (ML), both in theory and in practice, operates under two assumptions.
First, we have i.i d. samples from a distribution.
Second, we care only about a single averaged scalar(e.g.
error/loss/risk) over this distribution.
Under these assumptions, we have mature methods andtheory: Modern learning methods excel when trained on i.i.d.
data to directly optimize a scalarloss and there are many theoretical for reasoning about generalization which explain when doesoptimization of a scalar on the train dataset translates to similar values of this scalar at test time.
However, the focus on scalar metrics such as average error misses many theoretically, practically, andsocially relevant aspects of model performance.
For example, models with small average error oftenhave high error on salient minority subgroups [Buolamwini and Gebru, 2018, Koenecke et al., 2020].
In general, ML models are applied to the heterogeneous and long-tailed data distributions of the realworld [Van Horn and Perona, 2017].
Attempting to summarize their complex behavior with only asingle scalar misses many rich and important aspects of learning.
These issues are compounded for modern overparameterized networks, as their nuanced test-timebehavior is not reflected by their train sets.
This presents an obstacle for algorithm design, becauseinterventions which alter a network’s properties on its train set do not always transfer to the test set.
For example, consider the setting of importance sampling : suppose we know that a certainsubgroup of inputs is underrepresented in the train dataset compared to the test distribution (breakingthe “i.i.d. train/test” assumption).
For underparameterized models, we can simply up-sample thisunderrepresented group to account for the distribution shift [see, e.g., Gretton et al., 2009].
Thisapproach, however, is known to empirically fail for overparameterized models [Byrd and Lipton,2019].
Because “what you see” (on the train dataset) is not “what you get” (at test time), wecannot make principled train-time interventions to affect test-time behaviors.
This issue extendsbeyond importance sampling: for example, theoretically principled methods for distributionallyrobust optimization (e.g. Namkoong and Duchi [2016]) fail for overparameterized deep networks,and require ad-hoc modifications [Sagawa et al., 2019].
We develop a theoretical framework which (1) sheds light on these existing issues, and (2) leads toimproved practical methods in privacy, fairness, and distributional robustness.
The core object in ourframework is what we call the “What You See Is What You Get” (WYSIWYG) property.
A trainingprocedure with the WYSIWYG property does not exhibit the pathologies of standard stochasticgradient descent (SGD): all test-time behaviors will be expressed on the train dataset as well, andthere will be “no surprises” in generalization.
What You See Is What You Get (WYSIWYG) as a Design Principle.
The WYSIWYG property isdesirable for two reasons.
The first is diagnostic: as there are “no surprises” at test time, all propertiesof a model at test time are already evident on the train dataset.
It cannot be the case, for example,that a WYSIWYG model has small disparate impact on the train dataset, but large disparate impactat test time.
The second reason is algorithmic: to mitigate any unwanted test-time behavior, it issufficient to mitigate this behavior on the train dataset.
This means that algorithm designers can beconcerned only with achieving “good behavior” on the train dataset, and the WYSIWYG propertyguarantees good behavior at test time.
In practice, this enables the usage of many theoreticallyprincipled algorithms which were developed in the underparameterized regime to also apply in themodern overparameterized (deep learning) setting.
For example, we find that interventions suchas importance sampling, or algorithms for distributionally robust optimization, which fail withoutadditional regularization, work exactly as intended with WYSIWYG (See Figure 1 for an illustration).
As WYSIWYG is a high-level conceptual property, we have to formalize it to use in practice.
Wedo so using the recent notion of Distributional Generalization (DG), as introduced by Nakkiranand Bansal [2020], Kulynych et al.
If classical generalization ensures that the values of themodel’s loss on the training dataset and at test time are close on average [Shalev-Shwartz et al., 2010],distributional generalization ensures that values of any other bounded test function (not only loss) areclose on training and test time.
We say that a model that satisfies an appropriate level of distributionalgeneralization exhibits the WYSIWYG property.
Achieving DG via Differential Privacy (DP).
Our key observation is that distributional generalization(DG) is formally implied by differential privacy (DP)
[Dwork et al., 2006, 2014]).
The spirit of thisobservation is not novel: DP training is known to satisfy much stronger notions of generalization(e.g.
, robust
Cummings al.,
Bassily et al., and Zakynthinou, 2020].
Weshow that a similar connection holds for the notion of distributional generalization, and prove (and improve) tight bounds relating DP, stability, and DG.
In particular, we show that if a training proceduresatisfies DP, it also satisfies DG.
This guarantees the WYSIWYG property for any method that isdifferentially-private, including DP-SGD on deep neural networks [Abadi et al., 2016].
We demonstrate how DG can be a useful design principle in two concrete settings.
First, we showthat we can mitigate disparate impacts of DP training [Bagdasaryan et al., 2019, Pujol et al., 2020] byleveraging importance sampling.
Second, we study the setting of distributionally robust optimization.
We show how ideas from DP can be used to construct heuristic optimizers, which do not formallysatisfy DP, yet empirically exhibit high DG.
Our heuristics lead to competitive results with SOTAalgorithms in five datasets in the distributional robustness setting.
Our Contributions.
We develop the theoretical connection between Differential Privacy (DP) and
Distributional Generalization (DG), and we leverage our theory to improve empirical performance inprivacy, fairness, and robustness applications.
Theoretically (Sections 2 to 4):
We provide tighter bounds than previously known connecting DP and strong forms of
generalization, and show that DP training methods satisfy DG, thus the WYSIWYG property.
We introduce DP-IS-SGD, an importance-sampling version of DP-SGD, and show it satisfies
DP and DG.
Experimentally (Section 5):
We use our framework to shed light on disparate impact of DP [Bagdasaryan et al.
, 2019]:The disparity in accuracy across groups at test time are provably reflected by the accuracydisparity on the train dataset .2.
We use our DP-IS-SGD algorithm to largely mitigate the disparate impact of DP usingimportance sampling.3.
Based on our theoretical intuitions, we propose a DP-inspired heuristic: addition of gradientnoise.
We find this empirically achieves competitive and even
improved results in several
DRO settings.
Taken together, our results emphasize the central role of distributional generalization—and the
WYSIWYG property it implies—in designing machine learning algorithms which avoid the “pathologies” of standard SGD.
We also establish DP as a useful tool for achieving DG, and thus extend itsapplications further beyond privacy.
Theory of “What You See is What You Get” Generalization
We first review the notion of distributional generalization and demonstrate why it implies the WYSIWYG property.
Second, we show that stability implies distributional generalization.
Finally, weimprove on the known stability guarantees of differential privacy.
As a result, we extend the connections between differential privacy, stability, and generalization to distributional generalization,showing that stability and privacy imply the WYSIWYG property.
Notation.
We consider a learning task with a set of examples X and labels Y .
We assume that asource distribution of labeled examples z ≜ ( x, y ) ∼ D is defined over D = X
× Y .
Given ani.i.d.-sampled dataset S ∼ D n of size n , we use a randomized training algorithm θ ( S ) that outputs amodel’s parameter vector θ from the set Θ .
We denote by f θ ( x ) the resulting prediction function.
2.1 Distributional Generalization and WYSIWYG
If on-average generalization [Shalev-Shwartz et al., 2010] guarantees closeness only of loss valueson train and test data, distributional generalization (DG) also guarantees closeness of values of alltest functions ϕ
( z ; θ ) ∈ [0 , 1] beyond only loss:
Definition 2.1 (Based on Nakkiran and Bansal [2020]) .
An algorithm θ ( S ) satisfies δ -distributionalgeneralization if for all ϕ : D × Θ → [0 , 1] ,
[Equation]
[Equation]
By the variational characterization of the total-variation distance, Equation (1) is equivalent to thebound d TV ( P 1 , P 0 )
≤ δ, where P 1 and P 0 are both distributions of (cid:0) z, θ ( S ) (cid:1) over the randomnessof S ∼ D n and the training algorithm θ ( · ) , with the difference that z ∼ S in the case of P 1 (traindistribution), and z ∼ D in the case of P 0 (test distribution).
It might seem that DG only ensures average closeness of bounded tests on train and test data.
This isnot, however, the full picture.
Consider generalization in terms of a broader class of functions:
Definition 2.2 (Kulynych et al.
An algorithm θ ( S ) satisfies ( δ, π ) -distributional generalization if for a given property function π : D × Θ → R k it holds that d TV ( π♯P 1 , π♯P 0 ) ≤ δ, where π♯Pis the distribution of π ( T ) for T ∼ P .
Because TV distance is preserved under post-processing, we can see that δ -distributional generalization implies ( δ, π ) -distributional generalization for all property functions.
Informally, δ -DG meansthat for all numeric property functions π
( z ; θ ) of a model, the distributions of the property values areclose on the train and test data, on average.
This fact captures the high-level idea of the “What You
See is What You Get” (WYSIWYG) guarantee.
Some example property functions:
• Subgroup loss: π ( z ; θ ) = 1 { z ∈ G } · ℓ ( z ; θ ) , for some subgroup G ⊂ D .
• Counterfactual fairness: π (( x, y ); θ ) =
f θ ( CF ( x )) − f θ ( x ) , where CF ( x ) is a counterfactual
version of x had it had a different value of a sensitive attribute [Kusner et al., 2017].
• Robustness to corruptions: π ( z ; θ ) =
ℓ ( A ( z ); θ ) , where A ( x ) is a possibly randomized
transformation that distorts the example, e.g., by adding Gaussian noise.
• Adversarial robustness: π ( z ; θ ) =
ℓ ( A θ ( z ); θ ) , where A θ ( z ) is an adversarial example, e.g.generated using the PGD attack [Madry et al., 2018].
In the next sections, we show how a training algorithm can provably satisfy DG and therefore provide
WYSIWYG guarantees for all properties, including the ones above.
2.2 Distributional Generalization from Stability and Differential Privacy
The connections between privacy, stability, and generalization are well-known.
In particular, stabilityof the learning algorithm—its non-sensitivity to limited changes in the training data—implies generalization
[Bousquet and Elisseeff, 2002, Shalev-Shwartz et al., 2010].
In turn, differential privacyimplies strong forms of stability, thus ensuring generalization through the chain Privacy ⇒ Stability⇒
Generalization [Raskhodnikova et al., 2008, Dwork et al., 2015b,a, Wang et al., 2016].
Let us formally define differential privacy:
Definition 2.3 (Differential Privacy [Dwork et al., 2006, 2014]) .
An algorithm θ ( S ) is ( ϵ, δ ) -differentially private (DP) if for any two neighbouring datasets —differing by one example— S ,
S ′ of size n , for any subset K ⊆
Θ it holds that Pr[ θ ( S ) ∈ K ] ≤ exp( ϵ )
Pr[ θ ( S ′ )
∈ K ] + δ.
DP mathematically encodes a notion of plausible deniability of the inclusion of an example in thedataset.
However, it can also be thought as a strong form of stability [Dwork et al., 2015b].
As such,
DP implies other notions of stability.
We consider the following notion, which has been studied in the literature under multiple names andcontexts.
In the context of privacy, it is equivalent to (0 , δ ) -differential privacy, and has been calledadditive differential privacy [Geng et al., 2019], and total-variation privacy [Barber and Duchi, 2014].
In the context of learning, it has been called total-variation (TV) stability [Bassily et al., 2016].
Wetake this last approach and refer to it as TV stability:
Definition 2.4 (TV Stability) .
An algorithm θ ( S ) is δ -TV stable if for any two neighbouring datasets
S , S ′ of size n , for any subset T ⊆
Θ it holds that Pr[ θ ( S ) ∈ K ] ≤ Pr[ θ ( S ′ )
∈ K ] + δ.
It is easy to see that ( ϵ, δ ) -DP immediately implies (exp( ϵ ) − 1 + δ ) -TV stability.
From Classical to Distributional Generalization.
Similarly to the classical generalization, one wayto achieve distributional generalization is through strong stability:
Theorem 2.5.
Suppose that the training algorithm is δ -TV stable.
Then, the algorithm satisfies δ -DG.
We refer to Appendix B for the proofs of this and all other formal statements in the rest of the paper.
[Figure]
As DP implies TV-stability, by Theorem 2.5 we havethat DP also implies DG.
Next, we show that DP algorithms enjoy a significantly stronger stability guarantee than the previously known one, which meansthat the DG guarantee that one obtains from DP isalso stronger.
Proposition 2.6.
Suppose that the training algorithmis ( ϵ, δ ) -DP.
Then, the algorithm satisfies δ ′ -TV stability with δ ′ =
(exp( ϵ ) − 1+2 δ ) / (exp( ϵ )+1) .
Figure 2 shows that the known bound of exp( ϵ ) + 1 − δ quickly become vacuous, unlike the onein Proposition 2.6.
In Appendix A we discuss the relationship of this result to other works ingeneralization literature, and show that our bound is tight in Appendix B.
Stronger Distributional Generalization Guarantees.
Although DG implies generalization boundsfor all properties in a black-box way, it is sometimes possible to obtain tighter bounds for specificproperties.
For example, if we are concerned with accuracy on subgroups, then directly applying 183 δ -DG to the subgroup loss property yields a bound that decays with the size of the subgroup: accuracyon very small subgroups is not guaranteed to generalize well.
In Appendix C.1 we show that TVstability in fact implies “subgroup DG”, which guarantees that the accuracy on even arbitrarily smallsubgroups generalizes well in expectation.
As another example, in Appendix C.2 we show that TVstability also ensures the generalization of calibration properties of the learning algorithm.
Example Applications
To demonstrate that DG is a useful property in algorithm design, in the remainder of this paper weuse it to construct simple and high-performing algorithms for two example applications: mitigationof disparate impact of DP and ensuring group-distributional robustness.
Mitigating Disparate Impact of DP.
First, we consider applications in which learning presentsprivacy concerns, e.g., in the case that the training data contains sensitive information.
Using trainingprocedures that satisfy DP is a standard way to guarantee privacy in such settings.
Training with
DP, however, is known to incur disparate impact on the model accuracy: some subgroups of inputscan have worse test accuracy than others.
For example, Bagdasaryan et al.
[2019] show that using
DP-SGD—a standard algorithm for satisfying DP [Abadi et al., 2016]—in place of regular SGDcauses a significant accuracy drop on “darker skin” faces in models trained on the CelebA dataset ofcelebrity faces [Liu et al., 2015], but a less severe drop on “lighter skin” faces.
Our goal is to mitigatesuch disparate impact.
This issue—a quality-of-service harm [Madaio et al., 2020]—is but one ofmany possible harms due to ML systems.
We do not aim to mitigate any other fairness-related issues,nor claim this is possible within our framework, or with technical solutions in general.
Formally, we assume the data distribution D is a mixture of m groups indexed by set G = { 1 , . . .
, m } ,such that D = (cid:80) m i
=1 q i D i .
The vector ( q i , . . . , q m ) ∈ [0 , 1] m represents the group probabilities, with 205 (cid:80) mi =
q
i
= 1 .
For given parameters ( ϵ, δ ) , we want to learn a model θ that simultaneously satisfies( ϵ, δ ) -DP, has high overall accuracy, and incurs small loss disparity :
[Figure]
[Equation]
[Figure]
[Equation]
Group-Distributional Robustness.
Next, we consider a setting of group-distributionally robustoptimization [e.g., Sagawa et al., 2019, Hu et al., 2018].
If in the standard learning approach we wantto train a model that minimizes average loss, in this setting, we want to minimize the worst-case(highest) group loss .
This objective can be used to mitigate fairness concerns such as those discussedpreviously, as well as to avoid learning spurious correlations [Sagawa et al., 2019].
Formally, we want to learn a model θ that minimizes the worst-case group loss :
[Equation]
Unlike the previous application, in this setting, we do not require privacy of the training data.
We usetraining with DP as a tool to ensure the generalization of the worst-case group loss, as described next.
Input: Training dataset S , model f θ parameterized by θ , loss function ℓ ; max gradient norm C , noise parameter σ , number of epochs T , sampling rate ¯ p ; group probabilities ⃗q = ( q 1 , . . . , q m )
Output: θ T 1: for t = 1 , . .
, T do 2:
Batch t ← Sample p ( · ) ( S ) , with p ( z ) ≜ ¯ p / m · q g ( z ) 3: ∆ t θ ←4: for ( x i , y i )
∈
Batch t do 5: ∆ t θ ← ∆ t θ + Clip ( ∇ θ ℓ ( f θ ( x i ) , y i ) , C ) 6: end for 7:
Sample ξ t ∼ σC · N (0 , I d × d ) 8: ∆ t θ ← 1 | Batch t | ∆ t θ + ξ t 9: θ t ← Optimizer ( θ t − 1 , ∆ t θ ) 10: end for
The highlighted parts indicate the differences with respect to DP-SGD.
We obtain DP-SGD as a special case when we have a single group with q = 1 (implying p ( z )
=
¯ p ).
Algorithms which Distributionally Generalize
In this section, we construct algorithms for the applications in Section 3.
Our approach followsthe blueprint: First, we apply a principled algorithmic intervention that ensures desired behavioron the training dataset (e.g., importance sampling).
Second, we modify the resulting algorithm toadditionally ensure DG, which guarantees that the desired behavior generalizes to the test data .
4.1 DP Training with Importance Sampling
Our first algorithm, DP-IS-SGD (Algorithm 1), is a version of DP-SGD [Abadi et al., 2016] whichperforms importance sampling.
DP-IS-SGD is designed to mitigate disparate impact while retaining
DP guarantees.
The standard DP-SGD samples data batches using uniform Poisson subsampling:
Each example in the training set is chosen into the batch according to the outcome of a Bernoulli trialwith probability ¯ p ∈ [0 , 1] .
To correct for unequal representation and the resulting disparate impact,we use non-uniform Poisson subsampling :
Each example z ∈ S has a possibly different probability 227 p ( z ) of being selected into the batch, where p ( z ) does not depend on the dataset S otherwise, and isbounded: 0 ≤ p ( z ) ≤ p ∗ ≤ 1 .
We denote this subsampling procedure as Sample p ( · ) ( S ) .
We assume that we know to which group any z = ( x, y ) belongs, denoted as g ( z ) , e.g., the group isone of the features in x .
We choose p ( z ) to satisfy two properties.
First, to increase the samplingprobability for examples in minority groups: p ( z ) ∝ 1 / q g ( z ) .
Second, to keep the average batchsize equal to ¯ p · n as in standard DP-SGD.
In the rest of the paper, we assume that the groupprobabilities ( q 1 , . . .
, q m ) are known, but it is possible to estimate them in a private way usingstandard methods [Nelson and Reuben, 2019].
We present DP-IS-SGD in Algorithm 1, along with itsdifferences to the standard DP-SGD.
DP Properties of DP-IS-SGD.
Uniform Poisson subsampling is well-known to amplify the privacyguarantees of an algorithm [Chaudhuri and Mishra, 2006, Li et al., 2012].
For example, Li et al.[2012] show that if an algorithm θ ( S ) satisfies ( ϵ, δ ) -DP, then θ ◦ Sample ¯ p ( S ) provides approximately( O (¯ pϵ ) , ¯ pδ )
-DP for small values of ϵ .
We show in Appendix B that non-uniform Poisson subsamplingprovides the same amplification guarantee with ¯ p = p ∗ , where p ∗ is the maximum value of p ( · ) .
As this guarantee is independent of the internal workings of θ ( S ) , it is loose.
For DP-SGD, one wayof computing tight privacy guarantees of subsampling is using the notion of Gaussian differentialprivacy (GDP)
[Dong et al., 2019].
GDP is parameterized by a single parameter µ .
If an algorithmθ ( S ) satisfies µ -GDP, one can efficiently compute a set of ( ϵ, δ ) -DP guarantees also satisfied byθ ( S )
[Dong et al., 2019].
We show that we can use any GDP-based mechanism for computing theprivacy guarantee of DP-SGD to obtain the privacy guarantees of DP-IS-SGD in a black-box manner:
Proposition 4.1.
Let us denote by µ (¯ p, σ, C, T ) (see Algorithm 1) a function that returns a µ -GDPguarantee of DP-SGD.
Then, DP-IS-SGD satisfies a GDP guarantee µ ( p ∗ , σ, C, T ) .
4.2 Group-DRO with Noisy Gradients
We showed that DP-IS-SGD enjoys theoretical guarantees for both DP and DG.
However, DP modelsoften have lower test accuracy compared to standard training [Chaudhuri et al., 2011].
This can be anunnecessary disadvantage in settings where privacy is not required, such as in group-distributionalrobustness.
Thus, we explore non-DP algorithms which do not come with theoretical guarantees on
DG, but are inspired by our theory, and satisfy good empirical DG in practice.
DP-SGD uses gradient clipping (line 5 in Algorithm 1) and gradient noise (lines 7–8).
Individually,these are used as regularization methods for improving stability and generalization [Hardt et al.
,2016, Neelakantan et al., 2015], thus possibly improving DG in practice.
Following this, we relax
DP-IS-SGD to only use addition of noise to the gradient as a regularizer.
This sacrifices privacy inexchange for practical performance.
Specifically, we apply gradient noise to three standard algorithmsfor achieving group-distributional robustness: importance sampling (IS-SGD), importance weighting(IW-SGD)
[Gretton et al., 2009], and gDRO [Sagawa et al., 2019].
This results in the followingvariations: IS-SGD with noisy gradient (IS-SGD-n), IW-SGD with noisy gradient (IW-SGD-n), andgDRO with noisy gradient (gDRO-n).
See Appendix D for more details.
Experiments
We empirically study the distributional generalization in real-world applications.
The code for theexperiments is available in the supplementary material.
Datasets.
We consider five datasets, studied in the distributional robustness and fairness literature:
CelebA [Liu et al., 2015], UTKFace [Zhang et al., 2017], iNaturalist2017 (iNat)
[Horn et al., 2018],
CivilComments [Borkan et al., 2019], and MultiNLI
[Williams et al., 2018, Sagawa et al., 2019].
For every dataset, each example belongs to one group (e.g., CelebA) or multiple groups (e.g.,
CivilComments).
For example, in the CelebA dataset, there are four groups: “blond male”, “malewith other hair color”, “blond female”, and “female with other hair color”.
We present more detailson the datasets, their groups, and used model architectures in Appendix E.
5.1 Enforcing DG in Practice
[Figure]
We empirically confirm that a training procedure with DP guarantees also has a bounded DG gap.
In practice, it is not possible to compute the exact DG gap.
Asa proxy in the rest of the paper, we use the difference betweentrain-time and test-time worst-group accuracy.
This (1) followsthe empirical approach by Nakkiran and Bansal [2020] whichproposes to estimate the gap in Equation (1) using a finite setof test functions, and (2) measures the aspect of distributionalgeneralization that is relevant to our applications.
We providemore details on this choice of the proxy measure in Appendix E.2.
In Figure 3, we train a model on CelebA using DP-SGD fordifferent levels of privacy ϵ .
We see that the gap between trainingand testing worst-group accuracy increases as the level of privacy 288 ε is weakened, which is consistent with our theoretical bounds.
In Appendix E.3 we also explore how regularization methodswhich do not formally imply DG, can empirically improve DG.
5.2 Disparate Impact of Differentially Private Models
We evaluate DP-IS-SGD (Algorithm 1), and demonstrate that it can mitigate the disparate impact inrealistic settings where both privacy and fairness are required.
Figure 4 shows the accuracy disparity, test accuracy, and worst-case group accuracy, computed asin Equation (2), as a function of the privacy budget ϵ .
The models are trained with DP-SGD and
DP-IS-SGD.
When comparing DP-SGD and DP-IS-SGD with the same or similar ϵ , we observethat DP-IS-SGD achieves lower disparity on all datasets.
However, this comes with a drop in
20 Privacy budget (cid:15) 0.
A cc u r ac y d i s p a r it y
DP-SGD DP-IS-SGD (a) Accuracy disparity (lower is better)
20 Privacy budget (cid:15) 0.
W o r s t - g r oup acc
u r ac y
DP-SGD DP-IS-SGD (b) Worst-group accuracy (higher is better)    20 Privacy budget (cid:15) 0.
T e s t acc u r ac y
DP-SGD DP-IS-SGD (c)
Test accuracy (higher is better)
In summary we find that DP-IS-SGD achieves lower disparity at the same privacy budget comparedto standard DP-SGD, with only mild impact on test accuracy.
5.3 Group-DRO
We investigate whether our proposed versions of standard algorithms with Gaussian gradient noise(Section 4.2) can improve group-distributional robustness.
We evaluate empirical DG using worstgroup accuracy as a proxy for DG gap as in Section 5.1, following the evaluation criteria in priorwork [Sagawa et al., 2019, Idrissi et al., 2021].
State-of-the-art (SOTA) methods apply ℓ 2 regularization and early-stopping to achieve the best performance.
We compare three baselines with 313 ℓ 2 regularization, IS-SGD- ℓ 2 , IW-SGD- ℓ 2 , and gDRO- ℓ 2 to our three noisy-gradient variations:
IS-SGD-n, IW-SGD-n, and gDRO-n, as well as DP-IS-SGD.
We use the validation set to selectthe best-performing regularization parameter and epoch (for early stopping) for each method.
See
Appendix E for details on the experimental setup.
Table 1 shows the worst-group accuracy of each algorithm on five datasets.
When comparing IS-SGD,
IW-SGD, and gDRO with their noisy counterparts, we observe that the noisy versions in general havesimilar or slightly better performance compared to non-noisy counterparts.
For instance, IS-SGD-nimproves the SOTA results on CivilComments dataset.
This showcases that in terms of learningdistributionally robust models, noisy gradient can be potentially a more effective regularizer than thecurrently standard ℓ 2 regularizer .
We also find that DP-IS-SGD improves on baseline methods or even achieves similar SOTA performance on several datasets.
For instance, on CelebA and MNLI, DP-IS-SGD achieves betterperformance than IS-SGD- ℓ 2 , and achieves comparable performance to SOTA.
This is surprising, as
DP tends to deteriorate performance.
This suggests that distributional robustness and privacy mightnot be incompatible goals.
Moreover, DP can be a useful tool even when privacy is not required.
Related Work
DP and Strong Generalization.
DP is known to imply a stronger than standard notion of generalization, called robust generalization [Cummings et al., 2016, Bassily et al., 2016].
Robust generalizationcan be thought as a high-probability counterpart of DG: generalization holds with high probabilityover the train dataset, not only on average over datasets.
We focus on our notion of DG for bothconceptual and theoretical simplicity: DG formulated as in Equation (1) enables us to prove tight
[Table]
theoretical bounds that are independent of sample size.
A more comprehensive discussion of relationsto robust generalization is in Appendix A. It can also be useful to consider intermediary definitions,varying in strength from DG to robust generalization.
In Appendix C, we introduce such an notion(“strong DG”) and show its connections to DP.
Other than robust generalization, our connections in Section 2 can also be derived from weakergeneralization bounds that rely on information-theoretic measures [Xu and Raginsky, 2017, Russoand Zou, 2019, Steinke and Zakynthinou, 2020].
We detail this in Appendix A.
Disparate Impact of DP.
Bagdasaryan et al.
[2019], Pujol et al.
[2020] have shown that ensuring
DP in algorithmic systems can cause error disparity across population groups.
Concurrently, Xuet al.
[2021] have proposed DP-SGD-F, a variant of DP-SGD [Abadi et al., 2016] that dynamicallyadapts gradient-clipping bounds for different groups to reduce such disparate impact.
We show that asimpler approach based on importance sampling is also an effective way to reduce disparities, and weprovide a more general blueprint for designing algorithms with similar mitigations.
Group-Distributional Robustness.
Group-distributional robustness aims to improve the worst-casegroup performance.
Existing approaches include using worst-case group loss [Mohri et al., 2019,
Sagawa et al., 2019, Zhang et al., 2020], balancing majority and minority groups by reweighting orsubsampling [Byrd and Lipton, 2019, Sagawa et al., 2020, Idrissi et al., 2021], leveraging generativemodels [Goel et al., 2020], and applying various regularization techniques [Sagawa et al., 2019, Caoet al., 2019].
Although some work [Sagawa et al., 2019, Cao et al., 2019] discusses the importance ofregularization in distributional robustness, they have not explored potential reasons for this importance(e.g.
via the connection to distributional generalization).
Another line of work studies how to improvegroup performance without group annotations [Duchi et al., 2021, Liu et al., 2021a, Creager et al.
,2021], which is a different setting from ours as we assume the group annotations are known.
Conclusions
We argue that distributional generalization (DG) is a desirable property for learning algorithms,enabling more principled algorithm design in settings including deep learning.
DG captures a highlevel “What You See is What You Get” property, which ensures that any training-time behavior isreproduced at test time.
We show that this property is possible to achieve, by proving it is implied bydifferential privacy.
This enables us to leverage advances in differential privacy (such as DP-SGD forneural networks) to enforce DG in many applications.
We propose enforcing DG as a general design principle, and we use it to construct simple yeteffective algorithms in two settings.
In certain fairness settings, we largely mitigate the disparateimpact of differential privacy by using importance sampling and enforcing DG in our new algorithm
DP-IS-SGD.
In certain distributional robustness settings, inspired by DP-SGD, we propose usinga noisy gradient regularizer instead of the traditional ℓ 2 regularizer to enforce DG.
Compared to
SOTA algorithms, noisy gradient regularization can achieve competitive results across many standardbenchmarks.
We hope future work can explore extending this algorithmic design principle to othersettings, such as adversarial robustness, counterfactual fairness, and calibration.
A.1 Differential Privacy and Robust Generalization
DP is known to imply a stronger notion of generalization, called robust generalization , which is a “tailbound” version of DG [Dwork et al., 2015a, Cummings et al., 2016, Bassily et al., 2016].
The originalmotivations for robust generalization are slightly different, but in our notation, a training procedure θis said to satisfy ( γ, η ) -Robust Generalization
if and only if for any test ϕ : D × Θ → [0 , 1] , we have
[Equation]
Comparing this to the definition of DG (Definition 2.1), it is immediate that any procedure thatsatisfies ( γ, η ) -Robust Generalization, also satisfies ( γ + η ) -DG.
Any training method satisfying( ϵ, δ ) -DP also satisfies ( O ( ϵ ) , O ( δ ))
-robust generalization, as long as the sample size n is of sizeΩ(log(1 /δ )
/ϵ 2 )
[Bassily et al., 2016, Theorem 7.2], therefore it satisfies O ( ϵ + δ ) -DG by the previousimplication.
Thus, it is possible to recover the result that DP implies DG as a consequence of theseprevious works, although with looser bounds.
The difference between Distributional Generalization and Robust Generalization is that DG considersall quantities in expectation, while robust generalization considers tail bounds with respect to thetrain dataset.
We focus on DG for two reasons: First, we believe DG is conceptually simpler, as it canbe seen as simply the TV-distance between two natural distributions, and does not involve additionalparameters.
This simplicity is conceptually useful to the algorithm designer, but also enables us toprove simpler tight theoretical bounds that are independent of sample size.
Second, it is often possibleto lift results about DG to the stronger setting of robust generalization, with additional bookkeeping.
Thus, we focus on DG in this paper, with the understanding that stronger guarantees can be obtainedfor these methods if desired.
A.2 Information-Theoretic Generalization Bounds
Other than robust generalization, it is also possible to obtain bounds on generalization of arbitrarytest (loss) functions using information-theoretic measures [Xu and Raginsky, 2017, Russo and Zou,2019, Steinke and Zakynthinou, 2020].
Recently, Steinke and Zakynthinou showed that one suchinformation-theoretic measure — conditional mutual information (CMI) between training algorithmoutputs and the training dataset — is bounded if the training algorithm is DP or TV stable.
Thus, wecould relate stability and DG as in Section 2.2 using CMI as an intermediate tool.
In particular, Steinke and Zakynthinou show that for bounded ϕ : D × Θ → [0 , 1] :
[Equation]
where CMI D ( θ ) is the conditional mutual information of the training algorithm with respect to thedata distribution.
If the training algorithm satisfies ( ϵ, 0) -DP, they also show that CMI D ( θ ) ≤ n 2 ϵ 2 ,where n is the dataset size.
Plugging this into Equation (4), we can see that the generalizationupper bound (right-hand side) is ϵ .
This is significantly looser than the our bound in Section 2.2, asillustrated in Figure 5.
We can use our results to sharpen the CMI bounds of Steinke and Zakynthinou.
Using the fact that 683 δ -TV stability implies CMI D ( θ ) ≤ nδ , we can show a tighter CMI guarantee from DP:
Corollary A.1.
Suppose that the training algorithm θ : D n → Θ is ( ϵ, δ ) -DP.
Then,
[Equation]
This is immediate from Proposition 2.6.
As we illustrate in Figure 5, however, this sharper bound on
CMI from DP does not translate into a sharper bound on the generalization gap if we plug the resultinto Equation (4).
A recent line of work on information-theoretic bounds explored sharper bounds using individual-levelmeasures [Bu et al., 2020a, Haghifam et al., 2020].
As a direction for future work, it could also be
1 .ϵ
DG
RangeSteinke and Zakynthinou [2020] Corollary A.
Section 2.2 (Ours) possible to obtain tighter bounds on DG using individual-level notions of stability [Feldman and
Zrnic,
A.3 Tension between Differential Privacy and Algorithmic Fairness
Beyond empirical observations that training with DP results in disparate impact on performanceacross subgroups [Bagdasaryan et al., 2019, Pujol et al., 2020], Cummings et al.
[2019] and, morerecently, Sanyal et al.
[2022] theoretically analyze the inherent tensions between DP and algorithmicfairness.
It might appear that the tensions contradict our results in which we use DP or similar noise-addingalgorithms to reduce disparate impact (Section 4).
However, Cummings et al.
[2019] and Pujolet al.
[2020] discuss the relationship between privacy and disparate performance (accuracy or falsepositive/false-negative rates), whereas we discuss the relationship between privacy and generalization.
Even if a DP model has to incur at least a certain error on small subgroups on average [Sanyal et al.
,2022, Lemma 1], this error is guaranteed to be similar at train time and test time (from our theoreticalresults in Section 2.2).
In terms of empirical results, the lower bound on subgroup error in Lemma 1 from Sanyal et al.[2022] vanishes for subgroups of size greater than 100 even for small values of epsilon (e.g., 0.1).
The subgroups and values of epsilon in our experiments in Section 5 are all larger than this, thusin our regime we can achieve meaningful subgroup performance using the DP-IS-SGD algorithmdespite the fundamental trade-off.
B Proofs Omitted in the Main BodyB.1 TV-Stability implies Distributional Generalization
Proof of Theorem 2.5.
First, observe that the following distributions are equivalent as the dataset isan i.i.d.
sample:
[Equation]
It is thus sufficient to analyze the equivalent distributions instead.
By the post-processing property ofdifferential privacy, for any dataset S ∈ D n − 1 , any two examples z, z ′ ∈ D , and any set K ⊆ { 0 , 1 } :
[Equation]
as datasets S ∪ { z } and S ∪ { z ′ } are neighbouring.
Taking the expectation of both sides over 716 z, z ′ ∼ D and S ∼ D n − 1 , we get:
[Equation]
where the last equality is simply renaming of the variables for convenience.
Note that analogouslywe also can obtain a symmetric bound:
[Equation]
The total variation between these two distributions is bounded:
[Equation]
where the last inequality is by Equation (7).
Using the equivalences in Equation (5) we can see that:
[Equation]
which is the sought result.
B.2 Tight Bound on TV-Stability from DP
To prove Proposition 2.6, we make use of the hypothesis-testing interpretation of DP [Wassermanand Zhou, 2010].
Let us define the hypothesis-testing setup and the two types of errors in hypothesistesting.
For any two probability distributions P and Q defined over D , let ϕ : D → { 0 , 1 } be a 726 hypothesis-testing decision rule that aims to tell whether a given observation from the domain Dcomes from P or Q .
Definition B.1 (Hypothesis-testing FPR and FNR) .
Without loss of generality, the false-positive 729 error rate α ϕ (FPR, or type I error rate), and the false-negative error rate β ϕ (FNR, or type II errorrate) of the decision rule ϕ : D → [0 , 1] are defined as the following probabilities:
[Equation]
A well-known result due to Le Cam provides the following relationship between the trade-off betweenthe two types of errors and the total variation between the probability distributions:
[Equation]
DP is known to provide the following relationship between FPR and FNR of any decision rule:
Proposition B.2 (Kairouz et al.
Suppose that an algorithm θ ( S ) satisfies ( ϵ, δ ) -DP.
Then,for any decision rule ϕ : D → [0 , 1] :
[Equation]
We can now prove Proposition 2.6:
Proof.
Consider a hypothesis-testing setup in which we want to distinguish between the distributionsθ ( S ) and θ ( S ′ ) .
Let us sum the two bounds in Equation (10):
[Equation]
[Equation]
Let us take the optimal decision rule ϕ ∗ .
In this case, the bound in Equation (9) holds exactly:
[Equation]
Combining this with Equation (11), we get:
[Equation]
Next, we show that the upper bound is tight up to δ :
Proposition B.3.
There is an algorithm θ ( S ) satisfying ( ε, 0) -DP, such that d TV ( θ ( S ) , θ ( S ′ ))
=
exp( ε ) − 1 exp( ε )+1 for two neighbouring datasets S and S ′ .
Proof.
Consider two distributions P 0 and P 1 on a set { 0 , 1 } , with P 0 ( { 0 } )
=
P 1 ( { 1 } )
=
γ forsome γ to be chosen later, and P 0 ( { 1 } ) =
P 1 ( { 0 } )
= 1 − γ .
Those two distributions satisfy 747 d
TV ( P 0 , P 1 ) = 1 − 2 γ , as well as the closeness condition appearing in the definition of ( ε, 0) -DP
[Equation]
with exp( ε ) = 1 − γ γ .
Expressing now TV-distance in terms of ε , we get d TV ( P 0 , P 1 )
=
exp( ε ) − 1 exp( ε )+1 .
With those distributions in hand, it is easy to provide a mechanism θ : { 0 , 1 } → { 0 , 1 } satisfying the 750 desired property: on the input 0 , it generates output according to distribution P 0 , and on the input 1 ,it generates output according to distribution P 1 .
B.3 Privacy Analysis of DP-IS-SGD
First, we present a loose analysis of the privacy guarantees of non-uniform Poisson subsampling.
Lemma B.4.
Suppose that θ ( S ) satisfies ( ϵ, δ ) -DP
and Sample ( S ) is a Poisson sampling procedurewhere each of the sampling probabilities p
i depend on the element z i (but do not depend on the set Sotherwise) and is guaranteed to satisfy p i ≤ p ∗ .
Then θ ◦ Sample satisfies (ln(1 − p ∗ + p ∗ e ϵ ) , p ∗ δ ) -DP.
For small ϵ this can be bounded by ( O ( p ∗ ϵ ) , p ∗ δ ) -DP.
Proof of Lemma B.4.
Consider two neighboring datasets S and S ′
=
S ∪ { z 0 } for some z 0 ̸∈ S .
Wewish to show that for any set K , we have
[Equation]
and symmetrically for S and S ′ .
We will only prove first of those inequalities, as the second isanalogous.
Note that with probability p 0 ≤ p
the element z 0 is included in Sample ( S ′ ) and we have 763 Sample
[Equation]
S ′ )
= { z 0 } ∪ Sample ( S ) , otherwise the element z 0 is not included, and conditioned on z 0not being included Sample ( S ′ ) has the same distribution as Sample ( S ) .
Therefore,
Pr
[Equation]
Sample
[Equation]
Now for each realization Sample ( S ) =
˜ S , we have Pr( θ ( { z 0 } ∪ ˜ S ) ∈ K ) ≤ e ϵ Pr
( θ ( ˜ S ) ∈ K )
+ δby the assumed DP guarantee of the algorithm θ ( S ) .
We can average over all possible subsets ˜ S toget
[Equation]
Plugging this back to the inequality (12), we get
Pr( θ ( Sample ( S ′ ))
∈ K ) ≤
p 0 ( e ϵ Pr
( θ ( Sample ( S ))
∈ K )
+ δ )
+ (1 − p 0 )
Pr( θ ( Sample ( S ))
∈ K ) ≤ (1 − p ∗ + p ∗ e ϵ )
Pr( θ ( Sample ( S ))
∈ K )
+ p ∗ δ.
Finally, when ϵ ≤ 1 we have e ϵ ≤ (1 + 2 ϵ ) , and therefore (1 − p ∗ + p ∗ e ϵ ) ≤ 1 + 2 ϵp ∗ ≤ e 2 ϵp ∗ .
For the tight privacy analysis of non-uniform Poisson subsampling, we make use of the notion of 771 f -privacy:
Definition B.5 ( f -Privacy Dong et al.
An algorithm θ ( S ) satisfies f -privacy if for any twoneighbouring datasets S, S ′
the following holds:
[Equation]
where τ ( P, Q ) is a trade-off function between the FPR and FNR of distinguishing tests (see Appendix B.2):
[Equation]
and f ( α ) ∈ [0 , 1] is a convex, continuous, non-increasing function.
Bu et al.
[2020b] show that uniform Poisson subsampling (see Section 4.1) provides the followingprivacy amplification:
Proposition B.6 (Bu et al.
[2020b]) . Suppose that θ ( S ) satisfies f -privacy, and Sample ( S ) is auniform Poisson sampling procedure with sampling probability ¯ p .
The composition θ ◦ Sample ( S )satisfies f ′ -privacy with f ′ =
¯ pf
+ (1 − ¯ p ) I
d , where Id ( α ) =
1 − α is the trade-off function thatcorresponds to perfect privacy.
We show that a similar result holds for non-uniform Poisson subsampling:
Lemma B.7.
Suppose that θ ( S ) satisfies f -privacy, and Sample ( S ) is a non-uniform Poissonsampling procedure, where the sampling probabilities p i depend on the element
z i (but do not dependon the set S otherwise) and each is guaranteed to satisfy p i ≤ p ∗ .
The composition θ ◦ Sample ( S )satisfies f ′ -privacy with
[Equation]
To show this, we adapt the proof Proposition B.6, and make use of the following lemma:
Lemma B.8 (Bu et al.
[2020b]) .
Let { P i } i ∈
I and { Q i }
i ∈
I be two collections of probability distributions on the same sample space for some index set I .
Let ( λ i ) i ∈ I ∈ [0 , 1] | I | be a collection ofnumbers such that (cid:80) i ∈
I λ i = 1 .
If τ ( P i , Q i ) ≥ f for all i ∈ I , then for any p ∈ [0 , 1] :
[Equation]
Proof of Lemma B.7.
We can think of the result of the subsampling procedure as outputting a binaryvector ⃗b =
( b 1 , . . .
, b n ) ∈ { 0 , 1 } n , where each bit b i indicates whether an example z i ∈ S waschosen in the subsample or not.
We denote the resulting subsample as S ⃗b ⊆ S .
By definition of
Poisson subsampling, each bit b i is an independent sample b i ∼ Bern ( p i ) .
Let us denote by λ ⃗b thejoint probability of ⃗b .
The composition θ ( S ) ◦
Sample ( S ) can be expressed as a mixture distribution:
[Equation]
Analogously, for a neighbouring dataset S ′ ≜ S ∪{ z 0 } , with the sampling probability p 0 correspondingto z 0 , we have:
[Equation]
Applying Lemma B.8, we get f 0 -privacy with f 0 = p
f
+ (1 − p 0 ) Id .
Applying to an arbitrary other 800 z 0 ∈ D , we potentially get the worst-case privacy guarantee for the highest sampling probability, i.e., 801 f = p ∗ f
+ ( − p ∗ ) Id .
Proposition 4.1 is immediate from Lemma B.7 by the fact that GDP is a special case of f -privacy.
C Stronger Formalizations of Distributional Generalization
The main notion discussed in this paper is that of δ -Distributional Generalization.
To recap, a learningalgorithm θ ( S ) satisfies it if for any bounded test function ϕ
we have
[Equation]
In Appendix A we have compared it to a stronger property called ( γ, η ) -Robust Generalization;namely learning algorithm θ ( S ) satisfies it if for any bounded ϕ we have:
[Equation]
A natural middle ground between these two is the following notion.
Definition C.1.
We say that a learning algorithm θ ( S ) satisfies δ -strong Distributional Generalizationif and only if for all ϕ : D × Θ
→
[Equation]
S ∼D n
[Equation]
If DG as defined in Section 2 is a (supremum over all test functions) version of expected generalizationerror [Shalev-Shwartz et al., 2010, Xu and Raginsky, 2017, Russo and Zou, 2019], strong DG is ananalogous variant of expected absolute error
[Steinke and Zakynthinou, 2020].
Clearly any algorithmsatisfying ( η, γ ) -Robust
Generalization satisfies also ( η + γ ) -strong DG, and any algorithm satisfyingδ -strong-DG satisfies δ -DG — so in fact the δ -strong DG can be thought of as a strengthening of
Distributional Generalization as well as Robust Generalization.
One of the motivations behind the definition of δ -strong DG is the following issue of the Distributional
Generalization itself.
As it turns out, it is possible to artificially construct a mechanism for binaryclassification which satisfies 0 -DG, always outputs a classifier with test error
+ o (1) (almostcompletely useless), but with probability 1 / 2 outputs a classifier with training error
(and withprobability 1 / 2 the training error is 1 ).
That is problematic: it is natural to expect from a trainingprocedure with strong generalization guarantee, that if we used it to train a model up to a smalltraining error, we should have high confidence that the test error is also small — but this exampleshows that our confidence cannot be larger than 1 / 2 .
The δ -strong DG solves this problem: The following proposition is a formal statement of the desiredproperty, and can be easily deduced from Markov inequality.
Proposition C.2.
If the learning procedure θ satisfies δ -strong-DG, then for any bounded property ϕwe have Pr S,θ ( | E z ∼ S ϕ
( z, θ ( S ))
− E z ∼D ϕ
( z, θ ( S ))
| > δλ )
<
λ − 1 .
Equivalently, if a learning procedure satisfies δ -strong-DG it also satisfies ( δλ, λ − 1 )
-Robust Generalization for any λ > 1 .
The δ -strong DG defined this way is also a more direct strengthening of the classical notions ofgeneralization gap — which are usually defined as E S ∼D n ,θ | E z ∼ S ℓ ( z, θ ( S ))
− E z ∼D ℓ ( z, θ ( S ))
|where ℓ is the loss function.
We can follow the same proof strategy as in [Bassily et al., 2016, Theorem 7.2], simplifying itsignificantly, to show that TV stability (and in turn Differential Privacy) implies strong Distributional
Generalization, as soon as the training set is large enough.
Lemma C.3.
If a learning algorithm θ ( S ) satisfies δ -TV stability, and the sample size | S | ≥ 2 δ 2 , then 838 θ ( S ) satisfies 4 δ -strong DG.
Proof.
Let us assume first that we have a δ -TV stable mechanism M : D n → Q where Q is a set of 840 bounded queries q :
[Equation]
We wish to find another 2 δ -TV stable mechanism ˜ M : D n → Q such that if we use notation 842 q
S ≜ M ( S ) and ˜ q S ≜ ˜ M
( S ) , we have
[Equation]
Indeed, consider the following mechanism ˜ M ( S ) : on a sample S , get q S ≜ M ( S ) and evaluate˜ µ ≜ E z ∼ S q S ( z ) + 1 δn U where U is a uniform variable in [ − 1 / 2 , 1 / 2 ] .
Finally if ˜ µ ≥ E z ∼ D q S ( z )output q S , otherwise output − q S + 2 E
z
∼D q S ( z ) .
Since d TV ( η U , η U + γ ) ≤ γη , given q the mechanism µ q ( S ) ≜ E z ∼ S q ( z ) + 1 δn U is δ -stable.
Hence,by composition and post-processing this mechanism ˜ M is 2 δ -TV stable.
For n ≥ 2 δ − 2 the inequality (14) follows from Lemma C.4.
Indeed, let ˆ Z ≜ E z ∼ S q S
( z ) −
E z ∼D q S
( z ) , Z = | ˆ Z | , and σ = sign( ˆ Z + 1 δn U ) .
Note that, expanding the chosen notation, we have
E z ∼ S ˜ q S ( z ) − E z ∼D ˜ q S ( z ) = σZ , and E [ σ | Z ] = max( 2 Zδ , 1) .
The conclusion of the Lemma C.4 is 851 E
Z ≤ 2 E [ Zσ ] , which is exactly (14).
Since the mechanism ˜
M satisfies 2 δ -TV stability, by Theorem 2.5 applied to the mechanism ˜ M , wecan bound the right hand side of the inequality (14) by 4 δ .
Finally, for an arbitrary learning algorithm θ satisfying δ -TV stability, and any bounded test function ϕ
( z, θ ) , we can describe mechanism ˜ M which given a sample S ∈ D n outputs a functionq S ( z ) ≜
ϕ ( z, θ ( S )) .
By post-processing such a mechanism M also satisfies δ -TV stability, and theinequality (14) becomes
[Equation]
Since ϕ was chosen in an arbitrary way, this proves that mechanism θ indeed satisfies the conditionof 4 δ -strong-DG.
Lemma C.4.
Let Z be a non-negative random variable satisfying E Z ≥ δ and σ be a {± 1 } valuedrandom variable with E [ σ | Z ] = max( 2 Zδ , 1) .
Then E Z ≤ 2 E [ Zσ ] .
Proof.
This is just a calculation.
[Equation]
C.1 Subgroup-level Distributional Generalization from TV Stability
TV-stability implies a more granular, subgroup-level notion of distributional generalization:
Definition C.5.
Suppose that the data distribution D is a mixture of group-specific distributions D g ,for g ∈ G .
We define ( δ, G ) -subgroup-DG similarly to δ -DG as follows:
[Equation]
(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
where S g denotes a subset of examples in the dataset S that belong to the group g .
Subgroup DG is a stronger notion of DG which says that the model’s behavior on examples fromeach group in G distributionally generalizes in expectation, as long as the model encounters at leastone representative of the group in training.
In its definition, we explicitly prevent the case when thetraining dataset does not contain any group examples to avoid undefined behavior (what is the groupaccuracy on the training dataset if there are no group representatives in the dataset?).
We now show that TV-stability implies this granular notion of DG:
Proposition C.6.
δ -TV stability implies ( δ, G ) -subgroup-DG for any group partitioning G .
Proof.
Observe that the following distributions are equivalent:
[Equation]
[Equation]
The statement follows by applying each step from the proof of Theorem 2.5 to the equivalentdistributions in Equation (15).
Thus, the difference with the proof of Theorem 2.5 is that z, z ′ ∼ D g ,not z, z ′ ∼ D .
C.2 Calibration Generalization via Strong Distributional Generalization
We that the strong Distributional Generalization (Definition C.1) implies generalization of calibrationproperties.
For simplicity we consider binary classifiers, with Y = { 0 , 1 } and f θ : X → [0 , 1] .
Let usdefine the notion of calibration we consider:
Definition C.7.
The calibration gap , or Expected Calibration Error (ECE), of a classifier f θ isdefined as follows:
[Equation]
(where F ( p ) is the law of f θ ( x ) ) where all expectations are taken w.r.t.
( x, y ) ∼ D .
To estimate this quantity from finite samples, we will discretize the interval [0 , 1] .
The τ -binnedcalibration gap of a classifier f θ is:
[Equation]
So that CGAP τ → CGAP as the bin width τ → 0 .
The empirical version of this quantity, for train set 890 S , is:
[Equation]
Next, we show strong DG implies that the expected binned calibration gap is similar between trainand test.
Theorem C.8.
If the training method θ satisfies δ -strong DG, then
[Equation]
Proof.
First, let us define the family of tests:
[Equation]
The assumed δ -strong DG implies that each of these tests ϕ p have similar expectations between trainand test distributions:
[Equation]
[Table]
Thus, the difference in expected calibration gaps between train and test is:
[Equation]
where the last inequality follows from the definition of δ -strong DG applied to functions ϕ p .
D Additional Details on Algorithms
We define q g as the probability of group g , and m as the number of groups.
IS-SGD.
The weight for group g is w g = 1 / m · q g .
Let g i be the group that the i -th example belongs to.
We then sample (with replacement) from the training set with the i -th example having a w g i chanceof being sampled until we have b examples, where b is the batch size.
Finally, for each mini-batch,we optimize the standard cross-entropy loss with the sampled examples.
IW-SGD.
The weight for group g is w g = 1 / m · q g .
We optimize the following loss function:w
[Equation]
, y ) , where ℓ ( · , · ) is the cross-entropy loss and ( x, y ) ∈ S drawn uniformly random drawn from the dataset,and g is the group to which ( x, y ) belongs.
E Additional Experiment Details
E.1 Details on Datasets, Software, and Model Training
[Table]
[Table] [Table]
[Table] [Table]
All algorithms are implemented in PyTorch 1 [Paszke et al., 2019].
For DP-related utilities, weuse opacus 2 [Yousefpour et al., 2021].
Other packages, including numpy 3 [Harris et al., 2020],scipy
[Virtanen et al., 2020], tqdm 5 , and pandas 6 [pandas development team, 2020], are alsoused.
For gDRO [Sagawa et al., 2019], we use the implementation from wilds [Koh et al., 2021].
We use Nvidia 2080ti, 3080, and A100 GPUs.
Our experiments required approximately 400 hours of
GPU time.
Datasets.
For CelebA and CivilComments, we follow the training/validation/testing split in Kohet al.
For UTKFace and iNat, we randomly split the data into 17000/2000/4708 and550000/50000/75170 for training/validation/testing.
For MNLI, we use the same training/validation/testing split in Sagawa et al.
Tables 2 to 6 show the dataset statistics on each group.
All the datasets are publicly available for non-commercial use.
In our work, we adhere to additionalrules regulating the use of each dataset.
All datasets other than iNat could potentially contain personally identifiable information, and arelikely collected without consent, to the best of our knowledge.
They are all, however, collected from
[Table]
manifestly public sources, such as public posts on social media.
Thus, we consider the associatedprivacy risks low.
The data also contain offensive material (e.g., explicitly in the case of CivilComments dataset).
Weconsider the associated risks of reproducing the offensive behavior low, as we use the datasets only toevaluate our theoretical and theoretically-inspired results.
Models.
Similar to previous work [Sagawa et al., 2019], we use the ImageNet-1k pretrained
ResNet50 [He et al., 2016] from torchvision for CelebA, UTKFace, and iNat, and use the pretrained BERT-Base [Devlin et al., 2019] from huggingface [Wolf et al., 2019] for CivilCommentsand MNLI.
Hyperparameters.
We run 50 epochs for CelebA, 100 epochs for UTKFace, 20 epochs for iNat, and5 epochs for CivilComments and MNLI.
For image datasets (CelebA, UTKFace, and iNat), we use the
SGD optimizer and for NLP datasets (CivilComments and MNLI), we use the AdamW [Loshchilovand Hutter, 2019] optimizer.
We use opacus ’s
[Yousefpour et al., 2021] implementation on DP-SGDand DP-AdamW to achieve DP guarantees.
We fix the batch size for none-DP algorithms to 64 for CelebA and UTKFace, 256 for iNat, 16 for
CivilComments, and 32 for MNLI.
For DP-SGD and DP-IS-SGD, we set the sample rate to 0 .
0001for CelebA and iNat, 0 . 001 for UTKFace, and 0.00005 for CivilComments and MNLI.
E.2 Generalization of Worst-Case Group Accuracy as a Proxy for the DG Gap
Although generalization of worst-case group accuracy is not explicitly implied by DG, in ourexperiments it is practically equivalent to using the generalization gap of subgroup accuracy, which isbounded by TV stability.
Let us first concretely define the generalization gap of the worst-case groupaccuracy:
Definition E.1.
The on-average generalization gap of the worst-case accuracy is defined as thefollowing difference:
[Equation]
where we take ℓ (( x, y ) , θ ) ≜ 1 [ f θ ( x )
=
y ] to be the 0-1 loss.
In this definition we explicitly restrictthe datasets to include elements of each group g ∈ G , which is a technicality needed in order to avoidundefined behavior.
In all our experimental results, the worst-performing groups (the maximizers in Equation (17)) arealways the same on the training and test data.
As long as this holds—the worst-performing group isthe same on the train and test data—the generalization gap above simplifies to:
[Equation]
[Figure]
SGD
[Figure]
(cid:15) 0.
[Figure] [Figure] [Figure]
[Figure]
oup acc
[Figure]
ac y
[Figure] [Figure] regularization 10 −10 −10 −
Weight decay 0.
W o r s t - g r oup acc
u r ac y
Train Test (b) ℓ 2 regularization 10101010 4 Iterations 0.
W o r s t - g r oup acc
u r ac y
Train
[Figure]
where g ∗ ∈ G is the worst-performing group.
In Appendix C.1 we show that this simplified gap from
Equation (18) is bounded by TV stability.
Therefore, in practice the generalization gap in Equation (17) offers a lower bound on the DG gapin Equation (1).
Using it as a proxy for DG gap follows the spirit of the estimation approach by
Nakkiran and Bansal [2020] which proposes to estimate the DG gap by taking the maximum ofempirical generalization gaps for a finite set of relevant test functions (here, per-group accuracies).
Other Approaches to Estimate the DG Gap.
The generalization gap of worst-case group accuracycan be loose as a proxy.
Finding the worst-case test function is an object of study in the literature on 962 membership inference attacks [Shokri et al., 2017], because DG and the accuracy of such attacks areequivalent, as showed by Kulynych et al.
We avoid using accuracy of a membership inferenceattack as a proxy for DG gap in this work as the fact that differential privacy and regularizationimpacts vulnerability to these attacks is known and well-documented [Jayaraman and Evans, 2019,
Liu et al., 2021b].
This body of evidence from the field of membership inference offers an alternativesource of empirical support for our claims checked in Section 5.1.
E.3 Additional Details for Section 5.
As mentioned in Section 4, many regularization methods can be used to improve different generalization gaps.
For example, Sagawa et al.
[2019] show that strong ℓ 2 regularization helps withimproving group-distributional generalization, and Yang et al.
[2020] show that dropout helps withadversarial-robustness generalization.
However, these works do not have theoretical justification.
Our framework suggests a unifying reason why strong regularization is helpful in distributional robustness: because it enforces DG.
Following this theoretically-inspired intuition, other regularizationmethods beyond a combination of gradient noise and clipping (DP-SGD) can imply DG in practice.
We verify this hypothesis empirically.
Privacy, ℓ 2 Regularization, and Early Stopping.
In Figure 6, we train a neural network on CelebAusing DP-SGD, and decrease the “regularization strength” in several different ways: by increasingprivacy
budget
ϵ
(Figure 6a), decreasing the ℓ 2 regularization (Figure 6b), or increasing the numberof training iterations (Figure 6c).
7 We then measure the gap in worst-group accuracy on train vs. test(Appendix E.2).
We observe that for all regularizers, the gap between training and testing worst-groupaccuracy increases as the regularization is weakened.
Investigating ℓ 2 Regularization in Depth.
In Figure 7, we show the training and testing worstgroup accuracy with different strength of ℓ 2 regularization and on different epochs (w/ and w/o ℓ 2regularization).
We have three observations: (1) with properly tuned regularization parameter, the gapbetween training and testing worst-group accuracy can be narrowed, (2) the gap can start wideningin very early stage of training, and (3) the testing worst-group accuracy can fluctuate largely, whichhighlights the importance of using validation set for early stopping in this task.
[Figure]
)    17 Epoch 0.
W o r s t - g r oup acc
u r ac y
Train Test (i) iNat ( ℓ 2 reg.)
E.4 Additional Details for Section 5.2
Figure 8 shows the accuracy disparity, test accuracy, and worst-group accuracy for CelebA, UTKFace,and iNat on DP-SGD and DP-IS-SGD.
The reason that UTKFace has a similar disparity between DP-SGD and DP-IS-SGD is possiblybecause UTKFace has a relatively small difference in the number of training examples between thelargest group and the smallest group.
In UTKFace, the majority group has around seven times moreexamples than in the minority group, whereas in CelebA, this difference is 52 × .
E.5 Additional Details for Section 5.3
We compare different algorithms, including SGD- ℓ 2 and IW-SGD- ℓ 2 as baselines, and two otheralgorithms, IS-SGD- ℓ 2 [Idrissi et al., 2021] and gDRO- ℓ 2 [Sagawa et al., 2019] in terms of the grouprobustness.
We set the learning rate as 0 .
001 for CelebA, UTKFace, and iNat, 0 . 00002 for MNLI,and 0 .
00001 for CivilComments.
We use the validation set to select the hyperparameters:
For SGD- ℓ 2 , IW-SGD- ℓ 2 , IS-SGD- ℓ 2 , and gDRO- ℓ 2 , we select the weight decay from0 .
{+0001 , 0 . 01 , 0 . 1 , and 1 . 
+}
{+For DP-IS-SGD, we fix the gradient clipping to 1 . 
+}
(except for iNat, where we set the valueto 10 . 0 as 1 . 0 does not converge). We select the noise parameter from 1 . 0 , 0 . 1 , 0 . 01 , 0 . 001on CelebA and UTKFace, select the noise parameter from 0 . 0000001 , 0 . 000001 , 0 .
[Figure]
r oup acc
[Figure]
and
0001 on and select the noise parameter from on
and we select the standard deviation the
we
Statistical Concerns.
Although our results appear to be comparable to or better than SOTA, wecaution readers about the exact ordering of methods due to high estimation variance: these benchmarkshave small validation and test sets (e.g., CelebA has 182 validation examples) and so hyperparametertuning is subject to both overfitting and estimation error.
For example, we observe validationaccuracies which differ from their test accuracies by up to 5% in our experiments.
We attempt tomitigate this using three random train/val/test splits on CelebA, and avoid large hyperparametersweeps 8 , but this is not done in prior work.
E.6 Additional Experiments with Other Properties
As we detail in Section 2.1, a DG-inducing optimizer can ensure generalization of different typesof properties.
The focus of our applications in Section 3 is the subpopulation-performance property.
As a proof of concept that our WYSIWYG framework is practical for other properties,
we also
[Figure]
apply it to ensure generalization of adversarial robustness, performance under test-time adversarialperturbations [see, e.g., Madry et al., 2018].
For images, a standard approach to modeling the adversary is assuming they attack the target model θby feeding it adversarial examples which maximize the loss of the target model, yet are similar tosome natural examples ( x, y ) ∼ D in L p distance:
[Equation]
A common way to train models that are adversarially robust in this sense is adversarial training [Madry et al., 2018].
In adversarial training, we want to train models that minimize the followingobjective, called adversarial loss :
[Equation]
Generalization of adversarial loss is an issue, as adversarially trained models are known to exhibit“robust overfitting”: higher generalization gap of adversarial loss than that of the regular loss [Riceet al., 2020].
As in Section 4.2, inspired by our theoretical results, we apply the noisy-gradient regularizer tothe standard adversarial training [Madry et al., 2018] which uses the projected gradient descentalgorithm to solve the inner maximization in Equation (19).
As with DRO, we expect this toinduce stability, and ensure empirical distributional generalization, thus to reduce overfitting underadversarial perturbations.
Experimental Setup.
We use the CIFAR-10 dataset [Krizhevsky et al., 2009], and ResNet-18 [Heet al., 2016] as the network architecture.
We train the model to be robust against L ∞ perturbations ofat most γ = 8 / 255 bound, which is a standard setup for adversarial training on this dataset.
We varyσ (noise parameter) from 0.0 (regular adversarial training without gradient noise) to 0.01.
Results.
We summarize the results in Figure 9.
The left plot displays the gap in robust accuracy, andthe right plot displays the robust test accuracy.
As shown in Figure 9 (a), without gradient noise, instandard adversarial training the gap between robust training accuracy and robust test accuracy is largeat approximately 30 p.p., consistent with the prior observations of “robust overfitting”
[Rice et al.
By injecting noise into the gradient, our proposed approach decreases the adversarial accuracygeneralization gap by more than 3 × to less than 10 p.p.
More surprisingly, in our experiments,training with gradient noise achieves both a low adversarial accuracy gap and better robust accuracycompared to standard adversarial training, when using a low level of noise ( σ = 0 . 0005 ).
Theseexperimental results further demonstrate how DG can be a useful design principle in practice.