[Paper Review] Optimizing Data Usage via Differentiable Rewards

Xinyi Wang, Hieu Pham+4 more|arXiv (Cornell University)|Nov 22, 2019

Domain Adaptation and Few-Shot Learning60 references4 citations

TL;DR

This paper proposes Differentiable Data Selection (DDS), a reinforcement learning framework that trains a learnable scorer to dynamically weight training data based on gradient alignment with a development set, enabling efficient, end-to-end optimization of data usage. DDS achieves consistent performance gains over strong baselines in image classification and multilingual neural machine translation without heuristic design or expensive multi-pass training.

ABSTRACT

To acquire a new skill, humans learn better and faster if a tutor, based on their current knowledge level, informs them of how much attention they should pay to particular content or practice problems. Similarly, a machine learning model could potentially be trained better with a scorer that "adapts" to its current learning state and estimates the importance of each training data instance. Training such an adaptive scorer efficiently is a challenging problem; in order to precisely quantify the effect of a data instance at a given time during the training, it is typically necessary to first complete the entire training process. To efficiently optimize data usage, we propose a reinforcement learning approach called Differentiable Data Selection (DDS). In DDS, we formulate a scorer network as a learnable function of the training data, which can be efficiently updated along with the main model being trained. Specifically, DDS updates the scorer with an intuitive reward signal: it should up-weigh the data that has a similar gradient with a dev set upon which we would finally like to perform well. Without significant computing overhead, DDS delivers strong and consistent improvements over several strong baselines on two very different tasks of machine translation and image classification.

Motivation & Objective

• To address the challenge of suboptimal data usage in machine learning, where training data distribution often differs from test distribution, leading to poor generalization.
• To eliminate reliance on hand-designed heuristics or task-specific rules for data selection, weighting, or curriculum learning.
• To develop a generalizable, efficient method for optimizing data usage that adapts during training and improves performance on a development set.
• To enable automatic, adaptive data weighting that evolves with the model's learning state, avoiding the need for multiple training runs.

Proposed method

• Formulates data weighting as a differentiable function of training examples, parameterized by a scorer network trained alongside the main model.
• Uses gradient alignment between training examples and the development set as a reward signal to guide scorer optimization.
• Employs a bi-level optimization framework where the scorer is updated to minimize the model's loss on the development set through direct differentiation.
• Applies the reward signal via a differentiable update rule that efficiently propagates gradients through the scorer network.
• Introduces two concrete instantiations: one for image classification (CIFAR-10, ImageNet) and another for multilingual neural machine translation.
• Leverages recent advances in gradient-based meta-learning and auxiliary task learning to define a stable and effective reward function.

Experimental results

Research questions

• RQ1 Can a learnable scorer network be trained end-to-end to dynamically weight training data for improved development set performance?
• RQ2 How can a reward signal be designed to reflect data relevance without relying on handcrafted heuristics or multiple training runs?
• RQ3 Can the proposed differentiable framework generalize across diverse tasks such as image classification and neural machine translation?
• RQ4 Does gradient alignment between training and development set examples serve as a reliable and effective proxy for data utility?

Key findings

• DDS achieves consistent and significant improvements over strong baselines in both image classification and multilingual neural machine translation tasks.
• The method outperforms heuristic-based and meta-learning-based baselines, including Ren et al. (2018), by a large margin in image classification.
• In multilingual NMT, DDS learns to prioritize data from related languages (e.g., focusing on Czech when training for Slovak), adapting dynamically over training.
• The data weighting strategy evolves over time, with the model shifting focus to different language pairs or data types depending on the learning stage.
• The framework requires no additional training passes or expensive multi-run optimization, making it computationally efficient.