MixRL: Data Mixing Augmentation for Regression using Reinforcement Learning

1 INTRODUCTION . As machine learning ( ML ) becomes widely used in critical applications including manufacturing and finance , data augmentation for regression becomes essential as it provides an opportunity to improve model performance without additional data collection . In comparison to classification tasks like object detection in images , the goal of regression is to predict one or more real numbers . To emphasize the importance of data augmentation in regression , we provide a case study of semiconductor manufacturing . Here a common quality check is to measure the layer thicknesses of a 3-dimensional semiconductor and see if they are even . However , directly measuring each thickness results in destroying the semiconductor itself , so a recently-common approach is to take an indirect measurement by applying light waves on the semiconductor , measuring the spectrum of wavelengths that bounce back from all the layers , and use ML to predict the layer thicknesses from the spectrum data ( see Fig . 4 in Sec . 4 for an illustration ) . With enough spectrum data and thickness information , ML models can be trained to accurately predict thicknesses from a spectrum . The main challenge is that there is not enough training data , and the only cost-effective solution is to augment small amounts of data that exist . Even a small improvement in model performance from the data augmentation has significant impact in this industry . In general , any regression task that predicts real values like emissions , stock prices , or even someone ’ s salary can also benefit from data augmentation . Most data augmentation techniques are designed for image classification . In particular , Mixup ( Zhang et al. , 2018 ; Berthelot et al. , 2019 ; Yun et al. , 2019 ) is a popular data augmentation technique that is widely used for classification tasks , but is seldom used for regression because it assumes a distinct label space . The idea of Mixup is to mix pairs of examples using the key assumption that taking a linear interpolation between examples can be used to estimate the label of any examples in between . Mixup is known to effectively regularize the model . More recently , Manifold Mixup ( Verma et al. , 2019 ) has been proposed to improve the hidden representation and decision boundaries where two examples are mixed in multiple hidden layers of neural networks . However , the Mixup techniques are not readily applicable to a regression setting because the key linearity assumption does not necessarily hold . Since the label space is continuous , taking a linear interpolation of examples that are very different either data-wise or label-wise may result in arbitrarily-incorrect labels as shown in Fig . 1a . As a result , the linearity assumption only holds to a certain extent , and the degree may vary for each example . Moreover , other data augmentation techniques for classification including image processing ( e.g. , flipping or rotating ) and generative models ( e.g. , GAN ( Goodfellow et al. , 2014 ) and VAE ( Kingma & Welling , 2014 ) ) are even less applicable to a regression setting ( see Sec . 5 ) . We propose MixRL , a data mixing augmentation framework that is the first to tailor Mixup for regression tasks using reinforcement learning . MixRL uses a stricter linearity assumption where it only holds within a certain data or label distance . These distance limits may vary by example , and we formulate the problem of learning for each example how many nearest neighbors it should be mixed with . MixRL employs a meta learning framework that estimates how valuable mixing an example is for reducing the model loss on a small validation set using Monte Carlo policy gradient reinforcement learning . MixRL ’ s framework is inspired by the recent Data Valuation using Reinforcement Learning ( DVRL ) framework ( Yoon et al. , 2020 ) , which solves the different problem of measuring how individual examples contribute to model performance without any mixing involved . Fig . 1b shows how limiting the nearest neighbors to mix is better than mixing with all neighbors as in classification . To see if the augmentation is useful , we train simple models on the original four examples ( i.e. , no augmentation ) , the augmented data in Fig . 1a , and the augmented data in Fig . 1b . Evaluating the models on 20 random test examples results in Root Mean Square Error ( RMSE ; see Sec . 4 ) values of 0.2967 , 0.5615 , and 0.1834 , respectively , where a lower RMSE is better . We thus conclude that carefully mixing examples is important for improving regression performance . Experiments conducted on real and synthetic datasets show that MixRL shows better model performances relative to baselines , especially when the linearity is limited , and the mixing must be done selectively . In addition , MixRL only requires small validation sets and scales to large training sets . 2 LIMITED LINEARITY IN DATA AND LABEL SPACE FOR REGRESSION . We explain why the key linearity assumption used for Mixup in classification has limitations in a regression setting . In classification , the labels are discrete where many examples may have the same label . The original version of Mixup ( Zhang et al. , 2018 ) is to take a linear interpolation between any pair of examples xi and xj with the labels yi and yj to produce the new example λxi+ ( 1−λ ) xj with the label λyi + ( 1 − λ ) yj where λ ∼ Beta ( α , α ) . The linearity assumption turns out to be reasonable because the label difference between examples is only 0 or 1 and thus not that sensitive to the data difference . In contrast , the labels in regression are in a continuous space . Although there is still a many-to-one relationship where multiple examples may have the same label , the degree is much smaller than in classification . As a result , when two examples are mixed , the interpolated label can be arbitrarily different than the actual label , e.g. , mixing the points a and d in Fig . 1a results in a label nowhere near the actual label . In Sec . 4.1 , we also show empirical results where the label error increases for larger data or label distances . Furthermore , mixing examples with larger data or label distances tend to have increasingly-negative effects on the model trained on the augmented training set . Figs . 2a and 2b show the model accuracies using RMSE for the Product dataset ( described in Sec . 4 ) when mixing examples with different ranges of distances . Regardless of adjusting the label or data distance , there are diminishing benefits for larger distances . How do we limit the data and label distances to improve Mixup for regression ? One approach is to only limit the data distance , which limits the label distance as well . Suppose that the regression function f is continuous where limx→c f ( x ) = f ( c ) for any x and c within the domain of f . We show that a short-enough data distance sufficiently reduces the label distance as well . Given f ’ s domain D and limx→c f ( x ) = L , the following is known to hold : ∀ , ∃δ s.t . ∀x ∈ D , if |x − c| < δ , then |f ( x ) −L| < . We can use this result to prove that ∀ , ∃δ s.t . ∀xi , xj ∈ D , if αxi + ( 1− α ) xj = c , 0 ≤ α ≤ 1 , and |xi − xj | < δ then the absolute difference between the mixed example ’ s xj value and L is small where |αf ( xi ) + ( 1 − α ) f ( xj ) − L| = |α ( f ( xi ) − L ) + ( 1 − α ) ( f ( xj ) − L ) | ≤ α|f ( xi ) − L|+ ( 1− α ) |f ( xj ) − L| < α + ( 1− α ) = . We can not do the converse and limit the data distance by limiting label distance because there is a many-to-one mapping from data to labels , which means that two different examples may have identical labels ( e.g. , Fig . 1a ’ s a and d ) . However , limiting the data distance to also limit the label distance may be too restrictive because there is not much correlation between the data and label distances in real data . Fig . 2c shows how the data distance relates to the label distance for the Product dataset . Even for small data distances , the label distance has a large range , which means that the data distance would have to be extremely limited . Hence , our solution is to limit the data and/or label distance as needed instead of just the data distance . This approach turns out to be more practical as we demonstrate in Sec . 4.4 . 3 MIXRL . The goal of MixRL is to identify which examples to mix with which nearest neighbors . Instead of finding the actual distance limits themselves , we solve the identical problem of finding the number of nearest neighbors to mix per example for convenience . We use reinforcement learning because there is no training data on how mixing each example affects the model performance , and the reward function is non-differentiable as we explain below . MixRL ’ s framework is inspired by DVRL ( Yoon et al. , 2020 ) , but we solve the different problem of mixing examples and address new issues . 3.1 POLICY OPTIMIZATION REINFORCEMENT LEARNING . Finding the optimal policy involves taking the gradient of the objective function J ( θ ) = Eπθ [ R ] where π is the policy , θ is its parameters , and R is the reward function . For MixRL , we would like to minimize the regression model loss on a small validation set when mixing a batch of examples . However , the validation loss is computed using the regression model , which does not involve θ . Hence we can not analytically compute the differential of the reward function with respect to θ . Instead , we use the REINFORCE ( Williams , 1992 ) policy gradient algorithm , which is a Monte Carlo method that is also widely used for data valuation ( Yoon et al. , 2020 ) and neural architecture searching ( Zoph & Le , 2017 ) . It is known that∇θJ ( θ ) can be approximated ( see Sec . A.1 ) as : ∇θJ ( θ ) ≈ 1 m m∑ i=1 R ( τ i ) T−1∑ t=0 ∇θ log πθ ( ait|sit ) ( 1 ) where τ i is the ith state-action trajectory under policy πθ , m is the number of sample trajectories , T is the number of actions taken in a path , sit is a state at time t , and a i t is an action at time t. In our setting of minimizing loss , REINFORCE performs gradient descent for each example decreasing θ by α∇θJ ( θ ) where α is a learning rate . Although the estimated gradient is unbiased , it is known to have a high variance , which we reduce using baseline techniques ( Sutton & Barto , 2018 ) . 3.2 FRAMEWORK . We define notations used in MixRL ’ s framework shown in Fig . 3 . Let D = { ( xi , yi ) } Si=1 ∼ P be the training set where xi ∈ X is a d-dimensional input example , and yi ∈ Y is an e-dimensional label . Let Dv = { ( xvi , yvi ) } Vi=1 ∼ P t be the validation set , where P t is the distribution of the test set , which is not necessarily the same as P . Let fφ be a regression model , and L the loss function that returns a performance score comparing fφ ( xi ) with the true label yi using Mean Square Error ( MSE ) . We assume a list N of possible data and label nearest neighbors ( NNs ) that can be mixed with an example . For instance , N could contain the options “ 1 data NN ” , “ 2 data NNs ” , and “ 2 label NNs ” . The more fine-grained the NN options are , the more precisely MixRL can determine the optimal number of NNs to mix per example . We do not add a “ 0 NN ” option because selecting and not selecting it have identical effects , which makes the policy network training unstable because nothing can be learned . Instead , we support this option separately by excluding examples that are not worth mixing as we explain later in this section . The possible NNs can be represented as a one-hot encoding vector of |N | values . We now define the states and actions , which is an important design choice of MixRL . A state st is a batch of examples Db = { ( ( xi , yi ) , ki ) } Bi=1 where { ( xi , yi ) } Bi=1 ⊆ D and each ki is an index of an NN option in N that specifies how many data or label NNs xi should be mixed with . An action at is then choosing Dm ⊆ Db where each ( xi , yi ) in Dm is mixed with its N [ ki ] NNs in D. A policy πθ ( Dm|Db ) returns the probability of selecting Dm at state Db . For each episode , MixRL selects a batch ( state ) once and chooses a subset of the batch to obtain a reward ( action ) once . Since there is only one time step , the transition function does not play a role . A naı̈ve implementation of the policy network would have an input dimension of B× ( d+ e+ |N | ) and an output dimension of 2B , which may be too large to train for large batch sizes . For example , in a typical setting of B = 1000 , d = 100 , e = 1 , and |N | = 10 , the input and output dimensions Algorithm 1 : Pseudo code for Mixup value network training . Input : Training set D , validation set Dv , nearest neighbor options N , learning rate α , reward scaling constant C , moving average window W Output : Mixup value network hθ Initialize θ , Base = 0 ; while until convergence do Sample Db = { ( xi , yi ) , ki } B 1 ∼ P × Uniform ( N ) ; Dm = ∅ ; for ( ( xi , yi ) , ki ) ∈ D b do Add ( xi , yi , ki ) to D m with probability hθ ( xi , yi , ki ) ; Train regression model fφ with initialized φ on Db ∪Mix ( Dm , D , N ) ; Loss = 1 V ∑V i=1 L ( fφ ( x v i ) , y v i ) ; θ = θ − α · C · ( Loss−Base ) · ∇θ log ( πθ ( Dm , Db ) ) ; Base = W−1 W Base+ 1 W Loss ; return return hθ ; become 111 , 000 and 21000 , respectively . Instead , we can significantly reduce the network size by assuming independence among examples and decomposing the policy network ’ s prediction into two phases . First , a Mixup value network hθ ( x , y , k ) is used to estimate the probability of an example being mixed with k NNs . Next , we randomly select which examples to mix according to the probabilities . The value of πθ ( Dm|Db ) is thus ∏ ( x , y , k ) ∈Dm hθ ( x , y , k ) ∏ ( x , y , k ) ∈Db\Dm [ 1− hθ ( x , y , k ) ] . Also , hθ has input and output dimensions of only d+ e+ |N | and 1 , respectively , and is practical . The objective function J ( θ ) is the validation set loss lθ of fφ trained on Db ∪ Mix ( Dm , D , N ) ( either from scratch or from a pre-trained state ) where Mix ( Dm , D , N ) mixes each ( xi , yi ) in Dm with its N [ ki ] NNs in D : J ( θ ) = l ( φ ) = E ( xvi , y v i ) ∼P t L ( fφ ( xvi ) , y v i ) ( 2 ) Using Eq . 1 , the gradient of Eq . 2 can be approximated using the validation set Dvas : ∇θJ ( θ ) = ∇θl ( θ ) ≈ 1 V V∑ i=1 L ( fφ ( xvi ) , y v i ) · ∇θ log ( πθ ( Dm , Db ) ) ( 3 ) where∇θ log πθ ( Dm , Db ) =∇φ ∑ ( x , y , k ) ∈Dm log hθ ( x , y , k ) + ∑ ( x , y , k ) ∈Db\Dm [ 1− hθ ( x , y , k ) ] . We also use reward scaling ( Henderson et al. , 2018 ) to further improve the training . Algorithm 1 shows the pseudo code for training hθ . The computational complexity is not directly related to the training set size , but depends on the number of iterations needed to train hθ and how long each iteration takes . Another factor is the number of possible NNs |N | where a higher number may result in slower training . We show in Sec . 4.2 that MixRL scales to large training sets . Once we train hθ , we choose the ( x , y , k ) ’ s with the highest hθ ( x , y , k ) values of at least a threshold T . We optimize T ’ s value using the validation set ( see Sec . B.1 ) . We ignore any ( x , y , k ) whose ( x , y ) pair has already been chosen with a different k value for the same data or label distance . We apply T in order to exclude ( x , y , k ) ’ s that are not worth mixing ( i.e. , make them mix with 0 NNs ) . Although actor-critic methods ( Mnih et al. , 2016 ; Schulman et al. , 2017 ) improve on REINFORCE by training both value and policy networks , using the two networks may not be practical in our setting due to the high input dimensions . For a policy network , we can reduce its dimension using the fact that the log probabilities can be calculated by multiplying selection probabilities . Unfortunately , we can not use the same trick for a value network because its output values are not probabilities .

Summary: This paper improves on the idea of mixup by selecting samples for mixup via a model that selects suitable pairs found using knn in a batch. In order to provide a learning signal to this discrete process, the authors apply REINFORCE, using the loss of the downstream regressor (not applied to classification tasks). They show promising results on a number of regression tasks.