Measuring Generalization with Optimal Transport
Authors: Ching-Yao Chuang
Youssef Mroueh
Kristjan Greenewald
Antonio Torralba
Stefanie Jegelka
Venue: NeurIPS 2021 Spotlight (2021)
Keywords: generalization
optimal transport
representation

Abstract: Understanding the generalization of deep neural networks is one of the most important tasks in deep learning. Although much progress has been made, theoretical error bounds still often behave disparately from empirical observations. In this work, we develop margin-based generalization bounds, where the margins are normalized with optimal transport costs between independent random subsets sampled from the training distribution. In particular, the optimal transport cost can be interpreted as a generalization of variance which captures the structural properties of the learned feature space. Our bounds robustly predict the generalization error, given training data and network parameters, on large scale datasets. Theoretically, we demonstrate that the concentration and separation of features play crucial roles in generalization, supporting empirical results in the literature.

Measuring Generalization with Optimal Transport
Ching-Yan Yao Chuang', Youssef Mroueh', Kristjan Greenewald Antonio Torralba' Stefanic Jegelka
'MIT CSAIL, 'IBM Research AI, MIT-IBBM Watson Al Lab
{cychuang, torralba, stefje/@mit edu mrouehSus. ibm. com, kristjan. h htt.aaa.....eeeaeeeeee.. greenewald@ibm com
Abstract
Understanding the generalization of deep neural networks is one of the most im- portant tasks in deep learning. Although much progress has been made, theoretical error bounds still often behave disparately from empirical observations. In this work, we develop margin-bass -base generalization bounds, where the margins are nor- malized with optimal transport costs between independent random subsets sampled from the training distribution. In particular, the optimal transport cost can be inter- preted as generalization of variance which captures the structural properties of the learned feature space. Our bounds robustly predict the generalization error, given training data and network parameters, on large scale datasets. Theoretically, we demonstrate that the concentration and separation of features play crucial roles generalization, supporting empirical results in the literature. The code is available at pectps://githa com/ chingyaoc/ Berggyocc/V--eee -Vraann...
1 Introduction
Motivated by the remarkable empirical success of deep learning, there has been significant effort in statistical learning theory toward deriving generalization error bounds for deep learning, i.c complexity measures that predict the gap between training and test errors. Recently, substantial progress has been made, e.g., [3, 4, 7, 13, 21, 43, 56). Nevertheless, many of the current approaches lead to generalization bounds that are often vacuous or not consistent with empirical observations [14, 38].
In particular, Jiang et al. [25] present large scale study of generalization in deep networks and show that many existing approaches, e.g., norm-based bounds [4, 42, 43], are not predictive of generalization in practice. Recently, the Predicting Generalization in Deep Learning (PGDL) competition described in [26] sought complexity measures that are predictive of generalization error given training data and network parameters. To achieve a high score, the predictive measure of generalization had to be robust different hyperparameters, network architectures, and datasets. The participants [32, 40, 48] achieved encouraging improvement over the classic measures such as VC-dimension [54] and weight norm [4]. Unfortunately, despite the good empirical results, these proposed approaches are not yet supported by rigorous theoretical bounds.
In this work, attempt decrease this gap between theory and practice with margin bounds based on optimal transport. In particular, we show that the expected optimal transport cost of matching two independent random subsets of the training distribution is natural alternative to Rademacher complexity. Interestingly, this optimal transport cost can be interpreted as the k-variance [49], generalized notion of variance that captures the structural properties of the data distribution. Applied to latent space, it captures important properties of the learned feature distribution. The resulting k- variance normalized margin bounds can be easily estimated and correlate well with the generalization error on the PGDL datasets [26]. In addition, our formulation naturally encompasses the gradient
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

normalized margin proposed by Elsayed et al. [15], further relating our bounds to the decision boundary of neural networks and their robustness.
Theoretically, our bounds reveal that the concentration and separation of learned features are important factors for the generalization of multiclass classification. In particular, the downstream classifier generalizes well if (I) the features within a class are well clustered, and (2) the classes are separable in the feature space in the Wasserstein sense.
In short, this work makes the following contributions:
We develop new margin bounds based on k-variance [49], a generalized notion of variance based optimal transport, which better captures the structural properties of the feature distribution; We propose k-variance normalized margins that predict generalization error well on the PGDL challenge data;
We provide theoretical analysis shed light on the role of feature distributions in generalization, based on our k-variance normalized margin bounds.
2 Related Work
-asged Generalization Bounds Classic approaches in learning theory bound the generaliza- tion error with the complexity of the hypothesis class [5, 54J. Nevertheless, previous works show that these uniform convergence approaches are not able to explain the generalization ability of deep neural networks given corrupted labels [60] or on specific designs of data distributions as in [38]. Re- cently, substantial progress has been made to develop better dete-deenndent and iigorithm-dependen bounds [3, 6, 7, 44, 52, 56]. Among them, we will focus on margin- mragginbbased generalization bounds for multi-class classification [29, 31]. Bartlett al. [4] show that margin normalized with the product of spectral norms of weight matrices is able to capture the difficulty of the learning task, where the conventional margin struggles. Concurrently, Neyshabur et al. [43] derive opectrally-normalizee margin bounds via weight perturbation within a PAC-Bayes framework [34]. However, empirically, spectral norm-based bounds can correlate negatively with generalization [38. 25J. Elsayed et al. [15] margin, which can be interpreted as the first order approximation to
present gradient-normalizze the distance the decision boundary. Jiang al. [24] further show that normaltt- margins, when combined with the total feature variance, are good predictors of the generalization gap. Despite the empirical progress, gradient- normalized margins are not yet supported by theoretical bounds.
Empirical Measures of Generalization Large scale empirical studies have been conducted to study various proposed generalization predictors [24, 25J. In particular, Jiang et al. [25] measure the average correlation between various complexity measures and generalization error under different experimental settings. Building their study, Dziugaite et al. [14] emphasize the importance of the robustness of the generalization measures to the ex experimental setting. These works show that well-known complexity measures such as weight norm [37, 42], spectral complexity [4, 43], and their variants are often negatively correlated with the generalization gap. Recently, Jiang et al. [26] hosted the Predicting Generalization Deep Learning (PGDL) competition, encouraging the participants to propose robust and general complexity measures that can rank networks according to their generalization errors. Encouragingly. several approaches [32. 40, 48] outperformed the conventional baselines by large margin. Nevertheless, none of these are theoretically motivated with rigorous generalization bounds. Our k-variance normalized margins are good empirical predictors of the generalization gap, while also being supported with strong theoretical bounds.
3 Optimal Transport and k- V-riince
Before presenting our generalization bounds with optimal transport, first give brief introduction to the Wasserstein distance, a distance function between probability distributions defined via an optimal transport cost. Letting μ and Prob(1Rd) be two probability measures, the p- Wassersteii distance with Euclidean cost function is defined as
(ExxxalX-YVI 1/p
W) = inf
=Ell(n,e)
2

---

where I[ v) S Prob(R x R denotes the set of measure couplings whose marginals are μ and V. respectively. The I-Wasserstein distance is also known as the Earth Mover distance. Intuitively. Wasserstein distances measure the minimal cost to transport the distribution to v.
Based on the Wasserstein distance, Solomon al. [49] propose the k-variance, generalization of variance, to measure structural properties of distribution beyond variance.
Definition (Wasserstein--) k- k-variance). Given probability measure μ Prob(IRR) and parameter E N, the Wasserstein-p k-variance is defined as
par = (k, d) Bgg (Ws(
where us = Ei= 8z, for D2 lid. I and cp(k: d) is normalization term described in [49].
When A: and p = 2, the k-variance equivalent to the variance Var [X] of the random variable X J.. For k >> and d 2 3. Solomon et al. [49] show that when 2, the k-variance provides an intuitive way to measure the average intra-cluster variance of clustered measures. this work, we use the unnormalized version (c,(k, d) 1) of k-variance and p = 1, and drop the p in the notation:
(Wasserstein-1 k k-variance): Vara(p) $ [W (g)ss
Note that setting (k d) 1 is not an assumption, but instead an alternative definition of k-variance. The change in constant has no effect on any part of our paper, as we could reintroduce the constant of Solomon et al. [49] and simply include premultiplication term in the generalization bounds to cancel out. Section we will show that this unnormalized Wasserstein- k-variance captures the concentration of learned features. Next, we use it to derive generalization bounds.
4 Generalization Bounds with Optimal Transport
We present our generalization bounds in the multi-class setting. Let X' denote the input space and y = (1,...,
, K' denote the output space. We will assume a compositional hypothesis class F o $ where the hypothesis f ø can be decomposed as feature (representation) encoder & e 4 and predictor f E F. This includes dividing multilayer neural networks at an intermediate layer.
We consider the score-based classifier [f1,-- ..... fx), fc Fe, where the prediction for X' is given by arg max, yay fy(((z)). The margin of f for a datapoint (Ir,) y) is defined by
p)(((x), 4) == fy(((z)) max for (b(x))
(1)
v'
where / misclassifies if ()(())), S O. The dataset S = { is drawn i.i.d. from distribution j over X' x y. Define me as the number of samples in class c, yielding m >x me- We denote the marginal over class y as He and the distribution over classes by p(c). The
pushforward measure of μ with respect & is denoted as Φ#fl. We are interested in bounding
the expected zero-one loss of hypothesis f Φ: R o Φ) (eoa))/(eor(((c)a)ssss by the corresponding empirical 7-margin loss R.,. (f o )) = E (l
4.1 Feature Learning and Generalization: Margin Bounds with k- V-aainacee
Our theory is motivated by recent progress in feature learning, which suggests that imposing certain structure the feature distribution improves generalization [8, 28, 57, 59]. The participants [32, 40] of the PGDL competition [26] also demonstrate nontrivial correlation between feature distribution and generalization.
To study the connection between learned features and generalization, we derive generalization bounds based on the k-variance of the feature distribution. In particular, we first derive bounds for fixed encoder and discuss the generalization error of the encoder at the end of the section. Theorem provides a generalization bound for neural networks via the concentration of M.. in each class.
Theorem 2. Let f = [f...... fx] E F = Fi x ... x FK where F; : X > R. Fix y > 0. The following bound holds for all f E F with probability at least 1 - 8 0:
Lip(er(s Varm, (0#He) + log(1/5)
Rf x) < R, (f Φ) Boup
7
2m
3

---

where Lip(ps(;:)) suPo, 101( is the margin Lipschitz constant w.r.t
10())-0(<']l>=
We give a proof sketch here and defer the full proof to the supplement. For a given class and a given feature map Φ, let H. = (h(x)) = P (,(x), c)ls = (fi... fx): Jy E F). The last step deriving
Rademacher-baser generalization bounds [5] amounts to bounding for each class c:
c
E S
sup _ hEH. Me 1=1
h(T:)
Σ h((3)))
(2)
where S,S "" Typically would plug in the Rademacher variable and arrive the standard Rademacher generalization bound. Instead, our key observation is that the Kantorovich-Rubine duality [23] implies
Wi (i(A, v)
sup Ex))(())) Bex Lip(AjS1
where the supremum is over the 1-Lipschitz functions h : Rd + R. Suppose He is subset of L-Lipschitz functions. By definition of the supremum, the duality result immediately implies that (2) can be bounded with k-variance for k = mc:
A. S L = ( [Wi(o#Ms, 0+145]] = L (o=me(
(3)
This connection suggests that k-variance is a natural alternative to Rademacher complexity the margin is Lipschitz. The following lemma shows that this holds when the functioon, fn are Lipschitz: Lemma 3. The margin ps(- y) is Lipschitz in its first argument if each of the f, Lipschitz.
The bound in Theorem 2 is minimized when (a) the k-variance of features within each class is small, (b) the classifier has large margin, and (c) the Lipschitz constant of is small. In particular, (a) and (b) express the idea of concentration and separation of the feature distribution, which we will further discuss in Section 6.
Compared the margin bound with Rademacher complexity [31], Theorem 2 studies fixed encoder, allowing the bound to capture the structure of the feature distribution. Although the Rademacher- based bound is also data-dependent, it only depends on the distribution over inputs and therefore can neither capture the effect of label corruption nor explain how the structure of the feature distribution dg (p) affects generalization. Importantly, is also non-trivial estimate the Rademacher complexity empirically, which makes it hard to apply the bound in practice.
4.2 Gradient Normalized (GN) Margin Bounds with k- Variance
We next extend our theorem to the gradient- gradient-normalizen margin, variation of the margin (1) that empirically improves generalization and adversarial robustness [15, 24). Elsayed et al. [15) proposed it to approximate the minimum distance to decision boundary, and Jiang al. [24] simplified it to Py (6(z), y) == ()(((x) v) ((IV (6(z), v) ll2 +
where e is small value (10-6 in practice) that prevents the margin from going to infinity. The Vofyman (o(x)). where ties among the Vmax are broken arbitrarily in [15, 24] (ignoring subgradients). The normallzed margin py (z, v) can be interpreted as the first order Taylor approximation of the minimum distance of the input æ to the decision boundary for the class pair (v. s²) [15]. In particular, the distance is defined as the norm of the minimal perturbation in the input or feature space to make the prediction change. See also Lemma 20 in the supplement for interpretation of this margin in terms of robust feature separation. Defining the margin loss RV (f) (a.))) ( (()())))))) we extend Theorem 2 to the gradient-normalize margin.
gradient here is Vopy (0(x),)), y) == Vo.fv(á(z))
Theorem 4 (Gradient-oomalliza Margin Bound). Let f = [fi, ...
F= [F..... FF)) where F, : X + R. Fix > O. Then, for any 8 > O, the following bound holds for all f E F with probability least 8 0:
J
Liplōn)((.c) Varm. (6 (om
log(1/0) 2m
Rf o 6) < RV (f o Φ)) Epop
7
where Lip(ps = sup
By is the Lipschitz constant defined w.rlt th. ((02))

---

4.3 Estimation Error of k- Variance
So far, our bounds have used the v-variance, which is an expectation. Lemma 5 bounds the estimation error when estimating the k-variance from data. This may be viewed as generalization error for the learned features in terms of k-variance, motivated by the connections of k-variance test error. Lemma 5 (Estimation Error of k- Variance / "Generalization Error" of the Encoder). Given distribution and n empirical samples $" where each S , ~ μ², define the empirical Wasserstein-I k-variance: Vark,n(0=11) = E W1 9# ). Suppose the encoder satisfies supz,æ [1((z) 6(x')) S B, then with probability at least - 6 > 0, we have
2B log( (og(1/6)
Varx(0#11) S Vark,((0#11 +
nk
We can then combine Lemma 5 with our margin bounds to obtain full generalization bounds. The following corollary states the empirical version of Theorem 2:
Corollary 6. Given the setting in Theorem 2 and Lemma 5, with probability at least 1 - 6, for
m ===1 2n
me
Rf o 6) is upper bounded by
Lip(es(.c) 7
log(2K/0) 12 Y
log(
Ry (f x6 +
Eowpy
Var
1 (P+ + 2B, +2B,
the
+
2m
2n
Note that the same result holds for the gradient normalized margin Đy- The proof of this corollary simple application of a union bound on the concentration of the k-variance for each class and the concentration of the empirical risk. We end this section by bounding the variance of the empirical k-variance. While Solomon et al. [49] proved a high-probabillty concentration result using McDiarmid's inequality, we here use the Efron-Stein inequality to directly bound the variance.
Theorem 7 (Empirical variance). Given a distribution and an encoder Ø, we have
(Var,(((X))
Var Varu,n
nk
where (orx()(X)) = EE=uull((()) (x=))( is the variance of @
Theorem 7 implies that if the feature distribution / has bounded variance, the variance of the empirical k-variance decreases as l and n increase. The values of we used in practice were large enough that the empirical variance of k-variance small even when set 1.
5 Measuring Generalization with Normalized Margins
We now empirically compare the generalization behavior of neural networks to the predictions of our margin bounds. To provide unified view of the bound, we set the second term in the right hand side of the bound to a constant. For instance, for Theorem 2, we choose y = % Epep [Lip(ps(<,c) Varm. (0 [c)], yielding [ o Φ) & R. a 6) + 1/70 + ((m where
B( =)
Ry o0) =
(pr(o(s) BB [Li(pp((sc) Varm, (d+He)) S n0)
and [1() is the indicator function. This implies the model generalizes better if the normalized margin larger. We therefore consider the distribution of the k-variance normalized margin, where each data point is transformed into single scalar via
Py ()(()),,)) v)
مy(((x),)) v)
and
Egoup
Var me J.1 (b# Lip(er(.c)))
Eawp
Var me J.I (d (ip()r(.c))
respectively. For simplicity, we set l and as k: [ma/2] and n = 1. We refer to these normalized margins as k- Variance normalized Margin (A:V-Margin) and k- Variance Gradient Normalized Margin (#V-GN-Margin) respectively.
5

---

CIFAR
SVHN
CINIC
CINIC
Flowers
Pets
Fashion
CIEAR
VGG
NiN
FCN
FCN
NiN
NIN
VGG
NiN
Margin
13.59
16.32
2.03
2.99
0.33
1.24
0.45
5.45
SN-Margin
5.28
3.11
0.24
2.89
0.10
1.00
0.49
6.15
GN-Margin 1st [24]
3.53
35.42
26.69
6.78
4.43
1.61
1.04
13.49
GN-Margin 8th [24]
0.39
31.81
7.17
1.70
0.17
0.79
2.12
1.16
1st [24]
19.22
36,90
31.70
16.56
4.67
4.20
0.16
25,06
TV-GN-Margin 8th [24]
38.18
41.52
6.59
16.70
0.43
5.65
2.35
10.11
kV-Margin 1st
5.34
26.78
37.00
16.93
6.26
2.07
1.82
15.75
VV-Margin 8th
30.42
26.75
6.05
15.19
0.78
1.76
0.33
2.26
RV-GN-Margin' 1st
17.95
44.57
30.61
16.02
4.48
3.92
0.61
21.20
AV-GN-Margin' 8th
40.92
45.61
6.54
15.80
1.13
5.92
0.29
8.07
Table Mutual information scores on PGDL tasks. We compare different margins across tasks in PGDL. The first and second rows indicate the datasets and the architecture types used by tasks. The methods that are supported with theoretical bounds are marked with T. Our k-variance normalized margins outperform the baselines in 6 out of 8 tasks in PGDL dataset.
is NP-hard compute the exact Lipschitz constant of ReLU networks [47, 27]. Various approaches have been proposed to estimate the Lipschitz constant for ReLU networks [27, 16]. however they remain computationally expensive. As we show in Appendix B.4, naive spectral upper bound on the Lipschitz constant leads to poor results in predicting generalization. On the other hand, as observed by [27], a simple lower bound can be obtained for the Lipschitz constant of ReLU networks by taking the supremum of the norm of the Jacobian on the training set.' Letting y" arg max vsc fv(((z)), the Lipschitz constant of the margin can therefore be empirically approximated
Lip(ps(s, c)) == max IV ( zES.
(é(z)))
where S. = (2 9e) E 1/4 c} is the set of empirical samples for class (as noted in [27], although this does not lead to correct computation of Jacobians for ReLU networks, it empirically performs well).
practice, we take the maximum over samples in the training set. We refer the reader to [39] and [51] for an analysis of the estimation error of the Lipschitz constant from finite subsets. In the supplement (App. C.1), we show that for piecewise linear hypotheses such as ReLU margin is very close to 1 almost everywhere. We thus simply set the Lipschitz constant for the gradient-normalize margin.
networks, the norm of the Jacobian of the
5.1 Experiment: Predicting Generalization in Deep Learning
We evaluate our margin bounds on the Predicting Generalization in Deep Learning (PGDL) dataset [26]. The dataset consists 8 tasks, each task contains a collection of models trained with different hyperparameters. The models in the same task share the same dataset and model type, but can have different depths and hidden sizes. The goal correlates with their generalization error. In particular, the complexity measure maps the model and training dataset to real number, where the output should rank the models in the same order as the generalization error. The performance is then measured by the Conditional Mutual Information (CMI). Intuitively, CMI measures the minimal mutual information between complexity measure and generalization еrгor conditioned measure must robust all possible settings including different architectures, learning rates, batch sizes, etc. Please refer to [26] for details.
find complexity measure of networks that
different sets of hyperparameters. To achieve high CMI, the
Experimental Setup. We compare our k-variance normalized margins V-a-gnn and AV-GN- Margin) with Spectrally-Normalizee Margin (SN-Margin) [4], Gradient-Normalized Margin (GN- Margin) [15], and vaaa-varrinnee-nmmliieeedd GN-Margin (TV-GN-Margin) [24]. Note that the (TV-GN-Margin) of [24] corresponds to Pr(((z))))
/ /aar,., ([ Comparing to our kV-GN-Margin, our
normalization is theoretically motivated and involves the Lipschitz constants of well as the generalized notion of k-variance. As some of the margins are defined with respect to different layers
In general, the Lipschitz constant of smooth, scalar valued functions is equal to the supremum of the norm of the input Jacobian the domain [17, 27, 47].

---

of networks, we present the results with respect to the shallow layer (first layer) and the deep layer (8th layer the number of convolutional layers is greater than 8, otherwise the deepest convolutional layer). To produce a scalar measurement, we use the median to summarize the margin distribution, which can be interpreted as finding the margin y that makes the margin loss 0.5. We found that using expectation or other quantiles leads to similar results. The Wasserstein-1 distance in k-variance computed exactly, with the linear program the POT library [18]. All of our experiments are run on 6 TITAN X (Pascal) GPUs.
To ease the computational cost, all margins and k-variances are estimated with random subsets size min 220 x #classes, data size) sampled from the training data. The average results over 4 subsets are shown Table 1. Standard deviations are given in App. B.2, as well the effect of varying the size of the subset in App. B.3 Our v-variance normalized margins outperform the baselines in 6 out of tasks. Notably, our margins are the only ones achieving good empirical performance while being supported with theoretical bounds.
Margin Visualization. To provide qualitative comparison, we select four models from the first task of PGDL CUFARIO/VGGO) which have generalization error 24.9%, 26.2%, 28.6%, and 31.8%, respectively. We visualize margin distributions for each in Figure 1. Without proper normalization, Margin and GN-Margin struggle to discriminate these models. Similar to the observation in [25], SN-Margtn Margin even negatively correlates with generalization. Among all the apporaches. &V-GN VV-GN-Marrin is the only measure that correctly orders and distinguishes between all four models. This is consistent with Table I, where GN-Margin achieves the highest score.
Fm 0.349
Em: 249
Em 0.249
Er: 0.286
Ea: 256
Far: 0.286
Density
Density
Log Density
Ear: 0.315
Margin
GN-Margin 8th
SN-Margin
0.249
0.249
Er: 0.262
En: 0.262
Ixn 0.262
Em:
Density
Eim 0.318
Density
Em:
Density
Ea: 0.318
TV-GN-Margin 8th
kV-Margin 8th
kV-GN-Margin 8th
Figure 1: Margin Visualization of PGDL Models. From left to right, the correct order of the margin distributions should be red, green, orange, and blue. &V-GN-Margin the only measure that behaves consistently with the generalization error.
Next, we compare our approach against the winning solution of the PGDL competition: Mixup" DBI [40]. Mixup" DBI uses the geometric mean of the Mixup accuracy [61] and Davies Bouldin Index (DBI) to predict generalization. In particular, they use DBI to measure the clustering quality of intermediate representations of neural networks. For fair comparison, we calculate the geometric mean of the Mixup accuracy and the median of the k-variance normalized margins and show the results in Table Following [40], all approaches use the representations from the first layer. Our Mixup" & MV-G-nn outperforms the ofate-of-the-art [40] in 5 out of the 8 tasks.
5.2 Experiment: Label Corruption
sanity check proposed in [60] to examine whether the generalization measures are able to capture the effect of label corruption. Following the experiment setup in [60], we train two Wide-ResNets [58], one with true labels (generalization error = 12.9%) and one with random labels (generalization error 89.7%) CIFAR- 10 [30]. Both models achieve 100% training accuracy. We select the

---

CIFAR
SVHN NiN
CINIC FCN
CINIC
Flowers
Pets
Fashion
CIEAR
VGG
FCN
NiN
NiN
VGG
NiN
Mixup" 'DBI [40]
0.00
42.31
31.79
15.92
43,99
12.59
9.24
25.86
Mixup" VV-Margin
7.37
27.76
39.77
20.87
9.14
4.83
1.32
22.30
Mixup" &V-GN-NNG Margin
20.73
48.99
36.27
22.15
4.91
11,56
0.51
25.88
Table 2: Mutual information scores PGDL tasks with Mixup. We compare with the winner (Mixup" (MI) of the PGDL competition [26]. Scores of Mixup" DBI from [40].
DBBI
feature from the second residual block to compute all the margins that involve intermediate features and show the results in Figure Without k-variance normalization, margin and GN-Margin can hardly distinguish these two cases, while k-variance normalized margins correctly discriminate them.
Clazi0 Ciferi Randoes Label
Clari Ciarri Ciferi Raaion Latel
Clarto
Ciazi0 Citero Raalon Label
Cila Brabken Latel
Density
Margin
GN-Margin
VV-Margin
kV-GN-Margin
Figure 2: Margin distributions with clean or random labels. Without k-variance normalization, Margin and GN-Margin struggle to distinguish the models trained with clean labels or random labels.
5.3 Experiment: Task Hardness
Next, we demonstrate our margins are able to measure the "hardness" of learning tasks. We say that learning task is hard if the generalization error appears large for well-trained models. Different from the PGDL benchmark, where only models trained on the same dataset are compared, we visualize the margin distributions of NesttR trained on CIFAR-10 10 [30], SVHN [41], and MNIST [33], which have generalization error 12.9%, 5.3%, and 0.7%, respectively. The margins are measured on the respective datasets. In Figure 3, we again see that k-variance normalized margins reflect the hardness better than the baselines. For instance, CIFAR-10 and SVHN are indicated to be harder than MNIST the margins smaller.
CIPARIO
CIEAROO
CIFARIO
CIPARIO
SVIIN
SVIIN
Density
MNISU
MNNST
Margin
GN-Margin
kV-Margin
KV-GN-Margin
Figure 3: CIFAR-10, SVHN, and MNIST have different hardness. Although models achieve 100% training accuracy on each task, the test accuracy differs. With k-variance normalization, the margin distributions of the models are able to recognize the hardness of the tasks.
6 Analysis: Concentration and Separation of Representations
6.1
Concentration of Representations
In this section, we study how the structural properties the feature distributions enable fast learning. Following the Wasserstein-2 k-variance analysis in [49], we apply bounds by Weed and Bach [55] to demonstrate the fast convergence rate of Wasserstein- 1 v-variance when (1) the distribution has low intrinsic dimension (2) the support is clusterable.
8

---

Proposition 8. (Low-dimensional Measures, Informal) For 64 E Prob(R*), we have Varm (0+1) I (mm=1/u) for d' 2. Q441 is supported on an approximately d' oddimensional set where d' < d, we obtain betler rale: Varm(0g S O(m-//d).
We defer the complete statement to the supplement. Without any assumption, the rate gets significantly worse as the feature dimension d increases. Nevertheless, for an intrinsically d' d'-dimensional measure, the variance decreases with faster rate. For clustered features,
obtain even stronger rate:
Proposition (Clusterable Measures) A distribution y (n, clussterae if supp(u) lies in the union of n balls of radius at most A. If @ is (n, -)-cluterr,ble, then for all m S (n(2) Varm (0#H) S / m
We arrive at the parametric rate ((m-1/2) if the cluster radius A is sufficiently small. In particular, the fast rate holds for large when the clusters well concentrated. Different from conventional studies that focus on the complexity of a complete function class (such as Rademacher complexity [5]) our k-variance bounds capture the concentration of the feature distribution.
6.2 Separation of Representations
We showed the previous section that the concentration of the representations is captured by the k-variance and that this notion translates the properties of the underlying probability measures into generalization bounds. Next, we show that maximizing the margin sheds light on the separation of the underlying representations in terms of Wasserstein distance.
Lemma 10. (Large Margin and Feature Separation) Assume there exist fy.y 1.... K that are L-Lipschitz and satisfy the max margin constraint Py P)(((z),) y) 2 Y for all (I, y) ~ D i.e: fg(ó(z)) 2 fy (x(x)) + 7: V v y, VI supp(uy). Then Vy + y.. W1 (1 (Rg), Q# (Py') Z 7.
Lemma 10 states that large margins imply Wasserstein separation of the representations of each class. It also sheds light on the Lipschitz constant of the downstream classifier F: L Z 7/ming,y' (1 (uy), Φ (Hg/)). One would need more complex classifiers, i.e, those with larger Lipschitz constant, to correctly classify classes that are close to other classes, by Wasserstein distance, in feature space. We further relate the margin loss to Wasserstein separation:
Lemma 11. Define the pairwise margin loss Ray' for v' E y as
Re ofo = = (Em-m, - + (o(x))) + Earpy h (o(x) + 5,
Assume fo is L-Lipschitz for all e y. Given a margin 7 > 0, for all y + y², we have:
(1)) 2 (7
Ry.3/ (fo
We show similar relation for the gradient-normalizer margin [15] in the supplement (Lemma 20): gradient normalization results in a robust Wasserstein separation of the representations, making the feature separation between classes robust adversarial perturbations.
Example: Clean vs. Random Labels. Fi- nally, we provide an illustrative example on how concentration and separation are associated with generalization. For the label corruption setting from Section 5.2, Figure shows (-SNE visual- izations [53] of the representations learned with true or random labels on CIFAR-10. Training with clean labels leads to well-clustered repre- sentations. Although the model trained with random labels has 100% training accuracy, the resulting feature distribution is less concentrated and separated, implying worse generalization.
(a) Clean Labels
(b) Random Labels
Figure 4: t-SNE visualization of representations.
Classes are indicated by colors.

---

7 Conclusion
In this work, we present k-variance normalized margin bounds, new data-dependent generalization bound based on optimal transport. The proposed bounds predict the generalization error well on the large scale PGDL dataset [26]. We use our theoretical bounds to shed light on the role of the feature distribution in generalization. Interesting future directions include (1) trying better approximations to the Lipschitz constant such as [27, 47J, (2) exploring the connection between contrastive representation learning [8, 10, 22, 28] and generalization theory, and (3) studying the
generalization of adversarially robust deep learning akin to [12].
Acknowledgements This work was in part supported by NSF BIGDATA award IIS- 1741341, ONR grant N00014-20-1-2023 (MURI ML-SCOPE) and the MIT-MSR Trustworthy & Robust AI Collaboration.
References
[1] Martín Abadi, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dandelion Mané, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah, Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viégas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. TensorFlow: Large- scale machine learning on heterogeneous systems, 2015. URL https: //ww://wwwwwww.. tensorflow. org/ Software available from tensorflow.org.
[2] Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems in ai safety. arXiv preprint arXiv: (6066665655 2016.
[3] Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang. Stronger generalization bounds for deep nets via compression approach. In International Conference on Machine Learning. pages 254-263. PMLR, 2018.
[4] Peter Bartlett, Dylan Foster, and Matus Telgarsky. Specttally-normalizedd margin bounds for neural networks. In Advances in Neural Information Processing Systems, 2017.
[5] Peter L Bartlett and Shahar Mendelson. Rademacher and gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3(Nov):443-482, 2002.
[6] Peter L Bartlett, Olivier Bousquet, Shahar Mendelson, al. Local rademacher complexities. The Annals of Statistics, 33(4):1497-1537, 2005.
[7] Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1599-1607. PMLR, 2020.
[9] François Chollet. keras.
www./eeras-teeeeeeeeeeeeeeeeeeee 2015.
[10] Ching-Yas Yao Chuang, Joshua Robinson, Lin Yen-Chen, Antonio Torralba, and Stefanie Jegelka. Debiased contrastive learning. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
[11] Andrew Cotter, Maya Gupta. Heinrich Jiang, Nathan Srebro, Karthik Sridharan, Serena Wang, Blake Woodworth, and Seungil You. Training well-genert gell-generalizing classifiers for faimess metries and other data-dependent constraints. In International Conference on Machine Learning, pages 1397-1405. PMLR, 2019.
[12] John C. Duchi, Peter Glynn, and Hongseok Namkoong. Statistics robust optimization: generalized empirical likelihood approach. Mathematics of Operations Research.
10

---

[13] Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the International Conference on Uncertainty in Artificial Intelligence, 2017.
[14] Gintare Karolina Dziugaite, Alexandre Drouin, Brady Neal, Nitarshan Rajkumar, Ethan Ca- ballero, Linbo Wang. Ioannis Mitliagkas, and Daniel M Roy. In search of robust measures of generalization. Advances in Neural Information Processing Systems, 33, 2020.
[15] Gamaleldin Elsayed, Dilip Krishnan, Hossein Mobahi, Kevin Regan, and Samy Bengio. Large margin deep networks for classification. Advances in Neural Information Processing Systems, 3 31:842-852, 2018.
[16] Mahyar Fazlyab, Alexander Robey, Hamed Hassani, Manfred Morari, and George Pappas. Efficient and accurate estimation of lipschitz constants for deep neural networks. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
[17] Herbert Federer. Geometric measure theory. Springer, 2014.
[18] Rémi Flamary, Nicolas Courty, Alexandre Gramfort, Mokhtar Alaya, Aurélic Boisbunon, Stanislas Chambon, Laetitia Chapel, Adrien Corenflos, Kilian Fatras, Nemo Fournier, al. Pot: Python optimal transport. Journal of Machine Learning Research, 22(78))1-8, 2021.
[19] Nicolas Fournier and Amaud Guillin. On the rate of convergence in wasserstein distance of the empirical measure. Probability Theory and Related Fields, 162(3)7707-738, 2015.
[20] Rui Gao, Xi Chen, and Anton Kleywegt. Wasserstein distributionally robust optimization and variation regularization. arXiv preprint 1rXiv:1712.0666 2017.
[21] Noah Golowich, Alexander Rakhlin, and Ohad Shamir. Size-independent sample complexity of neural networks. In Conference On Learning Theory, pages 297-299. PMLR, 2018.
[22] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 9729-9738, 2020.
[23] F Hirzebruch N Hitchin L Hörmander, NJA Sloane B Totaro, and A Vershik M Waldschmidt. Grundlehren der mathematischen wissenschaften 332. 2006.
[24] Yiding Jiang, Dilip Krishnan, Hossein Mobahi, and Samy Bengio. Predicting the generalization gap in deep networks with margin distributions. In International Conference on Learning Representations, 2018.
[25] Yiding Jiang, Behnam Neyshabur, Hossein Mobahi, Dilip Krishnan, and Samy Bengio. Fantastic generalization measures and where find them. In International Conference on Learning Representations, 2019.
[26] Yiding Jiang, Pierre Foret, Scott Yak, Daniel M Roy, Hossein Mobahi, Gintare Karolina Dziugaite, Samy Bengio, Suriya Gunasekar, Isabelle Guyon, and Behnam Neyshabur. Neurips 2020 competition: Predicting generalization in deep learning, arXiv preprint arXiv:2012.07977, 2020.
[27] Matt Jordan and Alexandros Dimakis. Exactly computing the local lipschitz constant of relu networks. In Advances in Neural Information Processing Systems, volume 33, pages 7344-7353. Curran Associates, Inc., 2020.
[28) Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning. Advances in Neural Information Processing Systems, 33, 2020.
[29] Vladimir Koltchinskii, Dmitry Panchenko, et al. Empirical margin distributions and bounding the generalization error of combined classifiers. Annals of statistics, 30(1):1-50, 2002.
[30] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images.
2009.
11

---

[31] Vitaly Kuznetsov, Mehryar Mohri, and U Syed. Rademacher complexity margin bounds for learning with large number classes. In ICML Workshop on Extreme Classification: Learning with Very Large Number of Labels, 2015.
[32] Carlos Lassance, Louis Béthune, Myriam Bontonou, Mounia Hamidouche, and Vincent Gripon. Ranking deep learning generalization using label variation in latent geometry graphs. arXiv preprint arXiv:201 1. 121737, 2020.
[33] Yann LeCun and Corinna Cortes. MNIST handwritten digit database. 2010. URL http: //yann. lecun. exdm//xxb///ss..../
[34] David A McAllester. Pac-bayesian model averaging. In Proceedings of the twelfth annual conference on Computational learming theory, pages 164-170, 1999.
[35] Riccardo Miotto, Fei Wang, Shuang Wang, Xiaoqian Jiang, and Joel T Dudley. Deep learning for healthcare: review, opportunities and challenges. Briefings in bioinformatics, 1266:1133- 2018.
[36] Takeru Miyato, Toshiii Kataoka, Masanori Koyama, and Yuichi Yoshida. Spectral normalization for generative adversarial networks. In International Conference on Learning Representations, 2018.
[37] Vaishnavh Nagarajan and Zico Kolter. Generalization in deep networks: The role of distance from initialization. arXiv preprint arXiv:1901.01672 2019.
[38] Vaishnavh Nagarajan and Zico Kolter. Uniform convergence may be unable explain gener- alization in deep learning. In Advances in Neural Information Processing Systems (NeurIPS),
2019.
[39] Assaf Naor and Yuval Rabani. On lipschitz extension from finite subsets. Israel Journal of Mathematics, 219(1):11-161, 2017.
[40] Parth Natekar and Manik Sharma. Representation based complexity measures for predicting generalization in deep learning. arXiv preprini arXiv:2012.02775. 2020.
[41] Yuval Netzer, Tao Wang, Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. 2011.
[42] Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro. Norm-based capacity control in neural networks. In Conference on Learning Theory, pages 1376-1401. PMLR, 2015.
[43] Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro. A pac-bayesian approach to spectrally-normaiizedd margin bounds for neural networks. In International Conference Learning Representations, 2018
[44] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. Towards understanding the role of over- parammtttttinnnonnnn generalization of neural networks. In International Conference on Learning Representations, 2019.
[45] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary VeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In H. Wallach, H. Larochelle, A. Beygelzimer, F. d' d'Alché- Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems 32, pages 8024000440000 8035. Curran Associates, Inc., 2019. URL http: /tttt:r://ppprsseeeeeeeeeee neurips cc/paper/ 9015-pytorch-2r wewytorch-nn-immeenn www.iiipaaieettttllllae high- htw.r//eeeee........eooooooooeee -eeerorrranee-dee-eeeeii Learp--g--annnnnb-rrrrr library. pdf.
[46] Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and Duchesnay. Scikit-learn: Machine learning Python, Journal of Machine Learning Research, 12:2825-2834 2.1
12

---

[47] Kevin Scaman and Aladin Virmaux. Lipschitz regularity of deep neural networks: analysis and efficient estimation. Proceedings of the 32nd International Conference on Neural Information Processing Systems, pages 3839-3845 2018.
[48] Yair Schiff, Brian Quanz, Payel Das, and Pin- Yu Chen. Gi and pal scores: Deep neural network generalization statistics. arXiv preprint arXiv:2104.0346% 2021.
[49] Justin Solomon, Kristjan Greenewald, and Haikady N Nagaraja. k-variance: A clustered notion of variance. arXiv preprint arXiv:2012.06958 2020.
[50] Jonathan M Stokes, Kevin Yang, Kyle Swanson, Wengong Jin, Andres Cubillos-Ruiz, Nina M Donghia, Craig R MacNair, Shawn French, Lindsey A Carfrae, Zohar Bloom-Ackermane Ackermann, et al. deep learning approach antibiotic discovery. Cell, 180(4))688-702, 2020.
[51] Adrien Vacher, Boris Muzellee, Alessandro Rudi, Francis Bach, and Francois- Xavier Vialard.
A dimension-free computational upper-bound for smooth optimal transport estimation, arXiv preprint 101.05000, 2021.
[52] Guillermo Valle-Perez. Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter- funcction map is biased towards simple functions. In International Conference Learning Representatiins, 2019.
[53] Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine learning research, 9(1 1), 2008.
[54] Vladimir N Vapnik and A Ya Chervonenkis. On the uniform convergence of relative frequencies of events to their probabilities. In Measures of complexity, pages 11-30. Springer, 2015.
[55] Jonathan Weed and Francis Bach. Sharp asymptotic and finite-sample rates of convergence of empirical measures in wasserstein distance. Bernoulli, 25(44)::620-2644, 2019.
[56] Colin Wei and Tengyu Ma. Improved sample complexities for deep networks and robust classification via an all-layer margin. In International Conference on Learning Representations, 2020.
[57] Xueting Yan, Ishan Misra, Abhinav Gupta, Deepti Ghadiyaram, and Dhruv Mahajan. Clusterfit: Improving generalization of visual representations. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pages 6509-6518, 2020.
[58] Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedings of the British Machine Vision Conference (BMVC), pages 87.1-87.12 BMVA Press, September 2016.
[59] John Zarka, Florentin Guth, and Stéphane Mallat, Separation and concentration deep networks. In International Conference Learning Representations, 2021.
[60] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
[61] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
13


Reviews:
[{"review_id": "v6RdDxP544h", "reviewer": "Reviewer_tip4", "summary": "This paper proposes margin-based generalization bounds based on the k-variance, a generalization of the variance based on ideas from optimal transport. The k-variance, like the Rademacher complexity, is a data-dependent term, and captures certain properties of the data such as concentration and separation of features. Unlike the Rademacher complexity, the k-variance is easily empirically estimated. \n\nThe authors proceed to empirically evaluate how well their bounds can predict generalization, using the framework of recent work in the literature, namely the PGDL competition. The authors show competitive empirical results on these tasks. Although there are existing generalization measures which perform as well or better, this work is notable as it's bounds are theoretically supported. ", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This work addresses an important problem: finding theoretically motivated generalization bounds which actually reflect empirical observations in the literature. As the authors note in their review of related work, the current state of affairs for this problem is (1) many theoretically supported bounds do not reflect empirical generalization trends for deep models (2) the best empirical bounds (e.g. the Mixup*DBI winning solution of PGDL 2020) are heuristic in nature and are not grounded in theory. \n\nTheoretically grounded measures which have some empirical are of clear interest to the community: the theory provides insight into the provide to the extent it fits empirical observations. \n\nThe paper is clearly written, albeit quite dense with technical details. I would like to see a fuller conclusion section discussing the limitations of the work, e.g. what gaps remains between theory and practice, and avenues for future work.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I did not find significant discussion of this limitations of this work. Although this paper makes a contribution towards closing the gap between theory and practice, the gap is nevertheless not closed. I would like to see more discussion of how much more work there is left to do in this area.\n\n\nI found the discussion of potential negative societal impacts of this work in the appendix adequate. \n", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}, {"review_id": "s0HPpIBUJz", "reviewer": "Reviewer_C7Ku", "summary": "This paper develops generalization bounds based on class margins and the optimal transport cost (in Wasserstein-1 distance) between two independent random subsets sampled from the training distribution. The authors argue that the optimal transport cost term explains the correlation between clustered representations and good generalization. Empirically, they show that mixup+OT cost improves over mixup+DBI in predicting generalization error.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Explaining generalization in deep learning is an important problem. This paper nicely combined theory with empirical progress. \n\nHowever, there are a number of weaknesses:\n1. The main bound seems to be a minor variation of the standard margin-based Rademacher complexity bound; the technique isn't very novel.\n2. There is no justification for why Wasserstein 1 distance should be preferred over Wasserstein 2 distance, TV, or DBI for measuring intraclass clustering, other than that it is easy to prove a generalization bound for Wasserstein 1 as a variation of the Rademacher bound. It is possible that similar bounds can be shown for other metrics as well.\n3. Can the authors comment on why OT distance needs to be combined with mixup accuracy for SOTA prediction of generalization performance?\n4. Why not turn this OT cost into a regularization term in training? Will models have better generalization with this regularization?\n5. In table 1, why does kV work better than 1st layer, but kV-GN works better with 8th layer?\n6. Although some ablations are performed (in appendix), the results are not discussed. What are the conclusions?\n7. Why is Mixup + kV-GN-Margin in Table 2, column 1 smaller than kV-GN-Margin-8th in Table 1, column 1?\n\n\nMinor:\n1. Line 562 in appendix: Missing citation\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "ke8Ea3xzhQ2", "reviewer": "Reviewer_tYu6", "summary": "Firstly, the authors proposed new margin bounds based on k-variance which\ncan better capture the structural properties of the data distribution than other\nmargin based bounds. Secondly, they further developed a variant called gradient\nnormalized margin bounds which perform well empirically on the large scale\nPGDL challenge data. Moreover, they informally showed the k-variance bounds'\nability to capture the low-dimensional structure of data distribution. Last but\nnot least, they showed that maximizing the margin will give rise to the feature\nseparation. The main contribution of the paper is to introduce the idea of\nk-variance generalization bounds.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall, I think it is a good paper. The authors introduced the idea of k-variance\n(Solomon et al. 2020) to margin based generalization bounds, which is original\nand meaningful. But technically, it is not a hard paper. Most of the theory seem\nsimple corollaries of some existing ones, for example, proposition 8 and 9 can\nbe easily derived from Solomon et al. 2020. Furthermore, more interpretation\nof the k-variance bound is desirable.\n potentially major comment:\n- Proposition 8 needs formal de\fnition for the low dimensional structure of the\nmeasure.\n- Optimal transport is a much bigger topic than Wasserstein distance. The\nauthors are using the Wasserstein distance generalization bounds, so I suggest\nchange the title to  \"Measuring Generalization with k-variance\".\nminor comments:\n- line 640, n_k instead of m_k\"\n- line 681 and so forth, \\mathcal{N}_{\\epsilon^{\\prime}}", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There are no obvious limitations or potential negative societal impact of their\nwork.", "review_text": "", "extra_scores": {"ethical_concerns": "No.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "16"}, "extra_text": {}}, {"review_id": "czxALmXfca1", "reviewer": "Reviewer_j4gL", "summary": "### POST-REBUTTAL ###\nI would like to thank the authors for their reply. I am pretty satisfied with the justification for setting $c(k,d)=1$. The comparison between the different related works also looks great. I am not raising my score further to Accept as I still see the presence of the Lipschitz constant in the bound as a weakness but this is common to many works and hopefully its estimation will see some advances in the nearest future. \n### POST-REBUTTAL ###\n\nThis paper presents margin-based generalization bounds for deep neural networks based on the k-variance. This latter is a new way of measuring the variance in the observed data (or the probability distribution generating it) based on optimal transportation (OT) theory and the Wasserstein distance related to it. The obtained generalization bounds are claimed to have several important advantages as k-variance: 1) can be computed from the available data contrary to other complexity measures such as the popular Rademacher complexity; 2) portrays more truthfully the underlying generalization capacity of the considered model. In addition to k-variance, the obtained bound also relies on the Lipschitz constant of the margin function with respect to (wrt) the encoder and the desired threshold on the margin used to define the empirical risk. Experimental results on large-scale datasets suggest the usefulness of the introduced complexity measure in capturing the generalization of deep neural networks. ", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Originality**\n\nThis works tackles an important problem in machine learning which is to find complexity measures that reflect generalisation capacities of deep neural networks. As mentioned by the authors, some existing complexity measures have a solid theory behind them but fail in the situations of practical interest, while others work well in practice without any formal justification. The author’s attempt to provide a generalisation bound with a practically useful complexity term based on k-variance seems to be novel as, to the best of my knowledge, k-variance has never been used in generalisation bounds. Also, it indeed addresses several drawbacks of existing complexity measures. \n\n**Quality**\n\nOverall, the paper appears to be technically sound even though some parts of it need a more thorough discussion. Below, I formulate a couple of concerns that I had while reading the paper:\n\n1. p.3. l.94: Assuming $c(k,d) = 1$ is an important deviation from the original work of Solomon et al. but the implications of it are not discussed in the paper. Indeed, Solomon et al. say in their work that this factor is “the ambient scaling rate chosen to account for the rate at which the expectation approaches zero”. If the authors consider the case where $d>2$ (I believe that this is the case of interest when considering neural nets), it implies $c(k,d) = k^{\\frac{2}{d}} = 1$. What does it mean in the context of this work? My guess is that for reasonable values of $k$, one expects large values of $d$ to compensate for the lack of scaling here but I am not totally sure about it and would like the authors to discuss it. \n2. Given Lemma 3.1 in [4] and Eq. (3), why not using the former to derive Theorem 2 simply as its corollary by applying Eq. (3) to bound the Rademacher complexity? \n3. Solomon et al. derive several properties of k-variance that are used in Section 6 and appear to be useful in the considered context. However, some other more “puzzling” properties of k-variance are not mentioned in this work. For instance, there are distributions with d<=5 for which the variance estimate diverges ([49], Example 6.4) and thus cannot potentially be used as a valid complexity measure. I would suggest including this remark as an assumption to this work. \n\n**Clarity**\n\nThe paper is clearly written and is rather easy to follow. The experimental section is a little be compressed but that is to be expected given the existing page limit. Some minor remarks are:\n\n1. Section 2 can benefit a lot from presenting some generalisation bounds explicitly or at least by presenting a summary table showing other complexity measures used to derive margin-based generalisation bounds for neural nets. At the the very least, the work of [4] deserves to be presented with more details as it relies on Rademacher complexity to which the authors compare k-variance throughout the text. Also, the complexity terms from the two works can be compared as well as both have the ratio of Lipschitz constants/margin threshold terms in them. \n2. p.4 l.141: how do we know that label corruption is handled by k-variance solely from the shape of the generalisation bound? Maybe at this point it is worth referring to experimental results directly.\n3. In the experimental results, the authors use the lower bound on the Lipschitz constant to evaluate the considered complexity measure or set it to 1. In both cases, I think that one cannot really rely on the obtained values as they necessarily carry some bias (as mentioned in Appendix C.1) if the tightness of the approximation is not taken into account. This raises the question of whether the obtained theoretical results are really evaluated in the experimental study or whether they are very rough approximation of the latter. At the very least, the fact of evaluating Lipschitz constant may not be considered as an advantage when compared to [24] (p.9 l.224-225).\n\n**Significance**\n\nUnfortunately, I have several doubts regarding the significance of the obtained results after this initial round of review. K-variance is definitely an interesting measure of complexity to consider to assess the generalisation of neural nets, but its experimental evaluation and its introduction into the bounds raises several question to which I would like to hear the authors’ answers before revisiting my opinion. I am highly inclined to reconsider the given score if the authors manage to reply in a convincing manner to my remarks.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Regretfully, the authors only partially discuss the limitations of their work in the main paper and postpone the discussion regarding the Lipschitz constant to Appendix.  I think that discussion the drawbacks of the k-variance and the quality of approximation for the proposed complexity measure should appear in the main part of the paper. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper proposes a novel generalization bound that is both theoretically grounded and works in practice. This is an important achievement that has been acknowledged by all reviewers. The least convinced reviewer appreciated your answers and the consensus during the discussion was for accepting the paper. Still the discussion with the reviewer about the constant c(k,d) is important and should be added in the final version of the paper.
Decision: Accept (Spotlight)