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: The paper studies the role and implications of weighted empirical risk minimization, where one can also choose the weight of each sample instead of fixing it to be equal. The paper then shows that specific choices of weights lead to variance penalization and higher-order moments penalization (Theorems 1 and 3). Lemma 2 and 4 study the regimes for which these chosen weights also preserve the convexity of the original loss functions. This framework is then empirically validated through multiple experiments.