The weighted mean trick – optimization strategies for robustness

1 INTRODUCTION . The most common objective in machine learning is to minimize the average loss . However , by doing this it implies that all samples are equal . The goal of this work is to convince the reader that from the perspective of the model , not all samples are equal and optimizing a weighted mean is a more progressive approach . For instance , at the end of the training , few samples are left uncaptured by the model which yield large loss values . One might choose to put more weight on those samples to help the model learn ( individual fairness ) , or consider them noise and assign them less weight ( robustness ) . The decision will also be reflected by the loss variance ; it will decrease when hard samples are weighted more and increase otherwise . Before introducing more formally the weighted mean , we first justify theoretically the impact of variance penalization . Using the empirical Bernstein bound ( Maurer & Pontil , 2009 ) for i.i.d . loss values Z and bounded variance we have with probability 1− δ : E [ Z ] − 1 n n∑ i=1 Zi ≤ C1 √ 2Vn [ Z ] ln 2/δ n + C2 7 ln 2/δ 3 ( n− 1 ) ( 1 ) where C1 , C2 are problem depended constants . This inequality reveals two things . First , we have with high probability that the empirical mean is close to the theoretical value . This bound along with similar PAC-Bayes bounds ( Seldin et al. , 2012 ; Tolstikhin & Seldin , 2013 ) justifies the practical success of the empirical risk minimization ( ERM ) which has the objective to minimize the mean loss value ( Namkoong & Duchi , 2017 ) . Secondly , the difference between the two is bounded in terms of the empirical variance . This second implication led to a large and growing body of studies that investigate variance penalization ( Maurer & Pontil , 2009 ; Namkoong & Duchi , 2017 ; Duchi & Namkoong , 2019 ; Staib et al. , 2019 ; Lam , 2019 ; Heinze-Deml & Meinshausen , 2021 ; Hu et al. , 2018 ) . Moreover , by penalizing the variance , the intrinsic bias-variance tradeoff of the ERM can be controlled . Two studies influential for this paper are that of Duchi & Namkoong ( 2019 ) and Li et al . ( 2021 ) . Duchi & Namkoong ( 2019 ) proposed taking the expectation with respect to a different distribution which allowed penalizing variance while preserving the convexity of the objective . Similarly , Li et al . ( 2021 ) investigated optimizing a tilted empirical risk which is equivalent to penalizing all the higher-order moments simultaneously . In summary , previous methods either penalize only one moment or all the higher-order moments but are not flexible enough to penalize any desired combination of the higher-order moments . Inspired by the above works , we propose to optimize a weighted mean and prove that for certain weights , it is equivalent to optimizing any higher-order moments of the loss distribution such as variance , skewness , and kurtosis . Our approach generalizes that of Duchi & Namkoong ( 2019 ) ; Li et al . ( 2021 ) while simplifying the optimization procedure and enabling separate penalization of the higher-order moments . In particular , we will construct the weights w such that optimizing the weighted mean of the loss ` is equivalent to applying a variance penalization : E [ w ` ] = E [ ` ] + λV [ ` ] ( 2 ) or penalization of other central moments ( e.g. , skewness , kurtosis ) . Construction of these weights are treated by Theorems 1 and 3 from section 3 , which require minimal computational resources . Important note , penalizing directly the variance preserves convexity only for values of λ within a very narrow range ( Maurer & Pontil , 2009 ) . However , the weighted mean method yields a convex objective for any positive values of λ . In detail , our contributions are as follows : ( C1 ) We build upon the work of Duchi & Namkoong ( 2019 ) and prove that optimizing a weighted mean is equivalent to reducing or amplifying both variance ( Theorem 1 ) and higher-order moments ( Theorem 3 ) . ( C2 ) We derive the limits of the penalization interval which preserves convexity ( Lemma 2 ) and also show how to penalize with values outside this interval while still maintaining convexity ( Lemma 4 ) . ( C3 ) We connect the variance and higher-order moments penalization using weighted mean to the robustification of loss functions ( Lemma 5 and Lemma 6 ) . ( C4 ) We develop a convex version of the variance penalized cross-entropy loss which provides a higher accuracy in high noise scenarios with class dependent noise . ( C5 ) We show experimentally that a negative variance penalization improves model accuracy when training with noisy labels . The implications of the weighted mean are much broader than what we investigate in this work . We limit the scope of this paper to classification using deep neural networks trained with noisy labels . But the mathematical framework also covers the control of the bias-variance trade-off and can also be applicable to regression problems . In the sequel , we start by introducing the notation in section 2 , then in section 3 we present the moment penalization strategy and the weighted mean trick which is a computational technique to make moment penalization practical . Next , in section 4 , we illustrate the application of moment penalization using the weighted mean formulation to optimize for robustness when training with noisy labels . 2 NOTATIONS . In subsequent sections we will use the following notation . A training data set is defined as D = { ( xi , yi ) } ni=1 where xi ∈ X are the features and yi ∈ Y = { 1 , . . . , k } represent the class labels . A classifier f ( x ; θ ) is a mapping f : X ×Θ→ V from the feature space to the probability simplex V parameterized by θ . A loss function ` : V × Y → [ 0 , ∞ ) gives a penalty ` ( v , y ) when the model predicted the value v and label y was observed . A weight function w : V × Y → R assigns to each sample a weight . As we are interested in optimizing the model output value v that minimizes the penalty , we will focus our investigation on loss functions ` ( v , y ) that are convex in v and seek to preserve the convexity when optimizing the weighted mean . In ERM , we are interested with finding the model parameters θ that minimize the empirical risk calculated as ED [ ` ( f ( xi ; θ ) , yi ) ] and the weighted form ED [ w ( ` ( f ( xi ; θ ) , yi ) ) ` ( f ( xi ; θ ) , yi ) ] . To simplify the notation , we will drop the dataset D from the expectation subscript and the arguments of the loss and the weight function , i.e. , E [ ` ] ⇔ ED [ ` ( f ( xi ; θ ) , yi ) ] and E [ w ` ] ⇔ ED [ w ( ` ( f ( xi ; θ ) , yi ) ) ` ( f ( xi ; θ ) , yi ) ] . Similarly , the notation for minimum is simplified as min ` ⇔ minxi , yi∈D ` ( f ( xi ; θ ) , yi ) . −3 0 3 −3 0 3 0.0 0.0 0.2 0.2 1.0 1.0 2.02.0 3.9 3.9 −3 0 3 0.0 0.0 0.5 0.5 1.0 1.0 1.3 1.3 −3 0 3 λ < 0 λ = 0 λ > 0 λ > 0 λ < 0 1Figure 1 : Variance penalization for classification problems . Contour lines of the weights distribution for positive λ are shown on the left and for negative on the center plot . The right plot shows the use of variance penalization for outlier suppression or amplification ( robustification versus generalization ) . 3 PENALIZING MOMENTS . In ERM , the impact on the model parameters of a sample is determined by the balance between its loss value and the average of the training batch . Especially observed in the latter stages of the training when majority of samples have a small loss value except few which are left uncaptured by the model . Case in which we consider them as unlearned samples and amplify their impact or as outliers and suppress their impact on the model . The weighted mean trick consist in assigning weights to each sample based on the loss value to either amplify or suppress their impact on training . This is similar to dropout ( Hinton et al. , 2012 ) but for losses . However , weighted mean trick is a deterministic process and the weights are not restricted to 0 or 1 but can take any non-negative value . In what follows , we apply the weighted mean trick to extend the ERM framework to multi-objective optimization of the mean and higher-order moments of the loss function . First , we show that for some distinct weights , optimizing the weighted mean is equivalent to a simultaneous optimization of the mean and variance . Theorem 1 ( Variance Expansion ) . Let ` be a loss function with finite E [ ` ] and finite V [ ` ] and let w ( v , y ) = 1 + λ ( ` ( v , y ) − E [ ` ( v , y ) ] ) , then we have : E [ w ` ] = E [ ` ] + λV [ ` ] ( 3 ) Proof . Replacing w with the definition and using the linearity property of the expectation along with Proposition 7 we get : E [ w ` ] = E [ ` ] +λE [ ( ` −E [ ` ] ) ` ] = E [ ` ] +λE [ ( ` −E [ ` ] ) 2 ] = E [ ` ] +λV [ ` ] Thus , switching from mean to weighted mean allows us to control the bias-variance tradeoff through λ to improve the distributional robustness of the model ( Maurer & Pontil , 2009 ) . However , the range of λ values that preserve the convexity of the objective depends on the average and the minimum penalty returned by the loss function ` as shown by the next lemma . Lemma 2 . As introduced in Theorem 1 , the variance expansion of a convex loss function ` ( v , y ) in v yields a new objective that is also convex in v if λ ∈ [ 0 , λmax ] , where λmax = 1/ ( E [ ` ] −min ` ) . Proof is provided in Appendix A . This lemma precisely shows why directly penalizing the variance does not preserve convexity besides when λ takes values in a narrow interval . Moreover , directly penalizing the variance as part of a numeric optimization objective , the upper limit of the interval , λmax , is not constant and changes with each iteration of the optimization algorithm . Note that the further the minimum value min ` is from the average loss value E [ ` ] , the narrower the interval is . Conversely , when E [ ` ] = min ` results that ` is constant and V [ ` ] = 0 thus the objective is convex for λ > 0 . However , using the weighted mean trick the objective remains convex for any positive λ irrespective if V [ ` ] = 0 . Figure 1 shows the weights distribution for different values of λ . When λ is positive ( left plot ) samples closer to the decision boundary ( which also means with larger loss values ) receive more weight . On the other hand , when λ is negative ( center plot ) samples with larger loss values receive less weight . Of note , for λ < 0 the objective is not convex , however as we will show later will still converge to an optimal solution . The right plot shows a binary classification problem where each class consists of two clusters , squares and circles . The parameter λ controls the placement of the decision boundary with respect to these two clusters . Positive λ values place more weight on the cluster of squares which is closer to the decision boundary and as a result the boundary is horizontally aligned . On the other hand , negative λ values place more weight on samples farther from the boundary and in this case the cluster of circles . Therefore , this aligns the decision boundary with respect to the cluster of circles oriented diagonally . Variance is not the only central moment that can be penalized using the weighted mean . In fact , any combination of the moments can be penalized . The next result generalizes Theorem 1 . Theorem 3 ( Moments Expansion ) . Let ` be a loss function with finite first m central moments and define w ( v , y ) = ∑m i=1 λi ( ` ( v , y ) − E [ ` ( v , y ) ] ) i−1 , then we have : E [ w ` ] = λ1E [ ` ] + m∑ i=2 λ̃iE [ ( ` − E [ ` ] ) i ] ( 4 ) where λ̃i = λi + λi+1E [ ` ] for i < m and λ̃m = λm Proof is provided in Appendix A . We notice that penalizing moments higher than two incurs an additional penalization of the previous moment . For example , penalizing skewness by λ3 incurs a variance penalization of λ3E [ ` ] . Theorem 3 also has an algebraic interpretation . Note that the formula for weights is nothing more than a polynomial in ` ( v , y ) translated by E [ ` ] . When penalizing the variance , λ2 controls the slope of the linear equation , and when penalizing the skewness , λ3 controls the curvature of the quadratic equation . Moreover , the penalization factors λi also define the placement of the roots of the polynomial and the convexity of the weighted mean objective . Lemma 4 ( Convexity of Moments Expansion ) . Let ` ( v , y ) be a loss function convex in v and p : R→ [ 0 , ∞ ) be a non negative and differentiable convex function and M ≥ 0 , then the weighted objective w ( v , y ) ` ( v , y ) with w ( v , y ) = p ( ` ( v , y ) −M ) is convex in v if p is non decreasing . Proof is provided in Appendix A . The above result lists two requirements on the polynomial function p such that the weighted mean objective remains convex . First , the function must be non-negative , thus , the negative weights must be clipped to 0 . Secondly , when the function takes positive values , p must be non decreasing and convex . Of note , this result is similar in scope to Lemma 2 but does not generalize it as the weights are non negative . Figure 2 shows several examples of polynomials . Variance penalization implies p is an affine function , and thus , convex . However , for λ1 > 0 the polynomial is non decreasing ( left plot ) and for λ1 < 0 the polynomial is non increasing ( center plot ) . Thus , only positive variance penalization will result in a convex objective . Of note , Lemma 2 upper bounds λ1 if weights are not clipped to 0 , however , if clipping is used λ1 is not upper bounded . When penalizing skewness , p is a quadratic function ( right plot ) . In this case , if λ2 > 0 then p is convex and non decreasing only on a restricted interval instead of the entire real line . Thus , convexity can be preserved by appropriately adjusting the roots of the polynomial such that p is convex on [ min ` − E [ ` ] , +∞ ] . Clipping negative weights to 0 prevents the optimization objective to switch from minimization to maximization , usually an undesirable behavior . Consequently , the samples with a corresponding zero weight reach their maximum contribution in the moments penalization . Further increasing the factors λi , will have no effect on those samples and only samples with non-zero weights will participate in the training . As a result , the efficiency of the moments penalization will slightly fall . Moreover , the samples with zero weight will be excluded from training , technique used in the past by multiple studies . Rockafellar & Uryasev ( 2000 ) used a similar clipping technique to optimize only for samples that are part of the tail of the distribution . The Moments Expansion Theorem extends the variance expansion proposed by Duchi & Namkoong ( 2019 ) to include higher-order moments . However , even when only the variance is penalized , the two methods are still slightly different . Moments Expansion Theorem computes the weights directly from the loss values , whereas the variance expansion of Duchi & Namkoong ( 2019 ) solves a secondary optimization problem to find the weights . The use of a secondary optimization problem has the −1 0 1 −1 0 1 2 3 λ2 = 2 λ1 = 1 λ4 = 1 λ3 = 2.5 λ2 = 3 λ1 = 1 −1 0 1 λ2 = −2 λ1 = 1 −1 0 1 λ3 = 1 λ2 = 0.5 λ1 = −0.5 1Figure 2 : Polynomial functions for moments penalization . Dotted lines show the complete polynomialwhereas solid lines the clipped version . Left plot shows two convex and non decreasing polynomials . Center plot shows a convex but non increasing polynomial and thus will result in a non convex objective . Right plot shows a convex and non decreasing polynomial on the interval [ −1 , +∞ ] , and thus will yield a convex objective when min ( ` ) − E [ ` ] ≥ −1 . advantage of penalizing the variance more consistently despite being more computationally expensive . On the contrary , when using the moments expansion , the penalized variance can be slightly lower depending on the number of weights that are 0 . However , the advantage is that the direct computation of the weights makes it easier to include the method into existing analysis frameworks . A similar method that extends the optimization objective to include penalization factors for higher-order moments was proposed by Li et al . ( 2021 ) . The proposed method replaces the ERM objective with a tilted version calculated as 1tK ( t ) = 1 t logE [ e t ` ] where K ( t ) is the cumulant-generating function of the loss ` . The penalization of the higher order moments of the tilted ERM can be recovered from the power series expansion of K ( t ) . The distinction between the two is that the moments penalization introduced in this paper represents a generalization of the tilted ERM as it allows any combination of the higher order moments to be penalized whereas tilted ERM uses a single parameter that governs the penalization factors . In summary , the moments penalization implemented using the weighted mean trick is more flexible , however , it comes at a cost as there are more parameters to tune when penalizing multiple moments compared to tilted ERM of Li et al . ( 2021 ) . Convergence and convergence rates . The moments penalization problem along with variance expansion of Duchi & Namkoong ( 2019 ) fall under the class of distributional robust stochastic programsSun & Xu ( 2016 ) ( DRSP ) which is a subclass of ambiguity programsRoyset & Wets ( 2017 ) ( AP ) where the general objective is : AP : min θ∈Θ sup P∈D ( θ ) ϕ ( θ , P ) ( 5 ) where Θ is the set of model parameters and D ( θ ) is the ambiguity set . In DSRP , the bivariate function ϕ ( θ , P ) = EP [ ` ] where ` is the loss function and the ambiguity set D ( θ ) is a set of probability distributions . In the case of moments penalization problem , the ambiguity set D ( θ ) depends on the model parameters and is a singleton as the weights uniquely transform the empirical distribution . As a result , the optimal value of the inner maximization problem becomes supP∈D ( θ ) ϕ ( θ , P ) = EP [ ` ] . Intuitively , a model will converge if changes in its parameters will cause minor changes in the distribution P , and with each step the distribution will approach the optimum distribution P ∗ . Formally , to quantify the changes in the distribution , we would need a distance or a metric . Sun & Xu ( 2016 ) use the total variation metric and a pseudometric to prove uniform convergence , possibly at an exponential rate , if P converges to P ∗ under total variation metric and ` is uniformly bounded ( see Sun & Xu , 2016 , Th . 1 and Prop . 3 ) . Royset & Wets ( 2017 ) proposed a hypo-distance metric and proved lop-convergence given that the bivariate function ϕ ( θ , P ) satisfies some assumptions , ( see Royset & Wets ( 2017 ) Def . 4.1 ) . Duchi & Namkoong ( 2019 ) provide guarantees for a number of stochastic risk minimization problems when only the variance is penalized and P is in the local neighborhood of the empirical distribution defined using the χ2-divergence . We refer the reader to the works of Sun & Xu ( 2016 ) and Royset & Wets ( 2017 ) and the references therein for additional guarantees if more information about the problem structure is available , or if other metrics are used . For the moments penalization problem , the moments penalization factors λi for i ≥ 2 determine how much the distribution P changes when the loss changes . Small values of the penalization factors will keep P in the neighborhood of the empirical distribution , whereas large values will make the weights sensitive to changes in the loss values that can cause stability or convergence issues . The exact values depend on the empirical distribution of the data and the choice of the model and loss function . Weighted mean trick in practice . To apply the method in practice , the classical batch training algorithm must be extended to include an additional step , the weights calculation . Instead of directly calculating the average loss , the user will calculate the loss value for each sample in the batch and then use the expression from Theorem 3 to compute the weights and the weighted mean . The moments penalization factors λi are the hyper-parameters and are tuned in ascending order with respect to i . However , penalizing higher-order moments might affect the impact of the lower-order ones , and thus , it might require a few iterations to find the optimal combination . The implementation of this algorithm in Pytorch ( Paszke et al. , 2019 ) is available on GitHub.1 Since the gradient of the weighted mean is the weighted gradient of the elements , this allows weights to control the impact of each sample on the model parameters . Of note , the theoretical results hold when switching from expectation to sample expectation , En [ ` ] = 1n ∑n i=1 ` ( f ( xi , θ ) , yi ) . Algorithm 1 : Training with Moments Penalization input : f ( x ; θ ) – model to be trained { xi , yi } n1 – batch of training data { λ } m1 – penalization factors ` ( v , y ) – loss function while stopping criteria not reached do for i← 1 to n do zi ← ` ( f ( xi , θ ) , yi ) ; / * sample loss * / w ← [ ∑m j=1 λj ( z − En [ z ] ) j−1 ] + Lw ← 1n ∑n i=1 wizi ; / * weighted mean * / θ ← θ − γ∇θLw ; / * update model parameters * /

Summary: This paper studies weighted Empirical Risk Minimization (ERM), where the weight at data point $(x_i, y_i)$ is a (polynomial) function of the loss value $\ell(f(x_i), y_i)$. For affine functions, the authors show that it is equivalent to perform variance penalization (Theorem 1 and Lemma 2). Polynomials with higher degrees involve higher order moments of the losses (Theorem 3 and Lemma 4). The authors also propose an iterative algorithm to perform the weighted ERM (Algorithm 1). In Section 3, a different choice of weights is proposed, that ensure the convexity of the cross-entropy criterion despite a negative $\lambda_1$, and experiments are presented.