What You See is What You Get: Principled Deep Learning via Distributional Generalization
Abstract
Having similar behavior at training time and test time—what we call a “What You
See Is What You Get” (WYSIWYG) property—is desirable in machine learning.
Models trained with standard stochastic gradient descent (SGD), however, do not necessarily have this property, as their complex behaviors such as robustness or subgroup performance can differ drastically between training and test time.
In contrast, we show that Differentially-Private (DP) training provably ensures the high-level WYSIWYG property, which we quantify using a notion of distributional generalization.
Applying this connection, we introduce new conceptual tools for designing deep-learning methods by reducing generalization concerns to optimization ones: to mitigate unwanted behavior at test time, it is provably sufficient to mitigate this behavior on the training data.
By applying this novel design principle, which bypasses “pathologies” of SGD, we construct simple algorithms that are competitive with SOTA in several distributional-robustness applications, significantly improve the privacy vs. disparate impact trade-off of DP-SGD, and mitigate robust overfitting in adversarial training.
Finally, we also improve on 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 independent and identically distributed (i i d.) samples.
Second, we care only about a single averaged scalar metric (error, loss, risk).
Under these assumptions, we have mature methods and theory: Modern learning methods excel when trained on i.i.d.
data to directly optimize a scalar loss, and there are many theoretical tools for reasoning about generalization , which explain when does optimization of a scalar on the training data translates to similar values of this scalar at test time.
The focus on scalar metrics such as average error, however, misses many theoretically, practically, and socially relevant aspects of model performance.
For example, models with small average error often have 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 real world [Van Horn and Perona, 2017].
Attempting to summarize their complex behavior with only a single scalar misses many rich and important aspects of learning.
These issues are compounded for modern overparameterized networks, as their nuanced test-time behavior is not reflected at training time.
For example, consider the setting of importance sampling : suppose we know that a certain subgroup of inputs is underrepresented in the training data compared to the test distribution (breaking the i.i.d.
assumption).
For underparameterized models, we can upsample this underrepresented group to account for the distribution shift [see, e.g., Gretton et al., 2009].
This approach, however, is known to empirically fail for overparameterized models [Byrd and
Lipton, 2019].
Because “what you see” (on the training data) is not “what you get” (at test time), we cannot make principled train-time interventions to affect test-time behaviors.
This issue extends
[Figure]
[Figure]
W o r s t - g r oup acc
u r ac y
SGD (CelebA)
[Figure]
DP - SGD (CelebA) beyond importance sampling.
For example, theoretically principled methods for distributionally robust optimization (e.g. Namkoong and Duchi [2016]) fail for overparameterized deep networks, and require ad-hoc modifications [Sagawa et al., 2019a].
We develop a theoretical framework which sheds light on these existing issues, and leads to improved practical methods in privacy, fairness, and distributional robustness.
The core object in our framework is what we call the “What You See Is What You Get” (WYSIWYG) property.
A training procedure with the WYSIWYG property does not exhibit the “pathologies” of standard stochastic gradient descent (SGD): all test-time behaviors will be expressed on the training data as well, and there will be “no surprises” in generalization.
What You See Is What You Get (WYSIWYG) as a Design Principle.
The WYSIWYG property is desirable for two reasons.
The first is diagnostic: as there are “no surprises” at test time, all properties of a model at test time are already evident at the training stage.
It cannot be the case, for example, that a WYSIWYG model has small disparate impact on the training data, but large disparate impact at test time.
The second reason is algorithmic: to mitigate any unwanted test-time behavior, it is sufficient to mitigate this behavior on the training data.
This means that algorithm designers can be concerned only with achieving desirable behavior at train time, as the WYSIWYG property guarantees it holds at test time too.
In practice, this enables the usage of many theoretically principled algorithms which were developed in the underparameterized regime to also apply in the modern overparameterized (deep learning) setting.
For example, we find that interventions such as importance sampling, or algorithms for distributionally robust optimization, which fail without additional 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 computational practice.
We do so using the notion of Distributional Generalization (DG), as introduced by Nakkiran and Bansal [2020], Kulynych et al.
If classical generalization ensures that the values of the model’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—are close on training and test time.
We say that a model which satisfies an appropriately high level of distributional generalization exhibits the WYSIWYG property.
Achieving DG in Practice.
Our key observation is that distributional generalization is formally implied by differential privacy (DP)
[Dwork et al., 2006, 2014]).
The spirit of this observation is not novel: DP training is known to satisfy much stronger notions of generalization (e.g., robust generalization , see Section 6 for more details), and stability than standard SGD [Dwork et al., 2015a, Cummings et al., 2016, Bassily et al., 2016, Steinke and Zakynthinou, 2020].
We show that a similar connection holds for the notion of distributional generalization, and prove (and improve) tight bounds relating DP, stability, and DG.
This guarantees the WYSIWYG property for any method that is differentially-private, including DP-SGD on deep neural networks [Abadi et al., 2016].
We demonstrate how DG can be a useful design principle in three concrete settings.
First, we show that we can mitigate disparate impact of DP training [Bagdasaryan et al., 2019, Pujol et al., 2020] by leveraging importance sampling.
Second, we study the setting of distributionally robust optimization [e.g., Sagawa et al., 2019a, Hu et al., 2018].
We show how ideas from DP can be used to construct heuristic optimizers, which do not formally satisfy DP, yet empirically exhibit DG.
Our heuristics lead to competitive results with SOTA algorithms in five datasets in the distributional robustness setting.
Third, we show that the heuristic optimizer is also capable of reducing overfitting of adversarial loss in adversarial training [Madry et al., 2018, Zhang et al., 2019, Rice et al., 2020].
Our Contributions.
We develop the theoretical connection between Differential Privacy (DP) and
Distributional Generalization (DG), and we leverage our theory to improve empirical performance in privacy, fairness, and robustness applications.
Theoretically (Sections 2 to 4):
We provide tighter bounds than previously reported 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 : The disparity in accuracy across groups at test time are provably reflected by the accuracy disparity on the train dataset .
We use our DP-IS-SGD algorithm to largely mitigate the disparate impact of DP using importance sampling. 3.
Based on our theoretical intuitions, we propose a DP-inspired heuristic: addition of gradient noise.
We find this empirically achieves competitive and even improved results in several DRO settings, and reduces overfitting of adversarial loss in adversarial training.
Taken together, our results emphasize the central role of the WYSIWYG property in designing machine learning algorithms which avoid the “pathologies” of standard SGD.
We also establish DP as a useful tool for achieving WYSIWYG, thus extend its applications 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 captures the
WYSIWYG property.
Second, we show that strong stability notions imply distributional generalization.
Finally, we improve 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 a source distribution of labeled examples z ≜ ( x, y ) ∼ D is defined over D = X
× Y .
Given an i.i.d.-sampled dataset S ∼ D n of size n , we use a randomized training algorithm T ( S ) that outputs a model’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 values on train and test data, distributional generalization (DG) also guarantees closeness of values of all test functions ϕ
( z ; θ ) ∈ [0 , 1] beyond only loss:
Definition 2.1 (Based on Nakkiran and Bansal [2020]) .
An algorithm T ( S ) satisfies δ -distributional generalization if for all ϕ : D × Θ → [0 , 1] ,
[Equation]
By the variational characterization of the total-variation (TV) distance [see, e.g., Polyanskiy and Wu, 2014, Chapter 6.3], Equation (1) is equivalent to the bound d TV ( P 1 , P 0 ) ≤ δ, where P 1 and Pare both distributions of (cid:0) z, T ( S ) (cid:1) over the randomness of S ∼ D n and the training algorithm T ( · ) , with the difference that z ∼ S in the case of P 1 (train), and z ∼ D in the case of P 0 (test).
It might seem that DG only ensures average closeness of bounded tests on train and test data.
This is not, however, the full picture.
Consider generalization in terms of a broader class of functions:
Definition 2.2 (Kulynych et al.
An algorithm T ( S ) satisfies ( δ, π ) -distributional generalization if for a given property function π : D × Θ → R k it holds that d TV ( π ♯ P 1 , π ♯ P 0 ) ≤ δ, where π ♯ P is 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 means that for all numeric property functions π ( z ; θ ) of a model, the distributions of the property values are close 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 θ ( x ′ )
− f θ ( x ) , where 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, stability of 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 privacy implies 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 T ( 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[ T ( S ) ∈ K ] ≤ exp( ϵ [Equation] Pr [Equation]
DP mathematically encodes a notion of plausible deniability of the inclusion of an example in the dataset.
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.
In the context of privacy, it is equivalent to (0 , δ ) -differential privacy, and has been called additive 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].
We take this last approach and refer to it as TV stability:
Definition 2.4 (TV Stability) .
An algorithm T ( S ) is δ -TV stable if for any two neighbouring datasets
S , S ′ of size n , for any subset T ⊆
Θ it holds that Pr[ T ( S ) ∈ K ] ≤ Pr[ T ( S ′ )
∈ K ] + δ.
It is easy to see that ( ϵ, δ ) -DP immediately implies δ ′ -TV stability with:
[Equation]
From Classical to Distributional Generalization.
Similarly to the classical generalization, one way to 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.
As DP implies TV-stability, by Theorem 2.5 we have that DP also implies DG.
We show that DP algorithms enjoy a significantly stronger stability guarantee than previously known, which means that the DG guarantee that one obtains from DP is also stronger.
Proposition 2.6.
An algorithm which is ( ϵ, δ ) -DP is also δ ′ -TV
stable with:
[Equation]
[Figure]
In Appendix A.2, we discuss the relationship of this result to other works in the literature on informationtheoretic generalization.
In particular, to Steinke and
Zakynthinou [2020] whose results can also be used to relate DP and DG.
Figure 2 shows that the known bounds quickly become vacuous unlike the bound in
Proposition 2.6.
In fact, we show that our bound is tight in Appendix B.
Stronger Distributional Generalization Guarantees.
Although DG immediately implies generalization for all bounded properties, it is possible to obtain tighter bounds from TV stability.
For example, directly applying δ -DG to the subgroup loss property yields a bound that decays with the size of the subgroup: accuracy on very small subgroups is not guaranteed to generalize well.
In Appendix C. we show that TV stability in fact implies “subgroup DG”, which guarantees that the accuracy on even small subgroups generalizes well in expectation.
As another example, in Appendix C.2 we show that
TV stability also ensures the generalization of calibration properties of the learning algorithm.
Example Applications
To demonstrate that WYSIWYG is a useful property in algorithm design, in the remainder of this paper we use it to construct simple and high-performing algorithms for three example applications: mitigation of disparate impact of DP, ensuring group-distributional robustness, and mitigation of robust overfitting in adversarial training.
Mitigating Disparate Impact of DP.
First, we consider applications in which learning presents privacy concerns, e.g., in the case that the training data contains sensitive information.
Using training procedures 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 inputs can 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 SGD causes a significant accuracy drop on “darker skin” faces in models trained on the CelebA dataset of celebrity faces [Liu et al., 2015], but a less severe drop on “lighter skin” faces.
Our goal is to mitigate such disparate impact.
This issue—a quality-of-service harm [Madaio et al., 2020]—is but one of many possible harms due to ML systems.
We do not aim to mitigate any other broad fairness-related issues, nor claim this is possible within our framework.
Formally, we assume the data distribution D is a mixture of m groups indexed by set G = { 1 , . . .
, m } , such that D = (cid:80) mi
=1 q i D i .
The vector ( q i , . . . , q m ) ∈ [0 , 1] m represents the group probabilities, with (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 :
[Equation]
Group-Distributional Robustness.
Next, we consider a setting of group-distributionally robust optimization [e.g., Sagawa et al., 2019a, Hu et al., 2018].
If in the standard learning approach we want to 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 discussed previously, as well as to avoid learning spurious correlations [Sagawa et al., 2019a].
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 use training with DP as a tool to ensure the generalization of the worst-case group loss.
Mitigating Robust Overfitting.
Finally, we consider the setting of robustness to test-time adversarial examples through adversarial training [Goodfellow et al., 2014].
A common way to train robust models in this sense in image domains is to minimize robust (adversarial) loss [Madry et al., 2018]:
[Equation] Input: Dataset S , loss ℓ ( z ; θ ) , initial parameters θ 0 , learning rate η , maximal gradient norm C , noise parameter σ , number of epochs T , sampling rate ¯ p , group probabilities ( q 1 , . . . , q m ) .
for t = 1 , . . .
, T do
Sample
[Table]
[Table]
sampling
[Table]
where γ > 0 is a small constant.
Rice et al.
[2020] observed that adversarially trained models exhibit “robust overfitting”: higher generalization gap of robust loss than that of the regular loss.
In this application, we similarly aim to use a relaxed version of training with DP as a tool to ensure generalization of robust loss, thus mitigate robust overfitting.
Algorithms which Distributionally Generalize
In this section, we construct algorithms for the applications in Section 3.
Our approach follows the blueprint: First, we apply a principled algorithmic intervention that ensures desired behavior on the training data (e.g., importance sampling).
Second, we modify the resulting algorithm to additionally ensure DG, which guarantees that the desired behavior generalizes to test time .
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] which performs 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 trial with 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 p ( z ) of being selected into the batch, where p ( z ) does not depend on the dataset S otherwise, and is bounded: 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 is one of the features in x .
We choose p ( z ) to satisfy two properties.
First, to increase the sampling probability for examples in minority groups: p ( z ) ∝ 1 / q g ( z ) .
Second, to keep the average batch size equal to ¯ p · n as in standard DP-SGD.
In the rest of the paper, we assume that the group probabilities ( q 1 , . . .
, q m ) are known, but it is possible to estimate them in a private way using standard methods [Nelson and Reuben, 2020].
We present DP-IS-SGD in Algorithm 1, along with its differences to the standard DP-SGD.
DP Properties of DP-IS-SGD.
Uniform Poisson subsampling is well-known to amplify the privacy guarantees of an algorithm [Chaudhuri and Mishra, 2006, Li et al., 2012].
For example, Li et al.
[2012] show that if an algorithm T ( S ) satisfies ( ϵ, δ ) -DP, then T ◦ Sample ¯ p ( S ) provides approximately ( O (¯ pϵ ) , ¯ pδ ) -DP for small values of ϵ .
We show in Appendix B that non-uniform Poisson subsampling provides 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 T ( S ) , it is loose.
For DP-SGD, one way of computing tight privacy guarantees of subsampling is using the notion of Gaussian differential privacy (GDP)
[Dong et al., 2019].
GDP is parameterized by a single parameter µ .
If an algorithm
T ( S ) satisfies µ -GDP, one can efficiently compute a set of ( ϵ, δ ) -DP guarantees also satisfied by
T ( S )
[Dong et al., 2019].
We show that we can use any GDP-based mechanism for computing the privacy 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 µ -GDP guarantee of DP-SGD.
Then, DP-IS-SGD satisfies a GDP guarantee µ ( p ∗ , σ, C, T ) .
4.2 Gaussian Gradient Noise
We showed that DP-IS-SGD enjoys theoretical guarantees for both DP and DG.
DP models, however, often have lower test accuracy compared to standard training [Chaudhuri et al., 2011].
This can be an unnecessary disadvantage in settings where privacy is not required, such as in our robustness applications.
Thus, we explore training algorithms which are inspired by our theory yet do not come with generic theoretical guarantees of DG.
Note that DP-SGD uses gradient clipping and noise (see Algorithm 1).
Individually, these are used as regularization methods for improving stability and generalization [Hardt et al., 2016, Neelakantan et al., 2015].
Following this, we relax DP-IS-SGD to only use gradient noise.
This sacrifices privacy guarantees in exchange for practical performance.
Specifically, we apply gradient noise to three standard algorithms for achieving group-distributional robustness: importance sampling (IS-SGD), importance weighting (IW-SGD)
[Gretton et al., 2009], and gDRO [Sagawa et al., 2019a].
This results in the following variations: IS-SGD-n, IW-SGD-n, gDRO-n, respectively.
Similarly, we apply gradient noise to standard PGD adversarial training [Madry et al., 2018].
See Appendix D for details.
Experiments
We empirically study the distributional generalization in real-world applications.
The code for the experiments is available at https://github.com/yangarbiter/dp-dg .
Datasets.
We use the following datasets with group annotations: CelebA
[Liu et al., 2015], UTKFace
[Zhang et al., 2017], iNaturalist2017 (iNat)
[Horn et al., 2018], CivilComments [Borkan et al., 2019], MultiNLI
[Williams et al., 2018, Sagawa et al., 2019a], and ADULT [Kohavi et al., 1996].
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”, “male with other hair color”, “blond female”, and “female with other hair color”.
Additionally, we use the
CIFAR-10 [Krizhevsky et al., 2009] dataset for the adversarial-overfitting application.
We present more details on the datasets, their groups, and used model architectures in Appendix E.
5.1 Enforcing DG in Practice
We empirically confirm that a training procedure with DP guarantees also has a bounded DG gap.
[Figure]
In practice, it is not possible to compute the exact DG gap.
As a proxy in applications which concern subgroup performance in this section, and Sections 5.2 and 5.3, we use the difference between train-time and test-time worst-group accuracy.
This (a) follows the empirical approach by Nakkiran and Bansal [2020] which proposes to estimate the gap in Equation (1) using a finite set of test functions, and (b) measures the aspect of distributional generalization that is relevant to our applications.
We provide more details on this choice of the proxy measure in Appendix E.2.
We train a model on CelebA using DP-SGD for varying levels of ϵ .
Figure 3 shows that the gap between training and testing worstgroup accuracy increases as the level of privacy decreases, which is consistent with our theoretical bounds.
In Appendix E.3 we also explore how regularization methods which do not necessarily 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 in realistic settings where both privacy and fairness are required.
Figure 4 shows the accuracy disparity, test accuracy, and worst-case group accuracy, computed as in Equation (3), 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 observe that
DP-IS-SGD achieves lower disparity on all datasets.
However, this comes with a drop in average
[Figure]
In summary we find that DP-IS-SGD can achieve lower disparity at the same privacy budget compared to standard DP-SGD, with mild impact on test accuracy.
Comparison to DP-SGD-F [Xu et al., 2021].
DP-SGD-F is a variant of DP-SGD which dynamically adapts gradient-clipping bounds for different groups to reduce the disparate impact.
We did not manage to achieve good overall performance of DP-SGD-F on the datasets above.
In Appendix E.4, we compare it to DP-IS-SGD on the ADULT dataset (used by Xu et al.
[2021]), finding that DP-ISSGD obtains lower disparity for the same privacy level, yet also lower overall accuracy.
5.3 Group-Distributionally Robust Optimization
We investigate whether our proposed versions of standard algorithms with Gaussian gradient noise (Section 4.2) can improve group-distributional robustness.
To do so, we evaluate empirical DG using worst-group accuracy as a proxy for DG gap as in Section 5.1, following the evaluation criteria in prior work [Sagawa et al., 2019a, Idrissi et al., 2022].
State-of-the-art (SOTA) methods apply ℓregularization and early-stopping to achieve the best performance.
We compare three baselines with ℓ 2 regularization, IS-SGD- ℓ 2 , IW-SGD- ℓ 2 , and gDRO- ℓ 2 to our noisy-gradient variations as well as
DP-IS-SGD.
We use the validation set to select the best-performing regularization parameter and epoch (for early stopping) for each method.
See Appendix E.5 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 have similar or slightly better performance compared to non-noisy counterparts.
For instance, IS-SGD-n improves the SOTA results on CivilComments dataset.
This showcases that in terms of learning distributionally robust models, noisy gradient can be a more effective regularizer than the currently standard ℓ 2 regularizer .
We also find that DP-IS-SGD improves on baseline methods or even achieves
SOTA-competetitive performance on several datasets.
For instance, on CelebA and MNLI, DP-ISSGD achieves better performance than IS-SGD- ℓ 2 .
This is surprising, as DP tends to deteriorate performance.
This suggests that distributional robustness and privacy are not incompatible goals.
Moreover, DP can be a useful tool even when privacy is not required.
5.4 Mitigating Robust Overfitting
As in the previous section, we expect that a modification of a standard projected gradient descent (PGD) method for adversarial training [Madry et al., 2018] with added Gaussian gradient noise (Section 4.2) improves the generalization behavior of adversarial training. [Table] [Figure] [Table]
[Figure] [Table] [Figure] [Table]
[Figure]
[Table]
[Figure]
To verify this, we adversarially train models on the CIFAR-10 dataset with varying levels of the noise magnitude.
We provide more details on the setup in Appendix E.6.
Figure 5 shows that in standard adversarial training without noise the gap between robust training accuracy and robust test accuracy is large at approximately 30 p.p., which is consistent with the prior observations of Rice et al.
By injecting noise into the gradient, our proposed approach decreases the generalization gap of robust accuracy by more than 3 × to less than 10 p.p.
Surprisingly, in our experiments, training with gradient noise achieves both a small adversarial accuracy gap and better adversarial test accuracy compared to standard adversarial training, when using a small noise magnitude ( σ = 5 × 10 − 4 ).
In terms of resulting robust accuracy, the method’s performance is comparable to early stopping, identified as the most effective way to prevent robust overfitting by Rice et al.
These experimental results demonstrate how WYSIWYG can be a useful design principle in practice.
Related Work
DP and Strong Generalization.
DP is known to imply a stronger than standard notion of generalization, called robust generalization 2 [Cummings et al., 2016, Bassily et al., 2016].
Robust generalization can be thought as a high-probability counterpart of DG: generalization holds with high probability over the train dataset, not only on average over datasets.
We focus on our notion of
DG for both conceptual and theoretical simplicity.
A more comprehensive discussion of relations to robust generalization is in Appendix A.1.
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 weaker generalization bounds that rely on information-theoretic measures [Steinke and Zakynthinou, 2020].
We detail this in Appendix A.2.
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.
Xu et al.
[2021] proposed a variant of DP-SGD for reducing disparate impact.
We compare our method to DP-SGD-F in Appendix E.4.
In another line of related work, Sanyal et al.
[2022], Cummings et al.
[2019] show fundamental trade-offs between model’s loss and DP training.
As our theoretical results concern generalization, not loss per se, our results do not contradict these theoretical trade-offs.
We discuss the relationship in detail in Appendix A.3.
Group-Distributional Robustness.
Group-distributional robustness aims to improve the worst-case group performance.
Existing approaches include using worst-case group loss [Mohri et al., 2019,
Sagawa et al., 2019a, Zhang et al., 2020], balancing majority and minority groups by reweighting or subsampling [Byrd and Lipton, 2019, Sagawa et al., 2019b, Idrissi et al., 2022], leveraging generative models [Goel et al., 2020], and applying various regularization techniques [Sagawa et al., 2019a, Cao et al., 2019].
Although some work [Sagawa et al., 2019a, Cao et al., 2019] discusses the importance of regularization in distributional robustness, they have not explored potential reasons for this (e.g. via the connection to generalization).
Another line of work studies how to improve group performance without group annotations [Duchi et al., 2021, Liu et al., 2021, Creager et al., 2021], which is a different setting from ours as we assume the group annotations are known.
Robust Overfitting.
Rice et al.
[2020], Yu et al.
[2022] have shown that adversarially trained models tend to overfit in terms of robust loss.
Rice et al.
[2020] proposed to use regularization to mitigate overfitting, but the noisy gradient has not been explored for this.
We showed that the WYSIWYG framework can serve as an alternative direction for mitigating and explaining this issue.
Conclusions and Future Work
We argue that a “What You See is What You Get” property, which we formalize through the notion of distributional generalization (DG), can be desirable for learning algorithms, as it enables principled algorithm design in settings including deep learning.
We show that this property is possible to achieve with DP training.
This enables us to leverage advances in DP to enforce DG in many applications.
We propose enforcing DG as a general design principle, and we use it to construct simple yet effective algorithms in three settings.
In certain fairness settings, we largely mitigate the disparate impact of differential privacy by using importance sampling and enforcing DG in our new algorithm DPIS-SGD.
In our analysis, however, the privacy and DG guarantees of DP-IS-SGD deteriorate in the presence of very small groups.
Future work could explore individual-level accounting [Feldman and Zrnic, 2021] for a tighter analysis.
In certain worst-case generalization settings, inspired by DP-SGD, we propose using a noisy-gradient regularizer.
Compared to SOTA algorithms in DRO, noisy gradient achieves competitive results across many standard benchmarks.
In certain adversarialrobustness settings, our proposed noisy-gradient regularizer significantly reduces robust overfitting.
An interesting direction for future work would be to explore its effectiveness in large-scale settings, e.g., ImageNet [Croce and Hein, 2022].
We hope future work can explore extending this design principle to ensure generalization of other properties, such as calibration and counterfactual fairness.
Acknowledgements
We thank Kamalika Chaudhuri, Benjamin L. Edelman, Saurabh Garg, Gautam Kamath, Maksym
Andriushchenko, and Cassidy Laidlaw for useful feedback on an early draft.
BK acknowledges support from the Swiss National Science Foundation with grant 200021-188824.
Yao-Yuan Yang
thanks NSF under CNS 1804829 and ARO MURI W911NF
A.1 Differential Privacy and Robust Generalization
DP is known to imply a stronger notion of generalization, called robust generalization , which is a “tail bound” version of DG [Dwork et al., 2015a, Cummings et al., 2016, Bassily et al., 2016].
The original motivations for robust generalization are slightly different, but in our notation, a training procedure T 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 that satisfies ( γ, η ) -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 previous implication.
Thus, it is possible to recover the result that DP implies DG as a consequence of these previous works, although with looser bounds.
The difference between Distributional Generalization and Robust Generalization is that DG considers all quantities in expectation, while robust generalization considers tail bounds with respect to the train dataset.
We focus on DG for two reasons: First, we believe DG is conceptually simpler, as it can be seen as simply the TV distance between two natural distributions, and does not involve additional parameters.
This simplicity is conceptually useful to the algorithm designer, but also enables us to prove simpler tight theoretical bounds that are independent of sample size.
Second, it is often possible to 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 obtained for these methods if desired.
A.2 Information-Theoretic Generalization Bounds
Other than robust generalization, it is also possible to obtain bounds on generalization of arbitrary test (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 such information-theoretic measure— conditional mutual information (CMI) between training algorithm outputs and the training dataset—is bounded if the training algorithm is DP or TV stable.
Thus, we could relate stability and DG as in Section 2.2 using CMI as an intermediate tool.
In particular, Steinke and Zakynthinou show that for a given ϕ : D × Θ → [0 , 1] :
[Equation]
where CMI D ( T ) is the conditional mutual information of the training algorithm with respect to the data distribution.
If the training algorithm satisfies ( ϵ, 0) -DP, they also show that CMI D ( T ) ≤ n 2 ϵ 2 , where n is the dataset size.
Plugging this into Equation (6), we can see that the generalization upper bound (right-hand side) is ϵ .
This is significantly looser than our bound in Section 2.2, as illustrated in Figure 2.
A recent line of work on information-theoretic bounds explores sharper generalization bounds using individual-level measures [Bu et al., 2020a, Haghifam et al., 2020].
Analogously, as a direction for future work, it could also be possible to obtain tighter bounds on DG using per-instance notions of stability [Wang, 2017, Feldman and Zrnic, 2021].
A.3 Tension between Differential Privacy and Algorithmic Fairness
Beyond empirical observations that training with DP results in disparate impact on performance across subgroups [Bagdasaryan et al., 2019, Pujol et al., 2020], Cummings et al.
[2019] and, more recently, Sanyal et al.
[2022] theoretically analyze the inherent tensions between DP and algorithmic fairness.
It might appear that this trade-off contradicts our results in which we claim that using DP or similar noise-adding algorithms with additional train-time interventions can reduce disparate impact.
However, Cummings et al.
[2019] and Pujol et al.
[2020] discuss the relationship between privacy and disparate performance (accuracy or false-positive/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 theoretical results 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, thus in our regime we can achieve meaningful subgroup performance using the DP-IS-SGD algorithm despite the fundamental trade-off.
B Proofs Omitted in the Main Body B.1 TV-Stability implies Distributional Generalization
Proof of Theorem 2.5.
First, observe that the following distributions are equivalent as the dataset is an i.i.d.
sample:
[Equation]
It is thus sufficient to analyze the equivalent distributions instead.
By the post-processing property of differential 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 z
[Equation]
z ′ ∼ D and S ∼ D n −
[Equation]
we get:
[Equation]
where the last equality is simply renaming of the variables for convenience.
Note that analogously we also can obtain a symmetric bound:
[Equation]
The total variation between these two distributions is bounded:
[Equation]
where the last inequality is by Equation (9).
Using the equivalences in Equation (7) 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 [Wasserman and Zhou, 2010].
Let us define the hypothesis-testing setup and the two types of errors in hypothesis testing.
For any two probability distributions P and Q defined over D , let ϕ : D → { 0 , 1 } be a hypothesis-testing decision rule that aims to tell whether a given observation from the domain D comes from P or Q .
Definition B.1 (Hypothesis-testing FPR and FNR) .
Without loss of generality, the false-positive error rate α ϕ (FPR, or type I error rate), and the false-negative error rate β ϕ (FNR, or type II error rate) 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 between the 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 T ( 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
T ( S ) and T ( S ′ ) .
Let us sum the two bounds in Equation (12):
[Equation]
Let us take the optimal decision rule ϕ ∗ .
In this case, the bound in Equation (11) holds exactly:
[Equation]
Combining this with Equation (13), we get:
[Equation]
Next, we show that the upper bound is tight:
Proposition B.3.
There is an algorithm T ( S ) satisfying ( ε, δ ) -DP, such that d TV ( T ( S ) , T ( S ′ ))
=
exp( ε ) − 1+2 δ exp( ε )+ for any two neighbouring datasets S and S ′ .
Proof.
We use the construction of the reduced mechanism by [Kairouz et al., 2015].
Consider a mechanism T : { 0 , 1 } → { 0 , 1 , 2 , 3 } , defined as follows:
[Equation]
Observe that this mechanism satisfies ( ε, δ ) -DP, and d TV ( T (0) , T (1)) =
exp( ε ) − 1+2 δ exp( ε )+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 T ( S ) satisfies ( ϵ, δ ) -DP and Sample ( S ) is a Poisson sampling procedure where each of the sampling probabilities p i depend on the element
z i (but do not depend on the set S otherwise) and is guaranteed to satisfy p i ≤ p ∗ .
Then T ◦ 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 .
We wish 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 is analogous.
Note that with probability p 0 ≤ p
the element z 0 is included in Sample ( S ′ ) and we have Sample ( S ′ )
= { z 0 } ∪ Sample ( S ) , otherwise the element z 0 is not included, and conditioned on znot being included Sample ( S ′ ) has the same distribution as Sample ( S ) .
Therefore,
Pr( T ( Sample ( S ′ ))
∈ K ) =
p
Pr( T ( { z 0 }∪ Sample ( S ))
∈ K )
+(1 − p 0 )
Pr( T ( Sample ( S ))
∈ K ) .
Now for each realization Sample ( S ) =
˜ S , we have Pr( T ( { z 0 } ∪ ˜ S ) ∈ K ) ≤ e ϵ Pr( T ( ˜ S ) ∈ K ) + δ by the assumed DP guarantee of the algorithm T ( S ) .
We can average over all possible subsets ˜ S to get
[Equation]
Plugging this back to the inequality (14), we get
[Equation]
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 f -privacy:
Definition B.5 ( f -Privacy Dong et al.
An algorithm T ( S ) satisfies f
-privacy if for any two neighbouring 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]
(15) 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 following privacy amplification:
Proposition B.6 (Bu et al.
[2020b]) .
Suppose that T ( S ) satisfies f -privacy, and Sample ( S ) is a uniform Poisson sampling procedure with sampling probability ¯ p .
The composition T ◦ Sample ( S ) satisfies f ′ -privacy with f ′ =
¯ pf
+ (1 − ¯ p ) I
d , where Id ( α ) =
1 − α is the trade-off function that corresponds to perfect privacy.
We show that a similar result holds for non-uniform Poisson subsampling:
Lemma B.7.
Suppose that T ( S ) satisfies f -privacy, and Sample ( S ) is a non-uniform Poisson sampling procedure, where the sampling probabilities p i depend on the element
z i (but do not depend on the set S otherwise) and each is guaranteed to satisfy p i ≤ p ∗ .
The composition T ◦ Sample ( S ) satisfies f ′ -privacy with f ′ =
p ∗ + (1 − p ∗ ) Id .
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 of numbers 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 binary vector ⃗b =
( b 1 , . . .
, b n ) ∈ { 0 , 1 } n , where each bit b i indicates whether an example z i ∈ S was chosen 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 the joint probability of ⃗b .
The composition T ( 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 corresponding to 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 z 0 ∈ D , we potentially get the worst-case privacy guarantee for the highest sampling probability, i.e., f = p ∗ f +
[Equation]
1 − p
[Equation]
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 learning algorithm T ( 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 T ( 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 T ( S ) satisfies δ -strong Distributional Generaliza- tion if and only if for all ϕ : D × Θ → [0 , 1] :
[Equation]
error [Shalev-Shwartz et al.,
, Xu and Raginsky, 2017, Russo and Zou, 2019], strong DG is an analogous variant of expected absolute error [Steinke and Zakynthinou, 2020].
Clearly any algorithm satisfying ( η, γ ) -Robust
Generalization satisfies also ( η + γ ) -strong DG, and any algorithm satisfying δ -strong-DG satisfies δ -DG.
Thus, δ -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 binary classification which satisfies 0 -DG, always outputs a classifier with test error
+ o (1) (almost completely useless), but with probability 1 / 2 outputs a classifier with training error
(and with probability 1 / 2 the training error is 1 ).
That is problematic: it is natural to expect from a training procedure with strong generalization guarantee, that if we used it to train a model up to a small training error, we should have high confidence that the test error is also small — but this example shows that our confidence cannot be larger than 1 / 2 .
The δ -strong DG solves this problem: The following proposition is a formal statement of the desired property, and can be easily deduced from Markov inequality.
Proposition C.2.
If the learning procedure T ( S ) satisfies δ -strong-DG, then for any bounded ϕ
( z ; θ ) we have:
[Equation]
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 of generalization as expected absolute difference, which are usually defined as E S ∼D n , T | E z ∼ S ℓ ( z, T ( S ))
− E z ∼D ℓ ( z, T ( S ))
| where ℓ is the loss function.
We can follow the same proof strategy as in [Bassily et al., 2016, Theorem 7.2], simplifying it significantly, 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 T ( S ) satisfies δ -TV stability, and the sample size | S | ≥ 2 δ 2 , then
T ( 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 bounded queries q : D → R satisfying | sup x ( q ( x )) − inf x ( q ( x )) | ≤ 1 .
We wish to find another 2 δ -TV stable mechanism ˜ M : D n → Q such that if we use notation q S ≜ M ( S ) and ˜ q S ≜ ˜ M ( S ) , we have
[Equation]
Consider the following mechanism ˜ M ( S ) :
For an input dataset S , obtain q S ≜ M ( S ) , and evaluate ˜ µ ≜ E z ∼ S q S ( z ) + 1 δn U , where U is a uniform random variable in [ − 1 / 2 , 1 / 2 ] .
Finally, if ˜ µ ≥ E z ∼ D q S
( z ) , then output q S , otherwise output − q S + 2 E z ∼D q S ( z ) .
As d TV ( η U , γ + η U ) ≤ γη , a mechanism µ q ( S ) ≜ E z ∼ S q ( z ) + 1 δn U is δ -stable.
Hence, by composition and post-processing, the mechanism ˜ M is 2 δ -TV stable.
For n ≥ 2 δ − 2 the inequality (16) 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 E
Z ≤ 2 E [ Zσ ] , which is exactly (16).
As the mechanism ˜
M satisfies 2 δ -TV stability, by Theorem 2.5 applied to the mechanism ˜ M , we can bound the right hand side of the inequality (16) by 4 δ .
Finally, for an arbitrary learning algorithm T ( S ) satisfying δ -TV stability, and any bounded test function ϕ
( z ; θ ) , we can describe mechanism ˜ M which, given a sample S ∼ D n , outputs a function q S ( z ) ≜
ϕ ( z ; T ( S )) .
By post-processing, such a mechanism M also satisfies δ -TV stability, and the inequality (16) becomes
[Equation]
As ϕ ( z ; θ ) was chosen in an arbitrary way, this proves that mechanism T ( S ) indeed satisfies the condition of 4 δ -strong-DG.
Lemma C.4.
Let Z be a non-negative random variable satisfying E Z ≥ δ and σ be a {± 1 } valued random 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]
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 from each group in G distributionally generalizes in expectation, as long as the model encounters at least one representative of the group in training.
In its definition, we explicitly prevent the case when the training dataset does not contain any group examples to avoid undefined behavior (what is the group accuracy 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]
The statement follows by applying each step from the proof of Theorem 2.5 to the equivalent distributions in Equation (17).
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 show that strong Distributional Generalization (Definition C.1) implies generalization of calibration properties.
For simplicity we consider binary classifiers, with Y = { 0 , 1 } and f θ : X → [0 , 1] .
Let us define the notion of calibration we consider:
Definition C.7.
The calibration gap , or Expected Calibration Error (ECE), of a classifier f θ is defined 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 τ -binned calibration 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 S , is:
[Equation]
Next, we show strong DG implies that the expected binned calibration gap is similar between train and test.
Theorem C.8.
If the training method T ( S ) satisfies δ -strong DG, then
[Equation]
Proof.
First
[Equation]
let us define the family of tests:
[Equation]
The assumed δ -strong DG implies that each of these tests ϕ p have similar expectations between train and test distributions:
[Equation]
Thus
[Equation]
the difference in expected calibration gaps between train and test is:
[Equation]
[Table]
[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 chance of 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:
[Equation]
, where ℓ ( z, θ ) is the cross-entropy loss and z ∈ S drawn uniformly random drawn from the dataset, and g is the group to which z belongs
[Table]
E Additional Experiment Details E.1 Details on Datasets, Software, and Model Training
[Table]
[Table] [Table]
13013 Mammalia  3624 Mollusca  1057 Plantae  22117 Protozoa  37 Reptilia
[Table]
[Table] [Table]
All algorithms are implemented in PyTorch 4 [Paszke et al., 2019].
For DP-related utilities, we use opacus 5 [Yousefpour et al., 2021].
Other packages, including numpy 6 [Harris et al., 2020], scipy
[Virtanen et al., 2020], tqdm 8 , and pandas 9 [pandas development team, 2020], are also used.
For gDRO [Sagawa et al., 2019a], 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 Koh et al.
For UTKFace and iNat, we randomly split the data into 17000/2000/4708 and 550000/50000/75170 for training/validation/testing.
For MNLI, we use the same training/validation/testing split in Sagawa et al.
[2019a].
For Adult [Kohavi et al., 1996], we randomly split the data into 35000/3000/5842 for training/validation/testing.
Tables 2 to 7 show the dataset statistics on each group.
All the datasets are publicly available for non-commercial use.
In our work, we adhere to additional rules regulating the use of each dataset.
All datasets other than iNat could potentially contain personally identifiable information, and are likely collected without consent, to the best of our knowledge.
They are all, however, collected from manifestly public sources, such as public posts on social media.
Thus, we consider the associated privacy risks low.
The data also contain offensive material (e.g., explicitly in the case of CivilComments dataset).
We consider the associated risks of reproducing the offensive behavior low, as we use the datasets only to evaluate our theoretical and theoretically-inspired results.
Models.
Similar to previous work [Sagawa et al., 2019a], we use the ImageNet-1k pretrained
ResNet50 [He et al., 2016] from torchvision for CelebA, UTKFace, and iNat, and use the pre-
[Table]
[Devlin al., 2019]
For ADULT, we follow the setup in [Xu et al., 2021] and use logistic regression with standard optimization, and DP-based training methods.
We fix the batch size to 256 (for SGD), weight decay to 0 .
01 , and number of epochs to 20 .
For the DP algorithms, we use gradient norm clipping to 0 .
5 , and sampling rate of 0 .
For all training algorithms, we train five model times with different random seeds and we record the mean and standard error of the mean of our metrics.
The noise parameter σ for DP-SGD-F and DP-SGD is set to 1 . 0 , and we set the σ for DP-IS-SGD to 5 . 0 to achieve similar privacy budget ϵ ≈ 0 .
The additional noise parameter for DP-SGD-F σ 2 is set to 10 σ as in Xu et al.
Hyperparameters.
We run 50 epochs for CelebA, 100 epochs for UTKFace, 20 epochs for iNat, and 5 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 [Loshchilov and Hutter, 2019] optimizer.
We use opacus ’s
[Yousefpour et al., 2021] implementation of DP-SGD and 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 .for
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 our experiments it is practically equivalent to using the generalization gap of subgroup accuracy, which is bounded by TV stability.
Let us first concretely define the generalization gap of the worst-case group accuracy:
Definition E.1.
The on-average generalization gap of the worst-case accuracy is defined as the following difference:
[Equation]
where we take ℓ (( x, y ) , θ ) ≜ 1 [ f θ ( x )
=
y ] to be the 0/1 loss.
In this definition we explicitly restrict the datasets to include elements of each group g ∈ G , which is a technicality needed in order to avoid undefined behavior.
In all our experimental results, the worst-performing groups (the maximizers in Equation (19)) are always the same on the training and test data.
As long as this holds—the worst-performing group is the same on the train and test data—the generalization gap above simplifies to:
[Equation]
[Figure] [Figure]
[Figure]
oup
[Figure] [Figure]
[Figure]
W o r s t - g
[Figure]
oup acc u
[Figure]
ac y
Train
[Figure] [Figure]
[Figure]
10 4 Iterations 0.
[Figure] [Figure]
W o
[Figure]
s t -
[Figure]
oup acc
[Figure]
ac y
[Figure]
where g ∗ ∈ G is the worst-performing group.
In Appendix C.1 we show that this simplified gap from
Equation (20) is bounded by TV stability.
Therefore, in practice the generalization gap in Equation (19) offers a lower bound on the DG gap in 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 of empirical 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 accuracy can be loose as a proxy.
Finding the worst-case test function is an object of study in the literature on membership inference attacks [Shokri et al., 2017], because DG and the accuracy of such attacks in their standard formalization are equivalent, as showed by Kulynych et al.
In this work, we opt for a simpler and direct approach described above.
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.
[2019a] show that strong ℓ 2 regularization helps with improving group-distributional generalization, and Yang et al.
[2020] show that dropout helps with adversarial-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 regularization methods 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 CelebA using DP-SGD, and decrease the “regularization strength” in several different ways: by increasing privacy budget ϵ (Figure 6a), decreasing the ℓ 2 regularization (Figure 6b), or increasing the number of training iterations (Figure 6c).
10 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-group accuracy 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 ℓregularization).
We have three observations: (1) with properly tuned regularization parameter, the gap between training and testing worst-group accuracy can be narrowed, (2) the gap can start widening in very early stage of training, and (3) the testing worst-group accuracy can fluctuate largely, which highlights 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.
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 likely because
UTKFace has a relatively small difference in the number of training examples between the largest group and the smallest group.
In UTKFace, the majority group has around seven times more examples than in the minority group, whereas in CelebA, this difference is 52 × .
Comparison with DP-SGD-F [Xu et al., 2021].
We did not manage to obtain good performance from DP-SGD-F on CelebA, UTKFace, and iNat, possibly because of the different domain—images— than tabular data considered by Xu et al.
To proceed with the comparison, we evaluate the algorithms on the census data—ADULT dataset [Kohavi et al., 1996]
(see Table 7 for dataset statistics)—that
Xu et al.
[2021] used in their work.
As subgroups, we consider four intersectional groups composed of all possible values of the “sex” attribute and prediction class (an income higher/lower than 50k).
We show the results in Table 9.
For a comparable epsilon value (0.69 for DP-SGD-F, and
[Figure]
r oup acc
[Figure]
[Table]
E.5 Additional Details for Section 5.
We compare different algorithms, including SGD- ℓ 2 and IW-SGD- ℓ 2 as baselines, and two other algorithms, IS-SGD- ℓ 2 [Idrissi et al., 2022] and gDRO- ℓ 2 [Sagawa et al., 2019a] in terms of the group robustness.
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 from. 0001 , 0 . 01 , 0 . 1 , and 1 . 0 .
For DP-IS-SGD, we fix the gradient clipping to 1 .
(except for iNat, where we set the value to 10 .
as 1 . 0 does not converge).
We select the noise parameter from 1 . 0 , 0 . 1 , 0 . 01 , 0 .on
CelebA and UTKFace, select the noise parameter from 0 . 0000001 , 0 . 000001 , 0 .
and 0 .
on iNat and select the noise parameter from 0 .
01 and 0 .
001 on CivilComments and MNLI.
For IW-SGD-n, IS-SGD-n, and gDRO-n, we select the standard deviation of the random noise from 0 .
001 , 0 . 01 , 0 . 1 , and 1 .
on CelebA, UTKFace, and iNat, and we select standard deviation of the random noise from 0 .
00001 , 0 . 0001 , and 0 .
001 on CivilComments and MNLI.
Statistical Concerns.
Although our results appear to be comparable to or better than SOTA, we caution readers about the exact ordering of methods due to high estimation variance: these benchmarks have small validation and test sets (e.g., CelebA has 182 validation examples), and so hyperparameter tuning is subject to both overfitting and estimation error.
For example, we observe validation accuracies which differ from their test accuracies by up to 5% in our experiments.
We attempt to mitigate this using three random train/val/test splits on CelebA, and avoid large hyperparameter sweeps 11 , but this is not done in prior work.
E.6 Additional Details for Section 5.
We use the CIFAR-10 dataset [Krizhevsky et al., 2009], and ResNet-18 [He et al., 2016] as the network architecture.
We train the model to be robust against L ∞ perturbations of at most γ = 8 /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.
In addition, we compare the performance of noisy gradient to adversarial training with early stopping —a simple but effective approach for mitigating overfitting in adversarial training [Rice et al., 2020].
In this experiment, we measure robust accuracy and its respective generalization gap, thus setting ℓ (( x, y ) , θ ) ≜ 1 [ f θ ( x ) =
y ] to be the 0/1 loss.
F Author Contributions
BK proposed to leverage the connection between differential privacy and distribution generalization, and led the theory developments and algorithm analysis.
Yao-Yuan Yang led the experiments, designed the experimental settings, and implemented the algorithms.
Yaodong Yu contributed to the experiments, especially in the NLP settings, and helped benchmark baselines.
The noisy-gradient algorithm is proposed by Yao-Yuan Yang and Yaodong Yu.
JB contributed to the theory, in particular the algorithmic analysis, initial proof of the tightness of the bound, and formal connections to robust generalization and calibration.
PN organized the team and managed the project, contributing in parts to the theory and experiments.
All authors participated in framing the results and writing the paper.