FisherSFT: Data-Efficient Supervised Fine-Tuning of Language Models Using Information Gain

We wanted to thank the reviewer for positive evaluation and appreciating that we bring classic ideas from statistical learning theory to LLMs. Our rebuttal is below. We focus on major issues and will incorporate all comments of the reviewer in the next version of our paper. If you have additional concerns, please reach out to us to discuss them.

Other Values of $d$ and $L$ in Synthetic Experiments

We plot the error curves for the synthetic problem in Section 5.1, and various $d \in \\{5, 10, 20, 30, 40\\}$ and $L \in \\{20, 30, 40, 50\\}$, at anonymized link. TokenOD consistently performs better. We also show the trends for $L$ (anonymized link) and $d$ (anonymized link) when fixing the number of samples at $1000$.

Computational Gains in Fast Implementation of TokenOD

We compare the computation times of the fast and slow implementations of TokenOD on the problem in Section 5.3 at anonymized link. The number of sentences is $5000$ and we plot the computation time for various values of $n$. The fast implementation of TokenOD is about $4$ times faster.

More Details on Experimental Setup in Section 5.3

The experiment can be described in our notation as follows. $x_{i,j}$ is the transformer embedding at the $j$-th token of the $i$-th sentence in the text corpus and $y_{i, j}$ is the identity of that token. For each compared method, we fine-tune a GPT-2 model on sentences collected by that method. At test time, we prompt the fine-tuned model with a few words from the original text corpus (represented by tokens $y_{i, 1}, y_{i, 2}, \dots, y_{i, p}$) and the model completes the sentence by generating $y_{i, p + 1}, y_{i, p + 2}, \dots, y_{i, 1024}$. Finally, we compare the quality of the completed sentences generated by the various models using an LLM as a judge.

Fisher Information

The Fisher information matrix in TokenOD is an approximation derived in Lemma 3.1. This algebraic form, an outer product of feature vectors, is the true Fisher information matrix for least squares. How is it useful beyond least squares? Simply put, optimization of this matrix leads to choosing feature vectors that cover all directions in training data uniformly. Because of that, it is possible to bound a worst-case prediction error over all feature vectors and be robust.

LLM as a Judge

We improved the evaluation, including reporting biases. See the rebuttal for Reviewer M5jD for details.

We wanted to thank the reviewer for positive evaluation, recognizing our contributions, and bringing up the prompt bias issue to our attention. Our rebuttal is below. We focus on major issues and will incorporate all comments of the reviewer in the next version of our paper. If you have additional concerns, please reach out to us to discuss them.

LLM as a Judge

We completely reworked the LLM-as-a-judge evaluation. The new prompt is

You are a judge of Shakespeare text.
<tag1>text1</tag1>
<tag2>text2</tag2>
Respond 2 if the text inside <tag2> is more fluent Shakespeare text than the text inside <tag1>. Respond 1 otherwise.

The prompt got simplified, does not name the methods, and targets our perceived benefit (improved language). We use a state-of-the-art LLM GPT-4o to judge. The text generated by the compared methods is randomized: one randomly-chosen method replaces text1 and the other text2. We tested the LLM judge and it chooses the first position with probability 0.54, which is slightly higher than 0.5 for a position-unbiased judge.

In addition to improving evaluation, we added a new baseline Ask-LLM (Sachdeva et al., 2024). See Appendix A of our paper for more details on this method.

We report the win rates of TokenOD on the Shakespeare dataset, as a function of sample size $n$, below:

We observe that TokenOD consistently outperforms all baselines.

We also added a new experiment on the Sherlock dataset from Sherlock Holmes Next Word Prediction Corpus. The evaluation protocol is the same as in the Shakespeare dataset, except that "Shakespeare" is replaced with "Sherlock". The win rates of TokenOD on the Sherlock dataset are:

We observe again that TokenOD consistently outperforms all baselines.

We wanted to thank the reviewer for positive evaluation, and recognizing the practicality of our solution as well as the importance of the solved problem. Our rebuttal is below. We focus on major issues and will incorporate all comments of the reviewer in the next version of our paper. If you have additional concerns, please reach out to us to discuss them.

Active Learning Using the Last-Layer Embedding

Model gradients are often used in active learning and bandit exploration (Deep Batch Active Learning by Diverse, Uncertain Gradient Lower Bounds, Neural Contextual Bandits with UCB-based Exploration). We particularly note that such methods are computationally intensive since the parameter space of modern neural networks is huge in comparison to the last-layer embedding. Our method uses the last-layer embedding as a featurizer and is similar in spirit to Neural Contextual Bandits with Deep Representation and Shallow Exploration. This approach is known to be robust and has a much lower empirical regret than other uncertainty representation techniques (Deep Bayesian Bandits Showdown: An Empirical Comparison of Bayesian Deep Networks for Thompson Sampling).

Empirical Convergence Rate of TokenOD

We plot the error rate of TokenOD on the synthetic problem in Section 5.1 at anonymized link. Specifically, we take the logarithm of the error rate and plot it as a function of the logarithm of the sample size. We observe a slope of $-0.3$, which means that the error rate is $O(n^p)$. We believe that this is close enough to the expected $p = -0.5$, especially since other factors may have played a role at our sample sizes.

More Extensive Empirical Evaluation

Our computational resources are limited and therefore we did not experiment beyond GPT-2. To partially address your concern, we further expanded our evaluation as follows. In addition to improving the LLM-as-a-judge evaluation, which other reviewers asked for, we added a new baseline Ask-LLM (Sachdeva et al., 2024). See Appendix A of our paper for more details on this method.

We report the win rates of TokenOD on the Shakespeare dataset, as a function of sample size $n$, below:

We observe that TokenOD consistently outperforms all baselines.

We also added a new experiment on the Sherlock dataset from Sherlock Holmes Next Word Prediction Corpus. The evaluation protocol is the same as in the Shakespeare dataset, except that "Shakespeare" is replaced with "Sherlock". The win rates of TokenOD on the Sherlock dataset are:

We observe again that TokenOD consistently outperforms all baselines.

Bias Towards Longer Sentences

TokenOD is biased towards selecting longer sentences. This bias naturally arises because each token in a sentence is a training point in supervised fine-tuning, for the conditional probability of token $y_{i, j}$ given the history embedding $x_{i, j}$. In a sense, longer sentences can contribute more to information gain. However, the selection of the sentence depends on its total information gain in comparison to other available sentences. We also observe a similar bias towards longer sentences in the GPT-2 code on Hugging Face

We wanted to thank the reviewer for positive evaluation and clearly summarizing the main technical contributions of our work. Our rebuttal is below. We focus on major issues and will incorporate all comments of the reviewer in the next version of our paper. If you have additional concerns, please reach out to us to discuss them.

LLM as a Judge

We completely reworked the LLM-as-a-judge evaluation. The new prompt is

You are a judge of Shakespeare text.
<tag1>text1</tag1>
<tag2>text2</tag2>
Respond 2 if the text inside <tag2> is more fluent Shakespeare text than the text inside <tag1>. Respond 1 otherwise.

The prompt got simplified, does not name the methods, and targets our perceived benefit (improved language). We use a state-of-the-art LLM GPT-4o to judge. The text generated by the compared methods is randomized: one randomly-chosen method replaces text1 and the other text2. We tested the LLM judge and it chooses the first position with probability 0.54, which is slightly higher than 0.5 for a position-unbiased judge.

In addition to improving evaluation, we added a new baseline Ask-LLM (Sachdeva et al., 2024). See Appendix A of our paper for more details on this method.

We report the win rates of TokenOD on the Shakespeare dataset, as a function of sample size $n$, below:

We observe that TokenOD consistently outperforms all baselines.

We also added a new experiment on the Sherlock dataset from Sherlock Holmes Next Word Prediction Corpus. The evaluation protocol is the same as in the Shakespeare dataset, except that "Shakespeare" is replaced with "Sherlock". The win rates of TokenOD on the Sherlock dataset are:

We observe again that TokenOD consistently outperforms all baselines.

How Is the Fine-Tuned LLM Used in Data Selection

Our data selection procedure depends on the LLM through embeddings $x_{i, j}$. They arise in both the log-likelihood in (2) and algorithm TokenOD. Simply put, TokenOD selects diverse sentences, which cover the embeddings $x_{i, j}$ more uniformly. As you pointed out, we reduce the original $d L \times d L$ Fisher information matrix into a $d \times d$ matrix, which can be optimized efficiently. While this neglects token-level prediction accuracy, it incorporates the LLM through the embeddings. A more direct optimization of the original matrix is an interesting direction that would require addressing the computational challenge.

Different Weaker LLM for Optimal Design

We agree that this is feasible when the embeddings of the weaker and stronger LLMs can be related. We have not done this in our work because we wanted to start with a simpler problem, the same LLM is used for both the optimal design and fine-tuning. We will discuss this option in the paper.

Missing Reference

Thank you for the reference. We will include it in the next version of the paper.