Measuring Generalization with Optimal Transport

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.e 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 , 25 , 38 ] . In particular , Jiang et al . [ 25 ] present a 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 to 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 , we attempt to 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 a natural alternative to Rademacher complexity . Interestingly , this optimal transport cost can be interpreted as the k-variance [ 49 ] , a 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 kvariance 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 ( NeurIPS 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 ( 1 ) 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 on 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 a theoretical analysis to shed light on the role of feature distributions in generalization , based on our k-variance normalized margin bounds . 2 Related Work . Margin-based Generalization Bounds Classic approaches in learning theory bound the generalization error with the complexity of the hypothesis class [ 5 , 54 ] . 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 ] . Recently , substantial progress has been made to develop better data-dependent and algorithm-dependent bounds [ 3 , 6 , 7 , 44 , 52 , 56 ] . Among them , we will focus on margin-based generalization bounds for multi-class classification [ 29 , 31 ] . Bartlett et 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 spectrally-normalized margin bounds via weight perturbation within a PAC-Bayes framework [ 34 ] . However , empirically , spectral norm-based bounds can correlate negatively with generalization [ 38 , 25 ] . Elsayed et al . [ 15 ] present a gradient-normalized margin , which can be interpreted as the first order approximation to the distance to the decision boundary . Jiang et al . [ 24 ] further show that gradient-normalized 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 , 25 ] . In particular , Jiang et al . [ 25 ] measure the average correlation between various complexity measures and generalization error under different experimental settings . Building on their study , Dziugaite et al . [ 14 ] emphasize the importance of the robustness of the generalization measures to the 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 in 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 a 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-Variance Before presenting our generalization bounds with optimal transport , we first give a brief introduction to the Wasserstein distance , a distance function between probability distributions defined via an optimal transport cost . Letting µ and ⌫ 2 Prob ( Rd ) be two probability measures , the p-Wasserstein distance with Euclidean cost function is defined as Wp ( µ , ⌫ ) = inf ⇡2⇧ ( µ , ⌫ ) ⇣ E ( X , Y ) ⇠⇡kX Y kp ⌘1/p , where ⇧ ( µ , ⌫ ) ✓ Prob ( Rd ⇥ Rd ) denotes the set of measure couplings whose marginals are µ and ⌫ , respectively . The 1-Wasserstein distance is also known as the Earth Mover distance . Intuitively , Wasserstein distances measure the minimal cost to transport the distribution µ to ⌫ . Based on the Wasserstein distance , Solomon et al . [ 49 ] propose the k-variance , a generalization of variance , to measure structural properties of a distribution beyond variance . Definition 1 ( Wasserstein-p k-variance ) . Given a probability measure µ 2 Prob ( Rd ) and a parameter k 2 N , the Wasserstein-p k-variance is defined as Vark , p ( µ ) = cp ( k , d ) · ES , S̃⇠µk h Wpp ( µS , µS̃ ) i , where µS = 1 k Pk i=1 xi for xi i.i.d.⇠ µ and cp ( k , d ) is a normalization term described in [ 49 ] . When k = 1 and p = 2 , the k-variance is equivalent to the variance Var [ X ] of the random variable X ⇠ µ . For k > 1 and d 3 , Solomon et al . [ 49 ] show that when p = 2 , the k-variance provides an intuitive way to measure the average intra-cluster variance of clustered measures . In this work , we use the unnormalized version ( cp ( k , d ) = 1 ) of k-variance and p = 1 , and drop the p in the notation : ( Wasserstein-1 k-variance ) : Vark ( µ ) = ES , S̃⇠µk ⇥ W1 ( µS , µS̃ ) ⇤ . Note that setting cp ( 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 a premultiplication term in the generalization bounds to cancel it out . In Section 6 , we will show that this unnormalized Wasserstein-1 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 , where the hypothesis f can be decomposed as a feature ( representation ) encoder 2 and a predictor f 2 F . This includes dividing multilayer neural networks at an intermediate layer . We consider the score-based classifier f = [ f1 , . . . , fK ] , fc 2 Fc , where the prediction for x 2 X is given by argmaxy2Y fy ( ( x ) ) . The margin of f for a datapoint ( x , y ) is defined by ⇢f ( ( x ) , y ) : = fy ( ( x ) ) max y0 6=y fy0 ( ( x ) ) , ( 1 ) where f misclassifies if ⇢f ( ( x ) , y ) 0 . The dataset S = { xi , yi } mi=1 is drawn i.i.d . from distribution µ over X ⇥ Y . Define mc as the number of samples in class c , yielding m = PK c=1 mc . We denote the marginal over a class c 2 Y as µc and the distribution over classes by p ( c ) . The pushforward measure of µ with respect to is denoted as # µ . We are interested in bounding the expected zero-one loss of a hypothesis f : Rµ ( f ) = E ( x , y ) ⇠µ [ ⇢f ( ( x ) , y ) 0 ] by the corresponding empirical -margin loss R̂ , m ( f ) = E ( x , y ) ⇠S [ ⇢f ( ( x ) , y ) ] . 4.1 Feature Learning and Generalization : Margin Bounds with k-Variance Our theory is motivated by recent progress in feature learning , which suggests that imposing certain structure on 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 a fixed encoder and discuss the generalization error of the encoder at the end of the section . Theorem 2 provides a generalization bound for neural networks via the concentration of µc in each class . Theorem 2 . Let f = [ f1 , · · · , fK ] 2 F = F1 ⇥ · · · ⇥ FK where Fi : X ! R. Fix > 0 . The following bound holds for all f 2 F with probability at least 1 > 0 : Rµ ( f ) R̂ , m ( f ) + Ec⇠p Lip ( ⇢f ( · , c ) ) Varmc ( # µc ) + r log ( 1/ ) 2m , where Lip ( ⇢f ( · , c ) ) = supx , x02X |⇢f ( ( x ) , c ) ⇢f ( ( x 0 ) , c ) | || ( x ) ( x0 ) ||2 is the margin Lipschitz constant w.r.t . We give a proof sketch here and defer the full proof to the supplement . For a given class c and a given feature map , let Hc = { h ( x ) = ⇢f ( ( x ) , c ) |f = ( f1 . . . fK ) , fy 2 F } . The last step in deriving Rademacher-based generalization bounds [ 5 ] amounts to bounding for each class c : c = ES , S̃⇠µmcc 2 4 sup h2Hc 1 mc mcX i=1 h ( x̃i ) 1 mc mcX i=1 h ( xi ) 3 5 , ( 2 ) where S , S̃ ⇠ µmcc . Typically we would plug in the Rademacher variable and arrive at the standard Rademacher generalization bound . Instead , our key observation is that the Kantorovich-Rubinstein duality [ 23 ] implies W1 ( µ , ⌫ ) = sup Lip ( h ) 1 Ex⇠µh ( x ) Ex⇠⌫h ( x ) , where the supremum is over the 1-Lipschitz functions h : Rd ! R. Suppose Hc is a 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 : c L · ES , S̃⇠µmcc [ W1 ( # µS , # µS̃ ) ] = L ·Varmc ( # µc ) . ( 3 ) This connection suggests that k-variance is a natural alternative to Rademacher complexity if the margin is Lipschitz . The following lemma shows that this holds when the functions fj are Lipschitz : Lemma 3 . The margin ⇢f ( . , y ) is Lipschitz in its first argument if each of the fj is 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 f 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 to the margin bound with Rademacher complexity [ 31 ] , Theorem 2 studies a fixed encoder , allowing the bound to capture the structure of the feature distribution . Although the Rademacherbased 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 # ( µ ) affects generalization . Importantly , it is also non-trivial to 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 use the gradient-normalized margin , a 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 a decision boundary , and Jiang et al . [ 24 ] simplified it to ⇢̃f ( ( x ) , y ) : = ⇢f ( ( x ) , y ) / ( kr ⇢f ( ( x ) , y ) k2 + ✏ ) , where ✏ is a small value ( 10 6 in practice ) that prevents the margin from going to infinity . The gradient here is r ⇢f ( ( x ) , y ) : = r fy ( ( x ) ) r fymax ( ( x ) ) , where ties among the ymax are broken arbitrarily as in [ 15 , 24 ] ( ignoring subgradients ) . The gradient-normalized margin ⇢̃f ( x , y ) can be interpreted as the first order Taylor approximation of the minimum distance of the input x to the decision boundary for the class pair ( y , y0 ) [ 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 an interpretation of this margin in terms of robust feature separation . Defining the margin loss R̂r , m ( f ) = E ( x , y ) ⇠S [ ⇢̃f ( ( x ) , y ) ] , we extend Theorem 2 to the gradient-normalized margin . Theorem 4 ( Gradient-Normalized Margin Bound ) . Let f = [ f1 , · · · , fk ] 2 F = [ F1 , · · · , Fk ] where Fi : X ! R. Fix > 0 . Then , for any > 0 , the following bound holds for all f 2 F with probability at least 1 > 0 : Rµ ( f ) R̂r , m ( f ) + Ec⇠p '' Lip ( ⇢̃f ( . ) ( · , c ) ) Varmc ( # µc ) # + r log ( 1/ ) 2m , where Lip ( ⇢̃f ( . ) ( · , c ) ) = supx , x02X |⇢̃f ( ( x ) , c ) ⇢̃f ( ( x 0 ) , c ) | || ( x ) ( x0 ) || is the Lipschitz constant defined w.r.t . 4.3 Estimation Error of k-Variance So far , our bounds have used the k-variance , which is an expectation . Lemma 5 bounds the estimation error when estimating the k-variance from data . This may be viewed as a generalization error for the learned features in terms of k-variance , motivated by the connections of k-variance to test error . Lemma 5 ( Estimation Error of k-Variance / “ Generalization Error ” of the Encoder ) . Given a distribution µ and n empirical samples { Sj , S̃j } nj=1 where each Sj , S̃j ⇠ µk , define the empirical Wasserstein-1 k-variance : dVark , n ( # µ ) = 1n Pn j=1 W1 ( # µSj , # µS̃j ) . Suppose the encoder satisfies supx , x0 k ( x ) ( x0 ) k B , then with probability at least 1 > 0 , we have Vark ( # µ ) dVark , n ( # µ ) + r 2B2 log ( 1/ ) 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 , for m = PK c=1b mc 2n c , Rµ ( f ) is upper bounded by R̂ , m ( f ) + Ec⇠py 2 64 Lip ( ⇢f ( · , c ) ) 0 @ dVarbmc2n c , n ( # µc ) + 2B s log ( 2K/ ) nbmc2n c 1 A 3 75+ s log ( 2 ) 2m . Note that the same result holds for the gradient normalized margin ⇢̃f . The proof of this corollary is a 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-probability 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 h dVark , n ( # µ ) i 4Varµ ( ( X ) ) nk , where Varµ ( ( X ) ) = Ex⇠µ [ || ( x ) Ex⇠µ ( x ) ||2 ] is the variance of # µ. Theorem 7 implies that if the feature distribution # µ has bounded variance , the variance of the empirical k-variance decreases as k and n increase . The values of k we used in practice were large enough that the empirical variance of k-variance was small even when we set n = 1 .

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.