Malladi Lab — Writing

The Hidden Infinity in Preference Learning

2024-07-09T00:00:00-07:00

TL;DR: I demonstrate from first principles how offline preference learning algorithms (e.g., SimPO) can benefit from length normalization, especially when training on model-annotated preference data. The derivation also lends insight into a subtle but important challenge in training LMs: infinite strings.

Aligning language models using human feedback has become increasingly popular. The process of collecting this feedback (i.e., a human rater reading and ranking pieces of text) is non-differentiable, so people turned to reinforcement learning (RL).

Learning from Human Feedback as RL

It turns out that it is confusing to translate the problem of improving language models via human feedback to the RL setting. First question: how should we define the states and the actions?

Originally, Jaques et al., 2017 designed the states in the problem to be the partial context, usually some combination of pre-filled prompt tokens and generated tokens, so the action was to pick a single token to generate next. I’ll refer to this as the token-level formulation. On the other hand, later works (Ziegler et al., 2019; Stiennon et al., 2020) framed the problem in a bandit setting, where the state is the pre-filled prompt and the action is to select which sequence to generate out of a few candidates. I’ll refer to this as the bandit formulation. Both of these designs are not perfectly suited to our setting, because we get human feedback on a sequence level but ask our LM to act on (i.e., generate) individual tokens. More recently, some works have designed objectives that incorporate token-level rewards or constraints alongside sequence-level feedback (Zeng et al., 2024; Lai et al., 2024). On the more theoretical and conceptual side, Rafailov et al., 2024 has a great illustration of this topic and draws meaningful connections between the two settings, and I’m sure that many RL papers are dedicated to parsing this gap between trajectory-level rewards and step-level actions.

The reason I bring up these two settings is to see how each one deals with the case of a model producing infinite-length strings. And yes, I know that it’s impossible in practice for LMs to actually generate an infinite-length string. But they can certainly put a tiny amount of probability mass on an infinite-length string. And if they do so for many such strings, then this ends up being a huge amount of probability mass lost to sequences that we’ll never see.

In the token-level formulation, it’s clear that an infinite-length string corresponds to essentially an infinite-length trajectory rolled out from the policy. One may think the bandit formulation somehow circumvents this problem because it provides only two bounded-length sequences for the model to choose from. But, when deriving optimal solutions to the bandit setting, you still need to deal with the problem of infinite “arms” (i.e., possible responses sampled from your policy). For example, Appendix A.1 of the DPO paper glosses over the fact that the partition function $Z_\theta$ may not exist if you don’t bound the length of the generations or somehow assume that the reward decays with length. If you’re still not convinced that infinite-length strings matter, then stay tuned to see how this idea eventually plays out into a rigorous understanding of the role of length normalization.

Length Normalization

We can now derive how length normalization might make sense from the RL perspective. My disclaimer here is that I definitely worked backwards from the idea of length normalization to arrive at this justification, so there may be many other algorithms that are equally or better suited to operating in the setting I described. This derivation is also hand-wavy and meant more to provide intuition more than anything else. I’ll use the bandit-level formulation.

Suppose we have a prompt $x$ and two candidate completions $y_1$ and $y_2$. We will use $\Pr[y_1\succ y_2\mid x]$ to denote the proportion of people who prefer $y_1$ over $y_2$ as a completion to $x$.

The first thing we need to do is make some assumption about the reward structure. The Bradley-Terry (BT) model was derived to rigorously describe the distribution of human preferences:

$\Pr[y_1\succ y_2\mid x] = \frac{\exp(r_\text{h}(x, y_1))}{\exp(r_\text{h}(x, y_1)) + \exp(r_\text{h}(x, y_2))}$

where $r_\text{h}(x,y)$ is the human ground-truth reward function that we don’t have access to. We already arrive at our first complication. In practice, some datasets are constructed by using an auxiliary model (e.g., GPT-4 or PairRM by Jiang et al., 2023) to pick which response is preferred. So, let’s call the reward function of such a language model $r_\text{LM}(x,y)$ and now consider the case where our data is annotated by a model (i.e., “RLAIF”).

One thing we do kind of know about $r_\text{LM}(x,y)$ is that it usually grows with the length of the response $y$. That is, models tend to favor longer responses. And in fact, one popular paper (Singhal et al., 2023) showed that for some models, the reward grows linearly with the length of $y$. But in the case of alignment, it’s not clear if we always want our models to output longer sequences. For example, if we are trying to prevent harmful behavior, length is likely a spurious correlation. On the other hand, if we want to promote helpful behavior, increasing the length may somewhat correlate with providing useful information. Regardless, let’s disentangle length from the reward that we want to maximize and define $r^\star(x,y)$ such that

$r_\text{LM}(x,y) = |y| r^\star(x,y)$

Now, we want to learn a model $\pi^\star$ that is trained with access to only data annotated according to $r_\text{GPT}$ but maximizes the length-normalized reward $r^\star$. In other words, we can write:

$r^\star(x,y) = \frac{\log(\pi^\star(x,y))}{|y|}$

where I’m keeping with the standard ML practice of considering the log-likelihood of a softmax-parametrized model instead of the likelihood directly. We can now define a score $\Lambda$ using $r^\star$.

$\Lambda[y_1\succ y_2 \mid x] = \sigma(r^\star(x,y_1) - r^\star(x,y_2))$

where $\sigma$ is the sigmoid function. Note that the DPO derivation does not rely on pulling this score out of thin air because they used the Bradley-Terry assumption on $r_\text{h}$ to show that, if the data is human-annotated, then $\sigma(r_\text{h}(x, y_1) - r_\text{h}(x, y_2))$ is exactly equal to the ground-truth probability that $y_1$ is preferred over $y_2$. Here, however, I want to make it clear that we have no reason to believe that $r^\star$ obeys the Bradley-Terry assumption. But, optimistically, if it did, then $\Lambda$ would define a valid probability distribution capturing how much the model prefers $y_1$ over $y_2$ if we removed the bias that favored long responses.

Now, we want to train the model to maximize this score in the case that $y_1 = y_w$ is the winning response and $y_2=y_l$ is the losing one. So we can write our length-normalized objective as

$\mathcal{L}(\pi_\theta, \mathcal{D}) = \mathop{\mathbb{E}}_{(x, y_w, y_l)\sim\mathcal{D}} \left[\log \sigma \left(\frac{\log \pi_\theta (y_w\mid x)}{|y_w|} - \frac{\log \pi_\theta(y_l\mid x)}{|y_l|}\right)\right]$

And so we arrive at a length-normalized preference learning objective. There are some interesting takeaways from this derivation.

If we’re going to annotate preference data using other models, then we can actually adjust our objectives to counteract known biases in the scoring. Length normalization is one example of this strategy. If and when other ones come to light, we could feasibly use this approach to allow our models to learn from imperfectly annotated data. We can analogously use this approach to correct biases in human-annotated data as well, though we ought to exercise a lot more caution and thought when shifting or otherwise modifying human-expressed preferences.
Note that no matter what I do, this current strategy cannot link $r_{\text{LM}}$ to $r_\text{h}$ – that is, I can’t really describe how the rewards that the model provides relate to the ground-truth human rewards. Instead, all I can say is that length-normalized objectives provide one possible strategy for mitigating the absorption of biases from the models used to annotate the data. By the same token, it doesn’t matter if the annotator LM was trained on human data or not. These nuances are something I hope to explore in future work!
Connection to average reward maximization: Another way to see $r^\star(x,y)$ is that it’s the average reward of each token in the sequence. This links length normalization to a long line of RL algorithms that maximize the average reward instead of the total reward (see a survey in Mahadevan, 1996). This formulation is especially useful when dealing with cyclical or infinite-horizon problems. Note that another way to deal with infinite-horizon problems is to discount rewards that will arrive in the future – if you were to do this, then you would arrive at the length-regularized DPO objective (Park et al., 2024).
If we just look at the final length-normalized objective, then we can see that it explicitly requires $\pi_\theta(y_w\mid x)$ to increase much more than it would if there was no normalization. Other modifications to the objective, like target reward margins, do not explicitly ensure that this will happen, because they operate on the distance between the two likelihoods. After all is said and done, we usually just use $\pi_\theta$ to generate sequences, so directly modifying it seems like it would be more effective.

Discussion

The derivation above demonstrates how we might adjust our objective functions to enable learning from imperfectly (e.g., automatically) annotated data. I think that using other models to grade or annotate preference data will become increasingly prevalent, due to growing evidence that preference learning with on-policy data is more effective (Tajwar et al., 2024) and the rising cost of finding and hiring qualified human annotators. In this post, I considered the case where the preference data was model-annotated and the model had a bias favoring longer responses. The analysis provides some insight into one particular situation in which length normalization makes sense, but it by no means provides a general prescriptive guideline for objective design.

I was originally motivated to derive this after reading about SimPO (Meng et al., 2024), which proposes a length-normalized preference learning objective. If you are curious about the other modifications in SimPO – removing the reference model and setting a global target reward margin – then you might be interested in reading our recent paper on how the reference model can prevent DPO from aligning effectively.

The possibility of infinite-length strings has been explored in several prior works, mainly with a focus on characterizing how prone a particular model architecture (e.g., transformer or RNN) is to leaking probability mass to infinite-length strings and how certain objectives might exacerbate or mitigate this issue (Welleck et al., 2020; Du et al., 2023).

There is also a big leap missing from what I have derived so far. We may train a policy $\pi^\star$ to maximize rewards but when we sample from the model (e.g., via greedy decoding), we don’t actually obtain a sequence that maximizes the log likelihood prescribed by the model. This was, after all, the motivation for developing many different decoding strategies, including ones that implicitly plan or search. In the context of post-training, I find this observation fascinating, given how the field has primarily focused on next-token prediction objectives as a means to constrain or guide long-form generations.¹ Some of my research going forward will be on this topic of the gap between next-token prediction and long-form generation. If you are interested in this direction, I enjoyed reading several of Clara Meister’s papers.

Since this is a blog post and not a paper, I’m not doing a full literature survey on these topics. As I mentioned, there are many RL papers on these topics, and my main research area is not RL. That being said, if I didn’t mention your paper and you think it is relevant, kindly send me an email!

Acknowledgments: Thank you (in alphabetical order) to Angelica Chen, Xinyi Chen, Yu Meng, and Mengzhou Xia for discussions that helped me refine my derivation and for suggesting several relevant papers. And also thank you to Adithya Bhaskar, Tianyu Gao, Lucy He, and Kaifeng Lyu for helping me proofread this post. Thanks also to Ben Eysenbach for pointing me in the direction of average reward maximization.

Citation: If you find this blog post to be helpful in your work, please use the following bibtex citation.

            @misc{malladi2024hiddeninfinity,
  title   = {The Hidden Infinity in Preference Learning},
  author  = {Malladi, Sadhika},
  year    = {2024},
  month   = {July},
  url     = {https://www.cs.princeton.edu/~smalladi/blog/2024/06/27/dpo-infinity/},
}

          

OK, to be fair, there are papers that now train models to predict several tokens into the future (Gloeckle et al., 2024), but I would argue this is just a patch over the underlying problem. Similarly, there are papers that allow a model to access and train on its own generations (e.g., STaR, Zelikman et al., 2022), but I’m not fully convinced that this closes the gap between training and inference. ↩

Using LESS Data to Tune Models

2024-04-04T00:00:00-07:00

TL;DR: We describe how data selection for modern-day LLMs differs from prior settings and how our algorithm, LESS, effectively selects relevant data to cultivate specific capabilities in models during instruction tuning.

Paper: https://arxiv.org/abs/2402.04333

Code: https://github.com/princeton-nlp/LESS/

This post will take the following structure:

We introduce the motivation for data selection and describe how the criteria for “good” data depends heavily on the setting. One can either try to identify representative datapoints for the in-domain setting or relevant ones for the transfer setting.
Our algorithm, LESS, effectively selects relevant data to induce capabilities in the instruction tuning setting. LESS identifies 5% of the dataset that induces stronger performance than training on the full dataset.
We conduct an in-depth analysis of prior works on data selection for various settings and provide insights into their technical details, strengths, and limitations.
We conclude by identifying trends in data selection in the era of LLMs.

Motivation

The training dataset is a crucial design choice when building a machine learning model. Dataset choice can drive the capabilities of the resulting model in various ways (see, for example, CodeLLaMA, and DeepSeek-Math). Also, training models can be expensive, and the cost usually scales with the size of the dataset, so dataset selection offers one way to improve efficiency and reduce cost.

Coreset selection selects data such that the selected subset represents the full dataset.
Transfer data selection selects the subset that is closest to the target data points.

We distinguish two settings for data selection: in-domain data selection and transfer data selection. In the former, the selected data is drawn from the same distribution as the evaluation data¹ , whereas in the latter, evaluation is performed on different data.

In-domain data selection aims to identify the most representative subset of data, often referred to as a coreset (Sener & Severese 2018), from a large in-domain training dataset. The selection criteria include representativeness, coverage, correctness, and more. We talk more about this setting latter.
Transfer data selection seeks to choose the most relevant data from a broad training pool that closely align with the target examples. The selection criterion focuses on relevancy.

Our work focuses on transfer data selection for instruction tuning. Instruction tuning has proven to be a highly effective way to quickly adapt language models to follow human instructions. Depending on the data used, models can be tuned to be general-purpose instruction followers (e.g., Alpaca, Vicuna, Zephyr) or solve more structured tasks per human instructions (i.e., targeted instruction tuning). Our work focuses on selecting data for the latter case, where models are tuned to perform particular types of reasoning (e.g., using a passage to answer a question). In this case, the data selection problem can be understood as bootstrapping a few examples to identify relevant data to solve a task.

Selecting LESS Datapoints for Targeted Instruction Tuning

The targeted instruction tuning setting poses the following research question:

Given just a few handwritten examples of a query type and a particular pre-trained model, how can we identify the most relevant data to train on out of a large pool of available instruction tuning data?

Several data selection strategies have been developed for pre-training, such as continued training on domain-specific data (Gururangan et al., 2020) and using n-gram statistics to identify relevant data (Xie et al., 2023). However, the instruction tuning setting is unique in that using all of the available data can hurt the development of specific capabilities (Wang et al., 2023), and one wants to somehow account for the properties of the pre-trained model when selecting data. So, we avoid using heuristic definitions of useful data and instead frame data selection as a rigorous optimization problem. As we describe in the next section, LESS selects training data to minimize the loss on the target data (i.e., the few handwritten validation examples). Check out our paper here and play with the code on GitHub!

Conceptual Approach

Suppose we have a handwritten validation example $z$ and a huge dataset of candidate training points $\mathcal{D}$. At the heart of any transfer data selection algorithm is the same question: how does training on some point $x\in\mathcal{D}$ affect the model’s performance on $z$? We explicitly formulate this by approximating how the validation loss $\ell(z;\theta)$ changes when we take one training step (i.e., update the model from $\theta_t$ to $\theta_{t+1}$) on a candidate datapoint $x$:

$\ell(z;\theta_{t+1}) \approx \ell(z;\theta_t) + \langle \nabla \ell (z;\theta_t), \theta_{t+1} - \theta_t \rangle$

Assume that we were training with SGD with step size $\eta_t$, we can further derive the following formulation:

$\ell(z;\theta_{t+1}) -\ell(z;\theta_t) \approx \eta_t\langle \nabla \ell (z;\theta_t), \nabla \ell (x;\theta_t)\rangle$

We can see that selecting $x$ to maximize $\langle \nabla \ell (z;\theta_t), \nabla \ell (x;\theta_t)\rangle$ will maximally reduce the validation loss on $z$. The method was initially proposed and employed in TracIn (Pruthi et al., 2020) to gain insights into how training examples influence the model’s predictions. The formulation is also very suitable for transfer learning, because there is no need to assume any relationship between $x$ and $z$. But, there are a couple of modifications required to make it work for our setting (see our paper for details):

Adam: LLMs are generally tuned using Adam, which has a more complicated update formula involving the moving averages of the gradient moments. We can plug in that update, which we denote as $\Gamma(x;\theta_t)$ instead of $\nabla\ell(x;\theta_t)$. See the paper for details on the Adam formulation and the resulting technical complications.
Multi-Epoch: We usually train our models over several epochs, which means each candidate $x\in\mathcal{D}$ would be seen several times over the course of training. We would want our estimate of the influence of $x$ to take this into account, but we want to avoid the computational cost of training on the entire candidate dataset for several epochs. Instead, we approximate how the model would adapt to seeing the same data by performing a warmup training period on a randomly selected 5% of the data and aggregating the influence over the whole run. See the paper for how we aggregate influences over multiple epochs.
Variable-Length Instruction Data: Instruction tuning sequences have differing lengths. Experiments in our paper showed that shorter sequences exhibit much larger norms, so the inner product $\langle \nabla \ell (z;\theta_t), \nabla \ell (x;\theta_t)\rangle$ is much larger for shorter $x$. We therefore decide to instead use the cosine similarity instead of the inner product when measuring the influence. This is not a failure of the formulation described above (which indeed is quite simple); instead, it indicates that we ought to perform data selection for individual tokens instead of entire sequences. However, measuring the gradient of every token in the sequence is prohibitively expensive with today’s methods, so we stick to sequence selection and leave a token-level formulation to future work.

Altogether, once we make the necessary modifications, we arrive at the following formula for computing the influence of a training datapoint $x$ on a validation point $z$.

$\textrm{Influence}(x, z) = \sum_{i=1}^N \bar{\eta}_i \cos (\nabla \ell(z;\theta_i), \Gamma(x; \theta_i))$

where $\Gamma(x; \theta_i)$ is the Adam update mentioned above, $N$ is the number of epochs during warmup training (see next section), $\bar\eta_i$ is the average learning rate in the $i$th epoch, and $\theta_i$ is the model after the $i$th epoch. The above formula makes it clear that we need to handle model gradient vectors, which can be very large, and we need to aggregate influences over several model checkpoints. In the next section, we describe how we compute this formula efficiently.

Selecting LESS Data

An overview of the four steps of our algorithm, LESS.

LESS consists of four major steps to make influence estimation feasible and scalable:

Warmup LoRA training: We train the model with a warmup phase with a random subset of data and checkpoint the model $\theta_i$ and optimizer update $\Gamma$ over $N$ epochs. We choose to train with LoRA to operate on gradients in a much smaller space (i.e., ~100M parameters for a 7B model).
Compute gradient features: We acquire low-rank Adam gradients by further projecting the LoRA gradients down to a smaller dimension $d$ (i.e., 8192 in our experiments).²
Select data: Given a few instances from target tasks, we first acquire their compressed gradients. We then calculate the influence $\mathrm{Inf}$ and pick the examples with the highest scores.³
Train models: We train models on the selected data, using either full-parameter fine-tuning or efficient fine-tuning approaches like LoRA.

Note that the first and second steps are computed once per candidate training set and can be stored as a gradient datastore. The datastore can be reused to quickly select data for different validation tasks. Moreover, the model used to select data in steps 1-3 can be different from the model trained in step 4, and we call this setting LESS-T, where “T” stands for transfer.

Results

Training on 5% selected data often outperforms training on the full dataset

We construct our dataset pool to be a combination of subsets FLAN V2, COT, Open Assistant 1, and Dolly datasets. On Mistral-7B and Llama-2-13B, we find that the selected 5% of the data outperforms using the full dataset. Additionally, we find that the data selected with Llama-2-7B is also effective for instruction-tuning Llama-2-13B and Mistral-7B (LESS-T).

Training on just 5% of the data, selected by our algorithm LESS, outperforms training on the full dataset.

In-depth results of using LESS on three benchmarks.
LESS-T indicates that we used LLaMA-2-7B to select the data for training the model.

LESS outperforms the baselines

We compare our approach to several baselines with different relevancy criteria.

Random: random selection
BM25: computes lexical overlap as a relevancy score and selects the examples with the highest scores
DSIR (Xie et al., 2023): weight candidate training data with n-gram features, and resample data based on the weights
RDS: uses model’s hidden representations as features and selects the examples with the highest similarity scores

We surprisingly find that LESS is the only approach that consistently outperforms random selection. Other approaches, unfortunately, either provide minimal improvement over random selection (BM25), or underperform random selection. In the next section, we provide qualitative examples to have an in-depth understanding of why other approaches fail.

LESS outperforms all baselines.

LESS selects examples with similar underlying task structures

We provide top selected examples by BM25, RDS and LESS for a TydiQA example. The TydiQA example presents a paragraph and a question in Bengali, and the goal is to locate an answer to the question within the given paragraph. BM25 and RDS select examples in Bengali, but these examples are related to a different task. In contrast, LESS selects a question-answering example written in English, which is more relevant to the target task.

This pattern holds true for the top examples selected for other questions. Upon further investigation, we discovered that the top examples selected by BM25 and RDS are predominantly in Bengali, whereas LESS consistently chooses examples in English that are specifically related to question answering. This observation suggests that LESS prioritizes examples that share a similar reasoning process, while the other approaches place too much emphasis on superficial cues such as language or topic, rather than the underlying task structure.

Qualitative analysis of selected data.
LESS circumvents surface-form cues to instead select examples with similar reasoning types as the validation data.

More computation on data selection enhances performance

Our ablations in the paper show that spending more computation in any of the steps of LESS can improve the performance of the method at the cost of additional runtime. For example, using a longer warmup phase in step 1, increasing the projected dimension in step 2, and aggregating the influence estimate over more model checkpoints can all improve performance at the cost of runtime and/or memory. We report results in a setting where the data selection cost is reasonable: selecting and training on the data requires less time than training on all available data. However, our results show that training on LESS data can even improve performance, so it may be worthwhile to invest more resources into the selection stage.

Past Works on Data Selection

In this section, we discuss many related works in-domain data selection and transfer data selection.⁴ Our goal is to briefly cover the many broad intuitions that can inform different data selection algorithms.

In-Domain Data Selection

As we mentioned before, the goal of data selection depends heavily on the setting. When ample labeled in-domain data is available (e.g., image classification in vision is an obvious case, pre-training data selection is also one), the goal is to identify the most representative subset of data, so much of the past work has sought to define the notion of “representative” datapoints. Sener & Severese 2018 explicitly defines this problem to be a core-set problem (Bachem et al., 2017, Tremblay et al., 2019), which aims to choose a set of datapoints such that each data point is close to at least one selected data point. Subsequent research has linked specific attributes of datapoints, such as being “easy to forget” (Toneva et al., 2019), exhibiting “large prediction error” (Paul et al., 2021), being “hard to memorize” (Feldman et al., 2020), being “redundant in clustering” (Birodkar et al., 2019) and having “uncertain predictions” (Chitta et al., 2019) with their representativeness and importance. Sorscher et al., 2022 consolidate these attributes under the category of hard examples, proposing that these instances are the most critical for effective training. One can also use gradient-based features (Mirzasoleiman et al., 2020, Wang et al., 2020, Killamsetty et al., 2021) for data selection. This naturally lends itself to a more general meta-learning formulation, which we discuss more below.

Pre-Training

In recent years, the pre-training and fine-tuning paradigm has proven to be an effective way to build large-scale foundation models. Unlike models trained with curated task-specific datasets like CIFAR or ImageNet, these models are trained with massive Internet-scraped data consisting of trillions of tokens or images. These two settings differ in two ways: (1) quality: web-scale datasets are generally heterogeneous in quality; and (2) scale: data selection has strong implications on efficiency in the foundational model era. We consider general pre-training data selection to be a form of in-domain data selection, as it aims to choose data that covers all potential use cases.

Algorithmic Filtering: Early in 2019, researchers have found that excluding high-perplexity examples from CommonCrawl could significantly boost training efficiency, since such examples are often nonsensical or of low quality (Wenzek et al., 2019, Marion et al., 2023). More sophisticated processing identified high-quality data by filtering domains and URLs, ensuring diversity, and removing redundancy (e.g., via MinHashLSH or semantic deduplication (Abbas et al., 2023)). Many popular datasets were constructed in this way, including C4 (Raffel et al., 2021), RefinedWeb (Penedo et al., 2023), SlimPajama (Shen et al., 2023), and more. Similar efforts have been explored for pre-training VIT models (Abbas et al., 2024).

LLM-Aided Filtering: An extreme form of algorithmic filtering is to explicitly prompt LLMs to generate data satisfying certain properties. The approach’s efficacy is clearest in the Phi-series models (Gunasekar et al., 2023, Javaheripi et al., 2024), which achieved strong performance on math and coding benchmarks with only 1.3B model parameters. While the Phi-series models focus on generating textbook-style data, other recent work has shown that rephrasing pre-training data to mimic the stylistic and informational density of Wikipedia articles markedly enhances the data’s cost-effectiveness (Maini et al., 2024). Aside from generating data, LLMs can also be used to judge the quality of data. Recent works such as QuRating (Wettig et al., 2024) and Ask-LLM (Sachdeva et al., 2024) aim to exploit capable language models to directly provide quality scores to data instances, offering a metric for evaluating their potential impact on model training.

Meta-Learning Formulation: Instead of relying on human notions of quality, one can also phrase in-domain and transfer data selection as meta-learning problems (Nguyen et al., 2020). The outer loop selects data for models in the inner loop to train on. Meta-learning approaches have traditionally been very computationally expensive. Recent work, dubbed datamodels (Ilyas et al., 2022), seeks to make this bi-level optimization problem more tractable by directly training a model to predict the test performance that would result from excluding or including a particular datapoint. Subsequent work (Park et al., 2023) made this approach computationally efficient, and recently, Engstrom et al., 2024 used this approach to score pre-training examples based on how they affect performance on a target set of examples. Our work, LESS, can be interpreted as one such meta-learning formulation for selecting data in the instruction tuning setting.

Instruction Tuning

Instruction tuning stands as a pivotal process in unlocking capabilities of pre-trained base models by further training the models to make them follow human instructions. Many works have assembled massive instruction tuning datasets. Some early datasets are human annotated (e.g., Open Assistant and Dolly), though recent trends mostly use completions from GPT models (e.g., Orca, ShardGPT, UltraChat etc.). The queries in these datasets cover a broad spectrum of topics, and could be as diverse as pre-training datasets. They are mostly used to build general-purpose chatbots.

Recently, a lot of works have shown that high-quality data is essential for instruction tuning. The pioneering work LIMA (Zhou et al., 2023) illustrates that a mere 1,000 meticulously selected high-quality human-curated examples could lead to marked performance improvements. Numerous studies have thus endeavored to automate the data selection pipeline, including strategies for choosing examples based on their naturalness (Cao et al., 2023), employing GPT-4 for quality scoring (Chen et al., 2023), enhancing data diversity (Bukharin et al., 2023, Liu et al., 2023), and ensuring broad coverage (Du et al., 2023). Additionally, some research has explored the benefits of prioritizing longer examples (Zhao et al., 2024), but it remains uncertain whether if this simply aligns on a surface level with the tendency of GPT to favor longer outputs (Wang et al., 2023).

Transfer Data Selection

Transfer data selection and in-domain data selection differ in purpose. While in-domain data selection aims to cover the properties of the entire dataset, transfer data selection focuses on enhancing the performance of a specific subdistribution of data. This shift in focus leads to a change in the selection criterion from representativeness to relevance.

In transfer data selection, the goal is to identify the most relevant data points from a large pool of available data. This approach is particularly useful for building domain-specific models or improving the performance of specific tasks or queries. To achieve this, a subset of target data is typically required to serve as an anchor for the data selection process. Previous works in this area include the study by Gururangan et al. (2020), which demonstrates the effectiveness of continued training on topic-specific pre-training data to improve performance on domain-specific downstream tasks. Another notable work is by Xie et al. (2023), which introduces a data reweighting approach based on the n-gram similarity between the source data and the target distribution. While these two works focus on aligning data with surface form cues (i.e., topic and ngram matching), LESS selects data that matches the underlying task or reasoning type.

Conclusion and Future Directions

Data plays a crucial role in determining the capabilities of trained deep models. In the in-domain setting, data selection aims to select a small yet representative dataset. Reducing the dataset size directly reduces the time required for training, which can be extremely useful when pre-training LLMs. On other hand, in the transfer setting, one seeks to solve tasks that do not have much data associated with them, and this requires filtering the dataset to identify a relevant subset to train on. Our method, LESS, selects data in the transfer setting to perform targeted instruction tuning, and it identifies the most relevant 5% of the dataset that can induce better performance than training on the full dataset. More broadly, we see data as being an important area of study for improving LLMs. Spending more compute on data selection (or, more broadly, data generation) in many different settings has proven to be very valuable, though one has to ensure that the cost of selection does not skyrocket. Data selection can also go beyond improving or preserving performance to attribute particular model behaviors to the training data. For example, a recent follow-up to LESS (He et al., 2024) identifies seemingly benign data that somehow breaks the safety of models during fine-tuning. We’re excited to see where data selection goes next!

Acknowledgements: LESS is co-authored with Suchin Gururangan, Sanjeev Arora, and Danqi Chen. We thank (in alphabetical order) Dan Friedman, Tianyu Gao, Lucy He, Austin Wang, Alex Wettig, and Howard Yen for their helpful feedback on this post!

Footnotes:

In some cases, the distribution of the training data may not match that of the evaluation data. However, the training data could still serve as a good approximation and correlates with the performance on the evaluation data. For instance, the pre-training loss measured on a held-out dataset, typically provides a reliable indication of the model’s performance on downstream tasks. Therefore, we still consider pre-training data selection as in-domain data selection. ↩
We use the efficient random projection implementation used in TRAK (Park et al., 2023). See their amazing codebase here! ↩
According to the derivation, the training gradients are normalized using the Adam update rule, and the validation gradients for the target instances are used as-is. Both are compressed via random projection. ↩
A recent survey paper provides a more comprehensive and formal treatment of data selection methods in the context of language models. ↩

How to Scale Hyperparameters as Batch Size Increases

2024-01-22T00:00:00-08:00

TL;DR: Stochastic differential equations (SDEs) provide rigorous, empirically validated scaling rules that prescribe how to adjust hyperparameters when scaling training runs (Adam or SGD, language or vision) to highly distributed settings without sacrificing performance. This post focuses on empirically useful insights, and a subsequent post will describe the theoretical toolbox of SDEs in more detail.

Scaling Rules: Increase Batch Size without Hurting Performance

Even as early as 2011, researchers recognized that scaling training runs across GPUs can yield huge efficiency gains (HOGWILD! by Niu et al.). For example, loading a large minibatch using data parallel with 8x as many GPUs allows you to finish training nearly 8x as fast (modulo communication and latency). However, using a very large batch size naively hurts SGD performance and puts a limit on how much you can scale (Keskar et al., 2017). A few papers subsequently suggested that one needs a scaling rule to adjust the hyperparameters when increasing the batch size. There are different scaling rules for different optimizers.

Linear Scaling Rule (for SGD)

When scaling the batch size by $\kappa$, scale the learning rate also by $\kappa$.

Square Root Scaling Rule (for Adam)

When scaling the batch size by $\kappa$, scale the learning rate by $\sqrt{\kappa}$. Also, change the other hyperparameters, setting $\beta_1 = 1 - \kappa(1-\beta_1)$, $\beta_2 = 1-\kappa(1-\beta_2)$, and $\epsilon = \epsilon / \sqrt{\kappa}$.

Even so, without rigorous theory, it was unknown what the maximal performant batch size was, so scaling runs up was an expensive trial-and-error game. In 2014, Krizhevsky heuristically derived a square root scaling rule for SGD, stating that the learning rate should be scaled by $\sqrt{\kappa}$ when scaling the batch size by $\kappa$ (see the bottom of page 5). This ends up being incorrect! Even Krizhevsky notes that the linear scaling rule yields empirically stronger performance. But later work in (Hoffer et al., 2017) agreed with the square root scaling rule.

At the same time, an empirical work trying to train a ResNet-50 on ImageNet in 1 hour (i.e., in a highly distributed setting), derived the linear scaling rule under the assumption that the gradient doesn’t change much during training (Goyal et al., 2017). Despite a few other optimization tricks, the test accuracy still degraded at large batch size. (Smith et al., 2020) also derived the linear scaling rule using reasoning analogous to a central limit theorem but still noted the batch size needs to be small for the rule to hold.

To clear up which scaling rule is correct and when it will break, we can turn to SDEs! Our work in NeurIPS 2021 used SDEs to design a simple test (requiring just one baseline run!) for the largest batch size you can parallelize to via the linear scaling rule without sacrificing performance. See the figure below. We also designed an efficient simulation of the SDE that provided evidence that the SDE is the correct way to model many realistic SGD training settings. (The next section explains why SDEs are so useful in analyzing SGD.)

Figure from our work in NeurIPS 2021. Predicting the failure of the linear scaling rule with just one baseline run, using insights from SDEs! The red dashed line indicates where the necessary condition for the scaling rule fails, and the shaded red area covers the batch sizes at which the test error has increased by 20% or more from its baseline value. These experiments are on CIFAR-10 but there are plenty more settings in the paper!

In the meantime, language models started to grow popular. (You et al., 2020) designed a scheme to train a BERT model in 76 minutes with Adam, and they empirically discovered that scaling the learning rate by $\sqrt{\kappa}$ works well. Our work in NeurIPS 2022 approached the question theoretically and established new SDE approximations for Adam and RMSprop. These yielded a square root scaling rule for Adam, which requires scaling the other optimization hyperparameters in addition to the learning rate. Experiments showed that the square root scaling rule preserved test accuracy, perplexity, and even post-fine-tuning performance. Naturally, these scaling rules will break at some large batch size, but we couldn’t find any empirical setting at the time where this happened.

Various works have since used the SDE approximation to derive useful distributed training protocols when using EMA (Busbridge et al., 2023) and Local SGD (Gu et al., 2023).

Into the Weeds with SDEs

Now that I’ve described the importance of SDEs, I’ll describe what they are exactly and how they can be used. SDEs describe a continuous trajectory of the model parameters over the course of training. The hypothesis is that this continuous trajectory is a reasonable and easily analyzable approximation of the discrete parameter trajectory that SGD prescribes. So let’s see how it shakes out.

At each step, SGD samples a minibatch $B_t\subset \mathcal{D}$ and updates the parameters according to some loss function $\ell$: $\theta_{t+1} = \theta_t - \eta\nabla \ell(B_t;\theta_t)$. Everyone knows the crucial hyperparameter $\eta$, the learning rate, but people often overlook the batch size hyperparameter, $B$, which is usually set to be the largest allowable value on the GPU (for pre-training) and carefully grid-searched (in fine-tuning). These two hyperparameters together control how much noise there is in the gradient estimate and how it affects the trajectory. The gradient noise is what makes SGD different from GD, so we will try to study this property carefully to understand why SGD results in better generalization than GD.

One approach to studying SGD is to use a continuous approximation of the discrete trajectory. This admits the derivation of results about the exploration and convergence of the optimization. One way to get a continuous approximation is to take the learning rate to be very small (i.e., $\eta\to 0$), which gets us to the famous gradient flow (GF) limit: $dX_t = -\nabla\ell(\mathcal{D};X_t)dt$. I’ll start to use $X$ to denote the continuous trajectory of the parameters and $\theta$ to denote the discrete one. GF used to be and still somewhat is the tool of choice to analyze how optimization proceeds. But if we look closely, we can see that GF is agnostic to the batch size! So, even if I was doing full-batch GD, I would still end up with GF as the limit. This means that GF can’t tell us much about the benefit of SGD over GD.

OK, so we need to find a continuous process that can actually model the gradient noise. This is where SDEs come in handy! Here’s the SDE used to approximate SGD (Li et al., 2019):

$dX_t = \underbrace{-\nabla \ell(\mathcal{D};X_t)dt}_{\text{Drift: Gradient Flow}} + \underbrace{(\eta\Sigma(X_t))^{1/2} dW_t}_{\text{Diffusion: Brownian Motion}}$

The first term of our SDE is the drift, the deterministic term in the process, which is just the full-batch gradient in this case. The second term is the diffusion, which we traditionally model using Brownian motion $W_t$. So, the SDE is actually very intuitive: we believe that SGD is doing something like GD + noise and we can clearly see that intuition in these two terms. Another key feature is that the learning rate $\eta$ actually shows up here! As we increase the learning rate, the noise in the SDE increases, which also makes sense: when we take larger steps with a mini-batch gradient, the trajectory increasingly diverges from full-batch GD. The last remaining mystery in this equation is $\Sigma(X_t)$, which is the gradient noise covariance.

$\Sigma^{(B)}(X_t) = \mathbb{E}_\gamma [(\nabla \ell(\mathcal{D}; X_t) - \nabla \ell(B_t; X_t)(\nabla \ell(\mathcal{D}; X_t) - \nabla \ell(B_t;X_t))^\top]$

where $\gamma$ denotes the optimization seed (which determines the minibatch $B_t$). We can assume the gradient for a single datapoint is the full batch gradient plus some Gaussian noise (discussed more at the end of the post). Then, if we take a minibatch $B_t$ with size $B$, the minibatch gradient is:

$\nabla \ell(B_t;X_t) = \nabla \ell(\mathcal{D}; X_t) + \frac{1}{B} \sum_{i=1}^B z_i$

where $z_i\sim \mathcal{N}(0, \Sigma(X_t))$ is drawn i.i.d. per datapoint. Then, when we scale the batch size, we can see how the gradient noise, and thus the diffusion term in the SDE, changes.

$\Sigma^{(\kappa B)}(X_t) = \frac 1\kappa \Sigma^{(B)}(X_t)$

This equation also matches our intuitions: when we increase the batch size, the noise (i.e., diffusion) in the process should get smaller, since the trajectory is closer to full-batch GD. Now we have basically everything we need in order to derive some scaling rules!

Deriving Scaling Rules from SDEs

Now that we have the basics of SDEs down, we can go back to the scaling rules at the start of the post.

Linear Scaling Rule (for SGD)

When scaling the batch size by $\kappa$, scale the learning rate also by $\kappa$.

Square Root Scaling Rule (for Adam)

The reasoning to derive scaling rules from the SDE goes like this. Changing the batch size scales $\Sigma(X_t)$ by $1/\kappa$ so one should scaling the learning rate proportionally to preserve the scale of the diffusion term. Language models are usually trained with the Adam optimizer, which requires a different SDE approximation (see our NeurIPS 2022 paper here). That SDE, through similar reasoning, yields the square root scaling rule.

Keep in mind that the accuracy of the SDE approximation is a sufficient condition for the scaling rule: if the SDE is a good approximation, then the scaling rule should hold. So when does the SDE approximation break? Section 4.1 of our paper illustrates a good warmup setting to build intuition on when the gradient noise gets to be too small and the SDE approximation is no longer good. If you find your hyperparameters taken to extremes by these scaling rules, consider shifting your baseline run to a larger batch size.

What do Scaling Rules Preserve?

Traditionally, when training vision models with SGD, the empirical goal of using a scaling rule is to preserve the test accuracy of the model. But for language models trained with Adam, what metrics are we interested in? Perplexity? In-context learning ability? Or even more broadly, factuality? Truthfulness? Fortunately, when we use the formal language of SDEs, we can show results that conceptually encapsulate all of these possibilities (and ones we haven’t dreamed up yet).

To recap, the way that one gets from SDEs to scaling rules is as follows. Changing the batch size changes the equation of the SDE. Adjusting the optimization hyperparameters according to the scaling rule changes the equation of the SDE back to what it was originally with the baseline run. If the SDE is indeed a faithful approximation of SGD/Adam, then preserving the SDE equation under changing batch sizes will also preserve the optimization trajectory of the baseline run of SGD/Adam.

So now we enter the question of what a “faithful approximation” is. As I hinted before, we might only care about preserving certain features of our trained model (e.g., test accuracy or truthfulness) under changing batch sizes. And we have no idea how these features change with the actual parameters — for example, two models may have parameters that are very close to each other but have different test accuracies. Yet another complicating factor is that SGD and Adam are stochastic (i.e., depend on the optimization seed), so we have to figure out what it means for two stochastic trajectories to approximate each other.

Altogether, we need a more sophisticated notion of approximation than just saying the two trajectories are close to each other in $\ell_2$ norm. SDEs handle this issue using test functions, which compute something as a function of the model parameters. This notion of approximation, called a weak approximation (Li et al., 2019), gives a guarantee that the SDEs and SGD/Adam produce models that have close values of all possible test functions. This is a very broad class of functions, and it likely includes all of the ones we would care about. The tricky thing is, to make the theory work out, we have to assume something fairly general about these test functions — that they grow at most polynomially with the model parameters — and we can’t be confident that test accuracy, truthfulness, etc. satisfy this assumption. But, fortunately, extensive experiments in our papers (on SGD and on Adam) show that scaling rules preserve various test functions, including test accuracy, perplexity, and even post-fine-tuning performance (on GLUE tasks).

Figures from our NeurIPS 2022 paper. Training RoBERTa models on Wiki+Books (left) and GPT-2 on WikiText-103 (right) with different batch sizes. Adjusting the Adam optimizer hyperparameters per the square root scaling rule ensures performance is preserved!

Discussions and Extensions

Here we discuss some of the more nuanced points of SDEs for those who are interested. I will dive more into these ideas in a subsequent post, which will contain more of the theoretical underpinnings of SDEs.

Is SGD gradient noise actually Gaussian? In the SDE section, I assumed that the gradient of one datapoint is the full-batch gradient plus some Gaussian noise. It is generally agreed that gradient noise is additive, but a big topic of discussion in the SDE community is whether the gradient noise is actually Gaussian or if it’s “heavy-tailed” (i.e., third-and-higher moments are non-negligible). The SDE derived above, called an Ito SDE, uses Brownian motion for the diffusion term, but some argue that we should instead use a the more general Levy process for the diffusion, which yields a Levy SDE (Simsekli et al., 2019). These are usually harder to analyze, and much of the empirical evidence motivating the usage of a Levy SDE has been disproven (Xie et al., 2021). Moreover, our work in NeurIPS 2021 showed empirically that using Gaussian noise with the first and second moments matching the naturally occurring noise in SGD is sufficient to preserve the generalization performance of SGD in vision. For these reasons, and also because of the ease of using an Ito SDE for analysis, the Ito SDE remains the common choice for approximating discrete optimization trajectories. More recently, however, the SGD noise has been considered to be heavy-tailed when nodes in a distributed setting fail unexpectedly (Schaipp et al., 2023). Due to the computational resources to make these measurements, it’s hard to collect evidence of which setting is more faithful to language modeling training, but our insights derived using Ito SDEs (in our NeurIPS 2022 paper) generally preserve the performance of language models when performing pre-training and fine-tuning. Other analyses using similar noise assumptions, including our upcoming ICLR 2024 paper showing that using momentum does not change the SGD trajectory when using small learning rates, also empirically hold when training language models.
Implicit Bias of SGD: The SDE that I described doesn’t directly give much information about the implicit bias (i.e., the ability of SGD to choose generalizing solutions out of many possible empirical risk minimizers). Recent works have studied using the SDE to describe the late phase of training, where the training loss is nearly zero, and the primary term driving the trajectory is the diffusion (Li et al., 2022). Analyzing this SDE reveals that SGD tends towards flatter minima. It is not totally clear how flatness and generalization are related to each other (Jiang et al., 2020; Andriushchenko et al., 2023; Wen et al., 2023), but a careful analysis of the SDE may yield a more nuanced understanding than classical generalization measures.

Conclusion

SDEs are a powerful, generalizable tool for studying stochastic optimization. The results are mostly agnostic to the model architecture and dataset, so just a little bit of understanding goes a long way and provides useful empirical insights. In a separate post, I’ll talk a bit more about the proof techniques and rigorous considerations involved in using SDEs to approximate discrete optimization.

Acknowledgements: Thanks (in alphabetical order) to Tianyu Gao, Surbhi Goel, Bingbin Liu, Kaifeng Lyu, Abhi Venigalla, Mengzhou Xia, and Howard Yen for their feedback on this post! I’m deeply indebted to Zhiyuan Li for patiently teaching me about SDEs and their technical features. Feedback from Twitter prompted me to add a link to the section in our paper that provides an easy setting to work out the scaling rule in. The works I mention in this paper were co-authored with (in alphabetical order) Sanjeev Arora, Zhiyuan Li, Kaifeng Lyu, Abhishek Panigrahi, Runzhe Wang, and Tianhao Wang.