A Bayesian Approach to Data Point Selection

TL;DR

This paper introduces a Bayesian method for data point selection in deep learning, addressing computational and theoretical issues of existing approaches by jointly inferring instance weights and model parameters using stochastic gradient Langevin MCMC.

Contribution

It proposes a novel Bayesian model for DPS, enabling efficient joint inference of weights and parameters, and demonstrates scalability to large models and diverse domains.

Findings

Efficient joint inference with stochastic gradient Langevin MCMC.

Scalable to large language models and diverse tasks.

Improves over bi-level optimisation in computational efficiency.

Abstract

Data point selection (DPS) is becoming a critical topic in deep learning due to the ease of acquiring uncurated training data compared to the difficulty of obtaining curated or processed data. Existing approaches to DPS are predominantly based on a bi-level optimisation (BLO) formulation, which is demanding in terms of memory and computation, and exhibits some theoretical defects regarding minibatches. Thus, we propose a novel Bayesian approach to DPS. We view the DPS problem as posterior inference in a novel Bayesian model where the posterior distributions of the instance-wise weights and the main neural network parameters are inferred under a reasonable prior and likelihood model. We employ stochastic gradient Langevin MCMC sampling to learn the main network and instance-wise weights jointly, ensuring convergence even with minibatches. Our update equation is comparable to the widely used SGD and much more efficient than existing BLO-based methods. Through controlled experiments in both the vision and language domains, we present the proof-of-concept. Additionally, we demonstrate that our method scales effectively to large language models and facilitates automated per-task optimization for instruction fine-tuning datasets.