Efficiently Learning at Test-Time: Active Fine-Tuning of LLMs

## Abstract

Abstract Recent efforts in fine-tuning language models often rely on automatic data selection, commonly using Nearest Neighbors retrieval from large datasets.
However, we theoretically show that this approach tends to select redundant data, limiting its effectiveness or even hurting performance.
To address this, we introduce SIFT , a data selection algorithm designed to reduce uncertainty about the model’s response given a prompt, which unifies ideas from retrieval and active learning.
Whereas Nearest Neighbor retrieval typically fails in the presence of information duplication, SIFT accounts for information duplication and optimizes the overall information gain of the selected examples.
We focus our evaluations on fine-tuning at test-time for prompt-specific language modeling on the Pile dataset, and show that SIFT consistently outperforms Nearest Neighbor retrieval, with minimal computational overhead. Moreover, we show that our uncertainty estimates can predict the performance gain of test-time fine-tuning, and use this to develop an adaptive algorithm that invests test-time compute proportional to realized performance gains.
We provide the activeft (Active Fine-Tuning) library which can be used as a drop-in replacement for Nearest Neighbor retrieval.

## 1 Introduction

Figure: Figure 1: Selecting fine-tuning data using SIFT (red) robustly outperforms Nearest Neighbor retrieval (black) and avoids the failure-mode of Nearest Neighbor retrieval where the same data is selected repeatedly, which is a common result of information duplication.
Refer to caption: x1.png

The standard paradigm of machine learning separates training and testing.
Training aims to learn a model by *inductively* extracting general rules from data, and testing applies this model to new, unseen data.
We investigate an alternative *transductive* paradigm where the model is fine-tuned at test-time specifically to the given task.
Variations of this paradigm have been studied since the inception of machine learning as a field.
Early examples are local learning (Cleveland, ; Cleveland & Devlin, ; Atkeson et al., ) and local fine-tuning (Bottou & Vapnik, ).
More recently, with the advent of large pre-trained models which have good representations and are strong foundations for fine-tuning, the idea of *test-time fine-tuning* has re-gained attention (Krause et al., ; Sun et al., ).
Hardt & Sun () show that fine-tuning on data related to the prompt to a large language model (LLM) can significantly improve performance.
Also, test-time fine-tuning is the central component of state-of-the-art approaches to the ARC challenge (Chollet, ; Cole & Osman, ), a non-saturated benchmark which is intended to test reasoning capabilities based on “core knowledge” rather than mere memorization.

#### Active Fine-Tuning: Effective data selection for fine-tuning LLMs

Test-time fine-tuning demands automatic data selection since manually selecting data for each test instance is infeasible.
Moreover, the sample efficiency of test-time fine-tuning is a central bottleneck as the number of gradient steps is directly proportional to inference time.
Previous works on data selection for fine-tuning LLMs have fundamentally relied on Nearest Neighbor retrieval within some embedding space (Hardt & Sun, ; Xia et al., ).
We show theoretically and empirically that Nearest Neighbor retrieval is insufficient for fine-tuning LLMs since it can lead to the selection of redundant data.
Notably, recent works using influence functions for data selection such as Xia et al. () have pointed out this limitation.
In contrast, a large body of work on (inductive) active learning has studied non-redundant data selection (e.g., Sener & Savarese, ; Ash et al., ; Yehuda et al., ; Kirsch et al., ) that covers the data manifold well (cf. Figure 2).
Retrieval and active learning can be seen as two extreme ends of a spectrum: retrieval selects relevant but potentially redundant data, while active learning selects diverse but potentially irrelevant data.

Figure: Figure 2: We consider a scenario where we have a pre-trained language model capturing a latent manifold (red) in the large sequence space (white). We aim to improve the models performance on a given prompt (blue) by *efficiently* fine-tuning the model on *few* relevant and diverse data points (black) at test-time.

We bridge this gap by unifying ideas from retrieval and active learning in SIFT, an algorithm based on emerging literature on transductive active learning (Hübotter et al., ) that *Selects Informative data for Fine-Tuning* as illustrated in Figure 2.
Our results show that SIFT leads to substantial improvements in performance and efficiency.
Concretely, we show the following:

- 1.
Nearest Neighbor retrieval is insufficient (§2):
We prove that selecting the top-$N$ highest scoring points from a large dataset according to a fixed scoring function leads to the selection of redundant data.
- 2.
SIFT estimates uncertainty about responses (§3):
We develop the notion of *uncertainty about the response to the prompt*, and derive an anytime high probability bound to the total variation distance between the model’s distribution over responses and the ground truth which is governed by this uncertainty.

- 3.
SIFT provably reduces uncertainty (§4): We propose SIFT, an algorithm that selects data which reduces uncertainty about the response to the prompt. We prove statistical rates for the uncertainty reduction (§4.1) and show that SIFT is compute-efficient, with minimal overhead compared to Nearest Neighbor retrieval (§4.2).
- 4.
SIFT performs better and is more robust than Nearest Neighbor retrieval (§5): We find that fine-tuning an LLM on data selected by SIFT consistently and robustly improves performance, which is not the case with Nearest Neighbor retrieval.
Moreover, our results indicate that at test-time an LLM can learn more effectively by fine-tuning on few examples than from context.
- 5.
SIFT can invest test-time compute proportionally to performance gains (§6): We observe that our uncertainty estimates can accurately predict the performance gain of test-time fine-tuning.
Motivated by this, we dynamically adapt compute to the expected performance gain.

## 2 Test-Time Fine-Tuning

We define test-time fine-tuning of LLMs (Hardt & Sun, ) as follows.
We consider a domain $\mathcal{X}$ of token sequences and assume that we have access to a large dataset of examples $\mathcal{D}\subseteq\mathcal{X}$ which we call the *data space*.
We further assume that we have access to a pre-trained autoregressive language model that maps token sequences $\mathcal{X}$ to probability distributions over the next token from a vocabulary of size $V$.
Our work addresses the central question:

*Given a prompt $\boldsymbol{x}^{\star}\in\mathcal{X}$, how can we effectively select fine-tuning data
from the large dataset $\mathcal{D}$ such that the fine-tuned model performs well on the prompt?*

We then fine-tune the model for a single gradient step on each selected sequence.

Locally adjusting a model at test-time has gained popularity in the context of few-shot in-context learning (Brown et al., ; Wei et al., ; Bubeck et al., ; OpenAI, ) with retrieval augmented generation (RAG, Lewis et al., ; Guu et al., ; Borgeaud et al., ).
In contrast to this approach, test-time fine-tuning works by fine-tuning the parameters of a pre-trained model at test-time specifically to each prompt.
Notably, test-time fine-tuning takes time linear in the number of tokens whereas in-context learning with a transformer has quadratic complexity (Vaswani et al., ).
Next to this, Hardt & Sun () and other works have found (test-time) fine-tuning to perform substantially better than in-context learning (Hu et al., ; Mosbach et al., ).
This work further improves the performance of test-time fine-tuning.
Prior work has also studied how one can explicitly meta-learn the ability to perform test-time fine-tuning (Finn et al., ; Sun et al., ), though we find this capability to emerge even from models that are not explicitly trained in this way.

The central question studied in this work also arises when fine-tuning LLMs during post-training.
For example, in targeted instruction tuning, the goal is to fine-tune a model to obtain desired capabilities, which are commonly embodied by a set of examples $\boldsymbol{x}^{\star}$ (Xia et al., ).
The extension of our work to such a “batched” setting is straightforward.

### 2.1 Nearest Neighbor Retrieval is Insufficient

Figure: Figure 3: We retrieve two data points to answer the prompt. Nearest Neighbor selects redundant data, while SIFT yields maximal information (cf. §L).

Prior work on data selection for fine-tuning has relied on Nearest Neighbor retrieval.
The idea of making predictions on $\boldsymbol{x}^{\star}$ depending on its nearest neighbors has been around as long as machine learning itself (Fix, ; Cover & Hart, ).
Bottou & Vapnik () were the first to apply this idea to the fine-tuning of convolutional neural networks by selecting the nearest neighbors of a test image in pixel-space.
More recently, due to advances in representation learning (Devlin et al., ; Reimers & Gurevych, ) and efficiency (e.g., Johnson et al., ; Aumüller et al., ), Nearest Neighbor retrieval has regained attention and been applied to test-time fine-tuning (Hardt & Sun, ).

Xia et al. () use influence functions (Cook, ; Koh & Liang, ; Pruthi et al., ) to select data for fine-tuning LLMs.
This line of work aims to select data that reduces a first-order Taylor approximation to the test loss after fine-tuning, an approach that corresponds to Nearest Neighbor retrieval in a certain embedding space.
They highlight two main limitations of the use of influence functions and Nearest Neighbor retrieval for data selection:

- •
Nearest Neighbor retrieval leads to the selection of redundant data.
Figure 3 illustrates this limitation with a qualitative example.
We formalize this limitation in Proposition K.1, which we summarize here informally:
Informal Proposition 2.1.
Selecting the top-$N$ nearest neighbors from the data space (according to cosine similarity or Euclidean distance) may not reduce the uncertainty about the response to the prompt beyond fine-tuning on the closest neighbor.
Every additional passage may be redundant.
- •
Nearest Neighbor retrieval selects data with high positive cosine similarity to the prompt.
Yet, data with high *negative* cosine similarity can be equally informative as data with high positive cosine similarity (Xia et al., , Appendix K.2), but is ignored by standard Nearest Neighbor retrieval.

In this work, we propose SIFT and show that it naturally addresses both limitations.

###### Informal Proposition 2.1 .

## 3 Preliminaries: Uncertainty Estimation for Fine-Tuning

We suppose the assigned probability that $y\in[V]$ is the class label of an input $\boldsymbol{x}\in\mathcal{X}$ is given by $s_{y}(\boldsymbol{f}^{\star}(\boldsymbol{x}))$, where $s_{y}$ is the softmax $\smash{s_{y}(\boldsymbol{f})\operatorname*{\,\dot{=}\,}\exp(f_{y})/(\sum_{i=1}
^{V}\exp(f_{i}))}$.
That is, $\boldsymbol{f}^{\star}(\boldsymbol{x})$ denotes the “ground truth” logits for a given input $\boldsymbol{x}$.
In the context of language modeling, $V$ is the number of tokens in the vocabulary, and $y$ denotes the index of the next token.
We defer all proofs to Appendix K.

We use a surrogate model to quantify the informativeness of data, which we define next.

###### Assumption 3.1 ( Surrogate model: Linear model class within a known latent space) .

We assume ${\boldsymbol{f}^{\star}(\boldsymbol{x})=\boldsymbol{W}^{\star}\boldsymbol{\phi
}(\boldsymbol{x})}$ with ${\boldsymbol{W}^{\star}\in\mathbb{R}^{V\times d}}$ and where ${\boldsymbol{\phi}(\cdot)\in\mathbb{R}^{d}}$ denotes known embeddings.

The surrogate model uses the latent space induced by the pre-trained model to describe the data manifold. We emphasize that while SIFT relies on this surrogate model for data selection, it still fine-tunes the full pre-trained model, including latent features.
Surrogate dense embedding models of this kind have been used extensively for data selection via Nearest Neighbor retrieval (e.g., Lewis et al., ; Karpukhin et al., ; Borgeaud et al., ; Xia et al., ), and to understand the training dynamics and generalization of large neural networks (e.g., Jacot et al., ; Lee et al., ; Malladi et al., ; Templeton et al., ; Park et al., ).
Furthermore, a surrogate model that assumes linearity in some fixed latent space may be a reasonable approximation for test-time fine-tuning since the latent space of the unfrozen model is not expected to change substantially by a few gradient steps.

In this work, we explore a scenario where we have a pre-trained model $\smash{\boldsymbol{f}^{\mathrm{pre}}(\boldsymbol{x})={\boldsymbol{W}^{\mathrm{
pre}}}\boldsymbol{\phi}(\boldsymbol{x})}$.
We let ${\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W})\operatorname*{\,\dot{=}\,}
\boldsymbol{W}\boldsymbol{\phi}(\boldsymbol{x})}$ and denote by $\mathcal{L}(\boldsymbol{W};D)$ the negative log-likelihood loss of $\boldsymbol{f}(\cdot;\boldsymbol{W})$ on a dataset $D$ of inputs $\boldsymbol{x}$ with corresponding class labels $y$: $\smash{\mathcal{L}(\boldsymbol{W};D)\operatorname*{\,\dot{=}\,}-\sum_{(
\boldsymbol{x},y)\in D}\log s_{y}(\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}
))}$.

###### Assumption 3.1 ( Surrogate model: Linear model class within a known latent space) .

#### Uncertainty Estimation

Our first intermediate goal is to estimate the uncertainty about the response to a given prompt $\boldsymbol{x}^{\star}$ after having fine-tuned on selected data $D_{n}$ of size $n$.
To this end, we generalize prior work on confidence sets under categorical feedback (i.e., class feedback, Amani & Thrampoulidis, ; Zhang & Sugiyama, ) to our fine-tuning setting.
We consider the function class ${\mathcal{W}\operatorname*{\,\dot{=}\,}\{\boldsymbol{W}\in\mathbb{R}^{V\times d
}\mid\left\|\boldsymbol{W}-{\boldsymbol{W}^{\mathrm{pre}}}\right\|_{\mathrm{F}
}\leq B\}}$ where $\left\|\cdot\right\|_{\mathrm{F}}$ denotes the Frobenius norm and with $B$ a constant such that $\boldsymbol{W}^{\star}\in\mathcal{W}$.
Then given data $D_{n}$, we can refine the prior estimate ${\boldsymbol{W}^{\mathrm{pre}}}$ of $\boldsymbol{W}^{\star}$ by minimizing the regularized negative log-likelihood loss

$$ $\displaystyle\mathcal{L}^{\lambda}(\boldsymbol{W};D_{n})\operatorname*{\,\dot{ =}\,}\mathcal{L}(\boldsymbol{W};D_{n})+\frac{\lambda}{2}\left\|\boldsymbol{W}- {\boldsymbol{W}^{\mathrm{pre}}}\right\|_{\mathrm{F}}^{2}\vspace{-2pt}$ (1) $$

with regularization coefficient $\lambda>0$.
We write its minimizer as $\smash{\boldsymbol{W}_{\!n}\operatorname*{\,\dot{=}\,}\operatorname*{arg\,min}
_{\boldsymbol{W}\in\mathcal{W}}\mathcal{L}^{\lambda}(\boldsymbol{W};D_{n})}$.
We will further denote the ground truth probability distribution over the response to $\boldsymbol{x}$ by $\boldsymbol{s}^{\star}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}\boldsymbol{s
}(\boldsymbol{f}^{\star}(\boldsymbol{x}))$ and our approximation after selection of $n$ samples by $\boldsymbol{s}_{n}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}\boldsymbol{s}(
\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}_{\!n}))$.

We construct confidence sets of the form $[\boldsymbol{s}_{n}(\boldsymbol{x})\pm\beta_{n}(\delta)\sigma_{n}(\boldsymbol{
x})]$ centered around this prediction, and show their uniform anytime validity.
The width of these sets is characterized by our central quantity $\sigma_{n}(\boldsymbol{x})$ which we define next.
We consider the inner-product kernel $k(\boldsymbol{x},\boldsymbol{x^{\prime}})\operatorname*{\,\dot{=}\,}
\boldsymbol{\phi}(\boldsymbol{x})^{\top}\boldsymbol{\phi}(\boldsymbol{x^{
\prime}})$ and define for a set of inputs $X=\{\boldsymbol{x}_{1},\dots,\boldsymbol{x}_{n}\}\subseteq\mathcal{D}$:

$$ $\displaystyle\sigma_{X}^{2}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}k( \boldsymbol{x},\boldsymbol{x})-\boldsymbol{k}_{X}^{\top}(\boldsymbol{x})( \boldsymbol{K}_{\!X}+\lambda\kappa\boldsymbol{I}_{n})^{-1}\boldsymbol{k}_{X}( \boldsymbol{x})$ (2) $$

where $\boldsymbol{k}_{X}(\boldsymbol{x})=(k(\boldsymbol{x}_{1},\boldsymbol{x}),\dots
,k(\boldsymbol{x}_{n},\boldsymbol{x}))\in\mathbb{R}^{n}$, $\boldsymbol{K}_{\!X}\in\mathbb{R}^{n\times n}$ is the kernel matrix satisfying $(\boldsymbol{K}_{\!X})_{i,j}=k(\boldsymbol{x}_{i},\boldsymbol{x}_{j})$, and $\smash{\kappa\operatorname*{\,\dot{=}\,}\sup_{\boldsymbol{x}\in\mathcal{X},
\boldsymbol{W}\in\mathcal{W}}1/\lambda_{\min}(\boldsymbol{A}(\boldsymbol{x};
\boldsymbol{W}))}$.
Here, $\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W})\in\mathbb{R}^{V\times V}$ is the matrix satisfying ${(\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W}))_{i,j}\operatorname*{\,\dot{=}
\,}s_{i}(\boldsymbol{x};\boldsymbol{W})(\mathbbm{1}\{{i=j}\}-s_{j}(\boldsymbol
{x};\boldsymbol{W}))}$ which is the proper generalization of the derivative of the sigmoid function, standard in the analysis of binary feedback (Faury et al., ; Pásztor et al., ). We write $\smash{\sigma_{n}^{2}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}\sigma_{X_{n}}
^{2}(\boldsymbol{x})}$ where $X_{n}\subseteq\mathcal{D}\subseteq\mathcal{X}$ are the inputs in $D_{n}$.
With this we are ready to state our first result, namely that for careful choice of $\beta_{n}(\delta)$, the confidence sets contain $\boldsymbol{s}^{\star}(\boldsymbol{x})$ simultaneously for all $\boldsymbol{x}\in\mathcal{X}$ and $n\geq 1$ with probability at least $1-\delta$.

###### Theorem 3.2 (Confidence Sets) .

Let Assumption 3.1 hold and $\boldsymbol{W}^{\star}\in\mathcal{W}$.
Let $\delta\in(0,1)$ and set

$$ $\displaystyle\beta_{n}(\delta)\operatorname*{\,\dot{=}\,}2\sqrt{V(1+2B)}\left[ B+\frac{LV^{3/2}d}{\lambda}\log\left(\frac{2}{\delta}\sqrt{1+\frac{n}{d\lambda }}\right)\right]\in O\left(\log(n/\delta)\right)$ (3) $$

where ${L\operatorname*{\,\dot{=}\,}\sup_{\boldsymbol{x}\in\mathcal{X},\boldsymbol{W}
\in\mathcal{W}}\lambda_{\max}(\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W}))}$.
Then

$$ $\displaystyle\mathbb{P}\left(\forall n\geq 1,\boldsymbol{x}\in\mathcal{X}:d_{ \mathrm{TV}}(\boldsymbol{s}_{n}(\boldsymbol{x}),\boldsymbol{s}^{\star}( \boldsymbol{x}))\leq\beta_{n}(\delta)\sigma_{n}(\boldsymbol{x})\right)\geq 1-\delta$ $$

where $d_{\mathrm{TV}}\left(\boldsymbol{s},\boldsymbol{s^{\prime}}\right)
\operatorname*{\,\dot{=}\,}\frac{1}{2}\sum_{i}|s_{i}-s^{\prime}_{i}|$ is the total variation distance.

We use $\sigma_{n}(\boldsymbol{x})$ as a proxy to the *uncertainty about the response to $\boldsymbol{x}$* after having fine-tuned on the selected data $D_{n}$, since it directly governs the size of the confidence sets around our current estimate of response probabilities.
This uncertainty is a key quantity not just in classification:
In Section K.5, we state analogous confidence sets for regression with the standard squared error loss, building on results by Abbasi-Yadkori () and Chowdhury & Gopalan ().

###### Theorem 3.2 (Confidence Sets) .

#### The Close Relationship of Regularized Loss Minimization and Test-Time Fine-Tuning

Recall that test-time fine-tuning does not solve the regularized objective of Equation 1, but instead takes a single gradient step.
So why do we expect the surrogate model $\boldsymbol{f}(\cdot;\boldsymbol{W}_{\!n})$ be closely related to the fine-tuned $\boldsymbol{f}^{\mathrm{pre}}$?
To answer this question, we contrast two alternative models:

- •
$\boldsymbol{W}_{\!\lambda}\operatorname*{\,\dot{=}\,}\operatorname*{arg\,min}_
{\boldsymbol{W}}\mathcal{L}^{\lambda}(\boldsymbol{W})$, *(minimizer of regularized loss)*
- •
$\boldsymbol{\widehat{W}}_{\!\eta}\operatorname*{\,\dot{=}\,}{\boldsymbol{W}^{
\mathrm{pre}}}-\eta\boldsymbol{\nabla}\mathcal{L}({\boldsymbol{W}^{\mathrm{pre
}}})$ with any step size $\eta>0$, *(single gradient-step fine-tuning)*

where we keep the dataset $D$ fixed and omit the dependency on $D$.
Our following proposition shows that both models are close if the loss landscape is relatively smooth and for careful choice of $\lambda\approx\frac{1}{\eta}$.

###### Proposition 3.3 .

It holds that $\|\boldsymbol{W}_{\!1/\eta}-\boldsymbol{\widehat{W}}_{\!\eta}\,\|_{\mathrm{F}}
\leq\eta\,\|\boldsymbol{\nabla}\mathcal{L}(\boldsymbol{W}_{\!1/\eta})-
\boldsymbol{\nabla}\mathcal{L}({\boldsymbol{W}^{\mathrm{pre}}})\|_{\mathrm{F}}$.

Recent works have also observed $\boldsymbol{W}_{\!1/\eta}\approx\boldsymbol{\widehat{W}}_{\!\eta}$ empirically (Ali et al., ; ).
Intuitively, with a larger step size, $\smash{\boldsymbol{\widehat{W}}_{\!\eta}}$ is farther away from ${\boldsymbol{W}^{\mathrm{pre}}}$, and hence corresponds to the regularized estimate with less regularization.
This connection between regularized loss minimization and test-time fine-tuning is closely linked to the tight connection between regularization and early stopping (Morgan & Bourlard, ; Yao et al., ; Li et al., ).
We will use this connection in the following to derive SIFT in the context of fine-tuning.

###### Proposition 3.3 .

## 4 SIFT : Efficiently Reducing Uncertainty about the Response

We introduce SIFT, an algorithm for selecting data for fine-tuning that effectively reduces the uncertainty about the response to the prompt $\boldsymbol{x}^{\star}\in\mathcal{X}$.
Note that we can compute the uncertainty $\sigma_{X}(\boldsymbol{x}^{\star})$ about the response to the prompt $\boldsymbol{x}^{\star}$ for any selected data ${X\subseteq\mathcal{D}}$ in closed-form, since its definition (cf. Equation 2) depends only on the selected inputs $X$. SIFT minimizes this uncertainty about $\boldsymbol{x}^{\star}$:

$$ $\displaystyle\boldsymbol{x}_{n+1}$ $\displaystyle\operatorname*{\,\dot{=}\,}\operatorname*{arg\,min}_{\boldsymbol{ x}\in\mathcal{D}}\sigma_{X_{n}\cup\{\boldsymbol{x}\}}^{2}(\boldsymbol{x}^{ \star})=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\boldsymbol{k}_ {X_{n}\cup\{\boldsymbol{x}\}}^{\top}(\boldsymbol{x}^{\star})(\boldsymbol{K}_{ \!X_{n}\cup\{\boldsymbol{x}\}}+\lambda^{\prime}\boldsymbol{I}_{n+1})^{-1} \boldsymbol{k}_{X_{n}\cup\{\boldsymbol{x}\}}(\boldsymbol{x}^{\star}).$ ( SIFT ( $\lambda^{\prime}$ )) $$

SIFT selects data that minimizes a bound on the approximation error of the surrogate model, and then fine-tunes the full LLM using this data.
We discuss the design choices, including the choice of embeddings, that make SIFT efficient in §4.2.
In §C.1, we illustrate with an example of how SIFT balances relevance and diversity, where we also see that the free parameter $\lambda^{\prime}=\lambda\kappa$ controls this trade-off.
Larger $\lambda^{\prime}$ emphasize relevance of selected data, while smaller $\lambda^{\prime}$ emphasize diversity.
Probabilistically, SIFT can be interpreted as maximizing the information gain of the selected data $X_{n}$ on the response to the prompt $\boldsymbol{x}^{\star}$, that is, $\boldsymbol{x}_{n+1}=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}
\mathrm{I}\!\left(f(\boldsymbol{x}^{\star});y(\boldsymbol{x})\ \middle|\ y_{1:
n}\right)$ where ${y(\boldsymbol{x})}$ denotes a noisy observation of the response to $\boldsymbol{x}$.
We formally introduce this interpretation of SIFT in §G.

### 4.1 Uncertainty Provably Vanishes

We prove that unlike with Nearest Neighbor retrieval, the uncertainty about the response to the prompt vanishes if SIFT is used to select data for fine-tuning.
We give an informal overview here, and defer the formal treatment to §C.2.
Our theoretical analysis shows that test-time fine-tuning can fully reduce uncertainty only if the data space contains sufficient information to determine the correct response.
If the data space does not contain all relevant information, the remaining uncertainty is quantified by the limiting uncertainty after seeing “all data in the data space infinitely often”, which we call the *irreducible uncertainty* and denote by $\sigma_{\infty}(\boldsymbol{x}^{\star})$.
We provide the formal definition in §C.2, but intuitively, the irreducible uncertainty is the largest quantity satisfying $\sigma_{X}(\boldsymbol{x}^{\star})\geq\sigma_{\infty}(\boldsymbol{x}^{\star})$ for all $X\subseteq\mathcal{D}$.
We then specialize the result of Hübotter et al. () to show that the uncertainty about the response to the prompt shrinks at the rate $\smash{\widetilde{O}\left(1/\sqrt{n}\right)}$ until it reaches the irreducible uncertainty:

###### Informal Theorem 4.1 (Convergence Guarantee) .

Fix any $\lambda^{\prime}>0$ and let SIFT($\lambda^{\prime}$) select $X_{n}$ from the data space $\mathcal{D}$.
Then for all $n\geq 1$ and $\boldsymbol{x}^{\star}\in\mathcal{X}$,

$$ $\displaystyle\sigma_{n}^{2}(\boldsymbol{x}^{\star})-\sigma_{\infty}^{2}( \boldsymbol{x}^{\star})\leq\frac{O\left(\lambda^{\prime}\log(n)\right)}{\sqrt{ n}}.$ $$

Naturally, convergence is slower with a larger regularization parameter / smaller step size.
Notably, the irreducible uncertainty depends on the data space.
With a large and diverse data space, the irreducible uncertainty is typically negligible.
This statistical guarantee is a key property of SIFT.
As we show in Proposition K.1, Nearest Neighbor retrieval fails to satisfy a guarantee of this kind.

###### Informal Theorem 4.1 (Convergence Guarantee) .

### 4.2 Compute-Efficient Data Selection

We have established how to select informative data for fine-tuning.
Next to good statistical efficiency, good computational efficiency is key for selecting data at test-time.
In the following, we describe design choices such that SIFT has negligible overhead compared to Nearest Neighbor retrieval.

#### Sequence-Level Selection

In the self-supervised paradigm, each sequence of tokens $\boldsymbol{x}\in\mathcal{D}$ corresponds to a dataset of next-token predictions $x_{1:k}\mapsto x_{k+1}$.
Rather than selecting individual next-token predictions from the data space of all sub-sequences $x_{1:k}$, we select full sequences $\boldsymbol{x}$ from the significantly smaller data space $\mathcal{D}$, then fine-tune for a single gradient step on each sub-sequence within $\boldsymbol{x}$.
This is a common practice in prior works that use Nearest Neighbor retrieval for data selection (e.g., Xia et al., ; Hardt & Sun, ).

#### Surrogate Sequence Embedders

We use a surrogate sequence embedding model to generate embeddings of the data space and prompts.
We use the same embedding model as Hardt & Sun () which is a large Roberta model (Liu, ) with 355M parameters that was fine-tuned for one pass on the Pile training set.
The embedding dimension is $1024$.
Unlike Hardt & Sun (), we additionally normalize the embeddings to unit length, the reasons for which we discuss in §D.

We obtain decent performance with this surrogate model.
Nevertheless, our theoretical results indicate that using embeddings extracted from the LLM to be fine-tuned could further improve the performance of SIFT.
Empirical neural tangent embeddings (Wei et al., ; Holzmüller et al., ) and influence function embeddings (Xia et al., ) can be implemented efficiently and offer alternative latent spaces capturing the pre-trained model.
We hypothesize that the decent performance of the surrogate model is explained by the similarity of emergent latent spaces of language models that were trained on similar data.

Figure: Figure 4: The (multiplicative) computational overhead of SIFT compared to Nearest Neighbor retrieval is minimal. The compute overhead with a 1k data space is less than $1.05\times$.
Refer to caption: x2.png

#### Efficient Implementation of SIFT

In our experiments, we pre-select $200$ candidates via Nearest Neighbor retrieval with Faiss (Johnson et al., ) and then apply SIFT to select $50$ sequences from this smaller data space.
On the Pile dataset, we find that performance can be increased further by pre-selecting more candidates (cf. Figure 18 in §H) but the marginal gains diminish.
The precise performance benefit of pre-selecting more candidates may differ on other datasets.
We describe in §H how SIFT can be solved iteratively without computing the inverse in every iteration.
When a matrix of the size of the pre-selected data space fits in GPU memory, we find that SIFT has a negligible computational overhead compared to Nearest Neighbor retrieval.
We report results with an NVIDIA RTX 4090 GPU in Figure 4.(^1^11We use the client-server architecture described by Hardt & Sun () with CPU-only servers.)
While our main implementation of SIFT is fast if the data space is small, it does not scale linearly with the size of the data space $K$.
In §H, we describe how a priority queue can be used to achieve an almost-linear runtime in $K$.

## 5 Results

We focus on language modeling with causal language models.
Following Hardt & Sun (), we fine-tune a pre-trained LLM for a single gradient step each on $N=50$ selected data points in the order that they are selected, most to least relevant.
We use the Pile dataset (Gao et al., ) for evaluation, restricting our use to data which is obtained and used in compliance with the terms of service of the data host.
This version of the Pile contains a diverse set of 17 high-quality sub-datasets, ranging from Q&A to code, scientific publications, math, and more.
Concretely, we use the Pile training set containing 210M sequences of total size 1.3TB as data space for data selection, and we evaluate on the Pile test set.(^2^22We evaluate on 1% of the test set (0.1% with Phi-3), corresponding to $1`812$ sequences.)
We report the *bits per byte* metric as recommended by Gao et al. (), which is proportional to the negative log-likelihood loss normalized by a dataset-specific constant.
Error bars correspond to 90% confidence intervals computed via bootstrapping with $1`000$ samples.

#### Base Models and Baselines

We evaluate the GPT-2 model (Radford et al., ) with 124M parameters also evaluated by Hardt & Sun (), with the default learning rate of the transformers library (Wolf et al., ).
We obtain analogous results with GPT-2-large (774M parameters) and the state-of-the-art Phi-3 (3.8B parameters, Abdin et al., ).(^3^33We detail hyperparameter choices for larger models in §I.)
With Phi-3, we use Low-Rank Adaptation (LoRA, Hu et al., ), fine-tuning slightly less than 1% of the model’s total parameters. We compare SIFT with $\lambda^{\prime}=0.1$ to Nearest Neighbor retrieval (NN) and the failure mode of Nearest Neighbor retrieval that repeatedly selects the closest neighbor.
The failure mode of Nearest Neighbor retrieval (NN-F) corresponds to an extreme case of redundancy in the data space which we suspect to be a realistic scenario in larger or less curated datasets.
Finally, we compare to Uncertainty Sampling (US), which is a widely used active learning strategy (Lewis, ; Settles, ) that selects the data with the highest uncertainty in the model’s response by selecting according to $\smash{\boldsymbol{x}_{n+1}=\operatorname*{arg\,max}_{\boldsymbol{x}\in
\mathcal{D}}\sigma_{n}^{2}(\boldsymbol{x})}$.
We compare to the heuristic that uses US to choose from the $200$ nearest neighbors, in which case US can be understood as finding a diverse cover of this pre-selected data space (see, e.g., Holzmüller et al., ; Kirsch et al., ).
In contrast, SIFT *minimizes* the uncertainty in the model’s response to the prompt $\boldsymbol{x}^{\star}$, leading to a “denser” cover close to $\boldsymbol{x}^{\star}$ and a “coarser” cover further away from $\boldsymbol{x}^{\star}$ (cf. Figure 2).

Figure: Figure 5: Bits per byte (in % relative to the base model, $\downarrow$ better) after $50$ test-time iterations. Left: Performance gains of SIFT are consistent across models. The failure-mode of Nearest Neighbor consistently performs worse than the base model. Tables 4 and 5 in §F detail our results with GPT-2-large and Phi-3 analogously to Table 1. Right: Most choices of $\lambda^{\prime}$ lead to comparable performance. With $\lambda^{\prime}\to\infty$, SIFT($\lambda^{\prime}$) repeatedly selects the nearest neighbor.
Refer to caption: x3.png

**Table 1: Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. We only include datasets with at least $10$ examples in our test set. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of SIFT over the strongest baseline.**
|  | US | NN | NN-F | SIFT | $\Delta$ |
| --- | --- | --- | --- | --- | --- |
| NIH Grants | 93.1 (1.1) | 84.9 (2.1) | 91.6 (16.7) | 52.9 (9.0) | $\downarrow\!32.0$ |
| US Patents | 85.6 (1.5) | 80.3 (1.9) | 108.8 (6.6) | 62.2 (3.6) | $\downarrow\!18.1$ |
| GitHub | 45.6 (2.2) | 42.1 (2.0) | 53.2 (4.0) | 28.6 (2.2) | $\downarrow\!13.5$ |
| Enron Emails | 68.6 (9.8) | 64.4 (10.1) | 91.6 (20.6) | 52.4 (11.8) | $\downarrow\!12.0$ |
| Common Crawl | 92.6 (0.4) | 90.4 (0.5) | 148.8 (1.5) | 86.1 (0.9) | $\downarrow\!4.3$ |
| ArXiv | 85.4 (1.2) | 85.0 (1.6) | 166.8 (6.4) | 81.6 (1.9) | $\downarrow\!3.4$ |
| Wikipedia | 67.5 (1.9) | 66.3 (2.0) | 121.2 (3.5) | 63.7 (2.1) | $\downarrow\!2.6$ |
| PubMed Abstr. | 88.9 (0.3) | 87.2 (0.4) | 162.6 (1.3) | 84.8 (0.7) | $\downarrow\!2.4$ |
| Hacker News | 80.4 (2.5) | 79.2 (2.8) | 133.1 (6.3) | 77.8 (3.5) | $\downarrow\!1.4$ |
| Stack Exchange | 78.6 (0.7) | 78.2 (0.7) | 141.9 (1.5) | 77.0 (0.7) | $\downarrow\!1.2$ |
| PubMed Central | 81.7 (2.6) | 81.7 (2.6) | 155.6 (5.1) | 80.6 (2.7) | $\downarrow\!1.1$ |
| DeepMind Math | 69.4 (2.1) | 69.6 (2.1) | 121.8 (3.1) | 70.1 (2.1) | $\uparrow\!0.7$ |
| FreeLaw | 63.9 (4.1) | 64.1 (4.0) | 122.4 (7.1) | 65.5 (4.2) | $\uparrow\!1.6$ |
| *All* | 80.2 (0.5) | 78.3 (0.5) | 133.3 (1.2) | 73.2 (0.7) | $\downarrow\!5.1$ |

#### Insight 1: SIFT consistently selects better data for fine-tuning than Nearest Neighbor retrieval.

We show in Figure 1 that SIFT outperforms NN and avoids its failure mode where the same data point is selected repeatedly.
In Figure 5 (left), we show that the performance gains of SIFT are consistent across models.
Table 1 compares the performance of SIFT against NN across all datasets of the Pile, using GPT-2 as base model.
Overall, we find that SIFT improves performance both on datasets where NN already performs well, such as GitHub, and on datasets where NN performs poorly, such as NIH Grants.
On all datasets of the Pile, SIFT performs at least as well as the strongest baseline (within margin of error), suggesting that it is a robust method for data selection.
We observe the trend that relative performance gains of SIFT over Nearest Neighbor retrieval *increase with model capability*.
That is, with stronger base models, informativeness of selected data appears to become more important.

#### Insight 2: SIFT is robust to the choice of λ ′ superscript 𝜆 ′ \lambda^{\prime} italic_λ start_POSTSUPERSCRIPT ′ end_POSTSUPERSCRIPT .

We evaluate SIFT with varying choices of $\lambda^{\prime}$, and summarize the results in Figure 5 (right).
We include extended results in Table 11 of §J, showing that for all evaluated $\lambda^{\prime}$ between $1\mathrm{e-}8$ and $10$, SIFT performs at least on-par with Nearest Neighbor retrieval on *all* datasets of the Pile, often outperforming it.
This suggests that SIFT is robust to the choice of $\lambda^{\prime}$.
Nevertheless, there may be an advantage to adaptively tuning $\lambda^{\prime}$ (e.g., via cross-validation).
In particular, choosing the best $\lambda^{\prime}$ for each dataset, SIFT would outperform NN on every dataset.

#### Insight 3: SIFT selects data the “right” number of times.

Nearest Neighbor retrieval implicitly relies on non-redundancy within the data space to not select duplicate information, as illustrated in the example of Figure 3.
This is almost never the case in practice, and in the extreme case of duplicate data, Nearest Neighbor selects the same data point repeatedly.
SIFT does not rely on excluding previously selected data points.
Instead, SIFT may select the same data point any number of times, adaptively taking more than one gradient step on it, if beneficial.
To ensure that the selected data is maximally informative, SIFT takes into account the redundancy of data points explicitly.
This makes SIFT robust to information duplication by design.

Figure: Figure 6: Bits per byte (in % relative to NN / NN-F, $\downarrow$ better) after $50$ test-time iterations. Error bars correspond to standard errors. The left bars measure the performance gain over all of the Pile. The middle and right bars measure the performance gain for all prompts where SIFT selects $\#$ unique points.
Refer to caption: x5.png

We illustrate this in Figure 6 where we evaluate the performance gain of SIFT over Nearest Neighbor and its failure mode.
As expected, we find that on all test prompts where SIFT selects many unique points, SIFT outperforms repeatedly selecting the closest neighbor by a large margin.
Interestingly, we also find that on all test prompts where SIFT selects only a single point, SIFT outperforms Nearest Neighbor by a large margin.
This suggests that in some cases repeatedly taking gradient steps on the closest neighbor is beneficial, and SIFT identifies these cases.

#### Insight 4: Test-time fine-tuning can significantly improve language modeling ability.

Our results from Figure 7 indicate that test-time fine-tuning improves the performance of the base LLM substantially, surprisingly, even with a state-of-the-art model such as Phi-3.
Our Phi-3 with test-time fine-tuning and SIFT achieves $0.595$ bits per byte, outperforming the previous leader in the [Pile language modeling benchmark](https://paperswithcode.com/sota/language-modelling-on-the-pile), a $30\times$ larger model.(^4^44We compare to prior work in the Pile language modeling benchmark in Table 2 of §A.)
We compare test-time fine-tuning to the common in-context learning, where we include as much of the data as possible into the context window of the test instance, in addition to its original context, by concatenating text in order of selection.
While in-context learning tends to improve the performance of the base model, we find that fine-tuning at test-time systematically outperforms or performs on-par with in-context learning. Furthermore, the compute cost of in-context learning grows quadratically with the context window size, meaning that including long texts within large context windows is expensive.
Surprisingly, even at our small context windows of 4k tokens, fine-tuning is more efficient.
Remarkably, test-time fine-tuning consistently outperforms in-context learning by more than $25\%$ on math and coding, tasks that require more complex reasoning (§F).

Figure: Figure 7: Left: Bits per byte ($\downarrow$ better), comparing fine-tuning and in-context learning with $50$ test-time examples selected by SIFT. We find that fine-tuning systematically outperforms or performs on-par with in-context learning, even when fine-tuning only a LoRA adapter as with Phi-3. Test-time fine-tuning with Phi-3 (3.8B) surpasses the performance of the $5\times$ larger Phi-3 (14B). Right: Number of processed tokens per second ($\uparrow$ better) on an NVIDIA RTX 4090. The token throughput of fine-tuning is orders of magnitude larger than that of in-context learning. §I details our estimation of in-context learning’s compute cost. We report the total test-time compute in Figure 17 of §E.
Refer to caption: x6.png

#### Further Insights

In §D, we discuss additional findings on active fine-tuning such as that the performance gains of SIFT over Nearest Neighbor retrieval *grow with dataset size*, and that normalizing embeddings is important for the effectiveness of data selection.
In §E, we discuss additional findings on test-time fine-tuning, for example, the trend that *larger models learn faster at test-time*.

## 6 Compute-Proportional Test-Time Fine-Tuning

We have shown that test-time fine-tuning can improve language modeling ability and that SIFT is a robust method for data selection, outperforming Nearest Neighbor retrieval.
However, a key shortcoming of previous approaches to test-time fine-tuning is that they spend a fixed amount of test-time compute, regardless of the nature of the prompt, the available data, or the model.
This is not computationally scalable in many practical applications, since a fixed test-time compute budget leads to non-proportionate performance gains.
For example, for the prompt “Hello” to a chatbot we would not like to spend any test-time compute, while for a more complex prompt we would like to spend more compute.
In this section, we evaluate whether uncertainty estimates can be used to adaptively stop test-time fine-tuning such that the realized performance gain is proportional to the compute used.

Figure: Figure 8: Left: We visualize the empirical density of the uncertainty estimates $\hat{\sigma}_{n}$ wrt. the bits per byte $\mathrm{bpb}_{n}$. Brighter colors indicate higher density on a logarithmic scale. We observe a strong linear relationship between uncertainty estimates and bits per byte. Middle: We construct a “reliability diagram” of uncertainty estimates. Notably, since we evaluate with respect to bits per byte rather than an accuracy, canonical calibration plots are not applicable. In particular, it is well known that bits per byte do not go to zero for perfect models due to irreducible *aleatoric* uncertainty, which is not captured by our *epistemic* uncertainty estimates. Nevertheless, we observe that our epistemic uncertainty estimates are predictive of the model’s performance. The red line indicates a linear fit. Right: We visualize the bits per byte (in % relative to the base model, $\downarrow$ better) of all prompts whose model is fine-tuned at a given iteration. We find that by adaptively stopping with respect to the known uncertainties $\sigma_{n}$, we can spend test-time compute proportional to realized performance gains (see also Figure 27 in §J). *Remarks:* Results are with GPT-2. In the left and middle plots, we remove the lowest and highest 0.25% of uncertainty estimates (i.e., the outliers) for better visualization. In the left plot, we additionally remove the lowest and highest 0.25% of bits per byte.
Refer to caption: x7.png

#### Insight 5: The response uncertainty can predict performance gain.

We find that $\sigma_{n}(\boldsymbol{x}^{\star})$ is monotonically and linearly correlated at coefficient $\approx 0.4$ with the model error after $n$ test-time iterations, i.e., the bits per byte $\mathrm{bpb}_{n}(\boldsymbol{x}^{\star})$.
This is remarkable because $\sigma_{n}$ contains information only from the surrogate embedding model, and is normalized such that $\sigma_{0}(\boldsymbol{x}^{\star})=1$.
To determine the importance of the base model, we also evaluate the denormalized uncertainty estimate $\hat{\sigma}_{n}(\boldsymbol{x}^{\star})\operatorname*{\,\dot{=}\,}\sigma_{n}(
\boldsymbol{x}^{\star})\cdot\mathrm{bpb}_{0}(\boldsymbol{x}^{\star})$, which unlike $\sigma_{n}$ cannot be evaluated at test-time.
We multiply $\sigma_{n}$ by $\mathrm{bpb}_{0}$ to ensure that the uncertainty measure is in the same units as the performance metric, correcting for the use of normalized surrogate embeddings.
We find that $\hat{\sigma}_{n}(\boldsymbol{x}^{\star})$ is strongly correlated at coefficient $\gtrapprox 0.5$ with the bits per byte. We summarize correlations in Table 12 of §J and visualize the predictive capability of $\hat{\sigma}_{n}$ in Figure 8 (left) and Figure 8 (middle).
Our findings indicate that approximations of the base model’s uncertainty, before test-time fine-tuning, can be beneficial.
In future work, we intend to determine whether generating embeddings from the base model can provide such scale-correction.

Recall that SIFT minimizes the response uncertainty $\sigma_{n}$ to the given prompt.
The predictive ability of uncertainty estimates provides an intuitive explanation for the effectiveness of SIFT.

#### Compute-Proportional Performance Gains: Early stopping at the “right” time.

Motivated by the predictive power of uncertainty estimates, we evaluate whether they can be used to *adaptively stop* test-time fine-tuning such that the realized performance gain is proportional to the compute used.
In the following, we propose a such a stopping criterion for SIFT.
Using the approximation of the error via uncertainty estimates discussed above and that $\sigma_{0}(\boldsymbol{x}^{\star})=1$:

$$ $\displaystyle\text{performance gain}=\frac{\mathrm{bpb}_{0}(\boldsymbol{x}^{ \star})}{\mathrm{bpb}_{n}(\boldsymbol{x}^{\star})}\approx\frac{\sigma_{0}( \boldsymbol{x}^{\star})}{\sigma_{n}(\boldsymbol{x}^{\star})}=\frac{1}{\sigma_{ n}(\boldsymbol{x}^{\star})}.$ (4) $$

We would like to stop fine-tuning when further test-time compute does not yield proportional performance gain, i.e., when “${\text{performance gain}<\alpha\cdot n}$” with $n$ approximating the compute of $n$ iterations and $\alpha$ a constant comparing the units of compute and performance.
Plugging in our above approximation of the performance gain, we propose to stop test-time fine-tuning *before* iteration $n$ if

$$ $\displaystyle\sigma_{n}(\boldsymbol{x}^{\star})>(\alpha n)^{-1}.$ ( Adaptive SIFT ) $$

Intuitively, this stops fine-tuning the LLM when its progress in crafting a better response stalls.
For complex prompts that benefit from fine-tuning, Adaptive SIFT spends more test-time compute, whereas for prompts where the model is already strong or where the data space is not informative, Adaptive SIFT spends less test-time compute.
Figure 8 (right) shows that the performance gains of this approach are proportional to the compute used.

#### Towards Scaling Laws of Test-Time Fine-Tuning

Interestingly, our results bear resemblance to scaling laws of LLM pre-training (Kaplan et al., ; Henighan et al., ; Hoffmann et al., ).
These scaling laws express the performance of a model as a function of the compute used for pre-training (e.g., the number of parameters or training tokens).
Such scaling laws are crucial for determining how to optimally spend a fixed amount of compute.
Recently, scaling laws for “test-time inference” have gained attention, where test-time compute is usually spent on search (e.g., beam search) with a variable number of forward passes of a few-shot prompted base LLM (Brown et al., ; Snell et al., ).
Our results suggest that similar scaling laws exist for test-time fine-tuning, expressing the performance of a model as a function of the compute used for fine-tuning at test-time.
Such scaling laws can be an important tool to determine how to spend test-time compute.
There are many open questions in this direction, which we do not address in this work.
For example, how does model size affect the scaling laws of test-time fine-tuning? Or, can a model be fine-tuned at test-time to build reasoning chains?
Based on our results and previous evaluations of fine-tuning and in-context learning (e.g., Hu et al., ; Mosbach et al., ; Hardt & Sun, ), we conjecture that test-time fine-tuning may lead to a more efficient use of compute than repeatedly prompting a base LLM.
We believe that these open questions are exciting directions for future work.

## 7 Discussion and Future Work

We propose a data selection algorithm, SIFT, unifying ideas from retrieval and active learning.
SIFT estimates the uncertainty about the response to a given prompt after having been fine-tuned on some data (§3), and then selects the data that minimizes this uncertainty (§4).
This addresses the limitatzions of Nearest Neighbor retrieval (§2).
SIFT can be seen as a generalization of Nearest Neighbor retrieval from a search method to a learning method, which ensures explicitly that the retrieved data is maximally informative.
We show on the Pile dataset that SIFT consistently outperforms Nearest Neighbor retrieval in prompt-specific fine-tuning at test-time and that this kind of local learning can be more effective than locally learning from examples in-context (§5).
Finally, we observe that our uncertainty estimates can predict the performance gain of test-time fine-tuning, and use this to develop an adaptive algorithm which achieves compute-proportional performance gains (§6).

Test-time fine-tuning addresses a fundamental limitation of in-context learning, namely that in-context learning is typically limited to a fixed and finite context window.
In contrast, test-time fine-tuning allows the LLM to dynamically and effectively access a potentially unbounded non-parametric memory.
By improving the effectiveness of test-time fine-tuning, this work opens up several exciting directions for future research.
Test-time fine-tuning may be used to ground the model on a trusted dataset, mitigate biases against under-represented groups in the training data, or to dynamically include private data depending on user privileges.
Particularly interesting would be a broad evaluation on non-perplexity tasks such as code generation or in the life sciences with large-scale medical or protein data.
Unlike few-shot in-context learning which is limited in scope to autoregressive models, test-time fine-tuning and SIFT may be extended to other model classes such as diffusion models.
Furthermore, SIFT may be used effectively in other settings that require automatic data selection, such as targeted instruction tuning during post-training of LLMs.
Finally, our results suggest scaling laws for test-time fine-tuning and we outline several exciting open questions (§6).

## Contributions

JH conceived and led the project, being involved in all its components and leading the theory, implementation of the SIFT algorithm, design of experiments, and writing.
SB set up and ran the first experiments validating the approach, and contributed to running the final ablation studies.
IH ran additional experiments, especially those with larger models, and optimized the code.
AK advised.

## Acknowledgements

We would like to thank Armin Lederer, Vignesh Ram Somnath, Bhavya Sukhija, Scott Sussex, and Lenart Treven for feedback on early versions of the paper.
This project was supported in part by the European Research Council (ERC) under the European Union’s Horizon 2020 research and Innovation Program Grant agreement no. 815943, and the Swiss National Science Foundation under NCCR Automation, grant agreement 51NF40 180545.
Ido Hakimi was supported by an ETH AI Center Postdoctoral fellowship.

## Appendices

## Contents

0[2]

## Appendix A Comparison to the State-of-the-Art on the Pile Language Modeling Benchmark

We compare our results to the state-of-the-art in the [Pile language modeling benchmark](https://paperswithcode.com/sota/language-modelling-on-the-pile), GLM-130B (Zeng et al., ), Jurrassic-1 (Lieber et al., ), and GPT-3 (Brown et al., ), which are 130B, 178B, and 175B parameter models, respectively.
Our results are with the much smaller Phi-3 (3.8B parameters) and test-time fine-tuning with SIFT.

Due to our dataset to the non-copyrighted part of the Pile, the data distribution changes slightly.
To account for this, we take the reported results of prior work and exclude the datasets that have copyright restrictions from the evaluation.
We report results in Table 2.
Notably, prior reported results of state-of-the-art models miss evaluation of the Wikipedia dataset, which we therefore also exclude for a direct comparison.
To the best of our knowledge, our results achieve a new state-of-the-art on the Pile benchmark.

**Table 2: Evaluation of state-of-the-art models on the Pile language modeling benchmark, without copyrighted datasets. Results with GPT-3 are from Gao et al. (). Results with Jurassic-1 and GLM are from Zeng et al. () and do not report on the Wikipedia dataset. For a complete comparison, we also evaluate our Phi-3 with test-time fine-tuning when excluding the Wikipedia dataset.**
| Model | Bits per Byte | Bits per Byte (without Wikipedia) |
| --- | --- | --- |
| Jurassic-1 (178B, Lieber et al., ) | n/a | 0.601 |
| GLM (130B, Zeng et al., ) | n/a | 0.622 |
| GPT-2 (124M, Radford et al., ) | 1.241 |  |
| GPT-2 (774M, Radford et al., ) | 1.093 |  |
| Llama-3.2-Instruct (1B) | 0.807 |  |
| Llama-3.2-Instruct (3B) | 0.737 |  |
| Gemma-2 (2B, Team et al., ) | 0.721 |  |
| Llama-3.2 (1B) | 0.697 |  |
| Phi-3 (3.8B, Abdin et al., ) | 0.679 | 0.678 |
| Phi-3 (7B, Abdin et al., ) | 0.678 |  |
| Gemma-2 (9B, Team et al., ) | 0.670 |  |
| GPT-3 (175B, Brown et al., ) | 0.666 |  |
| Phi-3 (14B, Abdin et al., ) | 0.651 |  |
| Llama-3.2 (3B) | 0.640 |  |
| Gemma-2 (27B, Team et al., ) | 0.629 |  |
| *Test-Time FT with* SIFT + GPT-2 (124M) | 0.862 |  |
| *Test-Time FT with* SIFT + GPT-2 (774M) | 0.762 |  |
| *Test-Time FT with* SIFT + Phi-3 (3.8B) | 0.595 | 0.599 |

## Appendix B Extended Related Work

### B.1 Learning at Test-Time

The subject of learning at test-time has a rich history in statistics and machine learning.
By “learning at test-time” we refer to models that are constructed specifically for a given test instance, differing from the model used for other test instances.
The following discussion provides a brief overview with emphasis on the most recent developments.

#### k 𝑘 k italic_k -Nearest Neighbors (since 1950s)

One of the most basic forms of learning at test-time was developed by Fix () and Cover & Hart ().
Given the supervised data $\mathcal{D}\subseteq\mathcal{X}\times\mathcal{Y}$ with input domain $\mathcal{X}\subseteq\mathbb{R}^{d}$ and labels $\mathcal{Y}=\{0,\dots,K\}$, the $k$-NN algorithm predicts the label of a test instance $\boldsymbol{x}^{\star}\in\mathcal{X}$ by taking the majority vote of the $k$ nearest neighbors of $\boldsymbol{x}^{\star}$ in $\mathcal{D}$ according to some distance metric on $\mathcal{X}$ such as Euclidean distance.
In the case of regression, $\mathcal{Y}=\mathbb{R}$ and the prediction is the average of the labels of the $k$ nearest neighbors.
This is a simple and often effective method if the inputs are well-structured and low-dimensional, e.g., if $\mathcal{X}$ is a learned low-dimensional manifold (Geirhos et al., ).
When $K$ is large, as for example when $\mathcal{Y}$ is the set of all tokens in a language modeling task, naïve application of $k$-NNs is difficult, nevertheless they have been shown to be effective when mixed with parametric language models (Khandelwal et al., ).

#### Local Learning (since 1970s)

Local learning is the idea of using data “relevant” to the test instance $\boldsymbol{x}^{\star}$ to train a parametric model.
Formally, given a test instance $\boldsymbol{x}^{\star}$, conventually a model $f$ is used to predict $f(\boldsymbol{x}^{\star})$ where $f$ is trained to minimize the average loss over the training data.
Instead, local learning trains a model $f_{\boldsymbol{x}^{\star}}$ specifically for $\boldsymbol{x}^{\star}$ and predicts $f_{\boldsymbol{x}^{\star}}(\boldsymbol{x}^{\star})$.
Original works train a linear model by weighting data according to their proximity to $\boldsymbol{x}^{\star}$ (Cleveland, ; Cleveland & Devlin, ; Atkeson et al., ).
Here, each test instance trains a model from scratch since the optimal solution of linear regression is independent of initialization.
This perspective has regained interest recently in the context of neural networks, with Sun et al. () naming it *“test-time training”*.

#### Transductive Learning (since 1990s)

Vladimir Vapnik developed the general principle of *transduction* which he states in Vapnik () as follows:

> Vladimir Vapnik: *“When solving a problem of interest, do not solve a more general problem as an intermediate step. Try to get the answer that you really need but not a more general one.”*

This is perhaps the most general principle behind learning at test-time, and directly opposed to the principle of *induction* — extracting the most general rules from data — which has arguably dominated machine learning research over the last decades.
In a way, local learning is pushing the principle of transduction to the opposite extreme: Each test instance defines its own learning problem, with the test instance alone being the target of prediction.

#### Local Fine-Tuning (since 1990s)

Bottou & Vapnik () were the first to use local learning in conjunction with a *pre-trained* parametric model.
They train (i.e., “fine-tune”) the last layer of a convolutional neural network for handwritten digit classification based on the nearest neighbors to the test instance in pixel space.
Very recently, Hardt & Sun () applied the same idea to language models, showing that local fine-tuning can significantly improve the performance of large language models on standard benchmarks.
Previously, this idea has also been evaluated by Li et al. () and Basu et al. ().
*“Test-time fine-tuning”* (as well as “active inference”) has frequently been used to refer to this approach of locally fine-tuning a pre-trained model.
Within the last few years, test-time fine-tuning has regained substantial interest in the context of self-supervised learning, where the pre-trained model is fine-tuned on the *test instance itself*.
Notable applications of this approach are in vision (Jain & Learned-Miller, ; Shocher et al., ; Luo et al., ; Sun et al., ; Wang et al., ) and in language modeling (Krause et al., ), where it is called *dynamic evaluation*.
As one would also naïvely expect, test-time fine-tuning yields the largest improvements when the prompt is not (well-) represented in the pre-training data, e.g., due to a distribution shift (Gandelsman et al., ; Hardt & Sun, ).
Notably, test-time fine-tuning is the central component of the state-of-the-art approaches to the ARC challenge (Chollet, ; Cole & Osman, ), a non-saturated benchmark which is intended to test reasoning capabilities based on “core knowledge” rather than mere memorization.

#### (Few-Shot) In-Context Learning (since 2020s)

Very recently, with the advent of large language models (LLMs), learning at test-time has regained interest.
Brown et al. () showed that GPT-3 can *learn in-context* from input-label pairs that are appended to the prompt, an emergent phenomenon of LLMs that has been widely studied since (Von Oswald et al., ; Kossen et al., ; Bhattamishra et al., ).
In contrast to standard in-weights learning, in-context learning requires no parameter updates.
Interestingly, in-context learning adopts the same paradigm as local learning wherein a model is adapted specifically for the test instance $\boldsymbol{x}^{\star}$, here by skewing the autoregressive distribution towards the data included in the prompt.
This is often combined with the automatic sourcing of nearest neighbors to $\boldsymbol{x}^{\star}$ in an external dataset, which is known as *“retrieval augmented generation”* (RAG, Lewis et al., ; Borgeaud et al., ), and is akin to the other methods of test-time learning discussed above.
A crucial difference between test-time fine-tuning and in-context learning appears to be that learning from context works by *changing the test instance* (Bhargava et al., ) whereas in-weights learning works by *changing the model*.
With small datasets, in-context learning is therefore often more computationally efficient than test-time fine-tuning, however this ceases to be the case when the dataset grows since the complexity of transformers grows quadratically in the number of context tokens whereas the complexity of test-time fine-tuning grows linearly.

#### Linear Representations and Interpretability

Linear representations akin to ours from Assumption 3.1 have been used extensively in interpretability research (e.g., Templeton et al., ) and recently been explored in the context of LLMs more broadly (e.g., Park et al., ).
For example, Ribeiro et al. () learns linear approximations of a more complex model (such as an LLM) locally around each test instance.
SIFT can be understood as a method to determine which data to use for learning such linear approximations to get the best-possible interpretable model.

### B.2 Data Selection

Clearly, the choice of data to learn from at test-time is crucial for predictive performance.
Selecting uninformative data can increase inference time or even degrade performance (see, e.g., Kolossov et al., ).
Today, datasets for fine-tuning are often hand-designed, however, this is not possible in a test-time setting.
Automatic data selection has a rich history in machine learning, studied extensively in *search*, *experimental design* (Chaloner & Verdinelli, ), and *active learning* (Settles, ).
The following attempts to give a brief overview of the most recent developments.

#### (Document) Retrieval (since 1970s)

Retrieval methods aim to search a dataset $\mathcal{D}$ for the most relevant data to a given query/prompt.
The most classical methods such as TF-IDF (Sparck Jones, ) and BM25 (Robertson et al., ) are based on keyword matching, and were developed alongside the first search engines.
Due to their reliance on “bags of words”, i.e., sets of one-hot-encoded word vectors, they are known as *sparse retrievers*.
An alternative idea is to select the data $\boldsymbol{x}$ that maximizes the likelihood of the query $\boldsymbol{x}^{\star}$ given the data, i.e., $\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}p(\boldsymbol{x}^{\star
}\mid\boldsymbol{x})$, known as *query likelihood retrievers* (Ponte & Croft, ; Wang et al., ).
Here, the conditional probability can be a non-parametric term frequency or a parametric language model.
More recently, due to significant advances in representation learning (Devlin et al., ; Reimers & Gurevych, ), dense retrievers have become popular (e.g., Lewis et al., ; Karpukhin et al., ; Borgeaud et al., ).
A *dense retriever* embeds dataset and query into a metric vector space, and retrieves the nearest neighbors to the query.
Standard vector-based search methods use cosine similarity or (equivalently(^5^55Here we assume that vectors are normalized to unit length, cf. Section K.2.)) Euclidean distance.
Recent advances in algorithms and implementation mean that (approximate) nearest neighbor retrieval can be performed efficiently with databases of billions or even trillions of tokens (e.g., Johnson et al., ; Aumüller et al., ).
The most common metric is cosine distance, which coincides with Euclidean distance when vectors are normalized to unit length.
Nearest neighbor retrieval has been the de-facto standard for data selection in RAG and local learning.(^6^66There is substantial literature that investigates selection of “informative” data for RAG (e.g., Ye et al., ).)

#### Influence Functions (since 1970s)

Influence functions measure the change in a model’s prediction when a single data point is removed from the training data.
First proposed by Cook () for linear regression, they have since been used extensively to *interpret* predictions (Koh & Liang, ; Pruthi et al., ).
Very recently, Xia et al. () applied influence functions to select data that leads to the largest (approximate) reduction in test-loss.
Concretely, using a first-order Taylor approximation of the loss $\ell$ and if the model at time $t$ is updated via stochastic gradient descent with step size $\eta_{t}$ on data $\boldsymbol{x}$, the loss reduction can be approximated as

$$ $\displaystyle\ell(\boldsymbol{x}^{\star};\boldsymbol{\theta}_{t+1})-\ell( \boldsymbol{x}^{\star};\boldsymbol{\theta}_{t})\approx-\eta_{t}\langle \boldsymbol{\nabla}_{\!\!\boldsymbol{\theta}}\,\ell(\boldsymbol{x};\boldsymbol {\theta}_{t}),\boldsymbol{\nabla}_{\!\!\boldsymbol{\theta}}\,\ell(\boldsymbol{ x}^{\star};\boldsymbol{\theta}_{t})\rangle.$ $$

That is, the data $\boldsymbol{x}$ whose loss gradient is most aligned with the loss gradient of the test instance $\boldsymbol{x}^{\star}$, can be expected to lead to the largest loss reduction.(^7^77Xia et al. () normalize embeddings before computing the inner product (thus, maximizing cosine similarity) to account for varying gradient norms depending on sequence lengths.)
Note that this simply leads to nearest neighbor retrieval in an embedding space informed by the model at time $t$.
A major limitation of using influence functions for data selection is that they implicitly assume that the influence of selected data adds linearly (i.e., two equally scored data points are expected to doubly improve the model performance, Xu & Kazantsev, , Section 3.2).
This assumption does quite obviously not hold in practice as seen, e.g., by simply duplicating data.
The same limitation applies to the related approach of *datamodels* (Ilyas et al., ).
A recent line of work aims to address this limitation by designing simulators that can be probed with datasets to estimate their effect on a prediction requiring less compute than training the full model (Guu et al., ), yet, this does not address the data selection problem as the space of possible datasets is exponentially large.

#### Coverage & Inductive Active Learning

Next we discuss an orthogonal line of work, which takes into account the interaction between selected data, but not the interaction of that data with respect to a test instance.
Roughly speaking classical active learning studies how to most effectively select data from a domain $\mathcal{X}$ for learning a model over this domain $\mathcal{X}$.
Intuitively, this task can be thought of as selecting a subset $X\subseteq\mathcal{X}$ of fixed size that captures the most “information” about the target function $f$.
As such, this task is of an *inductive nature*: we aim to extract general rules from the data that can be applied to unseen data later, without concrete specification of the unseen data.
Approaches to (inductive) active learning are broadly aiming to select *diverse* data that covers the data manifold in $\mathcal{X}$ well.
Methods include those that maximize the mutual distances between selected data (e.g., CoreSet (Sener & Savarese, ), BADGE (Ash et al., ), and ProbCover (Yehuda et al., )) with respect to a latent distance metric and those “uncertainty sampling” methods that select data that the model is most uncertain about (e.g., *D-optimal design* (Wynn, ) and BatchBALD (Kirsch et al., )).(^8^88Section 5.2 of Holzmüller et al. () provides a comprehensive overview.)
Both families of methods can be seen as determining some decent covering of the data manifold in $\mathcal{X}$.
In a probabilistic sense, uncertainty sampling can be seen to minimize the “posterior predictive entropy” in expectation over the observed data.

### B.3 SIFT Unifies Work on Retrieval and Work on Coverage

In this work, we make the following central observation:

> *Learning and prediction is not a search problem; it requires synthesizing non-redundant relevant information.*

Current means of automatic data selection fall on to two extreme ends of a spectrum:
Retrieval methods search for relevant data without ensuring that data is non-redundant.
As such, naïve application of search methods is insufficient for a learning task since those generally do not take “distinctiveness” into account (cf. Section 2.1).
In contrast, coverage methods select non-redundant data without ensuring that data is relevant.

#### Transductive Active Learning: Unifying retrieval & coverage

Transductive active learning (Hübotter et al., ) unifies approaches to search and coverage.
In this work, we propose SIFT, an approach to test-time transductive active learning (i.e., transductive active learning with a single prediction target), which extends previously proposed algorithms (MacKay, ; Seo et al., ; Yu et al., ; Hübotter et al., ).
Similar algorithmic ideas have recently been evaluated empirically in a variety of other settings (Kothawade et al., ; Wang et al., ; Kothawade et al., ; Smith et al., ) such as Bayesian optimization (Hübotter et al., ), multi-task reinforcement learning (Bagatella et al., ), and the amortized fine-tuning of neural networks (Hübotter et al., ).
SIFT aims to select data that is both relevant and non-redundant with respect to the already seen data, whereby the hyperparameter $\lambda^{\prime}$ controls the trade-off between relevance and redundancy.
Hübotter et al. () introduce extensions of SIFT to more than one prediction target, i.e., amortizing learning across multiple prompts.
They show that if the prediction targets include *all of $\mathcal{X}$*, then the method reduces to a form of *inductive active learning*.

## Appendix C Further Details on SIFT

### C.1 How SIFT Balances Relevance and Diversity

Let us look more closely at the points selected by SIFT.
We will assume here for ease of notation that embeddings have unit length.(^9^99See Section K.4 for the expressions with non-normalized embeddings.)
The first point selected by SIFT has the largest (absolute) cosine similarity to the prompt within the latent space:

$$ $\displaystyle\boldsymbol{x}_{1}=\operatorname*{arg\,min}_{\boldsymbol{x}\in \mathcal{D}}\sigma_{\{\boldsymbol{x}\}}^{2}(\boldsymbol{x}^{\star})= \operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\frac{(\boldsymbol{\phi }(\boldsymbol{x}^{\star})^{\top}\boldsymbol{\phi}(\boldsymbol{x}))^{2}}{1+ \lambda^{\prime}}=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\Big{ (}\underbrace{\measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star}, \boldsymbol{x})}_{\text{cosine similarity of $\boldsymbol{\phi}(\boldsymbol{x} ^{\star}),\boldsymbol{\phi}(\boldsymbol{x})$}}\Big{)}^{\!\!2}.$ ( 1st point ) $$

This recovers the standard approach of Nearest Neighbor retrieval with respect to cosine similarity, provided cosine similarities are non-negative.
However, we show next that selecting more than one point, SIFT not only considers the relevance with respect to the prompt $\boldsymbol{x}^{\star}$, but also the redundancy with respect to the already seen data $\boldsymbol{x}_{1}$.

$$ $\displaystyle\boldsymbol{x}_{2}=\operatorname*{arg\,min}_{\boldsymbol{x}\in \mathcal{D}}\sigma_{\{\boldsymbol{x}_{1},\boldsymbol{x}\}}^{2}(\boldsymbol{x}^ {\star})=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\begin{bmatrix }\measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x}_{1}) \\ \measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x})\end{ bmatrix}^{\top}\!\!\begin{bmatrix}1+\lambda^{\prime}&\measuredangle_{ \boldsymbol{\phi}}(\boldsymbol{x}_{1},\boldsymbol{x})\\ \measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}_{1},\boldsymbol{x})&1+ \lambda^{\prime}\end{bmatrix}^{-1}\!\!\begin{bmatrix}\measuredangle_{ \boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\\ \measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x})\end{ bmatrix}\!\!.$ ( 2nd point ) $$

To illustrate how SIFT balances relevance and diversity, we compare the value of observing $\boldsymbol{x}_{1}$ twice to observing a different $\boldsymbol{x}$ with cosine similarity $\measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}_{1},\boldsymbol{x})=0$.
We show in Section K.4 that SIFT($\lambda^{\prime}$) prefers $\boldsymbol{x}$ over $\boldsymbol{x}_{1}$ for selecting $\boldsymbol{x}_{2}$ *if and only if*

$$ $\displaystyle\measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star}, \boldsymbol{x})^{2}>\frac{\lambda^{\prime}}{2+\lambda^{\prime}}\measuredangle_ {\boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})^{2}$ $$

**Table 3: The effect of $\lambda$ and its other interpretations on diversity of selected data (as the parameter is increased).**
| Parameter | Relation | Div. |
| --- | --- | --- |
| regularization $\lambda$ | $\lambda$ | ${\color[rgb]{0.94921875,0.328125,0.35546875}\definecolor[named]{pgfstrokecolor<br>}{rgb}{0.94921875,0.328125,0.35546875}\boldsymbol{\downarrow}}$ |
| step size $\eta$ | $1/\eta$ | ${\color[rgb]{0.16796875,0.3125,0.66796875}\definecolor[named]{pgfstrokecolor}{<br>rgb}{0.16796875,0.3125,0.66796875}\boldsymbol{\uparrow}}$ |
| noise $\rho$ (cf. §G) | $\rho^{2}$ | ${\color[rgb]{0.94921875,0.328125,0.35546875}\definecolor[named]{pgfstrokecolor<br>}{rgb}{0.94921875,0.328125,0.35546875}\boldsymbol{\downarrow}}$ |

The hyperparameter $\lambda^{\prime}$ controls the trade-off between relevance and diversity: if $\lambda^{\prime}=1$ then even if $\boldsymbol{x}$ has one third the relevance of $\boldsymbol{x}_{1}$, it is still preferred.
As $\lambda^{\prime}\to\infty$, SIFT($\lambda^{\prime}$) performs retrieval by repeatedly selecting the same point; and as $\lambda^{\prime}\to 0$, SIFT($\lambda^{\prime}$) aims only to select the most diverse points.
We observe the same relationship empirically on the Pile dataset (cf. Figure 9 (left)).
Table 3 summarizes the effect of the regularization parameter $\lambda$ and its interpretations.

### C.2 The Uncertainty of SIFT Provably Vanishes

We now formally prove that unlike with Nearest Neighbor retrieval, the uncertainty $\sigma_{n}^{2}(\boldsymbol{x}^{\star})$ about the response to the prompt vanishes if SIFT is used to select data for fine-tuning.
As discussed in §4.1, this requires that the data space contains sufficient information to determine the correct response.
In general, there might be an irreducible error remaining.
We will denote a basis of the embeddings $\{\boldsymbol{\phi}(\boldsymbol{x}):\boldsymbol{x}\in\mathcal{D}\}$ within the data space $\mathcal{D}$ by ${\boldsymbol{\Phi}\in\mathbb{R}^{m\times d}}$ with size $m$ and dimension $d$, and we denote by $\boldsymbol{\Pi}_{\boldsymbol{\Phi}}$ its orthogonal projection onto the orthogonal complement of the span of $\boldsymbol{\Phi}$.
Hübotter et al. () show that for all $X\subseteq\mathcal{D}$,

$$ $\displaystyle\sigma_{X}^{2}(\boldsymbol{x}^{\star})\geq\left\|\boldsymbol{\phi }(\boldsymbol{x}^{\star})\right\|_{\boldsymbol{\Pi}_{\boldsymbol{\Phi}}}^{2}$ (5) $$

where ${\left\|\boldsymbol{v}\right\|_{\boldsymbol{A}}=\sqrt{\boldsymbol{v}^{\top}
\boldsymbol{A}\boldsymbol{v}}}$ denotes the Mahalanobis distance.
We call $\smash{\sigma_{\infty}^{2}(\boldsymbol{x}^{\star})\operatorname*{\,\dot{=}\,}
\left\|\boldsymbol{\phi}(\boldsymbol{x}^{\star})\right\|_{\boldsymbol{\Pi}_{
\boldsymbol{\Phi}}}^{2}}$ the *irreducible uncertainty* about $\boldsymbol{x}^{\star}$.
It can be seen that $\smash{\sigma_{\infty}^{2}(\boldsymbol{x}^{\parallel})=0}$ for all $\smash{\boldsymbol{x}^{\parallel}\in\mathcal{X}}$ with $\smash{\boldsymbol{\phi}(\boldsymbol{x}^{\parallel})\in\operatorname*{span}{
\boldsymbol{\Phi}}}$.
That is, the irreducible uncertainty is zero for points in the span of the data space.
In contrast, for points $\smash{\boldsymbol{x}^{\perp}}$ with $\boldsymbol{\phi}(\boldsymbol{x}^{\perp})\in(\operatorname*{span}{\boldsymbol{
\Phi}})^{\perp}$, the irreducible uncertainty equals the initial uncertainty: $\sigma_{\infty}^{2}(\boldsymbol{x}^{\perp})=\sigma_{0}^{2}(\boldsymbol{x}^{
\perp})$.
The irreducible uncertainty of any prompt $\boldsymbol{x}^{\star}$ can be computed by simple decomposition of $\boldsymbol{\phi}(\boldsymbol{x}^{\star})$ into parallel and orthogonal components.
Hence, if the data space is large and includes all relevant information to answer the prompt, the irreducible uncertainty is negligible.

We will denote the *uncertainty reduction* about the prompt $\boldsymbol{x}^{\star}$ achieved by fine-tuning on $X$ by $\smash{\psi_{\boldsymbol{x}^{\star}}(X)\operatorname*{\,\dot{=}\,}\sigma_{0}^{
2}(\boldsymbol{x}^{\star})-\sigma_{X}^{2}(\boldsymbol{x}^{\star})}$ and note that SIFT selects $\boldsymbol{x}_{n+1}=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}
\psi_{\boldsymbol{x}^{\star}}(X_{n}\cup\{\boldsymbol{x}\})$.
Stating the convergence guarantee of SIFT requires one straightforward assumption.

###### Assumption C.1 .

The uncertainty reduction $\psi_{\boldsymbol{x}^{\star}}(X)$ is submodular.

Figure: Figure 9: Left: The parameter $\lambda^{\prime}$ controls the trade-off between relevance and diversity of the selected data. As $\lambda^{\prime}\to\infty$, SIFT selects the same point repeatedly whereas as $\lambda^{\prime}\to 0$, SIFT selects a diverse set of points. Middle: The irreducible uncertainty of test prompts from the Pile given neighbors selected from fractions of the Pile training dataset in the data space. The irreducible uncertainty captures how much information is available, and decays quickly. Right: We empirically observe that $\psi_{\boldsymbol{x}^{\star}}$ is monotone submodular, i.e., its “marginal gains” decrease as the number of iterations increases. The shaded region denotes the standard deviation, gray lines are from $10$ randomly selected prompts.
Refer to caption: x8.png

Intuitively, Assumption C.1 states that the marginal uncertainty reduction achieved by adding a point to the selected data (i.e., the ‘marginal gain’) decreases as the size of the selected data increases, which is a common assumption in prior work.(^10^1010Similar assumptions have been made by Bogunovic et al. () and Kothawade et al. ().)
Formally Assumption C.1 is satisfied if, for all $\boldsymbol{x}\in\mathcal{D}$ and $X^{\prime}\subseteq X\subseteq\mathcal{D}$,

$$ $\displaystyle\Delta_{\boldsymbol{x}^{\star}}(\boldsymbol{x}\mid X^{\prime}) \geq\Delta_{\boldsymbol{x}^{\star}}(\boldsymbol{x}\mid X)$ (6) $$

where $\Delta_{\boldsymbol{x}^{\star}}(\boldsymbol{x}\mid X)\operatorname*{\,\dot{=}
\,}\psi_{\boldsymbol{x}^{\star}}(X\cup\{\boldsymbol{x}\})-\psi_{\boldsymbol{x}
^{\star}}(X)$ is the *marginal uncertainty reduction* of $\boldsymbol{x}$ given $X$.

Though theoretically this assumption may be violated by some instances (Hübotter et al., , Example C.8), we observe that it is satisfied in practice (cf. Figure 9 (right)).
Under this assumption, $\psi_{\boldsymbol{x}^{\star}}(X_{n})\geq(1-1/e)\max_{X\subseteq\mathcal{D},|X|
\leq n}\psi_{\boldsymbol{x}^{\star}}(X)$ due to the seminal result on monotone submodular function maximization of Nemhauser et al. ().
That is, the iterative scheme of SIFT achieves a constant factor approximation of the optimal uncertainty reduction. Moreover, recent work on transductive active learning of Hübotter et al. () which we restate here shows that the uncertainty of SIFT converges to the irreducible uncertainty.
We assume w.l.o.g. that $\smash{\left\|\boldsymbol{\phi}(\boldsymbol{x})\right\|_{2}^{2}\leq 1}$ for all $\boldsymbol{x}\in\mathcal{X}$.

###### Theorem C.2 (Convergence Guarantee, formalization of 4.1 ) .

Let Assumption C.1 hold and $X_{n}$ be selected by SIFT($\lambda^{\prime}$) from the data space $\mathcal{D}$.
Then for all $n\geq 1$ and $\boldsymbol{x}^{\star}\in\mathcal{X}$,

$$ $\displaystyle\sigma_{n}^{2}(\boldsymbol{x}^{\star})\leq\sigma_{\infty}^{2}( \boldsymbol{x}^{\star})+\frac{d(1+2d\lambda^{\prime}\lambda_{\min}^{-1})\log(1 +\frac{\hat{\lambda}_{n}}{\lambda^{\prime}})}{\sqrt{n}}$ $$

where $\lambda_{\min}$ is the smallest eigenvalue of $\boldsymbol{\Phi}\boldsymbol{\Phi}^{\top}$ with $\boldsymbol{\Phi}\in\mathbb{R}^{m\times d}$ a basis of ${\{\boldsymbol{\phi}(\boldsymbol{x}):\boldsymbol{x}\in\mathcal{D}\}}$, and where $\hat{\lambda}_{n}\leq O\left(n\right)$ is the largest eigenvalue of $\boldsymbol{\Phi}_{n}\boldsymbol{\Phi}_{n}^{\top}$.

###### Proof.

Theorem C.2 follows from Theorem 3.2 of Hübotter et al. () noting that

- •
The SIFT objective is a special case of VTL (Variance-based Transductive Active Learning) with “target space” $\mathcal{A}=\{\boldsymbol{x}^{\star}\}$.
- •
Theorem 3.2 of Hübotter et al. () can be extended to finite-dimensional reproducing kernel Hilbert spaces (Hübotter et al., , Appendix C.6.4).
- •
The “maximum information gain of $n$ iterations”, $\gamma_{n}$, in the statement of Hübotter et al. () is bounded as follows (Srinivas et al., , Appendix C.3): $\smash{\gamma_{n}\leq d\log(1+\hat{\lambda}_{n}/\lambda^{\prime})}$.

∎

###### Assumption C.1 .

###### Theorem C.2 (Convergence Guarantee, formalization of 4.1 ) .

###### Proof.

## Appendix D Further Insights on Active Fine-Tuning

We expand the analysis of our results that we summarized in §5.
We analyze aspects of the two key contributions of our work separately: In the following, we analyze the performance of SIFT in active fine-tuning, and in §E, we analyze the performance of test-time fine-tuning more generally.

#### Insight 6: SIFT ’s improvement over NN grows with dataset size.

As shown in Figure 10, we find that the relative improvement of SIFT over Nearest Neighbor retrieval grows with dataset size.
We suspect that going from a small-size dataset to a medium-size dataset, the additional performance stems mainly from the ability of SIFT to adaptively select the same data for multiple gradient steps.
Going from a medium-size dataset to a large-size dataset, we suspect that the additional performance stems mainly from the ability of SIFT to select more diverse data points.

Figure: Figure 10: Bits per byte (in % relative to the Nearest Neighbor retrieval baseline, $\downarrow$ better). We evaluate data selection from 3%, 33%, and 100% of the Pile training dataset. We see a clear trend that SIFT’s improvement over Nearest Neighbor retrieval grows with dataset size — even from 33% to 100% with the highly curated Pile dataset.
Refer to caption: x9.png

#### Insight 7: Points with high negative cosine similarity may help.

With the Roberta embedding model, we find that there are no negative cosine similarities in the data (cf. Figure 21 in §J).
Choosing different embeddings such as influence embeddings can give negative cosine similarities (Xia et al., , Appendix K.2).
Inspection of those points found by Xia et al. () suggests that they can be equally informative as points with high positive cosine similarity.
Our derivation of SIFT naturally

Figure: Figure 11: Data selection via SIFT (red) and Nearest Neighbor (black) performs best with normalized embeddings.
Refer to caption: x10.png

addresses this by selecting points with large *absolute* cosine similarity.
Geometrically, points with positive or negative cosine similarity are both equally “parallel” to the test prompt.
Our theoretical results suggest that the informativeness of a data point is closely related to how parallel its embedding is to the test prompt.
We leave further investigation to future work.

#### Insight 8: Normalizing embeddings helps.

We evaluate the performance of Nearest Neighbor retrieval and SIFT with or without explicitly normalized embeddings in Figure 11.
We find that for both selection strategies, normalizing embeddings consistently improves performance.
Previously, Hardt & Sun () minimized the Euclidean distance between unnormalized embeddings, which we find to perform identically to maximizing cosine similarity.

## Appendix E Further Insights on Test-Time Fine-Tuning

#### Insight 9: Scaling pre-training compute is not all you need.​​

In Table 2 of §A, we compare state-of-the-art LLMs to our test-time fine-tuned models.
We show that our Phi-3 with test-time fine-tuning outperforms all evaluated base models, from a wide selection of state-of-the-art LLMs, by a large margin.
Notably, we see a clear advantage of using stronger base models, i.e., better initializations.
The leading base model Gemma-2 (27B, Team et al., ), which is $7\times$ larger and more recent than Phi-3, achieves $0.629$ bits per byte, whereas our test-time fine-tuned Phi-3 achieves $0.595$ bits per byte.
This indicates that scaling pre-training compute is not all you need to achieve state-of-the-art performance, and that test-time fine-tuning is an efficient method for improving the performance of a base LLM.

Figure: Figure 12: Improvement of $50$ test-time iterations over the base model (blue; $\downarrow$ better) with SIFT against the percentage of bytes occupied by the dataset in the Pile. Error bars correspond to standard errors. We observe the trend that test-time fine-tuning benefits prompts at the “boundary” of the data distribution most. The “outlier” GitHub dataset is highlighted in red.
Refer to caption: x11.png

#### Insight 10: Test-time fine-tuning yields largest gains at the boundary of the data distribution.

In Figure 12, we plot the improvement of test-time fine-tuning with SIFT over the base model against the weight of a dataset in the Pile.
We observe the trend that test-time fine-tuning yields largest performance improvements for datasets that have a smaller weight in the Pile.
We hypothesize that this trend occurs because the weight of a dataset in the Pile corresponds roughly to the weight of similar data in the pre-training dataset of GPT-2, in which case the performance gains would be largest for prompts that are at the “boundary” of the data distribution.
Notable is the outlier of the large GitHub dataset where test-time fine-tuning leads to large performance gains.
We hypothesize that this is because coding is relatively dissimilar to other data in the Pile, and therefore the GitHub dataset can be seen as “small” relative to the rest of the data.

We make the observation that if the problem domain is large (like general language modeling), almost every sub-task can be seen as at the “boundary” / as an “outlier”.
We see that datasets closest to the center of mass of the data distribution do not benefit as much from test-time fine-tuning as datasets that are further away from the center of mass.
Therefore, we expect test-time fine-tuning to benefit those models most that are learning a diverse data distribution as opposed to models that are learning a very concentrated data distribution.

#### Insight 11: The order of fine-tuning data does not matter.

In Figure 13, we evaluate the performance of test-time fine-tuning with Nearest Neighbor retrieval when taking gradient steps in the order of selected data compared to reversed order.
We find that the order of gradient steps does not affect the final performance.
This indicates that sequentially fine-tuning on selected data is not necessary, and that batched gradient steps can be used to further speed up test-time fine-tuning.
We leave a detailed exploration of batched updates to future work.

Figure: Figure 13: Taking gradient steps in order of selected data compared to reversed order. Data is selected using Nearest neighbor retrieval. We observe that the order of gradient steps does not affect the final performance.
Refer to caption: x12.png

#### Insight 12: Test-time fine-tuning works also with parameter-efficient fine-tuning.

In our experiments with Phi-3, we use Low-Rank Adaptation (Lora, Hu et al., ) with a rank of $64$.
We find that LoRA converges slower than fine-tuning the full model, and therefore use the learning rate $5\mathrm{e}{-4}$, which is a factor $10$ larger than the learning rate used for fine-tuning the full model.
In Figure 14, we evaluate the performance of LoRA compared to fine-tuning the full model.
On the smaller GPT-2 and GPT-2-large we use a rank of $32$.
We generally observe that fine-tuning with LoRA can recover roughly the same performance as fine-tuning the full model.
We expect that with more careful tuning of the learning rate, learning curves could be made more similar.

Figure: Figure 14: Bits per byte ($\downarrow$ better) against the number of test-time iterations. We compare parameter-efficient fine-tuning with LoRA and fine-tuning the full model. We use 0.1% of the full test set.
Refer to caption: x13.png

Figure: Figure 15: Bits per byte ($\downarrow$ better) with GPT-2-large and varying $\lambda^{\prime}$. A larger $\lambda^{\prime}$ can lead to overfitting in later iterations. We use 0.1% of the full test set.
Refer to caption: x14.png

#### Insight 13: Larger models appear to learn faster at test-time.

We find that with a larger model (e.g., GPT-2-large vs GPT-2), a smaller $\lambda^{\prime}$ tends to be more beneficial.
For example, keeping the learning rate fixed at $5\mathrm{e}{-5}$, using $\textsc{SIFT}(0.1)$ is the best choice for GPT-2, but leads to slight overfitting at later iterations for GPT-2-large as shown in Figure 15.
Recall that a smaller $\lambda^{\prime}$ leads to more diverse sampling of the data space.
Thus, this observed trend indicates that larger models learn faster, and therefore benefit more from less redundant training data.
The same trend can also be observed from the behavior of NN-F from Figure 16:
GPT-2-large overfits much faster with NN-F than GPT-2.
This offers a potential explanation why the advantage of SIFT over Nearest Neighbor retrieval grows with larger models (cf. §F.1).

Figure: Figure 16: Bits per byte ($\downarrow$ better) against the number of test-time iterations with various base models.
Refer to caption: x15.png

Figure: Figure 17: Total test-time compute ($\downarrow$ better) of test-time fine-tuning and the in-context baseline. As to be expected, we see that in-context learning becomes comparably less efficient with Phi-3 due to its larger context window. This figure complements Figure 7 which shows that test-time fine-tuning achieves higher token throughput than in-context learning.
Refer to caption: x16.png

## Appendix F Extended Results

This section includes additional per-dataset results to support our findings on active fine-tuning and test-time fine-tuning.

### F.1 Active Fine-Tuning

We compare SIFT against the data selection baselines Uncertainty Sampling (US), Nearest Neighbor retrieval (NN), and the failure-mode of Nearest Neighbor retrieval (with information duplication) that repeatedly retrieves the same point (NN-F).
Our results with GPT-2 as base model are summarized in the main text in Table 1.

- •
In Table 4, we include the comparison with GPT-2-large.
- •
In Table 5, we include the comparison with Phi-3.

We find that our results on GPT-2 are consistent across all models.
In particular, test-time fine-tuning with SIFT improves the base model on *all* datasets of the Pile, often significantly.
SIFT outperforms Uncertainty Sampling and Nearest Neighbor retrieval consistently.
Notably, we find that the improvement of SIFT over Nearest Neighbor retrieval is larger with stronger base models, indicating that informativeness of data becomes more important the stronger the base model.

### F.2 Test-Time Fine-Tuning

We compare the in-context baseline against test-time fine-tuning.

- •
In Table 6, we include the comparison with GPT-2.
- •
In Table 7, we include the comparison with GPT-2-large.
- •
In Table 8, we include the comparison with Phi-3.

We find that test-time fine-tuning consistently outperforms in-context learning with GPT-2 and GPT-2-large.
With Phi-3, in-context learning and test-time fine-tuning have roughly matching performance, though test-time fine-tuning is more computationally efficient (cf. Figure 7).
Interestingly, we observe that test-time fine-tuning leads to large gains on math (“DeepMind Math”) and coding (“GitHub”) on all models, two tasks that require more complex reasoning.

**Table 4: Results with GPT-2-large. Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. We only include datasets with at least $10$ examples in our test set. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of SIFT over the strongest baseline.**
|  | US | NN | NN-F | SIFT | $\Delta$ |
| --- | --- | --- | --- | --- | --- |
| NIH Grants | 96.6 (1.6) | 77.9 (4.8) | 107.6 (19.8) | 51.9 (9.3) | $\downarrow\!26.0$ |
| US Patents | 86.8 (2.3) | 78.9 (2.6) | 129.1 (7.7) | 64.7 (3.8) | $\downarrow\!14.2$ |
| Enron Emails | 73.9 (12.3) | 68.6 (13.6) | 102.9 (23.1) | 55.5 (12.2) | $\downarrow\!13.1$ |
| GitHub | 45.2 (2.4) | 42.8 (2.2) | 62.0 (4.5) | 31.0 (2.2) | $\downarrow\!11.8$ |
| Wikipedia | 71.0 (2.0) | 71.5 (2.0) | 141.3 (3.5) | 64.4 (2.2) | $\downarrow\!6.6$ |
| PubMed Abstr. | 94.5 (0.4) | 93.7 (0.6) | 202.6 (1.6) | 87.8 (0.7) | $\downarrow\!5.9$ |
| ArXiv | 90.6 (1.8) | 90.2 (2.0) | 175.8 (5.7) | 84.8 (2.1) | $\downarrow\!5.4$ |
| Hacker News | 79.4 (2.6) | 79.0 (2.9) | 138.7 (4.4) | 75.6 (3.6) | $\downarrow\!3.4$ |
| Stack Exchange | 84.1 (0.7) | 84.6 (0.8) | 165.2 (1.8) | 80.7 (0.9) | $\downarrow\!3.4$ |
| Common Crawl | 93.7 (0.6) | 89.9 (0.7) | 163.6 (2.1) | 87.1 (1.0) | $\downarrow\!2.8$ |
| PubMed Central | 87.9 (2.7) | 87.6 (2.7) | 157.8 (4.6) | 85.4 (3.1) | $\downarrow\!2.2$ |
| FreeLaw | 66.8 (4.2) | 67.4 (4.1) | 132.0 (6.4) | 68.3 (4.2) | $\uparrow\!1.5$ |
| DeepMind Math | 71.2 (2.2) | 72.2 (2.0) | 186.1 (4.1) | 74.2 (2.3) | $\uparrow\!3.0$ |
| *All* | 82.6 (0.6) | 80.6 (0.6) | 153.3 (1.4) | 74.9 (0.7) | $\downarrow\!5.7$ |

**Table 5: Results with Phi-3. Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of SIFT over the strongest baseline.**
|  | US | NN | NN-F | SIFT | $\Delta$ |
| --- | --- | --- | --- | --- | --- |
| GitHub | 80.6 | 80.8 | 105.2 | 46.5 | $\downarrow\!34.1$ |
| US Patents | 95.4 | 94.2 | 274.6 | 83.7 | $\downarrow\!10.5$ |
| Enron Emails | 113.6 | 86.6 | 319.9 | 78.7 | $\downarrow\!7.9$ |
| Wikipedia | 84.6 | 85.5 | 263.2 | 79.2 | $\downarrow\!5.4$ |
| PubMed Abstracts | 93.5 | 93.3 | 301.8 | 89.5 | $\downarrow\!3.8$ |
| NIH Grants | 100.4 | 100.1 | 327.6 | 98.6 | $\downarrow\!1.5$ |
| ArXiv | 95.5 | 96.5 | 282.4 | 94.3 | $\downarrow\!1.2$ |
| Common Crawl | 95.3 | 94.9 | 257.0 | 93.7 | $\downarrow\!1.2$ |
| PubMed Central | 80.3 | 82.1 | 204.9 | 79.7 | $\downarrow\!0.6$ |
| DeepMind Math | 76.4 | 75.5 | 221.4 | 75.3 | $\downarrow\!0.2$ |
| Hacker News | 95.1 | 94.8 | 243.8 | 95.0 | $\uparrow\!0.2$ |
| FreeLaw | 66.9 | 67.8 | 178.0 | 67.2 | $\uparrow\!0.3$ |
| Stack Exchange | 99.7 | 98.7 | 309.9 | 99.4 | $\uparrow\!0.7$ |
| *All* | 92.0 (1.1) | 91.6 (1.1) | 256.6 (7.1) | 85.7 (2.0) | $\downarrow\!5.9$ |

**Table 6: Comparison between the in-context baseline and test-time fine-tuning with GPT-2. Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. We only include datasets with at least $10$ examples in our test set. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of test-time fine-tuning over in-context learning.**
|  | Context | Fine-Tuning | $\Delta$ |
| --- | --- | --- | --- |
| GitHub | 74.6 (2.5) | 28.6 (2.2) | $\downarrow\!56.0$ |
| DeepMind Math | 100.2 (0.1) | 70.1 (2.1) | $\downarrow\!30.1$ |
| US Patents | 87.4 (2.5) | 62.2 (3.6) | $\downarrow\!25.2$ |
| FreeLaw | 87.2 (3.6) | 65.5 (4.2) | $\downarrow\!21.7$ |
| Stack Exchange | 97.2 (0.4) | 77.0 (0.7) | $\downarrow\!20.2$ |
| Enron Emails | 71.9 (10.2) | 52.4 (11.8) | $\downarrow\!19.5$ |
| ArXiv | 98.8 (1.8) | 81.6 (1.9) | $\downarrow\!17.2$ |
| Hacker News | 92.6 (2.7) | 77.8 (3.5) | $\downarrow\!14.8$ |
| NIH Grants | 67.7 (9.4) | 52.9 (9.0) | $\downarrow\!14.8$ |
| Wikipedia | 76.6 (2.1) | 63.7 (2.1) | $\downarrow\!12.9$ |
| PubMed Central | 92.3 (3.3) | 80.6 (2.7) | $\downarrow\!11.7$ |
| Common Crawl | 93.5 (0.7) | 86.1 (0.9) | $\downarrow\!7.4$ |
| PubMed Abstr. | 91.6 (0.6) | 84.8 (0.7) | $\downarrow\!6.8$ |
| *All* | 89.4 (0.5) | 73.2 (0.7) | $\downarrow\!16.2$ |

**Table 7: Comparison between the in-context baseline and test-time fine-tuning with GPT-2-large. Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. We only include datasets with at least $10$ examples in our test set. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of test-time fine-tuning over in-context learning.**
|  | Context | Fine-Tuning | $\Delta$ |
| --- | --- | --- | --- |
| GitHub | 74.6 (2.5) | 31.0 (2.2) | $\downarrow\!43.6$ |
| DeepMind Math | 100.2 (0.7) | 74.2 (2.3) | $\downarrow\!26.0$ |
| US Patents | 87.4 (2.5) | 64.7 (3.8) | $\downarrow\!22.7$ |
| FreeLaw | 87.2 (3.6) | 68.3 (4.2) | $\downarrow\!18.9$ |
| Hacker News | 92.6 (2.7) | 75.6 (3.6) | $\downarrow\!17.0$ |
| Stack Exchange | 97.2 (0.4) | 80.7 (0.9) | $\downarrow\!16.5$ |
| NIH Grants | 67.7 (9.4) | 51.9 (9.3) | $\downarrow\!15.8$ |
| Enron Emails | 71.9 (10.2) | 55.5 (12.2) | $\downarrow\!15.5$ |
| ArXiv | 98.8 (1.8) | 84.8 (2.1) | $\downarrow\!14.0$ |
| Wikipedia | 76.6 (2.1) | 64.4 (2.2) | $\downarrow\!12.2$ |
| PubMed Central | 92.3 (3.3) | 85.4 (3.1) | $\downarrow\!6.9$ |
| Common Crawl | 93.5 (0.7) | 87.1 (1.0) | $\downarrow\!6.4$ |
| PubMed Abstr. | 91.6 (0.6) | 87.8 (0.7) | $\downarrow\!3.8$ |
| *All* | 89.4 (0.5) | 74.9 (0.7) | $\downarrow\!14.5$ |

**Table 8: Comparison between the in-context baseline and test-time fine-tuning with Phi-3. Bits per byte (in % relative to the base model, $\downarrow$) after $50$ test-time iterations on individual datasets of the Pile. Bold numbers denote the best performing selected subset. Numbers in parentheses are standard errors. $\Delta$ denotes the performance gain of test-time fine-tuning over in-context learning.**
|  | Context | Fine-Tuning | $\Delta$ |
| --- | --- | --- | --- |
| DeepMind Math | 100.8 | 75.3 | $\downarrow\!25.5$ |
| GitHub | 71.3 | 46.5 | $\downarrow\!24.8$ |
| FreeLaw | 78.2 | 67.2 | $\downarrow\!11.0$ |
| ArXiv | 101.0 | 94.3 | $\downarrow\!6.4$ |
| Enron Emails | 81.8 | 78.7 | $\downarrow\!3.1$ |
| Hacker News | 97.6 | 95.0 | $\downarrow\!2.6$ |
| Stack Exchange | 100.9 | 99.4 | $\downarrow\!1.4$ |
| PubMed Central | 79.9 | 79.7 | $\downarrow\!0.2$ |
| US Patents | 83.3 | 83.7 | $\uparrow\!0.4$ |
| Wikipedia | 77.1 | 79.2 | $\uparrow\!2.1$ |
| NIH Grants | 95.1 | 98.6 | $\uparrow\!3.5$ |
| Common Crawl | 89.9 | 93.7 | $\uparrow\!3.8$ |
| PubMed Abstracts | 85.7 | 89.5 | $\uparrow\!3.8$ |
| *All* | 87.1 (1.7) | 85.7 (2.0) | $\downarrow\!1.4$ |

## Appendix G SIFT Maximizes Information Gain

We discuss here briefly that SIFT can be interpreted as maximizing the information gain of data $X_{n}$ on the response to the prompt $\boldsymbol{x}^{\star}$.

### G.1 Preliminaries: Information Theory and Gaussian Processes

#### Information Theory

We briefly recap several important concepts from information theory.
The (differential) entropy ${\mathrm{H}\left[\boldsymbol{f}\right]\operatorname*{\,\dot{=}\,}\mathbb{E}_{p
(\boldsymbol{f})}{}\left[-\log p(\boldsymbol{f})\right]}$ of a random vector $\boldsymbol{f}$ is one possible measure of uncertainty about $\boldsymbol{f}$.
Here, $-\log p(\boldsymbol{f})$ is also called the suprisal about an event with density $p(\boldsymbol{f})$.
The entropy can be interpreted as the expected suprisal about $\boldsymbol{f}$ upon realization.
The conditional entropy ${\mathrm{H}\!\left[\boldsymbol{f}\ \middle|\ \boldsymbol{y}\right]
\operatorname*{\,\dot{=}\,}\mathbb{E}_{p(\boldsymbol{f},\boldsymbol{y})}{}
\left[-\log p(\boldsymbol{f}\mid\boldsymbol{y})\right]}$ is the (expected) posterior uncertainty about $\boldsymbol{f}$ after observing the random vector $\boldsymbol{y}$.
The information gain ${\mathrm{I}\left(\boldsymbol{f};\boldsymbol{y}\right)=\mathrm{H}\left[
\boldsymbol{f}\right]-\mathrm{H}\!\left[\boldsymbol{f}\ \middle|\ \boldsymbol{
y}\right]}$ measures the (expected) reduction in uncertainty about $\boldsymbol{f}$ due to $\boldsymbol{y}$.
Refer to Cover () for more details.

#### Gaussian Processes

The stochastic process $f$ is a Gaussian process (GP, Williams & Rasmussen ()), denoted ${f\sim\mathcal{GP}({}\mu,k)}$, with mean function $\mu$ and kernel $k$ if for any finite subset ${X=\{\boldsymbol{x}_{1},\dots,\boldsymbol{x}_{n}\}\subseteq\mathcal{X}}$, ${\boldsymbol{f}_{\!X}\sim\mathcal{N}({}\boldsymbol{\mu}_{\!X},\boldsymbol{K}_{
\!X})}$ is jointly Gaussian with mean vector ${(\boldsymbol{\mu}_{\!X})_{i}=\mu(\boldsymbol{x}_{i})}$ and covariance matrix ${(\boldsymbol{K}_{\!X})_{i,j}=k(\boldsymbol{x}_{i},\boldsymbol{x}_{j})}$.

For Gaussian random vectors $\boldsymbol{f}$ and $\boldsymbol{y}$, the entropy is ${\mathrm{H}\left[\boldsymbol{f}\right]=\frac{d}{2}\log(2\pi e)+\frac{1}{2}\log
\det{\mathrm{Var}{}\left(\boldsymbol{f}\right)}}$ and the information gain is ${\mathrm{I}\left(\boldsymbol{f};\boldsymbol{y}\right)=\frac{1}{2}\left(\log
\det{\mathrm{Var}{}\left(\boldsymbol{f}\right)}-\log\det{\mathrm{Var}{}\left(
\boldsymbol{f}\mid\boldsymbol{y}\right)}\right)}$.

### G.2 Probabilistic Observation Model

We will focus in the following on the case of regression, which we introduced in Section K.5.
We suppose that observations of $f$ follow the probabilistic model

$$ $\displaystyle y_{\boldsymbol{x}}=f_{\boldsymbol{x}}+\varepsilon_{\boldsymbol{x }},$ $$

where we make the following assumptions about the prior distribution of $f$ and the noise $\varepsilon_{\boldsymbol{x}}$:

###### Assumption G.1 (Gaussian prior) .

We assume that ${f\sim\mathcal{GP}({}\mu,k)}$ with known mean function $\mu$ and kernel $k$.

###### Assumption G.2 (Gaussian noise) .

We assume that the noise $\varepsilon_{\boldsymbol{x}}$ is mutually independent and zero-mean Gaussian with known variance $\rho^{2}>0$.

Under Assumptions G.1 and G.2, the posterior distribution of $f$ after observing points $X$ with values $\boldsymbol{y}_{\!X}$ is $\mathcal{GP}({}\mu_{n},k_{n})$ with

$$ $\displaystyle\mu_{n}(\boldsymbol{x})$ $\displaystyle=\mu(\boldsymbol{x})+\boldsymbol{k}_{X}^{\top}(\boldsymbol{x})( \boldsymbol{K}_{\!XX}+\rho^{2}\boldsymbol{I})^{-1}(\boldsymbol{y}_{\!X}- \boldsymbol{\mu}_{\!X}),$ $\displaystyle k_{n}(\boldsymbol{x},\boldsymbol{x^{\prime}})$ $\displaystyle=k(\boldsymbol{x},\boldsymbol{x^{\prime}})-\boldsymbol{k}_{X}^{ \top}(\boldsymbol{x})(\boldsymbol{K}_{\!XX}+\rho^{2}\boldsymbol{I})^{-1} \boldsymbol{k}_{X}(\boldsymbol{x^{\prime}}),$ $\displaystyle\sigma_{n}^{2}(\boldsymbol{x})$ $\displaystyle=k_{n}(\boldsymbol{x},\boldsymbol{x}).$ $$

###### Assumption G.1 (Gaussian prior) .

###### Assumption G.2 (Gaussian noise) .

### G.3 The Probabilistic Interpretation of SIFT

Observe that the above definition of $\sigma_{n}^{2}$ matches the definition from Equation 2.(^11^1111Notably, it can also be shown that $\mu_{n}$ is the closed-form solution to the regularized loss from Equation 10.)
That is, under the above probabilistic model,

$$ $\displaystyle\sigma_{n}^{2}(\boldsymbol{x})=\mathrm{Var}{}\left(f(\boldsymbol{ x})\mid y_{1:n}\right).$ $$

As such, SIFT($\rho^{2}$) is minimizing the variance of the response to the prompt $\boldsymbol{x}^{\star}$ after observing the data $X_{n}$:

$$ $\displaystyle\boldsymbol{x}_{n+1}$ $\displaystyle=\operatorname*{arg\,min}_{\boldsymbol{x}\in\mathcal{D}}\mathrm{ Var}{}\left(f(\boldsymbol{x}^{\star})\mid y_{1:n},y(\boldsymbol{x})\right).$ By simple algebraic manipulation this can be seen to be equivalent to maximizing the information gain of the data on the response to the prompt $\boldsymbol{x}^{\star}$ : $\displaystyle\boldsymbol{x}_{n+1}$ $\displaystyle=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\frac{1}{ 2}\Big{(}\underbrace{\log\mathrm{Var}{}\left(f(\boldsymbol{x}^{\star})\right) missing}_{\mathrm{const}}-\log\mathrm{Var}{}\left(f(\boldsymbol{x}^{\star}) \mid y_{1:n},y(\boldsymbol{x})\right)\Big{)}$ $\displaystyle=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\mathrm{I }\!\left(f(\boldsymbol{x}^{\star});y(\boldsymbol{x})\ \middle|\ y_{1:n}\right).$ (7) $$

#### Discussion

The above offers a very intuitive probabilistic interpretation of SIFT($\rho^{2}$).
In this probabilistic interpretation, the regularization parameter $\lambda^{\prime}$ of SIFT is equal to the observation noise $\rho^{2}$.
Intuitively, larger observation noise leads to slower convergence of the estimate of $f$, analogously to our discussion of larger regularization parameter and smaller step size in Proposition 3.3.

The reason why SIFT($\rho^{2}$) can be interpreted *both* as minimizing the variance and as minimizing the entropy of the response to the prompt $\boldsymbol{x}^{\star}$ is that the variance is proportional to the entropy of the response to the prompt $\boldsymbol{x}^{\star}$.
As observed by Hübotter et al. (), if learning is amortized with respect to multiple prompts $\{\boldsymbol{x}^{\star}_{1},\dots,\boldsymbol{x}^{\star}_{m}\}=\mathcal{A}$, this ceases to be the case and the two objectives lead to different data selection schemes.
It appears to be a special property of non-amortized transductive active learning that measures of uncertainty and resulting data selection schemes are interchangeable.

Under Assumption C.1, the information gain $\mathrm{I}\left(f(\boldsymbol{x}^{\star});y_{1:n}\right)$, i.e., the “entropy reduction” of data $X_{n}$ selected by SIFT achieves therefore also a constant factor approximation of the maximum possible information gain $\max_{X\subseteq\mathcal{D},|X|\leq n}\mathrm{I}\left(f(\boldsymbol{x}^{\star}
);\boldsymbol{y}(X)\right)$.

### G.4 How SIFT Balances Relevance and Diversity

In §C.1, we discussed how chooses data that is both relevant and diverse.
The probabilistic interpretation offers a simple explanation for how this behavior naturally emerges from selecting the most informative data.
To this end, observe that the information gain from above can be expressed as

$$ $\displaystyle\mathrm{I}\left(f(\boldsymbol{x}^{\star});\boldsymbol{y}(X)\right )=\underbrace{\mathrm{H}\left[\boldsymbol{y}(X)\right]missing}_{\text{ diversity}}-\underbrace{\mathrm{H}\!\left[\boldsymbol{y}(X)\ \middle|\ f( \boldsymbol{x}^{\star})\right]}_{\text{relevance}}.$ (8) $$

The first term, the entropy $\mathrm{H}\left[\boldsymbol{y}(X)\right]$, measures the joint model uncertainty about the responses to the prompts in $X$ before the model is fine-tuned on them.
Maximizing this uncertainty encourages selection of diverse data since any redundant observations will not contribute to the model’s joint uncertainty.
The second term, the conditional entropy $\mathrm{H}\!\left[\boldsymbol{y}(X)\ \middle|\ f(\boldsymbol{x}^{\star})\right]$, measures the model uncertainty about the responses to the prompts in $X$ if the model were to know the response to the target prompt $\boldsymbol{x}^{\star}$.
This term can be interpreted as measuring how related the data in $X$ is to $\boldsymbol{x}^{\star}$.
If the data in $X$ was not related to $\boldsymbol{x}^{\star}$, this term would be large and the information gain would be small.

### G.5 The Perspective of Classification

The above interpretation takes the perspective of regression.
However, the above interpretation can be extended to classification.
We will focus here on the case of binary classification for notational convenience, but the same argument can be made for multi-class classification (Williams & Rasmussen, , Section 3.5).

In (binary) Gaussian Process Classification the logit $f\sim\mathcal{GP}({}\mu,k)$ is modeled as a Gaussian process, and the likelihood follows the model introduced in Section 3: $y(\boldsymbol{x})\sim\mathrm{Bern}(s(f(\boldsymbol{x})))$ where we have Bernoulli rather than categorical feedback and use the logistic function $s(a)\operatorname*{\,\dot{=}\,}1/(1+e^{-a})$ rather than the softmax by virtue of restricting to binary classification.

The standard approach (Williams & Rasmussen, , Section 3.4) is to approximate the posterior distribution of the latent function $f$ given observations $y_{1:n}$ by a Gaussian using Laplace’s method.
This Gaussian can be shown to have covariance $\smash{(\boldsymbol{K}_{\!X_{n}}^{-1}+\boldsymbol{W})^{-1}}$ with $\boldsymbol{W}\succeq\kappa^{-1}\boldsymbol{I}_{n}$ where $\smash{\kappa\operatorname*{\,\dot{=}\,}\sup_{a\leq B}1/\dot{s}(a)}$ and $\smash{\dot{s}(a)=s(a)(1-s(a))}$ denotes the derivative of the logistic function.(^12^1212In the binary case, this is equal to the more general $\kappa$ from the main text.)
It is then straightforward to derive that

$$ $\displaystyle\sigma_{n}^{2}(\boldsymbol{x}^{\star})$ $\displaystyle=k(\boldsymbol{x}^{\star},\boldsymbol{x}^{\star})-\boldsymbol{k}_ {X_{n}}^{\top}(\boldsymbol{x}^{\star})(\boldsymbol{K}_{\!X_{n}}+\boldsymbol{W} ^{-1})^{-1}\boldsymbol{k}_{X_{n}}(\boldsymbol{x}^{\star})$ $\displaystyle\leq k(\boldsymbol{x}^{\star},\boldsymbol{x}^{\star})-\boldsymbol {k}_{X_{n}}^{\top}(\boldsymbol{x}^{\star})(\boldsymbol{K}_{\!X_{n}}+\kappa \boldsymbol{I}_{n})^{-1}\boldsymbol{k}_{X_{n}}(\boldsymbol{x}^{\star})$ $$

Thus, SIFT minimizes a tight upper bound to the (approximate) posterior variance of the latent function $f$ at the prompt $\boldsymbol{x}^{\star}$.
The same relationship to maximizing information gain that was discussed above applies.

## Appendix H Efficient Computation of SIFT

In the following, we show how to select data via SIFT at low computational cost.
Our implementation extends the Faiss library (Johnson et al., ; Douze et al., ) for Nearest Neighbor retrieval.
We open-source the [activeft](https://github.com/jonhue/activeft) (Active Fine-Tuning) library which can be used as a drop-in replacement for Nearest Neighbor retrieval.

In our runtime analysis, we will denote by $K$ the size of the data space $\mathcal{D}$, and by $N$ the number of points to be selected.
We describe two implementations of SIFT:

- 1.
The first exact implementation has sequential computation cost $O\left(K^{2}N\right)$, however, computation can be effectively parallelized on a GPU.
- 2.
The second “fast” implementation assumes submodularity (i.e., Assumption C.1) and has computation cost $\smash{\widetilde{O}\left(K+N^{3}\right)}$ where $\smash{\widetilde{O}\left(\cdot\right)}$ suppresses log-factors.
This cost is only marginally above the cost of Nearest Neighbor retrieval.

Both implementations achieve virtually identical performance gains (cf. Figure 19 (right)), which is further evidence that Assumption C.1 is satisfied in our language modeling setting.

### H.1 Exact Implementation

The central object of the first implementation is the conditional kernel matrix of the data space given the selected points $X_{n}$:

$$ $\displaystyle\boldsymbol{K}_{\!n}\operatorname*{\,\dot{=}\,}\boldsymbol{K}_{\! \mathcal{D}}-\boldsymbol{K}_{\!\mathcal{D},X_{n}}(\boldsymbol{K}_{\!X_{n}}+ \lambda^{\prime}\boldsymbol{I}_{n})^{-1}\boldsymbol{K}_{\!X_{n},\mathcal{D}}.$ $$

The entries $k_{n}(\boldsymbol{x},\boldsymbol{x^{\prime}})$ of this matrix can be updated efficiently via the following relation (Chowdhury & Gopalan, , Appendix F) arising from properties of the Schur complement:

$$ $\displaystyle k_{n}(\boldsymbol{x},\boldsymbol{x^{\prime}})=k_{n-1}( \boldsymbol{x},\boldsymbol{x^{\prime}})-\frac{k_{n-1}(\boldsymbol{x}, \boldsymbol{x}_{n})k_{n-1}(\boldsymbol{x}_{n},\boldsymbol{x^{\prime}})}{k_{n-1 }(\boldsymbol{x}_{n},\boldsymbol{x}_{n})+\lambda^{\prime}}.$ (9) $$

The implementation is detailed in Algorithm 1.
The computation of the objective value in 4 and the kernel matrix update in 5 can be parallelized on a GPU.
Thus, the main bottleneck of this implementation is the requirement that the kernel matrix of size $K\times K$ fits onto a GPU.
In case this is not possible, such as with large data spaces, the following two sections detail methods to reduce the computational cost.

Figure: Algorithm 1 SIFT($\lambda^{\prime}$)

### H.2 Fast (Exact) Implementation

The following “fast” implementation of SIFT rests on the assumption that the objective function optimized by SIFT is submodular (cf. Assumption C.1).
Recall that this objective function can be expressed as $\boldsymbol{x}_{n+1}=\operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}
\psi_{\boldsymbol{x}^{\star}}(X_{n}\cup\{\boldsymbol{x}\})$ where $\psi_{\boldsymbol{x}^{\star}}(X)=\sigma_{0}^{2}(\boldsymbol{x}^{\star})-\sigma
_{X}^{2}(\boldsymbol{x}^{\star})$ denotes the *uncertainty reduction* about $\boldsymbol{x}^{\star}$ upon fine-tuning the model on data $X$.

The “trick” of the fast implementation is to use a max-heap (with $O\left(1\right)$ lookup and $O\left(\log K\right)$ insertion) to keep track of upper bounds of $\psi_{\boldsymbol{x}^{\star}}(X_{n}\cup\{\boldsymbol{x}\})$ for each $\boldsymbol{x}\in\mathcal{D}$.
The upper bounds come directly from the submodularity assumption:

$$ $\displaystyle\psi_{\boldsymbol{x}^{\star}}(X_{i}\cup\{\boldsymbol{x}\})\geq \psi_{\boldsymbol{x}^{\star}}(X_{j}\cup\{\boldsymbol{x}\})\quad\forall j\geq i.$ $$

At iteration $n$, we evaluate $\psi_{\boldsymbol{x}^{\star}}(X_{n-1}\cup\{\boldsymbol{x}\})$ for $\boldsymbol{x}$ in max-heap order.
As soon as we find a $\boldsymbol{x}$ whose re-computed upper bound is smaller than a previously re-computed upper bound, we stop the evaluation.
In the worst case, one might iterate through all $K$ points in each iteration, but in practice, it can sometimes be reasonable to assume that one only needs to consider $O\left(1\right)$ points per iteration.
This algorithm is known as the “lazy greedy algorithm” in submodular function maximization (Minoux, ) where it is typically seen to result in large speed-ups.

We summarize the fast implementation in Algorithm 2.
The kernel matrix $\boldsymbol{K}$ tracks the conditional kernel matrix of the prompt $\boldsymbol{x}^{\star}$ and the previously selected data $X_{n-1}$.
$\boldsymbol{\Lambda}$ tracks the (regularized) inverse of the kernel matrix of the previously selected data $X_{n-1}$.
Whenever necessary, the cached kernel matrix and cached inverse are updated.
We denote by $\boldsymbol{\Phi}\in\mathbb{R}^{(n-1)\times d}$ the matrix of embeddings of previously selected points and by $\smash{\tilde{\boldsymbol{\Phi}}\in\mathbb{R}^{n\times d}}$ the same matrix extended by $\boldsymbol{\phi}(\boldsymbol{x}^{\star})$ as the first row.

Initializing the max-heap takes time $\widetilde{O}\left(K\right)$ and is analogous to standard Nearest Neighbor retrieval.
Additionally, SIFT-Fast performs a data selection loop for $N$ iterations where each operation takes $O\left(N^{2}\right)$ time requiring persistent memory of size $O\left(N^{2}\right)$.
Notably, only the kernel matrix of the prompt and the previously selected data is kept in memory.

Figure: Algorithm 2 SIFT-Fast($\lambda^{\prime}$)

Figure: Algorithm 3 SIFT-Fast($\lambda^{\prime}$): Recompute

Figure: Algorithm 4 SIFT-Fast($\lambda^{\prime}$): UpdateState

### H.3 Pre-Selecting Data via Nearest Neighbor Retrieval

Figure: Figure 18: We run SIFT $\!(\lambda^{\prime}=1)$ with various values of $K$ and report the bits per byte ($\downarrow$ better) after $50$ test-time iterations. We find that performance on the Pile plateaus after $K=1`000$. Even at $K=50$, which equals the number of points selected, SIFT outperforms Nearest Neighbor retrieval due to being able to select the same points multiple times.
Refer to caption: x17.png

The reason for SIFT-Fast being so efficient is that it effectively “discards” all points in $\mathcal{D}$ that are completely irrelevant to the prompt.
Whereas SIFT recomputes the objective value of every point in $\mathcal{D}$ at each iteration, SIFT-Fast only reevaluates points that are potentially relevant.
An alternative to make SIFT fast is therefore simply to preemptively discard irrelevant points.
In our experiments we do so by pre-selecting a subset of size $K=200$ via Nearest Neighbor retrieval within $\mathcal{D}$ (cf. Appendix I for more details).
This step aims to eliminate all points from the data space that SIFT would not end up picking anyway while retaining a diverse set of relevant points.
Figure 18 shows the effect of $K$ on statistical performance and Figure 4 shows the effect on computational performance.

### H.4 Future Work: Improving GPU Utilization of SIFT-Fast

Figure: Figure 19: Left: Computational overhead of SIFT-Fast over Nearest Neighbor retrieval. This overhead is larger than the overhead of SIFT over Nearest Neighbor retrieval (cf. Figure 4). Right: SIFT-Fast achieves identical statistical performance to SIFT, which is further evidence that Assumption C.1 is satisfied in our language modeling setting.
Refer to caption: x18.png

In our experiments on the Pile dataset, we find that SIFT-Fast is less efficient than SIFT (cf. Figure 19 (left)).
We attribute this to the fact that for any given prompt, the closest neighbors in the data space are all relatively similar to the prompt (cf. Figure 20), meaning that each iteration of SIFT-Fast has to loop (sequentially) over the entire priority queue.
In contrast, SIFT performs this operation in parallel on a GPU.

Figure: Figure 20: Average cosine similarities of test prompts to closest $1`000$ neighbors in the data space of the Pile; with the Roberta embedding model.
Refer to caption: x19.png

We believe that a promising computational approach is to combine the advantages of the SIFT and SIFT-Fast implementations.
This could be achieved by keeping a large sub-selected kernel matrix on the GPU (akin to the SIFT implementation) and selectively using the SIFT-Fast implementation if points on the priority queue that are not in the sub-selected kernel matrix may be selected.
This would allow for a more efficient use of the GPU memory of SIFT-Fast, which we expect to yield comparable computational performance to the SIFT implementation in most cases, while still being able to handle large data spaces.

## Appendix I Experiment Details

We fine-tune the pre-trained model for a single gradient step each on $N=50$ selected data points.
We evaluate the performance on 1% of the test instances of the Pile.
We use the Pile training dataset as data space for data selection, which notably does *not* include data from the validation and test sets.

#### Evaluation

We use the standard implementation of the lm-evaluation-harness library (Gao et al., ) for computing the bits per byte.
This implementation computes the log-likelihood of a document using a rolling-window approach, which ensures that the model’s maximum context window is fully utilized.

#### Truncation of Long Sequences

Analogously to Hardt & Sun (), to generate embeddings, we naively truncate long sequences to the maximum sequence length of the embedding model, that is, we only consider the prefixes of long sequences for data selection.

#### Learning Rate and Optimizer

Following Hardt & Sun (), we use the Adam optimizer (Kingma & Ba, ) with $\epsilon$-value $1\mathrm{e}{-8}$.
We use the default learning rate $5\mathrm{e}{-5}$ of the transformers library (Wolf et al., ) unless noted otherwise.
Hardt & Sun () used a learning rate of $2\mathrm{e}{-5}$ for their experiments.
We show in Figure 24 that $5\mathrm{e}{-5}$ leads to strictly better performance of the Nearest Neighbor baseline.
In our ablation study over metrics for Nearest Neighbor retrieval (cf. Figure 11), which was conducted concurrently, we still used learning rate $2\mathrm{e}{-5}$ of Hardt & Sun ().

#### Low-Rank Adaptation (LoRA)

We use LoRA (Hu et al., ) for fine-tuning Phi-3, and also evaluate the performance of LoRA with GPT-2 and GPT-2-large (cf. §E).
We use LoRAs with rank $64$, output scaling $16$, without dropout and bias.
When fine-tuning with LoRA, we use the learning rate $5\mathrm{e}{-4}$.

#### Gradient Checkpointing

We additionally use gradient checkpointing (Chen et al., ) for fine-tuning Phi-3 to reduce memory footprint and allow fine-tuning on our hardware.
This suggests that our compute-time results with Phi-3 (cf. Figures 7 and 17) are conservative estimates.

#### Uncopyrighted Pile Dataset

We use only those datasets of the Pile where our use is in compliance with the terms of service of the data host (Gao et al., ).
This excludes the Books3, BookCorpus2, OpenSubtitles, YTSubtitles, and OWT2 datasets.

We provide an overview of all hyperparameters of test-time fine-tuning in Table 9.

**Table 9: Hyperparameters during test-time fine-tuning, unless noted otherwise.**
|  | GPT-2 | GPT-2-large | Phi-3 |
| --- | --- | --- | --- |
| $\lambda^{\prime}$ | $0.1$ | $0.01$ | $0.01$ |
| Learning rate | $5\mathrm{e}{-5}$ | $5\mathrm{e}{-5}$ | $5\mathrm{e}{-4}$ |
| Adam’s $\epsilon$-value | $1\mathrm{e}{-8}$ | $1\mathrm{e}{-8}$ | $1\mathrm{e}{-8}$ |
| Max. sequence length (in tokens) | $1024$ | $1024$ | $4096$ |
| LoRA | no | no | yes |
| Gradient checkpointing | no | no | yes |

### I.1 Properties of the Pile Dataset

Figure 21 shows the average cosine similarities of test prompts to neighbors in the data space of the Pile.
Table 10 shows the weight of each dataset in the Pile.

Figure: Figure 21: Average cosine similarities of test prompts to neighbors in the data space of the Pile; with the Roberta embedding model.
Refer to caption: x20.png

**Table 10: Overview of datasets in the (uncopyrighted) Pile. Weight is the percentage of bytes in the final dataset occupied by each dataset. Numbers are taken from Gao et al. () and renormalized.**
|  | Weight |
| --- | --- |
| Common Crawl | 24.14% |
| PubMed Central | 19.19% |
| ArXiv | 11.94% |
| GitHub | 10.12% |
| FreeLaw | 8.18% |
| Stack Exchange | 6.84% |
| US Patents | 4.87% |
| PubMed Abstracts | 4.09% |
| Project Gutenberg | 2.89% |
| Wikipedia | 2.04% |
| DeepMind Math | 1.65% |
| Ubuntu IRC | 1.17% |
| EuroParl | 0.97% |
| Hacker News | 0.83% |
| PhilPapers | 0.51% |
| NIH ExPorter Grants | 0.40% |
| Enron Emails | 0.19% |

### I.2 In-Context Baseline

In our evaluation of in-context learning, we use the following format to insert the selected data into the context of the model: We separate all retrieved token sequences with the string "\n\n" which can be seen as a paragraph separator, and additionally add this string between the data string and the prompt.

Notably, our results with in-context learning on GPT-2-large outperform the results previously reported by Hardt & Sun ().
We suspect that this is due to a combination of a more reasonable evaluation and using SIFT as opposed to Nearest Neighbor retrieval for data selection.

#### Evaluation of Inference Cost of In-Context Baseline

We estimate the inference cost of in-context learning as follows.
We evaluate the time it takes compute the rolling log-likelihood of the test instance with context included and subtract the time it takes to compute the rolling log-likelihood of the test instance without context.
This is a lower-bound of the inference cost of in-context learning, as unlike autoregressive generation, computing the log-likelihood is partially parallelized.

To compute the token throughput of the in-context baseline, we divide the total compute time by the number of tokens added to the context.

### I.3 Inference Cost with Test-Time Fine-Tuning

Figure 22 evaluates the inference cost when test-time fine-tuning is used with GPT-2.
Notably, we do not use a state-of-the-art GPU for this evaluation.
Our results show that test-time fine-tuning can even be computationally feasible in low-latency applications.

Figure: Figure 22: Cost of taking a single gradient step with GPT-2. Results are with an NVIDIA RTX 4090.
Refer to caption: x21.png

## Appendix J Ablations

This section summarizes ablations that we conducted to investigate test-time fine-tuning and SIFT.

- •
Hyperparameter $\lambda^{\prime}$: Table 11
- •
Learning Curves for Individual Datasets of the Pile: Figure 23
- •
Learning Rate: Figure 24
- •
Uncertainty Estimation:
–
Summary of correlations (Table 12)
–
Visualization of $\sigma_{n}$ (Figure 25)
–
We fit the power law $\sigma_{n}\simeq 0.244n^{-0.684}+0.123$, which accurately describes the scaling of $\sigma_{n}$ (Figure 26, Table 13)
- •
Compute-proportional Performance Gain:
–
Details on Adaptive SIFT (Figure 27)

Table: Table 11: Percentage of bits per byte after $50$ test-time iterations for varying $\lambda^{\prime}$, relative to the bits per byte of the base model. We only include datasets with at least $10$ examples in our test set. Bold numbers denote the best performing selected subset. Underlined numbers denote better or on-par performance with Nearest Neighbor retrieval. $\Delta$ denotes the performance gain of SIFT with the strongest $\lambda^{\prime}$ *per dataset* over Nearest Neighbor retrieval. Numbers in partentheses are standard errors. We remark that $\lambda^{\prime}$ is on a logarithmic scale. For any choice of $\lambda^{\prime}\in[1\mathrm{e-}8,10]$, SIFT *always* performs at least on-par with Nearest Neighbor retrieval.

Figure: Figure 23: Performance in some of the datasets of the Pile, with GPT-2 as base model. Error bars correspond to standard errors.
Refer to caption: x22.png

Figure: Figure 24: Ablation of the learning rate with data selected by Nearest Neighbor retrieval. We find that the default learning rate $5\mathrm{e}{-5}$ of the transformers library (Wolf et al., ) works best, and conduct our other experiments with this learning rate unless noted otherwise. Hardt & Sun () had previously used $2\mathrm{e}{-5}$ which we find to be suboptimal.
Refer to caption: x23.png

**Table 12: We find a strong / moderate correlation between the uncertainty estimates $\hat{\sigma}_{n}$ / $\sigma_{n}$ and bits per byte. We further consider the correlation at all test-time iterations (from 0 0 to $50$) as well as only at the final iteration. We report both the Spearman and Pearson correlation coefficients, measuring monotonic and linear relationships, respectively. Before determining the Pearson correlation, we exclude the 0.25% of the data points with the lowest and highest uncertainty estimates to avoid the influence of outliers. The p-value of all correlations is below $1\mathrm{e}{-5}$ due to the large sample size.**
|  |  | Spearman | Pearson |
| --- | --- | --- | --- |
| $\sigma_{n}$ | all steps | 0.485 | 0.421 |
| final step | 0.496 | 0.443 |  |
| $\hat{\sigma}_{n}$ | all steps | 0.722 | 0.581 |
| final step | 0.682 | 0.482 |  |
| $\log\sigma_{n}$ | all steps | 0.485 | 0.468 |
| final step | 0.496 | 0.466 |  |
| $\log\hat{\sigma}_{n}$ | all steps | 0.722 | 0.618 |
| final step | 0.682 | 0.526 |  |

Figure: Figure 25: We visualize the predictive ability of the uncertainty estimates $\sigma_{n}$ analogously to Figure 8.
Refer to caption: x24.png

Figure: Figure 26: We fit a scaling law to the uncertainty estimate $\sigma_{n}$. We report statistics of the fit in Table 13.
Refer to caption: x25.png

**Table 13: Power law fit to uncertainty estimate: $\sigma_{n}\simeq An^{-\beta}+B$, visualized in Figure 26. The model achieves an $R^{2}$ of $0.999$, indicating that $99.9\%$ of the variance in uncertainty estimates is explained by the model.**
| Parameter | Estimate | Standard Error | 95% Bootstrap Confidence Interval |
| --- | --- | --- | --- |
| $\beta$ | $0.684$ | $0.007$ | $[0.668,0.788]$ |
| $A$ | $0.244$ | $0.001$ | $[0.241,0.271]$ |
| $B$ | $0.123$ | $0.001$ | $[0.122,0.129]$ |

Figure: Figure 27: We evaluate Adaptive SIFT with the same choices of $\alpha$ as in Figure 8 (right). Left: Bits per byte of Adaptive SIFT ($\downarrow$ better) against test-time compute. Every marker corresponds to the performance of Adaptive SIFT with a given $\alpha$, where the associated test-time compute is the average number of test-time iterations on prompts. We compare Adaptive SIFT to SIFT, where we spend the same test-time compute on every prompt. We see a slight advantage of Adaptive SIFT over SIFT, due to adaptively stopping depending on the prompt. Our current experiment exhibits a bias as test-time compute approaches $50$, since we force-stop the compute at $50$ iterations. This biases Adaptive SIFT to perform similarly to SIFT. We hypothesize that the initial advantage of Adaptive SIFT over SIFT may grow with more test-time compute if compute is not force-stopped at $50$ iterations. Right: Frequency of stopping at a given iteration for given values of $\alpha$.
Refer to caption: x26.png

## Appendix K Proofs

This section provides the formal proofs of the results presented in the main text.

- •
§K.2 proves the insufficiency of Nearest Neighbor retrieval (2.1).
- •
§K.3 shows the close relationship of regularized loss minimization and test-time fine-tuning (Proposition 3.3).
- •
§K.4 details how SIFT balances relevance and diversity (§C.1).
- •
§K.5 states confidence sets for fine-tuning regression models that are analogous to the confidence sets for classification from the main text.
- •
§K.6 derives the confidence sets from the main text (Theorem 3.2).

### K.1 Notation

Throughout this work, $\log$ denotes the natural logarithm.
Unless noted otherwise $\{\dots\,\!\}$ denotes a multiset.
We define the feature map $\boldsymbol{\Phi}_{n}\operatorname*{\,\dot{=}\,}(\boldsymbol{\phi}(\boldsymbol
{x}_{1}),\dots,\boldsymbol{\phi}(\boldsymbol{x}_{n}))\in\mathbb{R}^{n\times d}$, which gives rise to the kernel matrix $\smash{\boldsymbol{K}_{\!n}\operatorname*{\,\dot{=}\,}\boldsymbol{K}_{\!X_{n}}
=\boldsymbol{\Phi}_{n}\boldsymbol{\Phi}_{n}^{\top}\in\mathbb{R}^{n\times n}}$ and the covariance operator $\smash{\boldsymbol{\Sigma}_{n}\operatorname*{\,\dot{=}\,}\boldsymbol{\Phi}_{n}
^{\top}\boldsymbol{\Phi}_{n}\in\mathbb{R}^{d\times d}}$.

### K.2 Insufficiency of Nearest Neighbor Retrieval ( 2.1 )

We refer to §C.2 for the formal definition of the irreducible uncertainty $\sigma_{\infty}(\boldsymbol{x}^{\star};\mathcal{D})$.

We remark that if embeddings are of unit length, the cosine similarity scoring function is equivalent to the (negative) Euclidean distance scoring function:

$$ $\displaystyle\left\|\boldsymbol{x}^{\star}-\boldsymbol{x}\right\|_{2}^{2}$ $\displaystyle=(\boldsymbol{x}^{\star}-\boldsymbol{x})^{\top}(\boldsymbol{x}^{ \star}-\boldsymbol{x})=\left\|\boldsymbol{x}^{\star}\right\|_{2}^{2}+\left\| \boldsymbol{x}\right\|_{2}^{2}-2{\boldsymbol{x}^{\star}}^{\top}\boldsymbol{x}= 2-2\cos(\boldsymbol{x}^{\star},\boldsymbol{x}).$ $$

We henceforth consider the Euclidean distance scoring function.

###### Proposition K.1 (Insufficiency of Nearest Neighbor Retrieval) .

Suppose w.l.o.g. that $\boldsymbol{\phi}(\boldsymbol{x})=\boldsymbol{x}$.
Consider the data space $\mathcal{D}=\bigcup_{i=1}^{d}\mathcal{D}_{i}$ where $\mathcal{D}_{i}=\{\boldsymbol{e}_{i}\mid j\in\mathbb{N}\}$ with $\boldsymbol{e}_{i}$ the $i$-th basis vector of $\mathbb{R}^{d}$.(^13^1313We remark that $\{\dots\,\!\}$ denotes a multiset.)
Let $\boldsymbol{x}^{\star}=\frac{1}{\sqrt{4+(d-1)}}(2,1,1,\dots,1)\in\mathbb{R}^{d}$.

Then, for all $n\geq 1$:

- 1.
If $X_{n}$ are the $n$ nearest neighbors of $\boldsymbol{x}^{\star}$ in $\mathcal{D}$, $\sigma_{n}^{2}(\boldsymbol{x}^{\star})\geq\sigma_{\infty}^{2}(\boldsymbol{x}^{
\star};\mathcal{D}_{1})\gg 0$.
- 2.
If $X_{n}$ is selected by SIFT, $\sigma_{n}^{2}(\boldsymbol{x}^{\star})\overset{n\to\infty}{\longrightarrow}
\sigma_{\infty}^{2}(\boldsymbol{x}^{\star};\mathcal{D})=0$.

###### Proof.

- 1.
Clearly, $\left\|\boldsymbol{x}^{\star}-\boldsymbol{e}_{1}\right\|_{2}^{2}<\left\|
\boldsymbol{x}^{\star}-\boldsymbol{e}_{i}\right\|_{2}^{2}$ for all $i>1$.
Hence, $X_{n}=\{\boldsymbol{e}_{1}\mid i\in[n]\}\subset\mathcal{D}_{1}$.
This is as if the data space was restricted to $\mathcal{D}_{1}$, and hence $\sigma_{n}^{2}(\boldsymbol{x}^{\star})\geq\sigma_{\infty}^{2}(\boldsymbol{x}^{
\star};\mathcal{D}_{1})$.
- 2.
This follows readily from Theorem C.2 and noting that $\operatorname*{span}{\mathcal{D}}=\mathbb{R}^{d}$, implying $\sigma_{\infty}^{2}(\boldsymbol{x}^{\star};\mathcal{D})=0$.

∎

###### Proposition K.1 (Insufficiency of Nearest Neighbor Retrieval) .

###### Proof.

#### Discussion

The setting examined in Proposition K.1 is an extreme case (where data exists with exact duplication), yet we deem that it illustrates a realistic scenario.
Particularly nowadays that similar information is accessible from many sources in different forms, it is crucial to explicitly select diverse data for fine-tuning.
We show here theoretically and in Section L.1 qualitatively that SIFT does not have this limitation.

### K.3 The close relationship of Regularized Loss Minimization and Test-Time Fine-Tuning ( Proposition 3.3 )

###### Proof.

We note that the regularized negative log-likelihood loss $\mathcal{L}^{\lambda}$ from Equation 1,

$$ $\displaystyle\mathcal{L}^{\lambda}(\boldsymbol{W};D)$ $\displaystyle=\underbrace{-\sum_{(\boldsymbol{x},y)\in D}\log s_{y}( \boldsymbol{W}\boldsymbol{\phi}(\boldsymbol{x}))}_{\mathcal{L}(\boldsymbol{W}; D)}+\frac{\lambda}{2}\left\|\boldsymbol{W}-{\boldsymbol{W}^{\mathrm{pre}}} \right\|_{\mathrm{F}}^{2},$ $$

is strictly convex in $\boldsymbol{W}$ and has a unique minimizer $\boldsymbol{W}_{\!\lambda}$ which satisfies

$$ $\displaystyle\boldsymbol{\nabla}\mathcal{L}^{\lambda}(\boldsymbol{W}_{\! \lambda};D)=\boldsymbol{\nabla}\mathcal{L}(\boldsymbol{W}_{\!\lambda};D)+ \lambda(\boldsymbol{W}_{\!\lambda}-{\boldsymbol{W}^{\mathrm{pre}}})= \boldsymbol{0}.$ $$

It follows that $\boldsymbol{W}_{\!\lambda}={\boldsymbol{W}^{\mathrm{pre}}}-\frac{1}{\lambda}
\boldsymbol{\nabla}\mathcal{L}(\boldsymbol{W}_{\!\lambda};D)$.

Finally, recall that $\boldsymbol{\widehat{W}}_{\eta}={\boldsymbol{W}^{\mathrm{pre}}}-\eta
\boldsymbol{\nabla}\mathcal{L}({\boldsymbol{W}^{\mathrm{pre}}};D)$.
We obtain

$$ $\displaystyle\|\boldsymbol{W}_{\!1/\eta}-\boldsymbol{\widehat{W}}_{\!\eta}\|_{ \mathrm{F}}$ $\displaystyle=\|\eta\boldsymbol{\nabla}\mathcal{L}({\boldsymbol{W}^{\mathrm{ pre}}};D)-\eta\boldsymbol{\nabla}\mathcal{L}(\boldsymbol{W}_{\!1/\eta};D)\|_{ \mathrm{F}}$ $\displaystyle=\eta\|\boldsymbol{\nabla}\mathcal{L}({\boldsymbol{W}^{\mathrm{ pre}}};D)-\boldsymbol{\nabla}\mathcal{L}(\boldsymbol{W}_{\!1/\eta};D)\|_{ \mathrm{F}}.$ $$

∎

###### Proof.

### K.4 How SIFT Balances Relevance and Diversity

#### 1st point

For non-unit length embeddings, the first selected point can be expressed as follows:

$$ $\displaystyle\boldsymbol{x}_{1}=\operatorname*{arg\,min}_{\boldsymbol{x}\in \mathcal{D}}\sigma_{\{\boldsymbol{x}\}}^{2}(\boldsymbol{x}^{\star})= \operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\frac{(\boldsymbol{\phi }(\boldsymbol{x}^{\star})^{\top}\boldsymbol{\phi}(\boldsymbol{x}))^{2}}{\left \|\boldsymbol{\phi}(\boldsymbol{x})\right\|_{2}^{2}+\lambda^{\prime}}= \operatorname*{arg\,max}_{\boldsymbol{x}\in\mathcal{D}}\begin{cases} \measuredangle_{\boldsymbol{\phi}}(\boldsymbol{x}^{\star},\boldsymbol{x})^{2}& \text{as $\lambda^{\prime}\to 0$}\\ (\boldsymbol{\phi}(\boldsymbol{x}^{\star})^{\top}\boldsymbol{\phi}(vx))^{2}& \text{as $\lambda^{\prime}\to\infty$}.\end{cases}$ $$

#### 2nd point

Next, we consider the second selected point.
We derive the results in terms of the dot product kernel $k(\boldsymbol{x},\boldsymbol{x^{\prime}})=\boldsymbol{\phi}(\boldsymbol{x})^{
\top}\boldsymbol{x^{\prime}}$ which is such that $k(\boldsymbol{x},\boldsymbol{x^{\prime}})=\measuredangle_{\boldsymbol{\phi}}(
\boldsymbol{x},\boldsymbol{x^{\prime}})$ for unit length embeddings.
Let $\boldsymbol{x}$ be such that $k(\boldsymbol{x}_{1},\boldsymbol{x})=0$.
We have

$$ $\displaystyle\psi_{\boldsymbol{x}^{\star}}(\{\boldsymbol{x}_{1},\boldsymbol{x} _{1}\})$ $\displaystyle=\begin{bmatrix}k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\end{bmatrix}^{\top}\begin{bmatrix }1+\lambda^{\prime}&1\\ 1&1+\lambda^{\prime}\end{bmatrix}^{-1}\begin{bmatrix}k(\boldsymbol{x}^{\star}, \boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\end{bmatrix}$ $\displaystyle=\frac{1}{(1+\lambda^{\prime})^{2}-1}\begin{bmatrix}k(\boldsymbol {x}^{\star},\boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\end{bmatrix}^{\top}\begin{bmatrix }1+\lambda^{\prime}&-1\\ -1&1+\lambda^{\prime}\end{bmatrix}\begin{bmatrix}k(\boldsymbol{x}^{\star}, \boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\end{bmatrix}$ $\displaystyle=\frac{2\lambda^{\prime}k(\boldsymbol{x}^{\star},\boldsymbol{x}_{ 1})^{2}}{(1+\lambda^{\prime})^{2}-1}$ $\displaystyle=\frac{2k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})^{2}}{2+ \lambda^{\prime}}.$ $$

For $\boldsymbol{x}$, we have

$$ $\displaystyle\psi_{\boldsymbol{x}^{\star}}(\{\boldsymbol{x}_{1},\boldsymbol{x}\})$ $\displaystyle=\begin{bmatrix}k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x})\end{bmatrix}^{\top}\begin{bmatrix}1+ \lambda^{\prime}&0\\ 0&1+\lambda^{\prime}\end{bmatrix}^{-1}\begin{bmatrix}k(\boldsymbol{x}^{\star}, \boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x})\end{bmatrix}$ $\displaystyle=\frac{1}{(1+\lambda^{\prime})^{2}}\begin{bmatrix}k(\boldsymbol{x }^{\star},\boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x})\end{bmatrix}^{\top}\begin{bmatrix}1+ \lambda^{\prime}&0\\ 0&1+\lambda^{\prime}\end{bmatrix}\begin{bmatrix}k(\boldsymbol{x}^{\star}, \boldsymbol{x}_{1})\\ k(\boldsymbol{x}^{\star},\boldsymbol{x})\end{bmatrix}$ $\displaystyle=\frac{k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})^{2}+k( \boldsymbol{x}^{\star},\boldsymbol{x})^{2}}{1+\lambda^{\prime}}.$ $$

We see that $\boldsymbol{x}$ is preferred over $\boldsymbol{x}^{\star}$ if and only if

$$ $\displaystyle\frac{k(\boldsymbol{x}^{\star},\boldsymbol{x}_{1})^{2}+k( \boldsymbol{x}^{\star},\boldsymbol{x})^{2}}{1+\lambda^{\prime}}>\frac{2k( \boldsymbol{x}^{\star},\boldsymbol{x}_{1})^{2}}{2+\lambda^{\prime}}\quad\iff \quad k(\boldsymbol{x}^{\star},\boldsymbol{x})^{2}>\underbrace{\frac{\lambda^{ \prime}}{2+\lambda^{\prime}}}_{c(\lambda^{\prime})}k(\boldsymbol{x}^{\star}, \boldsymbol{x}_{1})^{2}.$ $$

As $\lambda^{\prime}\to\infty$, $c(\lambda^{\prime})\to 1$; whereas as $\lambda^{\prime}\to 0$, $c(\lambda^{\prime})\to 0$.

We interpret the expressions extensively in Section 4.

### K.5 Confidence Sets for Regression

Before moving on to deriving confidence sets for the setting with categorical feedback, we state analogous results for the regression setting under the following standard assumptions.
For ease of notation, we consider the scalar case.

###### Assumption K.2 (Linear function in a known latent space) .

We assume $f^{\star}(\boldsymbol{x})=\boldsymbol{\phi}(\boldsymbol{x})^{\top}\boldsymbol{
w}^{\star}$ with ${\boldsymbol{w}^{\star}\in\mathbb{R}^{d}}$ and where ${\boldsymbol{\phi}(\cdot)\in\mathbb{R}^{d}}$ denotes known embeddings.
We assume that $\boldsymbol{w}^{\star}$ has bounded norm, i.e., $\left\|\boldsymbol{w}^{\star}-{\boldsymbol{w}^{\mathrm{pre}}}\right\|_{2}\leq B$ for some finite $B\in\mathbb{R}$.

###### Assumption K.3 (Sub-Gaussian Noise) .

We assume that the data follows

$$ $\displaystyle y_{n}=f^{\star}(\boldsymbol{x}_{n})+\varepsilon_{n}$ $$

where each $\varepsilon_{n}$ from the noise sequence $\{\varepsilon_{n}\}_{n=1}^{\infty}$ is conditionally zero-mean $\rho$-sub-Gaussian with known constant ${\rho>0}$.
Formally,

$$ $\displaystyle\forall n\geq 1,\lambda\in\mathbb{R}:\quad\mathbb{E}{}\!\left[e^{ \lambda\epsilon_{n}}\ \middle|\ D_{n-1}\right]\leq\exp\left(\frac{\lambda^{2} \rho^{2}}{2}\right)$ $$

where $D_{n-1}$ corresponds to the $\sigma$-algebra generated by the random variables $\{\boldsymbol{x}_{i},\epsilon_{i}\}_{i=1}^{n-1}$ and $\boldsymbol{x}_{n}$.

We consider the standard squared loss $\mathcal{L}(\boldsymbol{w};D)\operatorname*{\,\dot{=}\,}\frac{1}{2}\sum_{(
\boldsymbol{x},y)\in D}(f(\boldsymbol{x};\boldsymbol{w})-y)^{2}$ where we write $f(\boldsymbol{x};\boldsymbol{w})\operatorname*{\,\dot{=}\,}\boldsymbol{\phi}(
\boldsymbol{w})^{\top}\boldsymbol{w}$.
The regularized loss with minimizer $\boldsymbol{w}_{n}$ is then

$$ $\displaystyle\mathcal{L}^{\lambda}(\boldsymbol{w};D_{n})\operatorname*{\,\dot{ =}\,}\mathcal{L}(\boldsymbol{w};D_{n})+\frac{\lambda}{2}\left\|f-{\boldsymbol{ w}^{\mathrm{pre}}}\right\|_{2}^{2}$ (10) $$

where $\lambda>0$ is the regularization parameter.
In the following, we write $f_{n}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}f(\boldsymbol{x};\boldsymbol{w
}_{n})$ and $f^{\mathrm{pre}}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}f(\boldsymbol{x};{
\boldsymbol{w}^{\mathrm{pre}}})$.
The closed-form solution to the optimization problem from Equation 10 is well-known (see, e.g., Williams & Rasmussen, , Section 6.2.2) to be

$$ $\displaystyle f_{n}(\boldsymbol{x})=f^{\mathrm{pre}}(\boldsymbol{x})+ \boldsymbol{k}_{X_{n}}^{\top}(\boldsymbol{x})(\boldsymbol{K}_{\!X_{n}}+\lambda \boldsymbol{I}_{n})^{-1}(\boldsymbol{y}_{n}-\boldsymbol{f}^{\mathrm{pre}}_{n})$ $$

where $\boldsymbol{f}^{\mathrm{pre}}_{n}$ is the vector of predictions of $f^{\mathrm{pre}}$ at $X_{n}$ and $\boldsymbol{y}_{n}$ is the vector of observations in $D_{n}$.

The below result is an almost immediate consequence of the results of Abbasi-Yadkori () and Chowdhury & Gopalan ().

###### Theorem K.4 (Confidence Sets for Regression) .

Pick ${\delta\in(0,1)}$ and let Assumptions K.2 and K.3 hold.
Let

$$ $\displaystyle\beta_{n}(\delta)\operatorname*{\,\dot{=}\,}B+\rho\sqrt{2(\gamma_ {n}+1+\log(1/\delta))}$ $$

where $\gamma_{n}\operatorname*{\,\dot{=}\,}\max_{\boldsymbol{x}_{1},\dots,
\boldsymbol{x}_{n}}\frac{1}{2}\log\det{(\boldsymbol{I}_{n}+\lambda^{-1}
\boldsymbol{K}_{X_{n}})}$.
Then

$$ $\displaystyle\mathbb{P}\left(\forall n\geq 1,\boldsymbol{x}\in\mathcal{X}:|f^{ \star}(\boldsymbol{x})-f_{n}(\boldsymbol{x})|\leq\beta_{n}(\delta)\sigma_{n}( \boldsymbol{x})\right)\geq 1-\delta.$ $$

###### Proof.

Let us define the *residual* of the ground truth and pre-trained model as ${\tilde{f}^{\star}(\boldsymbol{x})\operatorname*{\,\dot{=}\,}f^{\star}(
\boldsymbol{x})-f^{\mathrm{pre}}(\boldsymbol{x})}$ with corresponding weight vector $\tilde{\boldsymbol{w}}$.
Analogously, let $\tilde{y}_{n}=\tilde{f}^{\star}(\boldsymbol{x}_{n})+\varepsilon_{n}$ be the observed error.
We have that $\tilde{\boldsymbol{w}}^{\star}\operatorname*{\,\dot{=}\,}\boldsymbol{w}^{\star
}-{\boldsymbol{w}^{\mathrm{pre}}}\in\mathbb{R}^{d}$ with norm $\left\|\boldsymbol{w}^{\star}-{\boldsymbol{w}^{\mathrm{pre}}}\right\|_{k}$.
The unbiased estimate of the remaining error is

$$ $\displaystyle\tilde{f}_{n}=\boldsymbol{k}_{X_{n}}^{\top}(\boldsymbol{x})( \boldsymbol{K}_{\!X_{n}}+\lambda\boldsymbol{I}_{n})^{-1}\tilde{\boldsymbol{y}} _{n}.$ $$

By Theorem 2 of Chowdhury & Gopalan (), for all $\boldsymbol{x}\in\mathcal{X}$ and $n\geq 1$, jointly with probability at least $1-\delta$, ${|\tilde{f}^{\star}(\boldsymbol{x})-\tilde{f}_{n}(\boldsymbol{x})|\leq\beta_{n
}(\delta)\sigma_{n}(\boldsymbol{x})}$.
It remains now only to observe that

$$ $\displaystyle|\tilde{f}^{\star}(\boldsymbol{x})-\tilde{f}_{n}(\boldsymbol{x})| =|f^{\star}(\boldsymbol{x})-f_{n}(\boldsymbol{x})|.$ $$

∎

###### Assumption K.2 (Linear function in a known latent space) .

###### Assumption K.3 (Sub-Gaussian Noise) .

###### Theorem K.4 (Confidence Sets for Regression) .

###### Proof.

### K.6 Confidence Sets for Classification ( Theorem 3.2 )

We begin by re-stating Corollary 1 of Amani & Thrampoulidis ().
Analogous results can be obtained from Theorem 1 of Zhang & Sugiyama ().
Substantial work has studied the special case of binary feedback, $K=2$ Faury et al. (); Pásztor et al. ().

Let $\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W})\in\mathbb{R}^{K\times K}$ be the matrix satisfying ${(\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W}))_{i,j}\operatorname*{\,\dot{=}
\,}s_{i}(\boldsymbol{x};\boldsymbol{W})(\mathbbm{1}\{{i=j}\}-s_{j}(\boldsymbol
{x};\boldsymbol{W}))}$.
Equivalently, ${\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W})=\mathrm{diag}{}\left\{
\boldsymbol{s}(\boldsymbol{x};\boldsymbol{W})\right\}-\boldsymbol{s}(
\boldsymbol{x};\boldsymbol{W})\boldsymbol{s}(\boldsymbol{x};\boldsymbol{W})^{
\top}}$.
Based on this matrix, we define ${L\operatorname*{\,\dot{=}\,}\sup_{\boldsymbol{x}\in\mathcal{X},\boldsymbol{W}
\in\mathcal{W}}\lambda_{\max}(\boldsymbol{A}(\boldsymbol{x};\boldsymbol{W}))}$ and ${\kappa\operatorname*{\,\dot{=}\,}\sup_{\boldsymbol{x}\in\mathcal{X},
\boldsymbol{W}\in\mathcal{W}}1/\lambda_{\min}(\boldsymbol{A}(\boldsymbol{x};
\boldsymbol{W}))}$.

###### Lemma K.5 (Corollary 1 of Amani & Thrampoulidis ( 2020 ) ) .

Assume $\boldsymbol{W}^{\star}\in\mathcal{W}$ and ${\boldsymbol{W}^{\mathrm{pre}}}=\boldsymbol{0}$.
Let $\delta\in(0,1)$ and set

$$ $\displaystyle\tilde{\beta}_{n}(\delta)\operatorname*{\,\dot{=}\,}\sqrt{\lambda }\left(B+\frac{1}{2\sqrt{K}}\right)+\frac{2K^{3/2}d}{\sqrt{\lambda}}\log\left( \frac{2}{\delta}\sqrt{1+\frac{n}{d\lambda}}\right).$ (11) $$

Then,

$$ $\displaystyle\mathbb{P}\left(\forall n\geq 1,\boldsymbol{x}\in\mathcal{X}:\| \boldsymbol{s}_{n}(\boldsymbol{x})-\boldsymbol{s}^{\star}(\boldsymbol{x})\|_{2 }\leq 2L\tilde{\beta}_{n}(\delta)\sqrt{\kappa(1+2B)}\left\|\boldsymbol{\phi}( \boldsymbol{x})\right\|_{\boldsymbol{V}_{n}^{-1}}\right)\geq 1-\delta,$ $$

where $\boldsymbol{V}_{\!n}\operatorname*{\,\dot{=}\,}\boldsymbol{\Sigma}_{n}+\kappa
\lambda\boldsymbol{I}_{d}$.

Our result follows from two auxiliary lemmas.

###### Lemma K.6 .

For any $\boldsymbol{s},\boldsymbol{s^{\prime}}\in\mathbb{R}^{K}$, $d_{\mathrm{TV}}\left(\boldsymbol{s},\boldsymbol{s^{\prime}}\right)\leq\frac{
\sqrt{K}}{2}\left\|\boldsymbol{s}-\boldsymbol{s^{\prime}}\right\|_{2}$.

###### Proof.

We have

$$ $\displaystyle d_{\mathrm{TV}}\left(\boldsymbol{s},\boldsymbol{s^{\prime}}\right)$ $\displaystyle=\frac{1}{2}\left\|\boldsymbol{s}-\boldsymbol{s^{\prime}}\right\| _{1}=\frac{1}{2}\sum_{i=1}^{K}|s_{i}-s^{\prime}_{i}|\leq\frac{1}{2}\sqrt{K} \sqrt{\sum_{i=1}^{K}(s_{i}-s^{\prime}_{i})^{2}}=\frac{\sqrt{K}}{2}\left\| \boldsymbol{s}-\boldsymbol{s^{\prime}}\right\|_{2}$ $$

where the inequality follows from Cauchy-Schwarz.
∎

The following lemma is a standard result in the literature (Srinivas et al., ; Chowdhury & Gopalan, ; Pásztor et al., ), which we include here for completeness.

###### Lemma K.7 .

Let $\sigma_{n}$ be as defined in Equation 2.
Then, $\sqrt{\kappa\lambda}\left\|\boldsymbol{\phi}(\boldsymbol{x})\right\|_{
\boldsymbol{V}_{\!n}^{-1}}=\sigma_{n}(\boldsymbol{x})$ for any $\boldsymbol{x}\in\mathcal{X}$.

###### Proof.

Note that $(\boldsymbol{\Sigma}_{n}+\kappa\lambda\boldsymbol{I}_{d})\boldsymbol{\Phi}_{n}
^{\top}=\boldsymbol{\Phi}_{n}^{\top}(\boldsymbol{K}_{\!n}+\kappa\lambda
\boldsymbol{I}_{n})$ which implies

$$ $\displaystyle(\boldsymbol{\Sigma}_{n}+\kappa\lambda\boldsymbol{I}_{d})^{-1} \boldsymbol{\Phi}_{n}^{\top}=\boldsymbol{\Phi}_{n}^{\top}(\boldsymbol{K}_{\!n} +\kappa\lambda\boldsymbol{I}_{n})^{-1}.$ (12) $$

Further, by definition of $\boldsymbol{k}_{n}$, $\boldsymbol{k}_{n}(\boldsymbol{x})=\boldsymbol{\Phi}_{n}\boldsymbol{\phi}(
\boldsymbol{x})$ which permits writing

$$ $\displaystyle(\boldsymbol{\Sigma}_{n}+\kappa\lambda\boldsymbol{I}_{d}) \boldsymbol{\phi}(\boldsymbol{x})=\boldsymbol{\Phi}_{n}^{\top}\boldsymbol{k}_{ n}(\boldsymbol{x})+\kappa\lambda\boldsymbol{\phi}(\boldsymbol{x})$ $$

and implies

$$ $\displaystyle\boldsymbol{\phi}(\boldsymbol{x})$ $\displaystyle=(\boldsymbol{\Sigma}_{n}+\kappa\lambda\boldsymbol{I}_{d})^{-1} \boldsymbol{\Phi}_{n}^{\top}\boldsymbol{k}_{n}(\boldsymbol{x})+\kappa\lambda( \boldsymbol{\Sigma}_{n}+\kappa\lambda\boldsymbol{I}_{d})^{-1}\boldsymbol{\phi} (\boldsymbol{x})$ $\displaystyle\overset{(\ref{eq:kernel_helper_lemma_helper1})}{=}\boldsymbol{ \Phi}_{n}^{\top}(\boldsymbol{K}_{\!n}+\kappa\lambda\boldsymbol{I}_{n})^{-1} \boldsymbol{k}_{n}(\boldsymbol{x})+\kappa\lambda(\boldsymbol{\Sigma}_{n}+ \kappa\lambda\boldsymbol{I}_{d})^{-1}\boldsymbol{\phi}(\boldsymbol{x})$ (13) $$

We have

$$ $\displaystyle k(\boldsymbol{x},\boldsymbol{x})$ $\displaystyle=\boldsymbol{\phi}(\boldsymbol{x})^{\top}\boldsymbol{\phi}( \boldsymbol{x})$ $\displaystyle\overset{(\ref{eq:kernel_helper_lemma_helper2})}{=}\left( \boldsymbol{\Phi}_{n}^{\top}(\boldsymbol{K}_{\!n}+\kappa\lambda\boldsymbol{I}_ {n})^{-1}\boldsymbol{k}_{n}(\boldsymbol{x})+\kappa\lambda(\boldsymbol{\Sigma}_ {n}+\kappa\lambda\boldsymbol{I}_{d})^{-1}\boldsymbol{\phi}(\boldsymbol{x}) \right)^{\top}\boldsymbol{\phi}(\boldsymbol{x})$ $\displaystyle=\boldsymbol{k}_{n}(\boldsymbol{x})^{\top}(\boldsymbol{K}_{\!n}+ \kappa\lambda\boldsymbol{I}_{n})^{-1}\boldsymbol{k}_{n}(\boldsymbol{x})+\kappa \lambda\boldsymbol{\phi}(\boldsymbol{x})^{\top}(\boldsymbol{\Sigma}_{n}+\kappa \lambda\boldsymbol{I}_{d})^{-1}\boldsymbol{\phi}(\boldsymbol{x})$ $\displaystyle=\boldsymbol{k}_{n}(\boldsymbol{x})^{\top}(\boldsymbol{K}_{\!n}+ \kappa\lambda\boldsymbol{I}_{n})^{-1}\boldsymbol{k}_{n}(\boldsymbol{x})+\kappa \lambda\boldsymbol{\phi}(\boldsymbol{x})^{\top}\boldsymbol{V}_{\!n}^{-1} \boldsymbol{\phi}(\boldsymbol{x}).$ $$

Reordering this equation, we obtain

$$ $\displaystyle\kappa\lambda\left\|\boldsymbol{\phi}(\boldsymbol{x})\right\|_{ \boldsymbol{V}_{\!n}^{-1}}^{2}=\kappa\lambda\boldsymbol{\phi}(\boldsymbol{x})^ {\top}\boldsymbol{V}_{\!n}^{-1}\boldsymbol{\phi}(\boldsymbol{x})=k(\boldsymbol {x},\boldsymbol{x})-\boldsymbol{k}_{n}(\boldsymbol{x})^{\top}(\boldsymbol{K}_{ \!n}+\kappa\lambda\boldsymbol{I}_{n})^{-1}\boldsymbol{k}_{n}(\boldsymbol{x})= \sigma_{n}^{2}(\boldsymbol{x}),$ $$

concluding the proof.
∎

We now proceed to prove a version of Theorem 3.2 with ${\boldsymbol{W}^{\mathrm{pre}}}=\boldsymbol{0}$.

###### Theorem K.8 .

Assume $\boldsymbol{W}^{\star}\in\mathcal{W}$ and ${\boldsymbol{W}^{\mathrm{pre}}}=\boldsymbol{0}$.
Let $\delta\in(0,1)$ and $\beta_{n}(\delta)$ as in Equation 3.
Then

$$ $\displaystyle\mathbb{P}\left(\forall n\geq 1,\boldsymbol{x}\in\mathcal{X}:d_{ \mathrm{TV}}(\boldsymbol{s}_{n}(\boldsymbol{x}),\boldsymbol{s}^{\star}( \boldsymbol{x}))\leq\beta_{n}(\delta)\cdot\sigma_{n}(\boldsymbol{x})\right) \geq 1-\delta.$ $$

###### Proof.

We have

$$ $\displaystyle d_{\mathrm{TV}}(\boldsymbol{s}_{n}(\boldsymbol{x}),\boldsymbol{s }^{\star}(\boldsymbol{x}))$ $\displaystyle\leq\frac{\sqrt{K}}{2}\left\|\boldsymbol{s}_{n}(\boldsymbol{x})- \boldsymbol{s}^{\star}(\boldsymbol{x})\right\|_{2}$ ( Lemma K.6 ) $\displaystyle{\leq}L\tilde{\beta}_{n}(\delta)\sqrt{K\kappa(1+2B)}\left\| \boldsymbol{\phi}(\boldsymbol{x})\right\|_{\boldsymbol{V}_{\!n}^{-1}}$ ( Lemma K.5 ) $\displaystyle=L\tilde{\beta}_{n}(\delta)\sqrt{\frac{K(1+2B)}{\lambda}}\sigma_{ n}(\boldsymbol{x}).$ ( Lemma K.7 ) $$

It remains to note that

$$ $\displaystyle L\tilde{\beta}_{n}(\delta)\sqrt{\frac{K(1+2B)}{\lambda}}$ $\displaystyle=L\sqrt{K(1+2B)}\left(B+\frac{1}{2\sqrt{K}}\right)+\frac{2LK^{2}d \sqrt{1+2B}}{\lambda}\log\left(\frac{2}{\delta}\sqrt{1+\frac{n}{d\lambda}}\right)$ $\displaystyle\leq 2\sqrt{K(1+2B)}\left[B+\frac{LK^{3/2}d}{\lambda}\log\left( \frac{2}{\delta}\sqrt{1+\frac{n}{d\lambda}}\right)\right]=\beta_{n}(\delta).$ $$

∎

With this we are ready to prove Theorem 3.2.

###### Proof of Theorem 3.2 .

We will proceed analogously to the proof of Theorem K.4.
That is, our objective will be to bound the deviation of our biased model, which we refer to as $\boldsymbol{W}_{\!n}=\operatorname*{arg\,min}_{\boldsymbol{W}\in\mathcal{W}}
\mathcal{L}^{\lambda}(\boldsymbol{W};D_{n})$, to $\boldsymbol{W}^{\star}$.
Let

$$ $\displaystyle\tilde{\mathcal{L}}(\boldsymbol{W}^{\prime};D)\operatorname*{\, \dot{=}\,}-\sum_{(\boldsymbol{x},y)\in D}\log s_{y}((\boldsymbol{W}^{\prime}+{ \boldsymbol{W}^{\mathrm{pre}}})\boldsymbol{\phi}(\boldsymbol{x}))\quad\text{ and}\quad\tilde{\mathcal{L}}^{\lambda}(\boldsymbol{W}^{\prime};D)\operatorname *{\,\dot{=}\,}\tilde{\mathcal{L}}(\boldsymbol{W}^{\prime};D)+\frac{\lambda}{2} \left\|\boldsymbol{W}^{\prime}\right\|_{\mathrm{F}}^{2}$ $$

with minimizer $\boldsymbol{W}_{\!n}^{\prime}\operatorname*{\,\dot{=}\,}\operatorname*{arg\,
min}_{\boldsymbol{W}^{\prime}:\left\|\boldsymbol{W}^{\prime}\right\|_{\mathrm{
F}}\leq B}\tilde{\mathcal{L}}^{\lambda}(\boldsymbol{W}^{\prime};D_{n})$.
We further define the residual weights $\smash{{\tilde{\boldsymbol{W}}^{\star}}\operatorname*{\,\dot{=}\,}\boldsymbol{
W}^{\star}-{\boldsymbol{W}^{\mathrm{pre}}}}$.

Next, we make the following observation:
In their proof of Lemma K.5, Amani & Thrampoulidis () bound

$$ $\displaystyle\|\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}_{\! n}^{\prime}))-\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};\tilde{\boldsymbol{ W}}^{\star}))\|_{2}\leq\mathrm{const}\cdot\|\mathrm{vec}(\tilde{\boldsymbol{W} }^{\star})-\mathrm{vec}(\boldsymbol{W}_{\!n}^{\prime})\|_{\tilde{\boldsymbol{G }}(\tilde{\boldsymbol{W}}^{\star},\boldsymbol{W}_{\!n}^{\prime})}$ (14) $$

where $\mathrm{const}$ is independent of $\boldsymbol{W}^{\star},{\boldsymbol{W}^{\mathrm{pre}}},\boldsymbol{W}_{\!n}^{\prime}$ and the matrix $\tilde{\boldsymbol{G}}(\tilde{\boldsymbol{W}}^{\star},\boldsymbol{W}_{\!n}^{
\prime})$ is invariant to a change of variables, i.e., $\tilde{\boldsymbol{G}}(\tilde{\boldsymbol{W}}^{\star},\boldsymbol{W}_{\!n}^{
\prime})=\boldsymbol{G}(\boldsymbol{W}^{\star},\boldsymbol{W}_{\!n}^{\prime}+{
\boldsymbol{W}^{\mathrm{pre}}})$ with $\tilde{\boldsymbol{G}}$ defined with respect to the loss $\tilde{\mathcal{L}}^{\lambda}$ and $\boldsymbol{G}$ defined with respect to the loss $\mathcal{L}^{\lambda}$.
Theorem K.8 applies to $\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}_{\!n}^{\prime}))$ and $\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}^{\star}-{
\boldsymbol{W}^{\mathrm{pre}}}))$ since the regularization of $\tilde{\mathcal{L}}^{\lambda}$ is unbiased and the residual weights satisfy $\smash{\|\tilde{\boldsymbol{W}}^{\star}\|_{\mathrm{F}}=\left\|\boldsymbol{W}^{
\star}-{\boldsymbol{W}^{\mathrm{pre}}}\right\|_{\mathrm{F}}\leq B}$ by assumption.

Since ${\tilde{\boldsymbol{W}}^{\star}-\boldsymbol{W}_{\!n}^{\prime}=\boldsymbol{W}^{
\star}-(\boldsymbol{W}_{\!n}^{\prime}+{\boldsymbol{W}^{\mathrm{pre}}})}$, the bounds of Equation 14 as well as Theorem K.8 then also apply to $\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};\boldsymbol{W}_{\!n}^{\prime}+{
\boldsymbol{W}^{\mathrm{pre}}})),\boldsymbol{s}(\boldsymbol{f}(\boldsymbol{x};
\boldsymbol{W}^{\star}))$.
Observing that ${\boldsymbol{W}_{\!n}=\boldsymbol{W}_{\!n}^{\prime}+{\boldsymbol{W}^{\mathrm{
pre}}}}$ as a direct consequence of the change of variables completes the proof.
∎

###### Lemma K.5 (Corollary 1 of Amani & Thrampoulidis ( 2020 ) ) .

###### Lemma K.6 .

###### Proof.

###### Lemma K.7 .

###### Proof.

###### Theorem K.8 .

###### Proof.

###### Proof of Theorem 3.2 .

## Appendix L Qualitative Examples

### L.1 Balancing Relevance and Diversity

The following details the data space and prompt used in the qualitative example of Figure 3.
We evaluate SIFT with $\lambda^{\prime}=1\mathrm{e}{-4}$ and normalized embeddings, using the same embedding model as in our main experiments.

**Table 14: Query and information about Michael Jordan within data space**
|  | Prompt |
| --- | --- |
|  | What is the age of Michael Jordan and how many kids does he have? |
|  | Data space |
| 1 | Michael Jordan was born on February 17, 1963, in Brooklyn, New York. |
| 2 | The age of Michael Jordan is 61 years. |
| 3 | Michael Jordan has five children. |
| 4 | Michael Jordan has 5 kids. |

**Table 15: Nearest Neighbor selects redundant data, while SIFT selects data with maximum information**
| Prompt | What is the age of Michael Jordan and how many kids does he have? |
| --- | --- |
|  | Nearest Neighbor |
| 1 | The age of Michael Jordan is 61 years. |
| 2 | Michael Jordan was born on February 17, 1963, in Brooklyn, New York. |
|  | SIFT (ours) |
| 1 | The age of Michael Jordan is 61 years. |
| 2 | Michael Jordan has five children. |

### L.2 Irreducible Uncertainty

Recall the definition of the irreducible uncertainty of a prompt $\boldsymbol{x}^{\star}$ provided the data space $\mathcal{D}$: $\smash{\sigma_{\infty}^{2}(\boldsymbol{x}^{\star})\operatorname*{\,\dot{=}\,}
\left\|\boldsymbol{\phi}(\boldsymbol{x}^{\star})\right\|_{\boldsymbol{\Pi}_{
\boldsymbol{\Phi}}}^{2}}$.
The projection can be computed as follows: $\smash{\boldsymbol{\Pi}_{\boldsymbol{\Phi}}=\boldsymbol{I}-\boldsymbol{\Phi}^{
\top}(\boldsymbol{\Phi}\boldsymbol{\Phi}^{\top})^{-1}\boldsymbol{\Phi}}$ where $\boldsymbol{\Phi}$ denotes a basis of the embeddings $\{\boldsymbol{\phi}(\boldsymbol{x}):\boldsymbol{x}\in\mathcal{D}\}$ (Hübotter et al., , Lemma C.22).

#### Interpretation of Irreducible Uncertainty in a Practical Example ( Figure 28 )

Figure: Figure 28: We plot the irreducible uncertainty $\smash{\sigma_{\infty}^{2}(\boldsymbol{x}^{\star})=\left\|\boldsymbol{\phi}( \boldsymbol{x}^{\star})\right\|_{\boldsymbol{\Pi}_{\boldsymbol{\Phi}}}^{2}}$ of the prompt $\boldsymbol{x}^{\star}$ “What is the age of Michael Jordan and how many kids does he have?” with respect to a data space that includes *only kids information*, *only age information*, and *both information*. The black line denotes the irreducible uncertainty about the prompt if additionally “Michael Jordan is age 61 and has 5 kids.” is included in the data space. The area underneath can be interpreted as the inherent uncertainty about the embeddings. We see that if some information is missing, the irreducible uncertainty is large. If all information is present, the irreducible uncertainty is close to the inherent uncertainty. We evaluate two embedding models, the recent Stella (Zhang, ) with state-of-the-art performance in the Massive Text Embedding Benchmark (Muennighoff et al., ) and a large Roberta model (Liu, ) that was fine-tuned on the Pile dataset (Hardt & Sun, ). The absolute values of the irreducible uncertainty differ between the two models due to the different dimensionalities of the latent spaces.
Refer to caption: x27.png

The irreducible uncertainty quantifies the “missing information” within the data space relative to a given prompt $\boldsymbol{x}^{\star}$.
Figure 28 provides evidence in the context of the practical example considered in Section L.1 that this interpretation is meaningful.
For completeness, the considered data space is provided in Table 16.

**Table 16: Information about Michael Jordan within data space considered in Figure 28**
|  | Data space |
| --- | --- |
|  | *age information* |
| 1 | Michael Jordan was born on February 17, 1963. |
| 2 | The age of Michael Jordan is 61 years. |
| 3 | Michael Jordan was born on the 17th of February, 1963. |
| 4 | On February 17, 1963, Michael Jordan came into the world. |
| 5 | Michael Jordan’s birth date is February 17, 1963. |
| 6 | Born on February 17, 1963, Michael Jordan is a legendary basketball player. |
| 7 | February 17, 1963, marks the birth of basketball icon Michael Jordan. |
| 8 | Michael Jordan was born on February 17th, in the year 1963. |
| 9 | The basketball star Michael Jordan was born on February 17, 1963. |
| 10 | Michael Jordan’s birthday is February 17, 1963. |
| 11 | On February 17th, 1963, Michael Jordan was born. |
| 12 | Michael Jordan’s date of birth is February 17, 1963. |
| 13 | Michael Jordan, born February 17, 1963, is now 61 years old. |
| 14 | As of today, Michael Jordan is 61 years old, having been born on February 17, 1963. |
| 15 | Michael Jordan, born on February 17, 1963, is currently 61 years old. |
| 16 | The current age of Michael Jordan, born February 17, 1963, is 61 years. |
| 17 | Michael Jordan, who was born on February 17, 1963, is 61 years old. |
| 18 | At 61 years old, Michael Jordan was born on February 17, 1963. |
| 19 | As of today, Michael Jordan, born on February 17, 1963, is 61 years old. |
| 20 | Michael Jordan was born on February 17, 1963, making him 61 years old. |
| 21 | Michael Jordan, born in 1963 on February 17, is now 61. |
| 22 | Born in 1963, Michael Jordan is 61 years old as of this year. |
| 23 | Michael Jordan is 61 years old, having been born in February of 1963. |
| 24 | Currently aged 61, Michael Jordan was born on February 17, 1963. |
| 25 | Michael Jordan, now aged 61, was born in February 1963. |
| 26 | As of now, Michael Jordan is 61 years old, born in February 1963. |
| 27 | The age of Michael Jordan is 61; he was born in February of 1963. |
|  | *kids information* |
| 28 | Michael Jordan has five children. |
| 29 | Michael Jordan has 5 kids. |
| 30 | Michael Jordan is the father of five children. |
| 31 | Michael Jordan has a total of five kids. |
| 32 | Five children belong to Michael Jordan. |
| 33 | Michael Jordan has raised five children. |
| 34 | Michael Jordan is a proud father of five kids. |
| 35 | There are five children in Michael Jordan’s family. |
| 36 | Michael Jordan’s family includes five children. |
| 37 | Michael Jordan is a parent to five kids. |
| 38 | The basketball legend Michael Jordan has five kids. |
| 39 | Michael Jordan has five kids in his family. |
| 40 | Michael Jordan is a father of five kids. |
| 41 | Michael Jordan is the father of 5 children. |
| 42 | In total, Michael Jordan has five children. |
| 43 | Michael Jordan’s family consists of five children. |
| 44 | Michael Jordan’s household includes five kids. |
| 45 | There are five kids in Michael Jordan’s family. |
| 46 | Michael Jordan has five kids in total. |
| 47 | Michael Jordan, the father of five children, is a family man. |
| 48 | Michael Jordan has five wonderful children. |
| 49 | The legendary Michael Jordan has five kids. |
| 50 | Michael Jordan has 5 children. |
| 51 | Michael Jordan’s family has 5 kids. |
| 52 | The father of five kids is none other than Michael Jordan. |
| 53 | Michael Jordan’s household consists of five children. |
| 54 | Michael Jordan has been blessed with five kids. |
|  | *joint information* |
| 55 | Michael Jordan is age 61 and has 5 kids. |

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Please refer to Weaknesses.
soundness: 3
strengths: Effective Data Selection with SIFT: SIFT combines uncertainty and diversity to select non-redundant data, enhancing test-time fine-tuning efficiency compared to traditional methods.
Comprehensive Experiments: Wide experiments on the Pile dataset demonstrate SIFT’s effectiveness, consistently outperforming traditional Nearest Neighbor and other baseline methods in fine-tuning efficiency and model performance.
summary: The paper proposes a novel data selection algorithm called SIFT (Select Informative data for Fine-Tuning) aimed at improving test-time fine-tuning for large language models (LLMs). The paper addresses the limitations of Nearest Neighbor retrieval in data selection by combining ideas from retrieval and active learning. Unlike traditional Nearest Neighbor methods that can select redundant data, SIFT optimizes for information gain, reducing redundancy in the selected examples and enhancing the performance of LLMs during test-time fine-tuning. Evaluated on the Pile dataset, SIFT shows consistent improvement over traditional methods with minimal computational overhead, achieving a robust balance between relevance and diversity in data selection.
weaknesses: This paper has the following drawbacks.

**Complexity and Clarity of Method Presentation**: The paper’s explanation of the SIFT algorithm, especially on pages 4 and 5, could benefit from clearer descriptions and simplification of symbols. The complex notation and detailed mathematical formulation may obscure understanding for readers, especially those less familiar with active learning or information-theoretic approaches. Providing a more accessible walkthrough or visual aids could improve clarity.

**Sensitivity to Hyperparameters**: The method relies on certain hyperparameters, such as the regularization parameter (λ′), which can significantly impact SIFT’s performance. Techniques for automatic tuning or guidelines for parameter selection would make the approach more user-friendly and robust.

**Relevance and Diversity**. It is strange to mention that "we provide an example of how SIFT balances relevance and diversity, where we also see that the parameter"

**Limited Novelty in Uncertainty-Based Data Selection**: The use of uncertainty as a criterion for data selection is not a novel concept and is a common technique in active learning. Could the author explain the difference with existing works such as [1,2]. Furthermore, could you discuss the novelty of utilizing uncertainty in your LLM background?


[1] Symmetric Uncertainty-Aware Feature Transmission for Depth Super-Resolution,
[2] LogitNorm: Mitigating Neural Network Overconfidence with Logit Normalization

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 4
questions: See weaknesses.
soundness: 3
strengths: 1. This paper verifies a meaningful yet little explored topic in LLM application, giving comprehensive and solid discussion on the challenge and the proposed solution. The whole paper is clearly organized and easy to follow.
2. The SIFT strategy is sopported with solid theoretical induction as well as experimental evidences.
3. The experimental study is convincing and covers a wide range of datasets.
summary: This paper introduces SIFT, a data selection algorithm aimed at improving the fine-tuning of language models by addressing the limitations of Nearest Neighbors retrieval, which often selects redundant data. SIFT reduces uncertainty about model responses by optimizing overall information gain and accounting for information duplication. The authors demonstrate that SIFT outperforms Nearest Neighbor retrieval in fine-tuning at test time in experiments, with minimal computational overhead. Additionally, they show that their uncertainty estimates can predict performance gains, leading to an adaptive algorithm that optimally allocates computational resources.
weaknesses: 1. The interpretability of this work might be improved by giving some instances of data selection. 
2. How is the possibility and gain of combining this strategy with orthogonal methods of LLM finetuning?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: See weaknesses.
soundness: 3
strengths: 1. Theoretical Foundation: The paper is theoretically robust and well-motivated, offering comprehensive analysis to demonstrate its effectiveness.
2. Organization and Insight: It is well-organized and self-consistent, providing thorough discussions on the research topic and outlining both current and future research directions.
summary: The paper introduces SIFT, a method for selecting informative data to fine-tune large language models (LLMs). It critiques the current Nearest Neighbors retrieval approach which selects redundant data and proposes SIFT, inspired by transductive active learning, to select relevant and diverse data for effective fine-tuning. Additionally, it offers a library as a drop-in replacement for Nearest Neighbor retrieval.
weaknesses: 1. Inference Cost: Comparison with Nearest Neighbor, this method may require more inference time. It would be beneficial to compare inference times with Nearest Neighbor across different datasets to quantify this.
2. Broader Evaluation: The paper could explore effectiveness on more datasets and larger models, such as LLaMA-3, to validate its scalability and generalizability.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: See weakness.
soundness: 3
strengths: - The paper introduces SIFT, a well-motivated algorithm that combines retrieval and active learning, effectively addressing data redundancy issues in traditional Nearest Neighbor retrieval for LLM fine-tuning.

- SIFT’s adaptive fine-tuning, which adjusts test-time compute based on predicted performance gains, is an efficient and practical innovation that can optimize resource use, especially in computationally constrained environments.

- The paper provides both theoretical proofs and empirical evidence to demonstrate that SIFT reduces model uncertainty and improves fine-tuning outcomes, adding robustness and credibility to the proposed approach.

- The release of the *activeft* library as a drop-in replacement for Nearest Neighbor retrieval supports transparency and facilitates future research in prompt-specific fine-tuning methods.
summary: This paper presents SIFT, a data selection algorithm aimed at improving test-time fine-tuning of LLMs by reducing redundancy in data selection. Unlike traditional Nearest Neighbor retrieval, which often selects redundant examples, SIFT combines retrieval and active learning principles to minimize model uncertainty for each prompt. This approach also enables adaptive fine-tuning, dynamically adjusting compute based on anticipated performance gains. Experiments on the Pile dataset indicate that SIFT achieves better efficiency and robustness than Nearest Neighbor retrieval, with minimal computational overhead. The authors provide an open-source library, activeft, for integration and reproducibility.
weaknesses: 1. The motivation and definition of the task lack clarity, particularly in distinguishing test-time fine-tuning from standard fine-tuning on selected data. It remains unclear why test-time fine-tuning is necessary in this context and how it fundamentally differs from simply fine-tuning on pre-selected data, which may impact understanding of the novelty and importance of SIFT’s approach.

2. The paper does not fully specify scenarios where test-time fine-tuning with SIFT would be most beneficial. This omission makes it difficult to assess the generalizability and practical applications of the method, particularly for those who unfamiliar with prompt-specific fine-tuning needs.

3. Although SIFT is described as having minimal overhead, the adaptive fine-tuning process could introduce additional complexity in real-time settings. A more thorough breakdown of the computational costs associated with adaptive adjustments would improve understanding of its efficiency.

Meta-review: This paper received four positive ratings, with all reviewers generally inclined to accept it. The paper introduces SIFT, a data selection algorithm designed to enhance the fine-tuning of language models by addressing the limitations of Nearest Neighbors retrieval, which often selects redundant data. SIFT reduces uncertainty in model responses by optimizing overall information gain and mitigating information duplication. The paper is well-motivated, providing a comprehensive analysis to demonstrate its effectiveness. It offers both theoretical proofs and empirical evidence to show that SIFT reduces model uncertainty and improves fine-tuning outcomes, thereby adding robustness and credibility to the proposed approach. The SIFT strategy is supported by solid theoretical foundations and experimental evidence. The experimental study is convincing and covers a wide range of datasets. The authors have addressed the concerns raised, resolving most of the doubts. Therefore, the Area Chair (AC) recommends accepting the paper.

Decision: Accept (Poster)