Optimizing Data Usage via Differentiable Rewards

1 INTRODUCTION . While deep learning models are remarkably good at fitting large data sets , their performance is also highly sensitive to the structure and domain of their training data . Training on out-of-domain data can lead to worse model performance , while using more relevant data can assist transfer learning . Previous work has attempted to create strategies to handle this sensitivity by selecting subsets of the data to train the model on ( Jiang & Zhai , 2007 ; Wang et al . ; Axelrod et al. , 2011 ; Moore & Lewis , 2010 ) , providing different weights for each example ( Sivasankaran et al. , 2017 ; Ren et al. , 2018 ) , or changing the presentation order of data ( Bengio et al. , 2009 ; Kumar et al. , 2019 ) . However , there are several challenges with the existing work on better data usage strategies . Most work data filtering criterion or training curriculum rely on domain-specific knowledge and handdesigned heuristics , which can be sub-optimal . To avoid hand designed heuristics , several works propose to optimize a parameterized neural network to learn the data usage schedule , but most of them are tailored to specific use cases , such as handling noisy data for classification ( Jiang et al. , 2018 ) , learning a curriculum learning strategy for NMT ( Kumar et al. , 2019 ) , and actively selecting data for annotation ( Fang et al. , 2017 ; Wu et al. , 2018 ) . Fan et al . ( 2018 ) proposes a more general teacher-student framework that first trains a teacher network to select data that directly optimizes development set accuracy over multiple training runs . However , because running multiple runs of training simply to train this teacher network entails an n-fold increase in training time for n runs , this is infeasible in many practical settings . In addition , in preliminary experiments we also found the single reward signal provided by dev set accuracy at the end of training noisy to the extent that we were not able to achieve results competitive with simpler heuristic training methods . In this paper , we propose an alternative : a general Reinforcement Learning ( RL ) framework for optimizing training data usage by training a scorer network that minimizes the model loss on the development set . We formulate the scorer network as a function of the current training examples only , making it possible to re-use the model architecture which is designed and trained for the main task . 1We will make the code publicly available upon acceptance . Thus , our method requires no heuristics and is generalizable to various tasks . To make the scorer adaptive , we perform frequent and efficient updates of the scorer network using a reward function inspired by recent work on learning using data from auxiliary tasks ( Du et al. , 2018 ; Liu et al. , 2019b ) , which use the similarity between two gradients as a measure of task relevance . We propose to use the gradient alignment between the training examples and the dev set as a reward signal for a parametric scorer network , as illustrated in Figure 1 . We then formulate our framework as an optimization problem found in many prior works such as meta-learning ( Finn et al. , 2017 ) , noisy data filtering ( Ren et al. , 2018 ) , and neural architecture search ( Liu et al. , 2019a ) , and demonstrate that our proposed update rules follow a direct differentiation of the scorer parameters to optimize the model loss on the dev set . Thus we refer to our framework as “ Differentiable Data Selection ” ( DDS ) . We demonstrate two concrete instantiations of the DDS framework , one for a more general case of image classification , and the other for a more specific case of neural machine translation ( NMT ) . For image classification , we test on both CIFAR-10 and ImageNet . For NMT , we focus on a multilingual setting , where we optimize data usage from a multilingual corpus to improve the performance on a particular language . For these two very different and realistic tasks , we find the DDS framework brings significant improvements over the baselines for all settings . 2 DIFFERENTIABLE DATA SELECTION . 2.1 RISK , TRAINING , AND DEVELOPMENT SETS . Commonly in machine learning , we seek to find the parameters θ∗ that minimize the risk J ( θ , P ) , the expected value of a loss function ` ( x , y ; θ ) , where 〈x , y〉 are pairs of inputs and associated labels sampled from a particular distribution P ( X , Y ) : θ∗ = argmin θ J ( θ , P ) where J ( θ , P ) = Ex , y∼P ( X , Y ) [ ` ( x , y ; θ ) ] ( 1 ) Ideally , we would like the risk J ( · ) to be minimized over the data distribution that our system sees at test time , ie . Ptest ( X , Y ) . Unfortunately , this distribution is unknown at training time , so instead we collect a training set Dtrain = { ( xi , yi ) : i = 1 , ... , Ntrain } with distribution Ptrain ( X , Y ) = Uniform ( Dtrain ) , and minimize the empirical risk by taking 〈x , y〉 ∼ Ptrain ( X , Y ) . Since we need a sufficiently large training set Dtrain to train a good model , it is hard to ensure that Ptrain ( X , Y ) ≈ Ptest ( X , Y ) . In fact , we often accept that training data comes from a different distribution than test data . The discrepancy between Ptrain ( X , Y ) and Ptest ( X , Y ) manifests itself in the form of problems such as overfitting ( Zhang et al. , 2017 ; Srivastava et al. , 2014 ) , covariate shift ( Shimodaira , 2000 ) , and label shift ( Lipton et al. , 2018 ) . However , unlike the large training set , we can collect a relatively small development set Ddev = { ( xi , yi ) : i = 1 , ... , Ndev } with distribution Pdev ( X , Y ) which is much closer to Ptest ( X , Y ) 2 . Since Ddev is a better approximation of our test-time scenario3 , we can use Ddev to get reliable feedback to learn to better utilize our training data from Dtrain . In particular , we propose to train a scorer network , parameterized by ψ , to provide guidance on training data usage to minimize J ( θ , Ddev ) . 2.2 REINFORCEMENT LEARNING FOR OPTIMIZING DATA USAGE . We propose to optimize the scorer ’ s parameters ψ in an RL setting . Our environment is the model state θ and an example 〈x , y〉 . Our RL agent is the scorer network ψ , which optimizes the data usage 2As is standard in machine learning experiments , we make sure that Ddev has no overlap with Dtrain or Dtest . Details of how we construct the Ddev can be found in Appendix A.2 and A.4 . 3For example , in Section 3.2 we would like to use training data from many different languages to improve the performance of a particular low-resource language . Here we can gather a small set of Ddev from the low resource language , even if we can ’ t gather a large training set in the language . Moreover , in a domain adaptation setting we can obtain a small dev set in the target domain , or in a setting of training on noisy data we can often obtain a small clean dev set . for the current model state . The agent ’ s reward on picking an example approximates the dev set performance of the resulting model after the model is updated on this example . Our scorer network is parameterized as a differentiable function that only takes as inputs the features of the example 〈x , y〉 . Intuitively , it represents a distribution over the training data where more important data has a higher probability of being used , denoted P ( X , Y ; ψ ) . Unlike prior methods which generally require complicated featurization of both the model state and the data as input to the RL agent ( Fan et al. , 2018 ; Jiang et al. , 2018 ; Fang et al. , 2017 ) , our formulation is much simpler and generalizable to different tasks . Since our scorer network does not consider the model parameters θt as input , we update it iteratively with the model so that at training step t , P ( X , Y ; ψt ) provides an up-to-date data scoring feedback for a given θt . Although the above formulation is simpler and more general , it requires much more frequent updates to the scorer parameter ψ . Existing RL frameworks simply use the change in dev set risk as the regular reward signal , which makes the update expensive and unstable ( Fan et al. , 2018 ; Kumar et al. , 2019 ) . Therefore , we propose a novel reward function as an approximation to ∆Jdev ( x , y ) to quantify the effect of the training example 〈x , y〉 . Inspired by Du et al . ( 2018 ) ( which uses gradient similarity between two tasks to measure the adaptation effect between them , we use the agreement between the model gradient on data 〈x , y〉 and the gradient on the dev set to approximate the effect of 〈x , y〉 on dev set performance . This reward simply implies that we prefer data that moves θ in the direction that minimizes the dev set risk : R ( x , y ) = ∆Jdev ( x , y ) ≈ ∇θ ` ( x , y ; θt−1 ) > · ∇θJ ( θt , Ddev ) ( 2 ) According to the REINFORCE algorithm ( Williams , 1992 ) , the update rule for ψ is thus ψt+1 ← ψt +∇θ ` ( x , y ; θt−1 ) · ∇θJ ( θt , Ddev ) ︸ ︷︷ ︸ R ( x , y ) ∇ψlog ( P ( X , Y ; ψ ) ) ( 3 ) The update rule for the model is simply θt ← θt−1 −∇θJ ( θt−1 , P ( X , Y ; ψ ) ) ( 4 ) For simplicity of notation , we omit the learning rate term . Full derivation can be found in Appendix A.1 . By alternating between Eqn . 4 and Eqn . 3 , we can iteratively update θ using the guidance from the scorer network , and update ψ to optimize the scorer using feedback from the model . Our formulation of scorer network as P ( X , Y ; ψ ) has several advantages . First , it provides the flexibility that we can either sample a training instance or equivalently scale the update from the training instance based on its score . Specifically , we provide an algorithm under the DDS framework for multilingual NMT ( see Sec . 3.2 ) , where the former is more efficient , and another more general algorithm for image classification ( see Sec . 3.1 ) , where the latter choice is natural . Second , it allows easy integration of prior knowledge of the data , which is shown to be effective in Sec . 4 .

Summary: This paper introduces a simple idea to optimize the weights of a weighted empirical training distributions. The goal is to optimize the population risk, and the idea is to optimize a distribution over the training examples to maximize the cosine similarity between training set gradients and validation set gradients. The distribution over the training set is parameterized by a neural network taking as arguments the