A Bayesian Approach to Data Point Selection

Authors: xinnuo  xu
minyoung  kim
brais  martinez
timothy  hospedales
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:59:24

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.

Main paper: 
 introduction: 
 Practical machine learning efficacy is heavily dependent on the choice, quality and quantity of training data, especially so in the case of neural networks that can easily fit every detail of the training set. This leads to challenges from how to learn reliably with imbalanced data [22] , noisy data, noisy labels [46] , and so on. Similarly there is often a key subset of data, which is most informative for a given learning problem, but buried among a much larger set of less relevant data. If the most salient data could be efficiently identified, learning could potentially be accelerated [15] . All these challenges are only growing in the era of large scale training on web-scraped data, where curation and gold-standard quality control are not feasible.Data Point Selection (DPS) algorithms aim to address these challenges by filtering or re-weighting the training data to reduce noise, imbalance, irrelevant background data and so on. The most established family of approaches [19, 15, 41, 45, 62] to DPS falls under the bi-level optimization or meta-learning umbrella, where one wraps the conventional learning problem with an outer loop that optimizes the dataset itself, so as to maximise performance on some validation set. These methods vary in the choice of their outer optimization variable (e.g., data point weights [41, 45, 19] vs mini-batch sampler [15] ), the method of computing meta-gradients (e.g., reverse mode differentiation [45, 62] or reinforcement learning [15] ), and the customization of their losses and other design parameters for the different scenarios (e.g., label-noise [45, 41] , etc). However, all the BLO approaches are quite expensive in computation and/or memory, which limits their applicability to the most salient use case of large models trained on large web data.In this paper we revisit the DPS problem from the perspective of Bayesian learning. Rather than constructing expensive nested optimization problems whose convergence is hard to analyse, we treat it as a problem of inferring the joint posterior over the main neural network parameters and instance-wise weights induced by a second weight-estimation neural network. This framework has several advantages in terms of being more efficient and scalable than typical BLO competitors and having a clear convergence guarantee. It is also able to address a variety of DPS-related problemsfrom noise and imbalance to data curation -within a single framework.Our empirical results present proof of concepts for all these capabilities on a variety of learning tasks in vision and language. We also show a use case of automating Large Language Models (LLMs) instruction fine-tuning (IFT) data curation for specific downstream tasks. The available IFT datasets are large, diverse, and of varying quality. This means that a key activity for natural language processing (NLP) researchers and developers is often finding the right composition of IFT data sources to optimize particular use cases. Our framework can automatically resample and curate the wide array of available auxiliary IFT data to optimize performance for each NLP task of interest. To our knowledge, no BLO alternative has been demonstrated on billion-parameter scale LLMs. We name our approach BADS (Bayesian Data Point Selection). 32 Our ApproachQuestion: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? Answer: [Yes] Justification: We have included our main claims in the abstract and introduction that accurately reflect the paper's contributions and scope. Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
problem setup: 
 For the DPS problem, we assume that we are given two datasets: the train set D t = {z t i } Nt i=1 and the meta set D m = {z m i } Nm i=1 . Each data point z i can be an input-target pair z i = (x i , y i ) in the traditional (class-)labeled data scenarios. But in the autoregressive generative model scenarios (e.g., LLM), z i is simply a sequence of tokens in which the inputs and targets are rather implicitly defined (e.g., all the tokens up until current time as input and the next token as target). We will use the notation l(z i ; θ) for the loss of the model θ on the data point z i , which must be well-defined in both scenarios.The meta dataset D m is considered as in-domain, meaning that the distribution of D m matches that of the downstream test task of interest. The size of D m , denoted by N m , is typically small, due to the cost of curation/annotation processes in practice. The train dataset D t consists of out-of-domain samples, possibly noisy, imbalanced, and uncurated, but the size N t is usually large. The goal of DPS is to select a (soft weighted) subset of the train set D t with the guidance of the meta set D m , so that the model trained on the selected train and meta dataset points performs well. Typical baselines include training with the meta-set alone, or union of meta and train sets.Perhaps one of the most widely adopted DPS techniques is the bi-level optimisation (BLO) formulation of the problem [41, 58] . Letting w i (≥ 0) be the weight (or importance) variable associated with the training data point z t i , the main intuition is to find the weights w ∈ R Nt + such that the model θ trained with the weighted train data with weights w yields the best performance on the meta set. More formally, this leads to the following BLO problem: EQUATIONHowever, a critical drawback is that solving this difficult problem is costly in computation and/or memory, and unrealible due the practical heuristics required. Typical BLO solutions to obtain the hypergradient dl/dω rely on approximate Hessian estimation or reverse mode differentiation with few-step SGD approximation of the inner optimisation. Aside from cost, for practical neural network implementations computed over minibatches, there is no theoretical guarantee for convergence. We view the DPS problem from a completely different perspective, and tackle it via Bayesian learning. Our model's generative process, that is, the graphical model representation, is depicted in Fig. 1 . The main neural network model parameters set θ is a random variable, which can generate data points in the meta dataset D m (precisely speaking, the backbone θ generates the target part of each data point). To make use of the train set D t in an appropriate way, we constrain θ to follow a prior distribution governed by the weighted train data with weights w ∈ R Nt + which are also random variables. Before observing the meta set D m , the weight vector w follows a prior distribution p(w) -a specific distributional choice for p(w) will be discussed later. Given w and D t , our backbone θ has to be compatible with the weight data {(w i , z t i )} Nt i=1 . This can be interpreted as placing a weighted-data-driven prior on θ, more specifically,EQUATIONwhere p(θ) is a base prior (e.g., 0-centered Gaussian that amounts to weight decay regularisation), and p(w i , z t i |θ) can be defined from the loss, e.g., exp(-w i •l(z t i ; θ)), following the conventional tricks [36, 25] . Then given θ, the meta data are generated following the likelihood defined as:EQUATIONThe equations ( 2) and (3) fully constitute the prior and likelihood for our Bayesian model. Our ultimate goal is to describe the distributions of θ and w after observing all evidences D t and D m , which boils down to the posterior inference p(θ, w|D t , D m ). Formally, we have:EQUATIONThe detailed derivations for Eq. ( 4) can be found in Appendix A. However, it is widely known that (4) does not admit any closed-form expressions. One main difficulty arises from the intractable normalizing constant in (4).Stochastic Gradient Langevin Dynamic Sampling. For computationally efficient posterior inference, we adopt the stochastic-gradient MCMC technique, specifically the Stochastic Gradient Langevin Dynamic (SGLD) sampling [55] . Applied to our model, we can obtain samples from the posterior p(θ, w|D t , D m ) by running the Langevin dynamic system (until convergence, i.e., mixing):EQUATIONwhere η is a small (constant) step size. There are two critical benefits: i) Since we differentiate the log-posterior, the difficult normalizing constant in (5) will disappear; ii) The update ( 5) is essentially gradient descent with additive Gaussian noise, leading to a computationally efficient update.Going one step further, even though the log-posterior involves the entire train data (and entire meta data), it is shown in [55] that the stochastic-gradient version (SGLD) that replaces the whole batch likelihood with a minibatched one, theoretically guarantees that the SGLD update converges to the posterior samples. More specifically, the SGLD update equations (one for θ and the other for w) can be written as follows:EQUATIONwhere B t and B m are minibatches from D t and D m , respectively, and ϵ θ , ϵ w ∼ N (0, I) are independent Gaussian samples.Repeating ( 6) and ( 7) for a sufficient amount of iterations (until we reach good mixing) leads us to posterior samples (θ, w). There are several options to take these samples for a final model for test prediction. One option is to collect latest M samples (either consecutive collection or thinning to take every kth samples) from the iterations, and either take the average as posterior means or perform full Bayesian treatment with the collected samples. Alternatively, we can just take the last single iterate (θ, w) as a point representative for the posterior distribution. For simplicity, we take the latter approach, which also works well empirically.We discuss several intuitions and implications of our proposed approach (Eq. 6-7). First, looking at the θ update Eq. ( 6), our model essentially updates θ in a way that it decreases the loss on the combined data of the whole meta data points and the weighted train data points with the current weights. This is a fairly intuitive strategy provided that the weights are properly determined. Then the next question is how the weights are determined. If we inspect the w update (Eq. 7), and take the gradient of the expected loss term with respect to w directly, we see that: i) those train data points z t i s with smaller losses at current backbone θ will get higher weights w i s; ii) those train data points z t i s with larger losses at current backbone θ will get lower weights w i s. This essentially means that our model performs loss alignment for DPS -In the course of training/update, once the backbone θ enters a good regime in the parameter space such that θ can assign (valid) low loss values on the in-domain meta data points, then it starts putting high weights on those train data points that have low losses under the current backbone. In other words, the model will assign high weights to those train data points that are well-aligned with the meta data points in terms of loss.Benefits over BLO Our Bayesian approach provides several benefits over BLO: (1) Efficiency. Our SGLD is efficient so does not require computationally demanding Hessian computations like implicit function theorem based methods (cf: [19] ) or huge memory demand like reverse-mode differentiation methods [41, 19] . (2) Sparsity. Our method straightforwardly achieves sparsity on the w weights allowing efficient sample selection unlike [19] . (3) Reliability. BLO-based methods rely on approximations (truncation, or Hessian approximations) for practical feasibility that make finding optimal solutions unreliable. Our straightforward Bayesian approach has reliable convergence properties thanks to being a standard application of SGLD.In Appendix C we provide a theorem showing that our SGLD algorithm converges to the true posterior. Our analysis is based on [64] where we make some adjustments for our case.Choice of Priors For the base prior p(θ), i.e., the prior before being driven by the weighted data, we adopt 0-mean Gaussian, which amounts to adding the weight decay regularisation for θ. For the weight prior p(w), we have made a careful design effort to come up with a viable sparsity inducing prior. Although encouraging sparsity in learned weights is ideal to avoid overfitting, during our initial experiments we have found that most of the weights eventually tend to vanish to 0, which is not what we actually want. We need to be able to impose both sparsity and a certain level of non-zero weights. To this end, we first introduce a hyperparameter β (e.g., 0.01) for the target sparsity level that we want to attain. Roughly saying, among the N t training data points, we aim to select ⌊N t •β⌋. For the (soft) weights, we impose i∈Dt w i ≈ ⌊N t •β⌋, which can be encoded in the prior form as:p(w) ∝ e -i wi-⌊Nt•β⌋ 2 /2σ 2 (8)where σ controls the strength of the regularisation. One technical difficulty in directly plugging (8) into the weight update (7) is that we have to load the whole {w i } Nt i=1 in memory for backprop. To avoid this issue, we use the following fact:EQUATIONwhere w represent historic running average of the entire weights. We basically build a computation graph only for the first term of batch B t weight sum, and regard the (historic) running average of the entire weights as constant during backprop. After each SGLD iteration, we update the running weight average with the new updated weights on the recent batch. We use the simple averaging scheme for the running average. To approximate the average weight w precisely, we only conduct the average over the most recent s avg step.Introducing impact constants In the SGLD principle, we have the log-likelihood terms that are proportional to the sizes of the datasets. In particular, we have N t and N m in (6). However, this scheme does not properly capture our preference to the in-domain meta data set in contrast to the noisy, out-of-domain train data set. To this end, we introduce the impact constants (hyperparameters) in the update equations where we downweigh or upweigh the loss terms of train and meta sets.Weight Network Instead of directly optimising individual weights w i , we can consider a weight network, w i = w(z t i ; ϕ), a neural network with parameters ϕ that takes the train data point z t i as input and returns its weight w i as output. We can then regard ϕ as random variables and the w update equation can be modified accordingly for ϕ update straightforwardly. This weight network approach can be useful for smoothing/regularising the output weights thanks to the smooth functional property of neural networks. Furthermore, if one needs to supplement the train dataset with extra new train samples after the model training, the learned weight network can be used for assigning weights or selecting samples from the new set, without retraining the whole model from the scratch. The posterior distribution similar to Eq. ( 4) is derived in full detail in Appendix B.Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: We have discussed several limitations of the current work, together with workarounds and our plans to address them in our future work. Please see the Limitations section right before the References. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
experiments: proof of concept: 
 In this section, we assess the effectiveness of the proposed BADS method in three critical scenarios where DPS is essential: Data Balancing, Data Denoising, and Efficient Learning. We begin by introducing the baseline systems and then present the experimental results for each of these scenarios.There are three types of DPS setups with different supervision signals:• Unsupervised DPS selects data without the guidance of a held-out meta set [47, 42] . Instead, it is guided by human-defined hypotheses, such as "challenging examples improve model performance". This approach aligns with curriculum learning. We include the online variant of AskLLM [42] , i.e. AskLLM-O, in our baseline comparisons. It selects examples from training set by querying a pretrained OpenLLaMA 3B to obtain the sampling score for each training sample.4 • Self-supervised DPS selects data with the guidance of a held-out meta set. However, the meta set does not share the same data distribution as the targeted test set [13, 4, 14, 49] . Typically, the examples in the meta set are selected from the training set based on specific hypotheses, such as "learnable examples enhance model performance". We include two approaches in our baseline comparisons: Contrastive Data Selection (CDS) [49] is tailored for data denoising. The algorithm assigns weights to each data point in D t according to the difference between the denoised and the noisy log probability, predicted using a denoised and a noisy model trained on a clean dataset and an uncurated dataset, respectively. These weights are then used to sample data points from minibatches in the training of LM. 5 Similar to CDS, ClassAct utilizes small proxy models trained on a limited portion of D t to calculate learnability scores for the remaining training data points.6 • Meta-set guided DPS selects data with the guidance of a small meta set that shares the same distribution as the test set, aiming to train a model that excels specifically on the target test set. The test set may encompass one or multiple downstream domains or tasks. This DPS is closely related to meta learning, domain adaptation, and transfer learning. Current methods primarily rely on Bilevel Optimization (BLO) for purposes such as data denoising [19, 37] , data balancing [41] , and efficient learning [28, 59, 29] . Considering both performance and code availability, we use the online BLO7 [41, 19] as our baseline. Our approach, BADS, also falls under this category.To ensure a fair comparison, we also introduce several baselines that train the backbone models In this scenario, we assess the model's capability to manage an imbalanced train set, D t . Despite being trained on this imbalanced dataset, the models are expected to perform effectively on a balanced test set. Following the setup in [41] , we use the standard MNIST handwritten digit classification dataset [33] to create a class-imbalanced binary classification task. A total of 5,000 images from classes 4 and 9 were selected as the train set In this scenario, the loss function is defined using the standard binary cross-entropy loss. Following [41] , all classification models are LeNet5 [32] . The training is conducted on a single GPU, using SGD with a fixed learning rate of 1e-3 and a mini-batch size of 100, over a total of 15,000 steps. In BADS, the weight network is implemented as a single-layer Feedforward Neural Network (FNN) with a sigmoid activation function. It takes the top-layer image embeddings from LeNet5 as input and outputs a weight w i ∈ [0, 1] for each image. The learning rate for the weight network is 1e-3 and the target sparsity level β is 0.005. Other hyperparameters, including those in the baselines, are detailed in Table 3 (Appendix E).The classification accuracy is presented in the top-left plot in We further evaluate the models' performance using meta sets D m of various sizes, with 5, 10, 25, and 50 examples per class included. 9 As illustrated in Figure 3 , with a very limited number of meta examples (5 per class), only BADS and BLO achieve over 90% accuracy on the balanced test set.As the number of available meta data increases, BADS consistently leads in performance. However, when sufficient meta data is provided, the gap between BADS and the other approaches narrows. We use ResNet32 [23] as the backbone classification model. Training is performed on a single GPU using SGD with a fixed learning rate of 1e-1 and a mini-batch size of 120 over 20,000 steps. The loss function is the standard multi-class cross-entropy loss. For BADS, the weight network structure retains the same as in scenario 1. However, in this scenario, data point weighting takes into account both the image and its associated label. We represent each label using a one-hot embedding, which is then concatenated with the top-layer image embeddings from ResNet32 and fed into the weight network. The learning rate for the weight network is set to 1e-4. Note that in Eq 7, the gradient of θ become large if N t is big. Therefore, we reduce the learning rate of the backbone classification model to 1e-4. The target sparsity level β is set to 0.8. 10 .The classification accuracy is shown in the top-middle plot in Figure 2 . All methods, except for the BLO approach 11 , manage to achieve over 60% accuracy, even using a train set contaminated by 50%.Notably, CDS, ClassAct and BADS deliver the highest three performances, with BADS surpasses all other methods by a noticeable margin. The bottom row of Figure 13 and the middle plot of Figure 14 (Appendix F) visualizes the weights allocated to the examples in each mini-batch by DPS approaches. If we raise the noise label ratio in the train set to 80% (Figure 4 ), both BLO and BADS still lead the performance, with BADS exceeds BLO and non-DPS approaches by 15% and 20% in classification accuracy, respectively. CDS does not exceed the performance of non-DPS approaches and starts to overfit on the noisy data after 15,000 training steps.When we lower the noise label ratio to 20% (Figure 5 ), all methods achieve a classification accuracy of around 80%. Although BADS continues to outperform both DPS and non-DPS approaches, the lead is narrower. 123.4 Scenario 3: Efficient Learning (WebNLG)In this scenario, we assess the model's ability to adapt to new domains with very few data points. To demonstrate the robustness of BADS across various research topics and backbone models, we focus on the Natural Language Generation (NLG) task in this section. NLG [18, 12] aims to accurately generate textual descriptions from input tuples. An example is shown in Figure 7 (Appendix D). The English benchmark introduced in WebNLG 2020 [6] All backbone models are the encoder-decoder T5-small [26] . Training is on a single GPU using Adam with a fixed learning rate of 3e-5 and a mini-batch size of 20 over 40,000 steps. The loss function is the standard negative log likelihood loss. In BADS, the weight network structure remains the same as in scenario 1. Its input is the embedding of the input sequence, represented by the contextual embedding of the "[EOS]" token, and the learning rate is set to 1e-4. The sparsity level β is 0.05. The BLEU scores are displayed in the top-right plot of Figure 2 . BADS leads the performance, achieving a 2 BLEU score advantage over the second-best system, Dupli-cate_Meta, and surpassing the remaining systems by more than 5 BLEU scores. The other three DPS approaches do not distinguish themselves from the non-DPS methods.In this scenario, we opt for a more controlled examination of data selection effectiveness due to the big amount of domains in the train set 13 . We select three occupationcentric domains-Athlete, Politician, and Astronaut-as CDS incorrectly weights the opposite domain higher. This illustrates that the effectiveness of DPS is linked to the overall performance of the models.ourWe evaluate the average GPU memory and time consumption using the WebNLG task. Instruction Fine-tuning (IFT) for LLMs is a practical application where all three mentioned scenarios are encountered simultaneously. IFT data can be acquired through prompting LLMs [52, 48, 56] , gathering existing Natural Language Processing (NLP) benchmarks [35, 54, 50, 43, 53, 38, 56] , or employing human annotation [52, 63, 48] . Noise is likely to accumulate during each of these data collection methods. Furthermore, since NLP benchmarks often vary greatly in size, the IFT data typically lacks balance. Additionally, the IFT data does not include data points from the downstream tasks, leading to domain shift in testing.We use the same IFT data as [57, 51] as our train set D t , which is a mix of FLAN V2 [35] , COT [54] , DOLLY [10] , and OPEN ASSISTANT 1 [31] . Following [57, 5] , we focus on four downstream tasks: MMLU [24] , which consists of multiple-choice questions across 57 sub-tasks, ARC-challenge/-easy [9] , and HellaSwag [61] . Following [57] , 5 examples were selected from each sub-task to create the meta set D m for MMLU, totaling 285 examples. Additionally, following [5] , for the other tasks, we randomly chose 25 examples from their validation set to create the respective meta sets. To facilitate checkpoint selection, we additionally create a validation set of equivalent size to the meta sets.Due to limited computational resources, we use OpenLLaMA 3B 14 as the backbone model. Training uses one A40 GPU utilizing Adam with a fixed learning rate of 3e-5 and a mini-batch size of 3.In BADS, while the weight network remains the same as described in scenario 3, we modify the input to be the average contextual embedding of all tokens in the sequence. The sparsity level β is 0.05. Results are presented in Table 2 . We excluded BLO from this experiment due to their prohibitive GPU memory usage. Given that the training data for the Mixing baseline predominantly consists of IFT data points, it is considered as a standard IFT process. Table 2 indicates that while individual methods may excel in specific tasks, none of the non-DPS baselines nor the CDS and ClassAct consistently surpass the others across all downstream tasks. However, BADS stands out by consistently outperforming all other baselines across every task, except for AskLLM-O, which, as indicated in Table 1 , demands significantly more computational resources. According to Table 4 (Appendix F) BADS prefers the human-written-from-scratch, OPEN ASSISTANT 1 and DOLLY, over the rest two created using existing benchmarks.To facilitate the use of our proposed datapoint selection method, we provide comprehensive guidelines for hyperparameter tuning in Appendix E. We also include an in-depth discussion and ablation study on the influence of hyperparameters on the method's effectiveness.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We plan to release code, with sufficient instructions to faithfully reproduce the main experimental results. Guidelines:• The answer NA means that paper does not include experiments requiring code. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
related work: 
 Recent works on DPS broadly fall into four categories: i) Approaches based on meta learning (or BLO), ii) Gradient-based methods, iii) Methods based on domain adaptation and transfer learning methods, and iv) General sample reweighing strategies.• Meta learning (BLO) approaches. The DPS can be formulated as a meta learning BLO problem where the outer loss is defined with the training data selection variables, and the inner optimisation is to minimize the model's loss with the selected data [19, 15, 41, 45, 62] . These methods vary in the choice of their outer optimization variables: either directly using the data point weights [41, 45, 19] or minibatch samplers [15] . To solve the BLO, most approaches rely on computing meta-gradients via reverse mode differentiation [45, 62] while some works utilised reinforcement learning techniques [15] . All these BLO-based methods are computationally demanding with large memory footprint, hindering them from being applied to large-scale models/data.• Gradient-based methods. The key idea is to measure the importance of training data points based on the alignment scores between the loss gradients on training and meta data points. The rationale behind the gradient alignment can be theoretically underpinned by the BLO perspective. As shown in [41] , the one-step inner loss update with zero initial weights in BLO reduces to the gradient alignment (cosine angle) between the train and meta data points [57] . However, the method in [57] requires evaluating and storing gradients of the entire training data points, thus computationally expensive. Furthermore, the final solution may not be optimal since the gradient computation is done with an initial network that is warm-up trained with a random subset of the training data. In [16] the weights are optimised by the expected gradient norms during the network training, which serves as a proxy for the importance of data points. In [28] they find the subsets that closely match the gradient of the training and meta sets using an orthogonal matching pursuit algorithm. Under online continual learning setup [2] , they formulate sample selection as a constraint reduction problem based on the constrained optimization view of continual learning.• Domain adaptation and transfer learning methods. Selecting a subset of the train set that are best aligned with the in-domain meta set can be naturally seen as domain adaptation and transfer learning. In the NLP community, [1] finds that the large language models implicitly learn sentence representations that cluster by domains without supervision, while in [20] , they show the effectiveness of domain-adaptive pre-training with data augmented using simple data selection strategies. In [39] they proposed domain adaptive transfer learning which computes the importance weights on the target dataset from the ideas in domain adaptation and estimation of the label distribution. In multi-task transfer learning community, some previous works identify the detrimental effects of the gradients from different tasks, and propose strategies to mitigate the conflicting gradients issue [60, 11, 34] .• General sample reweighing strategies. Since DPS essentially involves finding the optimal reweighed distribution with the underlying domain itself unchanged, several approaches aim to tackle the problem via importance sampling techniques. In [27] they derive a tractable upper bound to the per-sample gradient norm, and derive an estimator of the variance reduction achieved with importance sampling. The curriculum learning [3] is also closely related as one can design an optimal schedule of the successive training distributions. For instance, in [44] they introduce the so-called data parameters, associating samples and classes in the train data with learnable parameters, which governs their importance in the learning process within the curriculum learning framework. 
conclusion: 
 We revisited the DPS problem from the perspective of Bayesian learning. By treating the neural network of interest, as well as an auxiliary weighting neural network as random variables and inferring their joint posterior using SGLD, we achieve a simple and effective approach to data reweighting that is more reliable and scalable than BLO alternatives. Our framework is straightforwardly applicable to learning with data imbalance, label noise, and automating auxiliary data curation. We demonstrate our framework can apply to automating curation of the wide variety of auxiliary instruction fine-tuning data available for billion-scale language models. Overall this demonstrates a promising new kind of approach to the growing need for data optimization in neural network learning. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: The authors cite the original papers that produced the code package or dataset under proper licenses. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Our proposed BADS algorithm has the following limitations, which we plan to address as our future work.1. There are several hyperparameters involved, which need to be carefully tuned for best performance. These include: the sparsity level parameter β, the relative impact constants ρ's of the objective terms in our SGLD update (detailed in App. E), and the standard hyperparameters (e.g., batch size, learning rates, weight decay). Some of these hyperparameters may be optimised via Bayesian model selection in a principled manner. For instance, the sparsity level β can be regarded as a latent random variable (a part of the model) with a proper prior distribution imposed on it, and we can do posterior inference of β (or marginalise it out) together with w and θ in our SGLD update equations. We will pursue this in our future study. 2. Although our method is computationally far more efficient than BLO and some other DPS approaches, its GPU memory footprint is demanding compared to non-DPS algorithms. This mainly originates from the entire weight variables or the whole weight network parameters loaded/maintained in the memory for frequent updates. One possible workaround is to load only the weight variables that are assoicated with the current minibatches. We will be investigating this code optimisation further in our ongoing future study.A Posterior Derivation for Eq. 4From the graphical model in Fig. 1 , we exploit the conditional independenceEQUATIONEQUATIONIn ( 13) we use D m ⊥ (w, D t ) | θ, and in ( 14)1 p(Dm|Dt) is regarded as constant for it has nothing to do with θ and w.Here we provide posterior derivation for the weight network case w i = w(z t i ; ϕ). First, the weights become a deterministic function of D t and ϕ (here ϕ = weight network parameters) as shown in Fig. 6(a) . But since w is a deterministic function of ϕ and D t , we can simplify it by having w absorbed into ϕ while introducing conditional dependence (an arrow) from D t to ϕ, as depicted in Fig. 6(b ). Note that D t is always given, we treat ϕ as a random variable, and wherever w i appears, we replace it by w(z t i ; ϕ). More specifically, we make the following changes to the equations in the weight net scenario:Eq. 2: p(θ|ϕ, D t ) ∝ p(θ) • Nt i=1 p(w(z t i ; ϕ), z i t |θ)Eq. 8: p(ϕ|D t ) ∝ e -( i w(z t i ;ϕ)-⌊Ntβ⌋) 2 /2σ 2Eq. 4 (with detailed derivations):EQUATIONEq. ( 5):[θ, ϕ] ← [θ, ϕ] + η 2 ∇ θ,ϕ log p(θ, ϕ|D t , D m ) + ϵ √ η, ϵ ∼ N (0, I)Eq. ( 6):θ ← θ + η 2 ∇ θ log p(θ) -N t •E i∼Bt w(z t i ; ϕ)•l(z t i ; θ) -N m •E j∼Bm l(z m j ; θ) + ϵ θ √ ηEq. ( 7): EQUATIONAssumption C.1 (Adjusted from Assumption 4.3 in [64] ). There exists m > 0 and b ≥ 0 such thatEQUATIONholds for any θ, w.Assumption C.2 (Adjusted from Assumption 4.4 in [64] ). Any minibatch gradient of the logposterior is Lipschitz continuous. That is, there exists a constant L such that for anyEQUATION)holds for any θ, w, θ ′ , w ′ .Theorem C.3 (Adjusted from Theorem 4.5 in [64] ). Let d = dim(θ) + N t , B be the batch size, and ρ be the Cheeger constant. For any ϵ ∈ (0, 1), with the initial iterate satisfying p θ init w init ≤ R/2 ≤ ϵ/16 for R = R(ϵK -1 /12), and step size η = Õ(min{ρ 2 d -2 , B 2 ρ 2 d -4 }), the distribution µ SGLD K of the K-th iterate in our SGLD iterations Eq. (6-7) satisfies:EQUATIONfor some constant λ > 0,C 0 = Õ(ρ 2 ), C 1 = Õ(Rdρ -1 ), C 2 = Õ(dρ -1). Here ∥ • ∥ T V stands for the total variation distance, and R is defined as: EQUATIONTo enhance flexibility in managing the training process, we incorporate ρ t θ , ρ m θ , and ρ t w into Eq 7 as follows:EQUATIONThe training hyperparameters for BADS and CDS are shown in the table below. Other hyperparameters have been shown in the main paper. Table 3 : Hyperparameters for all experiments.There are two primary sets of hyperparameters in BADS:• hyperparameters for model's parameters update: η, ϵ θ , ϵ w , ρ t θ , ρ m θ , and ρ t w (Eq. 24-25), • hyperparameters for the prior distribution of the example weights, which are σ, β, s avg (Eq. 8-9).In general,• we set η to 1 and kept the Gaussian noise small with ϵ θ and ϵ w equal to 1e-5; • ρ t w is set to 1; • we set ρ m θ to 1 and primarily adjust ρ t θ ; • in most cases, ρ t θ is simply set to 1. However, if the training set contains noise-where the ground truth labels might be incorrect-the loss from the training examples, particularly in the early stages of training, may be unreliable. In such cases, we decrease ρ t θ ; • s avg should not be too large, as it may incorporate outdated weights from earlier training steps.We set it to 10; • we select β based on the proportion of training data we consider beneficial for the downstream tasks. For the LLMs Instruction Fine-tuning experiments, we adopt the same ratio used in the previous studies [52, 46] . • σ controls how tightly the weights should match β. We set σ based on our confidence in the selection ratio β: a smaller σ indicates greater confidence in β.• data denoising scenario is a special, Eq 7 shows that high losses from the training batch push the example weights w toward 0. With a high noise rate (50% and 80%), the training batch losses remain high throughout the training process. To prevent the weights w from collapsing to 0, we use a high selection ratio β and a low σ.E.3 Ablation study: Influence of the impact constants σ and sparsity level β.Impact constant σ Higher σ causes the example weights w to drift away from β, occasionally collapsing to 0, and may result in incorrect example weights (see Figure 8 ). The right column in Figure 11 shows that in all three proof-of-concept scenarios, the models achieve similarly good performance when σ is reduced to around 10 -5 .In both WebNLG and MNIST scenarios, high β leads to incorrect example weights (see Figure 9 ). Conversely, due to the reason we mentioned above, in data denoising (CIFAR) scenario, lower β leads to incorrect weights. The left column in Figure 11 shows that the backbone models' performance significantly declines in both the WebNLG and MNIST scenarios when the example weights are incorrect. In the CIFAR scenario, the impact is less pronounced. As β decreases from 0.8 to 0.05, the weighting network behaves as follows: at β = 0.8, it assigns higher weights to the City domain; at β = 0.2, it assigns similar weights to both domains; and at β = 0.05, it assigns significantly higher weights to the Artists domain.The experimental results discussed in Section 3.3.1 demonstrate that BLO converges more slowly compared to other approaches. In this section, we aim to examine the learning curves of DPS approaches, as illustrated in Figure 11 . The trends for each approach appear consistent across both 10: Model's performance in the three proof-of-concept scenarios with different β and σ.symmetric and asymmetric noise experiments. The convergence rate of BADS aligns with that of non-DPS approaches, whereas CDS converges faster than BADS, and BLO converges significantly slower.In the asymmetric noise experiments, overfitting is less pronounced compared to the symmetric noise experiments, where after 50,000 training steps, the test accuracy for all approaches decreases. BLO and BADS exhibit a notably slower rate of overfitting compared to other methods. Conversely, CDS overfits more quickly than non-DPS approaches, resulting in significantly lower test accuracy even compared to the non-DPS methods.In [41] , each training example is associated with a learnable scalar weight. The scalability issues of this method compared to the weight network used in BADS are detailed in Section 2.4. In this section, we examine the CIFAR-10 denoising task to assess performance differences that result from replacing the BADS weight network with the individual weights strategy described in [41] . From the top row of Figure 11 , it is evident that the performance of BADS with scalar weight (referred to as Bayesian_per_sample in the legend) slightly surpasses that of the original BADS (labeled as Bayesian in the legend).IFT sets Avg Example Weights OPEN ASSISTANT 1 [31] 0.0063 DOLLY [10] 0.0025 Flan-V2 [35] 0.0023 CoT [54] 0.0005 12: Scenario 3, domain adaptation using WebNLG benchmark. This plot shows the BLEU scores on WebNLG benchmark. All DPS models are trained on 2 domains, Artist and City, and tested on other 3 domains -Athlete, Politician, and Astronaut. The x-axis represents the training steps, while the y-axis shows the evaluation metric, BLEU. The code for this paper is available at https://github.com/XinnuoXu/BADS.Since[47] showed that unsupervised DPS can amplify class imbalances, and open-source LLMs generally do not accommodate vision input, we compare to AskLLM-O only in the LLM fine-tuning use case in Section 4.In our experiments, we use Mixing and Meta_only as the noisy and denoised model.To fairly compare the selection mechanism, we replaced their meta set using our meta set.https://github.com/danieltan07/learning-to-reweight-examples.gitWe acknowledge the existence of recent work on DPS with[57,7] or without[59,40] meta-data alignment. However, all these studies use two-stage pipelines, where data points are selected offline and then used in the final model training. Since our approach employs an online selection method, i.e. dynamically selecting data while the model is under training, we have only chosen baselines that follow the same style.Since ClassAct uses a similar approach to CDS but performs worse, we exclude it from this ablation study.Besides the primary results, we also present results for asymmetric noise in the top right plot of Figure(see Appendix F.1). Additionally, we explore performance variations by substituting the weight network with the individual weights strategy outlined in[41] (refer to Appendix F.2).11 The convergence of BLO is slower than that of other approaches. To further explore this, we conduct an ablation study to compare the learning curves for all methods in Appendix F.1.ClassAct uses a similar approach to CDS and has similar results. We exclude it from this ablation study.Given that the train set spans 16 domains, it is difficult to assess the adequacy of the data selection behavior. We still show the weighting plots for the 16 domain in the bottom row of Figure13in Appendix F.https://huggingface.co/openlm-research/open_llama_3b_v2 


Decision: All reviewers unanimously voted for acceptance. Weaknesses like missing baselines in the experiments and missing theoretical results, were outbalanced during the rebuttal period. Authors are strongly encouraged to import the additional results into the paper.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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.
limitations: limitations: The authors have adequately addressed the limitations proposed in their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Please see the weaknesses
limitations: limitations: No potential negative societal impact.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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)?
limitations: limitations: - Limitations are discussed, but after the page limit

Comments: title: RE: Official Comment by Reviewer 5MKv
comment: Regarding the embedding functions: In practice, using a fixed embedding function requires an additional embedding model, which increases memory usage. Additionally, the embedding model may require tuning, which would add extra time. We could consider including experiments to empirically understand the gap between using a changing function approximation and a fixed embedding function.

We sincerely appreciate the reviewer’s insightful and valuable feedback, as well as their time. We will incorporate the information provided in the rebuttal into the revised paper.
We sincerely appreciate the reviewer’s insightful and valuable feedback, as well as their time. We will incorporate the information provided in the rebuttal into the revised paper.
Dear reviewer oLJs,

We appreciate your thoughtful and detailed feedback once again.

Since today is the final day for the rebuttal discussion, we wanted to check if you have any remaining questions or suggestions that we can address before the deadline.

Thank you for your time.
Thanks, I understand now the discrepancy between Figure 2 in the paper and the ablation in the rebuttal pdf. I definitely think some discussion around the hyper parameter sensitivity based on the ablations in the rebuttal pdf should be part of the main paper, as this is important to guide future work around your method. 

Regarding the circular dependence when using embeddings I am not sure if I follow the authors arguments and would encourage them to think about the effects of using a changing embedding function as approximation of some fixed embedding function $f^*$. 

But overall many of my doubts or concerns have been answered and I definitely think the Bayesian perspective on DPS offers a novelty that makes this work stand out, so I have increased my original score from 6 to 7.
I would like to thank the authors for thoroughly addressing my comments and quickly incorporating the revelant comparisons to related work and baseline comparisons where relevant. This new information has improved my understanding and confidence in the impact of the paper.
We appreciate the reviewer’s engaging discussion and sharp questions.

**[RE: Circular Dependence]**

Strictly speaking, if we use the embedding $f_\theta(z^t)$ in the weightnet, then weight $w$ becomes a function of both $\theta$ and $D_t$, thus we can have circular dependency as the reviewer pointed out.

However, we would rather see this case in the following way: Suppose we were able to access the ideal oracle embedding function $f^*$. Then ideally the weightnet would use this $f^*$ internally to output the weight $w$, ie, $w = weightnet(f^*(z^t); \phi)$. This way, as $f^*$ is constant, our graphical model in Fig.1(b) of the attached PDF is still valid, and there is no circular dependency. In practice, however, since we cannot access $f^*$, one reasonable way to approximately mimic it is to use the current estimate $f(z^t; \theta)$ as an *online plugin proxy for* $f^*$. We hope this answers your question. FYI, in the implementation, we do stop-gradient $\theta$ in the weightnet update, which is in line with our argument of *online plugin proxy for* $f^*$. 

**[RE: Test/BLEU]**

For WebNLG, we did experiments on two different setups:
* **The main setup** (the performance is shown in *top-right plot in Figure 2 of the paper*): We use the original WebNLG domain adaptation set up (16 domains for training and 3 domains for testing) to report the models generation performance. Details are explained in *Appendix A*. 
* **Ablation setup (called 2-domain)** (the performance is shown in *Figure 8 in the Appendix*): The purpose of this setup is to easily assess the effectiveness of data selection, as visualizing the data selection across 16 domains can be challenging. Details and the rationale for this setup are explained in lines 291-301 of the paper. 

Generally, the Ablation setup is somewhat simpler than the main setup, resulting in a higher test BLEU score.

In the Hyperparameter Analysis section of the rebuttal, we did experiments using the Ablation setup to visualize how hyperparameters affect data selection behavior (shown in Figure 3 in the rebuttal PDF). So, reviewers comparing Figure 1 in the rebuttal PDF to Figure 8 in the paper will observe that the test BLEU score for $\beta=0.05$ consistently hovers around 44.

If the reviewer has any further questions, please feel free to reach out.
I thank the authors for answering some of my questions (non-negativity of weights, different weight NN inputs for MNIST and CIFAR, why using the weight NN). 

I also acknowledge that they added some ablation on WebNLG about the prior parameters $\beta$ and $\sigma$ as well as a derivation on the inference when a weight network is being used. Regarding this however I have 2 questions:
- How does inference change if instead of $z^t$ model embeddings $f_{\theta}(z^t)$ are used like in the experiments of the paper? Does that not cause a circular dependence between $\theta$ and $w$?
- Why is the Test/BLEU for $\beta=0.05$ in the ablation at 44 after 40k vs. only 36 in Figure 2 of the paper?