A Bayesian Approach to Data Point Selection

Summary

The paper proposes a method for data point selection (DPS) in the setting where only little data from the target (or meta) distribution is available, together with a lot of data from the training distribution. DPS aims to find the weights of data points in the training distribution such that the model parameters that minimise the loss of the therewith weighted training distribution also minimise the loss on the target distribution. The method proposed for DPS in this paper consists in treating the problem in a Bayesian way and doing posterior inference on both (data) weights and model parameters simultaneously, using stochastic gradient Langevin dynamic sampling (SGLD) to approximate the posterior. The crucial trick is to view the data from the target distribution as generated by the model (parameters), which allows for the inference. They show the effectiveness of their approach on 4 experiments and compare against 2 DPS baselines (bilevel optimisation (BLO) and contrastive data selection (CDS)) as well as a range of naive baselines (mixing available training and target distributions to different extents). - The first experiment is on data balancing for MNIST, where the class balance differs significantly between training (imbalanced) and target (balanced) distribution. - The second experiment is on data denoising for CIFAR, where the training distribution (noisy) differs from the target distribution (clean) in the degree of noisiness - The third experiment is on efficient learning on WebNLG, where there are differences in the topic or domain of the data between training distribution (Artist, City) and target distribution (Athlete, Politician, Astronaut) - The fourth experiment is on instruction finetuning (IFT) for LLMs, using an OpenLlama3B model. The proposed method compares favourably in all experiments in the main part of the paper. In the appendix results for the second experiment are additionally shown for longer training times, in which case BLO outperforms the proposed method.

Strengths

- The idea of framing and solving DPS in a Bayesian way seems novel and innovative - The paper is well written in general and is easy to read and follow - The proposed method is getting good empirical results as displayed in the paper, in particular also on the relevant use case of IFT for LLMs - Compared to BLO, the proposed method seems to have lower memory and compute requirements

Weaknesses

- For the most part, the method proposed is explained well and in detail. However, when it comes to a crucial aspect in the implementation that is subsequently used in all experiments (the weight network), details are brushed under the carpet as a 'straightforwardly modification'. However, for me the causal graph in Figure 1 seems to be different if the weights become a function of $z^t$, let alone $\theta$ if model embeddings $f_{\theta}(z^t)$ are used instead of raw data $z^t$ like in the experiments. This should be explained in more detail (see also Question about this below). - In the main part of the paper, performance of BLO is portrayed as clearly falling behind the proposed method and other methods in terms of accuracy on the CIFAR experiment (Fig 2, Fig 3, line 266). However in Fig 7 in the Appendix, it can be seen that with longer training, BLO is at least on par (symmetry noise) or even slightly better (asymmetry noise) in terms of accuracy. It would be better if this would already be mentioned in the main part of the paper - Assuming a non-negativity constraint on the weights and the loss, in Eq (4) and subsequently Eq (7) the gradient wrt $w$ will be negative for the data dependent term and thus each gradient update is only pointing away from $0$ through the gradient of the prior term $\log p(w)$ (and the noise term at the end). Therefore, preventing collapse to $0$ weights (as described in line 148) seems to crucially depend on a strong enough prior as well as the impact constant and $N_t$. At least for the not so compute intensive experiments on MNIST and CIFAR, it would be good to see an ablation on the effect of these impact constants and the prior governing sparsity parameter $\beta$. - The limitations are only spelled out after the page limit

Questions

- Which part in the prior formulation or in the optimisation ensures that the weights are non-negative? Or are they intentionally not constrained by the method in that sense? - How does the causal graph in Figure 1 change if a weight network that is using embeddings of training data as inputs is being used like in the experiments? How does posterior inference change in this case? - Why is the weight network being used in the experiments if in the ablation in the appendix, the element wise weights version of the proposed approach is outperforming it? - Why are the inputs to the weight network different for the MNIST experiment (image embeddings) and the CIFAR experiment (image embedding + one-hot encoded label)? - When comparing to BLO in terms of compute efficiency, have the authors taken into account proposals of making the Hessian vector product calculations more efficient for BLO (see https://iclr-blogposts.github.io/2024/blog/bench-hvp/ for example)?

Rating

7

Confidence

3

Soundness

2

Presentation

3

Contribution

4

Limitations

- Limitations are discussed, but after the page limit

Summary

This work proposes a Bayesian approach to the data point selection task. Specifically, the authors introduce the important weights to each training point and derive the posterior joint probability of the instance-wise weights and network parameters. The parameters and weights are then sampled iteratively based on the stochastic-gradient MCMC technique. Experiments on computer vision and natural language processing have been conducted to show the efficacy of the proposed method.

Strengths

1. This work views the data point selection problem from the perspective of Bayesian theory and introduces the posterior distributions of the instance-wise weights and network parameters. 2. Experimental results demonstrate the efficacy of the proposed method.

Weaknesses

1. I suggest the authors provide a detailed derivative for Eq. 4. 2. Eq. 5 still requires to compute the normalizing constant in Eq. 4, which serves as the bias of the gradient. 3. The authors claim that “Our method straightforwardly achieves sparsity on the w weights allowing efficient sample selection unlike”. I hope the authors provide more analysis and evidence about the sparsity of instance-wise weight.

Questions

Please see the weaknesses

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

No potential negative societal impact.

Summary

This paper proposes a new Bayesian method for Data Point Selection (DPS), called BADS. DPS aims to select training data points that optimize performance on downstream tasks. Instead of relying on bi-level optimization (BLO), BADS frames DPS as a posterior inference problem. The method uses a weight network to learn the importance of each data point and proposes a novel update rule for jointly updating model parameters and data point weights using Stochastic Gradient Langevin Dynamic sampling. Experiments on MNIST, CIFAR, and WebNLG datasets show that BADS outperforms or matches existing DPS methods, including BLO, in terms of accuracy and efficiency.

Strengths

- Originality: This paper introduces a novel perspective on DPS by using a Bayesian framework instead of the BLO approach. This offers a new way to think about data selection and its connection to posterior inference. The paper also leverages Langevin sampling for optimizing the derived posterior inference which adds to the novelty. - Clarity: The paper is well-written and easy to understand. The authors clearly explain their method and the motivations behind their design choices. The final derived method is simple and intuitive.

Weaknesses

- Limited theoretical analysis: The paper focuses on the empirical performance of BADS and lacks a deep theoretical analysis. Providing more theoretical insights on the convergence and optimality properties of BADS would strengthen the paper, especially since the authors claim that their method, unlike BLO, has theoretical convergence guarantees. - Hyperparameter sensitivity: The paper mentions that BADS has several hyperparameters that require careful tuning. A more in-depth discussion on the impact of these hyperparameters and potential strategies for automatic tuning would be beneficial. - Comparison and claims: The paper primarily compares BADS with BLO, but there are other methods that deal with selecting data to learn on, such as offline pruning methods [1], and online and offline batch selection [2, 3], among others. In addition to that, some of these methods have already been used for large scale model training, such as [4, 5]. Some important equations and their corresponding assumptions have been omitted. In particular, equation (4) is not derived and on first glance seems to assume independence of $w$ and $\mathcal{D}$, but it is not clear if this still holds after the introduction of the weight network. 1. https://arxiv.org/abs/2206.14486 2. https://arxiv.org/abs/2312.05328 3. https://arxiv.org/abs/2406.10670 4. https://arxiv.org/abs/2402.09668 5. https://arxiv.org/abs/2406.17711

Questions

Can the authors provide any theoretical guarantees on the convergence of BADS? What are the conditions under which BADS finds an optimal solution? What is the impact of different hyperparameters, especially the sparsity level parameter and impact constants? How robust is BADS to hyperparameter variability.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors have adequately addressed the limitations proposed in their work.