11footnotetext: Equal Contribution22footnotetext: Equal Advising Contribution

# Repurposing Language Models into Embedding Models: Finding the Compute-
Optimal Recipe

Alicja Ziarko*1 2 3
&Albert Q. Jiang*4
&Bartosz Piotrowski1
Wenda Li5
&Mateja Jamnik$\dagger$4
&Piotr Miłoś$\dagger$1 2 3 6
1IDEAS NCBR 2University of Warsaw 3Institute of Mathematics, Polish Academy of
Sciences
4University of Cambridge 5University of Edinburgh 6deepsense.ai

###### Abstract

Text embeddings are essential for many tasks, such as document retrieval,
clustering, and semantic similarity assessment. In this paper, we study how to
contrastively train text embedding models in a compute-optimal fashion, given
a suite of pre-trained decoder-only language models. Our innovation is an
algorithm that produces optimal configurations of model sizes, data
quantities, and fine-tuning methods for text-embedding models at different
computational budget levels. The resulting recipe, which we obtain through
extensive experiments, can be used by practitioners to make informed design
choices for their embedding models. Specifically, our findings suggest that
full fine-tuning and low-rank adaptation fine-tuning produce optimal models at
lower and higher computational budgets respectively.

##  1 Introduction

Figure 1:  The optimal loss achieved using four different fine-tuning methods
(full fine-tuning, only tuning the bias, low-rank adaptation, and freezing
transformer blocks) at given budgets. The horizontal axis is the computational
budget in floating point operations (FLOP) and the vertical axis is the
contrastive loss. The X marks are datapoints and dotted lines are fitted
linear trends for different methods. The solid black line is the “optimal
frontier,” i.e., the optimal loss achievable with a fixed budget and the best
method.

Text embeddings are vector representations of text sequences (Devlin et al.,
2019). A desired property of text embeddings is that their distribution in the
embedding space reflects the semantic relations of the embedded texts. This
property benefits multiple practical tasks (Muennighoff et al., 2023b), such
as document retrieval (where documents are ranked based on the distance
between their representations and that of a query), document clustering,
sentiment analysis, or textual similarity measurement.

In the era of large language models (LLMs) that are pre-trained on vast
textual data, it is a natural idea to capitalise on the language
representations learned by these models to achieve high-quality embedding
models. A common and effective approach has been to initialise an embedding
model from a pre-trained LLM and fine-tune it with a contrastive loss
(Neelakantan et al., 2022; Izacard et al., 2022). This phase is of crucial
importance, as the quality of the embeddings extracted directly from a pre-
trained LLM (e.g., by averaging hidden states of the last layer) is not
sufficient for most tasks.

However, the best performing modern LLMs are usually massive in terms of the
parameter counts (e.g., 175B parameters for GPT3 (Brown et al., 2020) and 340B
for PaLM2 (Anil et al., 2023)), and are therefore difficult to train with
limited resources. This motivates a practical question, that to the best of
our knowledge has not been addressed systematically, and thus we address it in
this article: _what is the best embedding model one can train from a backbone
decoder-only LLM with a fixed budget_?

To answer this, we performed an extensive empirical study. We started by
identifying design choices one can make when fine-tuning language models into
embedding models, including: model size, data volume, parameter-efficient
fine-tuning technique, and the hyperparameters of the chosen technique. Then,
we performed a grid search over the pre-defined design choices, and found the
best performing configuration under each computational budget. Using these
findings, we established scaling laws for contrastive fine-tuning, which
enabled us to build an algorithm that produces a general recipe for efficient
fine-tuning of a decoder-only LLM that obtains a high-quality embedding model.

#### Contribution

We comprehensively experimented with contrastively fine-tuning language models
into text embedding models. We analysed how the choice of model sizes, data
quantities, and fine-tuning methods affect the performance in the resource-
constrained regime. We compiled the findings into an algorithm that, given a
fixed computational budget, predicts the optimal network architecture, data
quantity, and parameter-efficient fine-tuning hyperparameters. We open-source
the code to train and evaluate our models at:
https://github.com/SeqDM/Efficient-Embeddings.

##  2 Related work

#### Embedding models

Training neural networks to represent text as continuous vectors was
popularised by word2vec (Mikolov et al., 2013), which produced semantically
meaningful embeddings of words.

BERT (Devlin et al., 2019) and the further contrastively trained SimCSE (Gao
et al., 2021c) quickly established encoder-only transformers as the go-to
architecture for embedding text.

More recently, decoder-only transformers have become increasingly more
powerful and efficient at the same time (Brown et al., 2020; Touvron et al.,
2023; Jiang et al., 2023). Building an embedding model based on them utilises
the knowledge from their pretraining, and thus it is a natural step.
Neelakantan et al. (2022) set a successful precedent in this paradigm by
initialising embedder training with decoder-only GPT models.

Notably, the current open state-of-the-art (the top open model on the MTEB
(Muennighoff et al., 2023b) leaderboard, as of 18 May 2024), SFR-Embedding-
Mistral (Meng et al., 2024), is also fine-tuned from the decoder-only Mistral
7B (Jiang et al., 2023).

#### Benchmarks for embedding models

Embedding models are versatile in their applications. Therefore, a broad and
diverse benchmark to provide a robust measure of their performance is called
for. The first such benchmark was BEIR introduced by Thakur et al. (2021). It
has nine different information retrieval tasks (e.g., duplicate-question
retrieval or citation-prediction) on 18 datasets. Recently, Muennighoff et al.
(2023b) introduced MTEB (Massive Text Embedding Benchmark), which is
substantially larger than BEIR and, in addition to information retrieval, has
seven other types of tasks (e.g., clustering, summarisation, and pair
classification) using 58 datasets and covering 112 languages. The authors also
provide a leaderboard that currently contains more than 220 entries.

#### Scaling laws

Scaling laws predict the performance of models at a larger scale (e.g., more
parameters, more data) from experiments at smaller scales. Kaplan et al.
(2020) and Hoffmann et al. (2022) used empirically fitted laws to predict the
performance of neural language models at different model sizes and data
scales, making it possible to calculate the optimal configuration before
training. Muennighoff et al. (2023a) inspected model training in data-
constrained, multi-epoch regime. Frantar et al. (2023) looked at scaling laws
for language models in the context of weight sparsity. Zhang et al. (2024)
investigates the scaling laws for model fine-tuning, including parameter
efficient methods: LoRA (Hu et al., 2022) and prompt tuning (Lester et al.,
2021). They show that the scaling laws and the best fine-tuning method is task
dependent. Biderman et al. (2024) investigates the properties of fine-tuning
models with LoRA in details, showing that LoRA underperforms full fine-tuning
when fine-tuning the models on mathematical or code datasets.

However, none of these works focus on scaling laws for training _embedding
models_. Our scaling law is targeted at the process of contrastive fine-tuning
for achieving a model yielding good embeddings. Moreover, our main goal is to
find which strategy for achieving good embeddings is optimal in a budget-
restricted settings, while taking into account the popularity of applying
parameter efficient methods for fine-tuning, like LoRA or partial model
freezing. In spirit, our work is similar to AutoML (He et al., 2021b), which
aims to automatically find the optimal ML architecture. However, our method
goes beyond and targets the data and the training method as well.

#### Parameter-efficient fine-tuning

Modern language models have a lot of parameters, and fine-tuning them using
consumer-grade hardware is difficult in general. Notwithstanding the recent
turn to in-context learning, fine-tuning is still necessary for most
applications. A range of techniques have been developed recently, to make the
fine-tuning more efficient (Houlsby et al., 2019; Li and Liang, 2021; Qi et
al., 2022; He et al., 2021a). These approaches are all applicable to embedding
models. Muennighoff (2022) explored a simple parameter-efficient approach to
repurpose GPT models into embedding models where only the bias tensors of the
transformer model are updated (Zaken et al., 2022). However, a systematic
study of what parameter-efficient methods are optimal under what scenarios has
not been performed, which is the aim of our work.

##  3 Preliminaries

###  3.1 Scaling laws and compute-optimal models

The constrained optimisation problem that we aim to solve is the following:
given a fine-tuning corpus and a family of pre-trained decoder-only language
models of different sizes, minimise the contrastive loss subject to a fixed
computational budget.

Following Hoffmann et al. (2022, Section 3.2), we perform a hyperparameter
search at different computational budget levels to obtain IsoFLOP profiles of
models. We then find the optimal models at each level, fit a loss-size scaling
law, and extrapolate it to unobserved FLOP budgets. We only focus on
contrastive fine-tuning since state-of-the-art text embedder models are all
trained in such fashion (Gao et al., 2021c; Neelakantan et al., 2022; Wang et
al., 2023).

We refer to the models with the lowest contrastive loss achievable with fixed
computational costs as compute-optimal models. We focus on optimising three
components specifically: model size, data quantity, and fine-tuning method.
For each fine-tuning method, we find the optimal model configuration at each
FLOP level and establish the relationship between the computational budget and
optimal loss for that method. Then, for any unobserved computational budget,
we use the scaling laws to predict the loss, and pick the method that gives
the lowest loss and its corresponding configuration.

###  3.2 Extracting representations from transformers

We start with a standard decoder-only transformer model architecture, such as
GPT architectures (Radford et al., 2019; Brown et al., 2020), or Pythia
(Biderman et al., 2023). Given a sequence of tokens $S=(s_{1},\ldots,s_{n})$,
we pass them through the transformer model and average the last layer
representations as the embedding vector
$\texttt{emb}(S)=\frac{1}{n}\sum^{n}_{i=1}h_{i}\in\mathbb{R}^{m}$, where
$h_{i}$ is the last layer hidden state of the token $s_{i}$.

Some works apply the last token pooling approach (Neelakantan et al., 2022)
instead of using mean pooling described above. Here, however, we default to
mean pooling since it is a recently more popular method overall (Wang et al.,
2022; Li et al., 2023; Günther et al., 2023).

###  3.3 Contrastive fine-tuning

We use the contrastive loss as in (Neelakantan et al., 2022), which has been
widely used in fine-tuning pre-trained language models (Wang et al., 2022,
2023; Günther et al., 2023).

For a batch of $n$ examples, that is, $n$ pairs of texts
$(x_{1},y_{1}),\ldots,(x_{n},y_{n})$, we only treat examples
$(x_{i},y_{i})_{i=1}^{n}$ as positive matches, and all the pairs
$(x_{i},y_{j})$, where $i\neq j$ as negatives. Concretely, we calculate the
batch contrastive loss in two steps. First, we calculate the logits:

$\texttt{logits[i,j]}=\texttt{sim}(\texttt{emb}(x_{i}),\texttt{emb}(y_{j}))\cdot\exp(\tau),$

where sim is the cosine similarity function between two vectors, and
$\tau\in\mathbb{R}$ is the temperature. The total loss is the sum of the cross
entropy losses on both the row and the column directions:111We adopt the
PyTorch notation here, where the cross entropy loss takes two arguments: the
logits of the probability distribution, and the indices of the true
categories.

                      labels = [0, 1, ..., n - 1]
                      l_r = cross_entropy(logits, labels)
                      l_c = cross_entropy(logits.T, labels)
                      loss = (l_r + l_c) / 2

#### Hard negatives in contrastive learning

In our contrastive learning setting, we use in-batch examples as negatives.
Some datasets additionally include tailored negative examples (_hard
negatives_) for each positive datapoint, which can be incorporated into the
contrastive loss to promote learning more precise representations.

However, works that do rely on hard negatives often arrive at embeddings by
having two distinct training phases: the expensive “unsupervised” phase one
involving vast data from the internet, and then a “supervised” phase two,
often targeted towards a specific downstream task, where training data of
lower quantity – but containing high-quality hard negatives – is used (Li et
al., 2023; Wang et al., 2022; Xiao et al., 2023). Since the first phase is
more expensive, it will benefit more from scaling laws such as the ones we
derive.

Moreover, the usefulness of hard negatives highly depends on their quality,
which may vary wildly between datasets. Therefore, conclusions reached with
hard negatives are more closely tied to the datasets the models are trained
on. Since we develop scaling laws agnostic to the dataset, we abstain from
using hard negatives in our experiments.

###  3.4 Fine-tuning methods

The most basic and straightforward method of fine-tuning is _full fine-tuning_
, where _all_ the weights of the model are updated under the contrastive
objective.

We also study parameter-efficient fine-tuning (PEFT) methods, which are
popular, especially in academic settings, because they have lower memory
requirements than full fine-tuning. Since the PEFT methods update fewer
parameters in the network than there are in total, the relationship between
their computational cost, model size, and data quantity is different from full
fine-tuning. Hence, we individually study the IsoFLOP profiles for four fine-
tuning methods, namely, _full fine-tuning_ , _block freezing_ ,_bias-only
tuning_ , and _Low-Rank Adaptation (LoRA)_. Each of the last three non-
standard methods is briefly characterised below.

#### Block freezing

In this approach, we _freeze_ the LLM token-embedding layer as well as the
first $k$ transformer blocks so that their parameters stay fixed during fine-
tuning. Only the latter part of the model is back-propagated through and
updates its parameters. To process the same amount of data, this is more
computationally efficient than full fine-tuning. By varying the number of
frozen blocks, one trades-off between the computational efficiency and
flexibility of the model.

#### Low-rank adaptation (LoRA)

This method was introduced by Hu et al. (2022) to significantly reduce the
number of trainable parameters of a neural network while maintaining high
performance. LoRA introduces a small number of weights into the model and only
trains those. It can be applied to a subset of dense layers of a model. It
works by parameterising each dense layer $W$ from a subset of linear layers of
a model by two low-rank matrices $A$ and $B$, and using $W+AB$ instead of $W$,
while training only $A$ and $B$.

#### Bias-only tuning

In this method, only bias parameters are fine-tuned, and the rest of the model
remains frozen. This approach was proposed by Zaken et al. (2022) as a
parameter-efficient fine-tuning method for BERT encoders, and recently applied
by Muennighoff (2022) to contrastively fine-tune GPT models.

###  3.5 Calculating computational cost

We denote the number of non-token-embedding parameters used in the forward
pass to process each token as $N_{F}$, the number of non-token-embedding
parameters in backward-propagation as $N_{B}$, the number of non-token-
embedding parameters updated during backward-propagation as $N_{U}$, the total
number of tokens processed as $D$, and the computational cost as $C$ floating
point operations (FLOP). Following the calculation by Kaplan et al. (2020), we
derive the relationship between the variables to be

$C=2N_{F}D+2N_{B}D+2N_{U}D.$

In case of full fine-tuning, every parameter is trainable, so
$N_{F}=N_{B}=N_{U}$, and therefore

$C=6N_{F}D.$

This is the same as the formula for the computational cost used by Kaplan et
al. (2020) for pre-training.

##  4 Experiments

Here, we first specify the relevant details of our experimental setup (Section
4.1). Next, we present the results of our experiments where we contrastively
train a grid of models of different sizes, using different computational
budgets, and apply different compute-optimal fine-tuning methods with varying
hyperparameters (Section 4.2). Based on the collected data, we fit a
mathematical formula describing the observed scaling laws (Section 4.3).
Finally, we provide general observations and recommendations for efficient
contrastive fine-tuning (Section 4.5).222The data collected in our experiments
and their interactive visualisations are available in this Colab notebook:
https://colab.research.google.com/drive/1EC2QdVrOIXhamLJI51CUywWJ33sApXX8

###  4.1 Experimental setup

We use eight decoder-only models from the Pythia suite (Biderman et al.,
2023), which have sizes 14M, 31M, 70M, 160M, 410M, 1B, 1.4B, and 2.8B of
parameters. All the models have been pre-trained on the Pile dataset (Gao et
al., 2021a) for 300 billion tokens. We consider six computational budgets
equally spaced on the log-scale: $1.51015$, $61015$, $2.41016$, $9.61016$,
$3.81017$, and $1.51018$ FLOP.

We fine-tune our models on the English partition of the BAAI BGE dataset (Xiao
et al., 2023), which contains 200 million semantically related (_query_ ,
_value_) pairs from various internet sources such as Wikipedia and Stack
Exchange. Note that since the BAAI BGE dataset is massive, for all our
experiments we fine-tune for less than one epoch on this corpus, which allows
us to avoid the effect of diminishing returns on subsequent training epochs
(Muennighoff et al., 2023a).

We use the AdamW optimiser (Loshchilov and Hutter, 2019) and a cosine learning
rate scheduler during training. The learning rate first goes through a linear
warm-up phase of $1/10$ of the total steps to a peak that is $1/10$ of the
pre-training peak learning rate, that is, to learning rates between $1.2\times
10^{-5}$ and $6\times 10^{-5}$. Then, it decays in a cosine schedule to $1/10$
of the maximum at the end of training. We employ gradient checkpointing (Chen
et al., 2016) and GradCache (Gao et al., 2021b) to limit the peak GPU RAM
usage during training. We use the Transformers (Wolf et al., 2020) library for
the training of the models and use the Accelerate (Gugger et al., 2022)
library to facilitate multi-GPU training. The batch size in our main line of
experiments is 1024 sequences and the context length is 75 tokens. We give the
full list of training hyperparameters in Appendix D.

We also perform a less extensive grid of experiments to confirm that our
findings are robust to change in hyperparameters. We find that our scaling law
formula fitted to losses from models with batch size 1024 and context length
of 75 tokens describes the loss for models trained with batch size 2048 and
context length of 512 tokens very well. The detailed description and plots are
shown in Appendix F.

In addition to controlling the final training contrastive loss achieved by the
models, we also measure downstream performance by evaluating the models on a
representative subset of the MTEB benchmark (Muennighoff et al., 2023b). The
benchmark defines eight categories of tasks in total (e.g., retrieval,
semantic text similarity). We select one task from each category by
determining which ones are the most correlated with the performance for the
whole category (see Appendix B for details).

###  4.2 Experimental results for different methods

#### Full fine-tuning

With full fine-tuning, the numbers of forward and backward parameters are the
same, and the computational cost is straightforwardly $C=6N_{F}D$. In Figure
2(a), we present the final contrastive loss vs. the number of parameters in
the network on a log-log scale.

(a)

(b)

Figure 2:  (a) IsoFLOP profiles for _full fine-tuning_. The horizontal axis is
the number of parameters in the model, and the vertical axis is the achieved
loss. Both axes use log-scale. The optimal model size tends to increase as the
computational budget increases. (b) IsoFLOP profiles for _block freezing_. The
axes are the same as for full fine-tuning. Each data point denotes the optimal
choice with respect to the fraction of active blocks during training, which is
noted above the points. The optimal model size tends to increase as the
computational budget increases, while the optimal active block fraction tends
to slightly decrease as the model size gets larger.

As expected, larger computational budgets consistently improve the loss. The
IsoFLOP profiles resemble the classical ones from Hoffmann et al. (2022).
However, they tend to be flatter when the model size and computational budget
are both small or both large. From the IsoFLOP profiles we can see what is the
optimal model size that minimises the loss at each computational budget.

#### Block freezing

Since the Pythia models we experimented with have $6,12,16,24$ or $32$ blocks,
we freeze $\frac{0}{6},\frac{1}{6},\ldots,\frac{5}{6}$ of their blocks if
their number is divisible by $6$, and
$\frac{0}{8},\frac{1}{8},\ldots,\frac{7}{8}$ if it is divisible by $8$. In
this setup, unlike for full fine-tuning, the token-embedding parameters always
remain frozen even when all the transformer blocks are active. We have
$N_{B}=N_{U}<N_{F}$, and hence cost $C=2N_{F}D+4N_{B}D$.

In Figure 2(b), we present IsoFLOP profiles for optimal loss at given model
sizes. The trend of the profiles is quite similar to that of full fine-tuning
presented in Figure 2(a).

We plot the loss with varying model sizes, computational budgets, and
fractions of frozen blocks in Figure 3. For smaller models ($\leq 160$M
parameters), a lower fraction of frozen blocks leads to a lower final loss
across all computational budgets. However, for larger models, using more than
$70\%$ of transformer often only marginally improves the model. For instance,
for the model with $11009$ parameters, it is visible that freezing $50-70\%$
of the blocks gives the optimal results for all budgets except for $3.81017$
FLOP.

Figure 3:  The effect of block freezing across all model sizes. Different
colours signify different computational budgets. Unless the model is large and
the computational budget small, it is always better to update all the (non-
embedding) weights of the model.

#### Bias-only tuning

Here we simply update only the bias parameters and not the weights of the
model. Figure 4(a) shows the final training loss across all the model sizes
and computational budgets. Increasing the computational budget decreases the
optimal loss that is achievable, although the absolute values of the losses
are high compared to other fine-tuning methods (the lowest achievable loss is
above 0.4 with bias-only tuning).

(a)

(b)

Figure 4:  (a) IsoFLOP profiles for _bias-only tuning_. The horizontal axis is
the number of parameters in the model, and the vertical axis is the achieved
loss. Both axes use log-scale. The optimal model size increases as the
computational budget increases, but the achievable loss is higher than for
other fine-tuning methods. (b) IsoFLOP profiles for _LoRA_ fine-tuning. The
axes are the same as for bias-only tuning. Each data point denotes the optimal
choice of the rank of LoRA matrices given the size of the model and the
computational budget.

#### Low-rank adaptation

We use the LoRA implementation from the PEFT library (Mangrulkar et al.,
2022). We attach the adapter to all the dense layers of the network, and vary
the adapter rank from $8$ up to $2048$, since even lower ranks resulted in an
unstable training.

We compare the loss obtained for each model size and computational budget for
different LoRA rank values in Figure 5. In almost all the combinations of the
model size and computational budget, the rank of 32 or 128 (occasionally 512)
is optimal. Wang et al. (2023) also use LoRA in the context of contrastive
learning, but they fix the rank of 16 for their experiments. Our findings
suggest that this choice should be reconsidered as being potentially
suboptimal.

Figure 5:  The effect of different LoRA ranks across all model sizes.
Different colours signify different computational budgets. The inflected
curves indicate that it is less beneficial to use a rank from either extremes
of the spectrum (8 or 2048). The detrimental effect of the high rank of 2048
is stronger for lower computational budgets. Ranks of 32 and 128 result in the
lowest loss overall.

In Figure 4(b), we present IsoFLOP profiles for LoRA, where each data point
represents the best LoRA rank setting for a given model size and computational
budget. We see that the relationship between the optimal loss and the LoRA
rank is quite consistent across model sizes and computational budgets, meaning
that the rank hyperparameter is not very sensitive with respect to the two
variables. In practice, this means that setting the LoRA rank to 32 or 128 are
likely to be good default choices.

#### Lookup table of optimal model sizes for different methods

In Figure 6(a) (for full fine-tuning), Figure 6(b) (for LoRA), and Figure 12
(for freezing and bias tuning, in Appendix A), we visualise the relationship
between the optimal model size and the computational budget. This makes
looking up the optimal configuration for a given method at a given budget very
convenient.

#### Downstream performance

In Figures 7, 8, 9, and 10, we present the IsoFLOP plots of the downstream
performance for all the four fine-tuning methods on the subset of the MTEB
benchmark (described in Appendix B). Based on the performance results, we also
present the optimal size choices per budget in Figure 13 (which is a
performance analogue of Figure 12) and the optimal performance plot in Figure
11 (a performance analogue of Figure 1). These are mostly consistent with the
loss results, albeit more noisy. In Appendix C, we demonstrate the significant
correlation between the final training loss and the downstream task
performance.

###  4.3 Scaling laws for embeddings

We aim to model the behaviour of the loss $L$ with an analytical formula as a
function of the number of parameters $N$ and the number of training tokens
$D$. Following Hoffmann et al. (2022), we start with the following formula:

$L(N,D)=C+\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}},$

where $A,B,C,\alpha,\beta\in\mathbb{R}_{+}$ are real-valued parameters to be
fitted. We found that this formula reasonably describes the loss behaviour for
each of our methods. However, the fraction of _trainable parameters_ is an
important factor that is not considered in the model above. Following the
modelling of Frantar et al. (2023), we propose a modified formula:

$L(S,N,D)=C+\frac{a_{d}\log(D)+b_{d}}{N^{\alpha}}+\frac{a_{s}(1-S)^{b_{s}}+c_{s}}{D^{\beta}},$

where $S$ is an additional variable representing the trainable fraction of
parameters, and $a_{d},b_{d},a_{s},b_{s},c_{s}\in\mathbb{R}$ are additional
coefficients to be fitted. We check that this new formula fits the data better
than the original one, and conjecture that it can be used for larger models
and larger computational budgets. More details about the derivation are
presented in Appendix E.

###  4.4 Compute-optimal frontier and recipe

In Figure 1, we plotted the best achievable losses against the computational
budgets for the four fine-tuning methods considered. We then fitted a linear
trend for each method on a log-log scale (equivalent to a power law
relationship on a normal scale), and highlighted the lowest loss achievable at
given budgets using the best corresponding fine-tuning method. We call the
contour of this function that predicts the optimal loss across methods the
compute-optimal frontier.

The equation describing the full fine-tuning and the LoRA linear fits are

$\displaystyle\ln(\mathrm{loss})=-0.21\cdot\ln(\mathrm{budget})+8.39,\text{
and }\ln(\mathrm{loss})=-0.22\cdot\ln(\mathrm{budget})+8.93,$

respectively. The equations suggest that at a budget lower than $9.061016$
FLOP, the optimal model is achieved with full fine-tuning, and at a higher
budget, the optimal model is achieved with LoRA.

(a)

(b)

Figure 6:  Optimal model size vs. computational budget for full fine-tuning
(a), and LoRA (b).

With this, one becomes fully equipped to deduce the optimal recipe for text
embedding model training for a given budget, by carrying out the procedure
detailed in Algorithm 1. Note that although the block freezing approach is not
predicted to be optimal at any budget, its performance is quite close to
optimal. It also has the advantage of being easy to implement and has lower
memory requirements, hence might be a good method to choose in practical use
cases for larger budgets. Its optimal model size and fraction of active blocks
hyperparameter can be worked out in the same way as for the other methods,
with the help of Figure 3 and Figure 12 (left pane) in Appendix A.

Algorithm 1 Recipe for compute-optimal embedding model

Input: compute budget $C$.

Output: fine-tuning method, model size, data quantity, and (optionally) the
method’s hyperparams.

if $B\leq$ $9.061016$ FLOP then

Use full fine-tuning.

Go to Figure 6(a) to find the optimal model parameters $N$ given budget $C$.

Calculate the data quantity $D$.

return Full fine-tuning, $N$, $D$, ().

else

Use LoRA.

Go to Figure 6(b) to find the optimal model parameters $N$ given budget $C$.

Go to Figure 5 to find the LoRA rank $R$ according to $C$ and $N$.

return LoRA, $N$, $D$, ($R$).

###  4.5 Takeaways

We compile the takeaways from the extensive experiments, by analysing each
fine-tuning method individually and contrasting them afterwards.

  * •

Bias-only tuning is not a good fine-tuning strategy for embedding model
training, as it is consistently worse than other strategies under IsoFLOP
comparisons.

  * •

LoRA and block freezing are both effective approaches that improve the
performance of the trained embedding models at higher budgets. For LoRA, the
rank hyperparameter is not very sensitive to model size or computational
budgets, with an optimum at around 128.

  * •

Finding the optimal recipe for fine-tuning embedding models requires a non-
trivial effort, and the resulting optimal configurations are sometimes
counter-intuitive (e.g., not using the largest model size possible). Carefully
tuning hyperparameters is crucial and should be practised systematically.

##  5 Limitations and future work

#### Other families of models

In this paper, we experimented with the Pythia (Biderman et al., 2023) suite
of pre-trained models, which are of different sizes but have all been trained
on 300B tokens. This means that they have been over-trained with respect to
the Chinchilla scaling law (Hoffmann et al., 2022) to different degrees. The
relationship between the optimal model size and the computational budget might
differ for a suite of pre-trained models with different over-training ratios.
However, we notice that prior efforts (Dettmers et al., 2024) have found that
fine-tuning conclusions reached with Pythia do generalise to other model
families such as OPT (Zhang et al., 2022), LLaMA (Touvron et al., 2023), and
BLOOM (Scao et al., 2022), in the quantised setting.

#### Averaging

Due to constraints on our computational resources, we ran each experiment
exactly once. Averaging over multiple random seeds for the experiments would
reduce the variance involved in each, and potentially result in less noisy
conclusions. This is a straightforward extension that we leave to carry out in
future work, should we have more resources. We also plan to evaluate the
trained models on the entirety of the MTEB benchmark (Muennighoff et al.,
2023b), instead of using the losses from the last steps of training, to better
reflect the performances of embedding models.

#### Other forms of embedding readout

We only experimented with mean pooling as a way of extracting embeddings from
the transformer models. Other readout methods could be used, like the max
pooling or extracting the hidden state corresponding to the end-of-sequence
token, which might result in different scaling laws and optimal frontiers.

##  6 Conclusions

In this paper, we systematically investigated the influence of pre-trained
model size, fine-tuning data quantity, and fine-tuning method in the final
performance of embedding models repurposed from language models. As a result,
we devised an algorithm that takes in the computational budget and returns the
optimal configuration for the embedding model to train. This enables
researchers who wish to adapt language models to embed their own data,
especially those with limited computational resources, to obtain optimal text
embedding models with greater efficiency in time and resources.

##  7 Acknowledgements

The authors would like to thank Denis Kuznedelev, Soroush Tabesh and Konrad
Staniszewski for helpful discussions and feedback. AQJ is supported by a
Peterhouse graduate studentship. PM was supported by National Science Center
Poland under the grant agreement 2019/35/O/ST6/03464. We gratefully
acknowledge Polish high-performance computing infrastructure PLGrid (HPC
Centers: ACK Cyfronet AGH) for providing computer facilities and support
within computational grant no. PLG/2023/016717.

## References

  * Anil et al. (2023) R. Anil, A. M. Dai, O. Firat, M. Johnson, D. Lepikhin, A. Passos, S. Shakeri, E. Taropa, P. Bailey, Z. Chen, E. Chu, J. H. Clark, L. E. Shafey, Y. Huang, K. Meier-Hellstern, G. Mishra, E. Moreira, M. Omernick, K. Robinson, S. Ruder, Y. Tay, K. Xiao, Y. Xu, Y. Zhang, G. H. Ábrego, J. Ahn, J. Austin, P. Barham, J. A. Botha, J. Bradbury, S. Brahma, K. Brooks, M. Catasta, Y. Cheng, C. Cherry, C. A. Choquette-Choo, A. Chowdhery, C. Crepy, S. Dave, M. Dehghani, S. Dev, J. Devlin, M. Díaz, N. Du, E. Dyer, V. Feinberg, F. Feng, V. Fienber, M. Freitag, X. Garcia, S. Gehrmann, L. Gonzalez, and et al.  Palm 2 technical report.  _CoRR_ , abs/2305.10403, 2023.  doi: 10.48550/ARXIV.2305.10403.  URL https://doi.org/10.48550/arXiv.2305.10403.
  * Biderman et al. (2024) D. Biderman, J. G. Ortiz, J. Portes, M. Paul, P. Greengard, C. Jennings, D. King, S. Havens, V. Chiley, J. Frankle, C. Blakeney, and J. P. Cunningham.  LoRA learns less and forgets less, 2024.
  * Biderman et al. (2023) S. Biderman, H. Schoelkopf, Q. G. Anthony, H. Bradley, K. O’Brien, E. Hallahan, M. A. Khan, S. Purohit, U. S. Prashanth, E. Raff, A. Skowron, L. Sutawika, and O. van der Wal.  Pythia: A suite for analyzing large language models across training and scaling.  In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett, editors, _International Conference on Machine Learning, ICML 2023, 23-29 July 2023, Honolulu, Hawaii, USA_ , volume 202 of _Proceedings of Machine Learning Research_ , pages 2397–2430. PMLR, 2023.  URL https://proceedings.mlr.press/v202/biderman23a.html.
  * Brown et al. (2020) T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei.  Language models are few-shot learners.  _CoRR_ , abs/2005.14165, 2020.  URL https://arxiv.org/abs/2005.14165.
  * Chen et al. (2016) T. Chen, B. Xu, C. Zhang, and C. Guestrin.  Training deep nets with sublinear memory cost.  _CoRR_ , abs/1604.06174, 2016.  URL http://arxiv.org/abs/1604.06174.
  * Dettmers et al. (2024) T. Dettmers, A. Pagnoni, A. Holtzman, and L. Zettlemoyer.  QLoRA: Efficient finetuning of quantized LLMs.  _Advances in Neural Information Processing Systems_ , 36, 2024.
  * Devlin et al. (2019) J. Devlin, M. Chang, K. Lee, and K. Toutanova.  BERT: pre-training of deep bidirectional transformers for language understanding.  In J. Burstein, C. Doran, and T. Solorio, editors, _Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers)_ , pages 4171–4186. Association for Computational Linguistics, 2019.  doi: 10.18653/V1/N19-1423.  URL https://doi.org/10.18653/v1/n19-1423.
  * Frantar et al. (2023) E. Frantar, C. Riquelme, N. Houlsby, D. Alistarh, and U. Evci.  Scaling laws for sparsely-connected foundation models.  _CoRR_ , abs/2309.08520, 2023.  doi: 10.48550/ARXIV.2309.08520.  URL https://doi.org/10.48550/arXiv.2309.08520.
  * Gao et al. (2021a) L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, S. Presser, and C. Leahy.  The pile: An 800GB dataset of diverse text for language modeling.  _CoRR_ , abs/2101.00027, 2021a.  URL https://arxiv.org/abs/2101.00027.
  * Gao et al. (2021b) L. Gao, Y. Zhang, J. Han, and J. Callan.  Scaling deep contrastive learning batch size under memory limited setup.  In A. Rogers, I. Calixto, I. Vulic, N. Saphra, N. Kassner, O. Camburu, T. Bansal, and V. Shwartz, editors, _Proceedings of the 6th Workshop on Representation Learning for NLP, RepL4NLP@ACL-IJCNLP 2021, Online, August 6, 2021_ , pages 316–321. Association for Computational Linguistics, 2021b.  doi: 10.18653/V1/2021.REPL4NLP-1.31.  URL https://doi.org/10.18653/v1/2021.repl4nlp-1.31.
  * Gao et al. (2021c) T. Gao, X. Yao, and D. Chen.  SimCSE: Simple contrastive learning of sentence embeddings.  In M. Moens, X. Huang, L. Specia, and S. W. Yih, editors, _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021_ , pages 6894–6910. Association for Computational Linguistics, 2021c.  doi: 10.18653/V1/2021.EMNLP-MAIN.552.  URL https://doi.org/10.18653/v1/2021.emnlp-main.552.
  * Gugger et al. (2022) S. Gugger, L. Debut, T. Wolf, P. Schmid, Z. Mueller, S. Mangrulkar, M. Sun, and B. Bossan.  Accelerate: Training and inference at scale made simple, efficient and adaptable.  https://github.com/huggingface/accelerate, 2022.
  * Günther et al. (2023) M. Günther, L. Milliken, J. Geuter, G. Mastrapas, B. Wang, and H. Xiao.  Jina embeddings: A novel set of high-performance sentence embedding models.  _CoRR_ , abs/2307.11224, 2023.  doi: 10.48550/ARXIV.2307.11224.  URL https://doi.org/10.48550/arXiv.2307.11224.
  * He et al. (2021a) J. He, C. Zhou, X. Ma, T. Berg-Kirkpatrick, and G. Neubig.  Towards a unified view of parameter-efficient transfer learning.  _CoRR_ , abs/2110.04366, 2021a.  URL https://arxiv.org/abs/2110.04366.
  * He et al. (2021b) X. He, K. Zhao, and X. Chu.  AutoML: A survey of the state-of-the-art.  _Knowledge-Based Systems_ , 212:106622, 2021b.
  * Hoffmann et al. (2022) J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. de Las Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre.  Training compute-optimal large language models.  _CoRR_ , abs/2203.15556, 2022.  doi: 10.48550/ARXIV.2203.15556.  URL https://doi.org/10.48550/arXiv.2203.15556.
  * Houlsby et al. (2019) N. Houlsby, A. Giurgiu, S. Jastrzebski, B. Morrone, Q. de Laroussilhe, A. Gesmundo, M. Attariyan, and S. Gelly.  Parameter-efficient transfer learning for NLP.  _CoRR_ , abs/1902.00751, 2019.  URL http://arxiv.org/abs/1902.00751.
  * Hu et al. (2022) E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen.  LoRA: Low-rank adaptation of large language models.  In _The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022_. OpenReview.net, 2022.  URL https://openreview.net/forum?id=nZeVKeeFYf9.
  * Izacard et al. (2022) G. Izacard, M. Caron, L. Hosseini, S. Riedel, P. Bojanowski, A. Joulin, and E. Grave.  Unsupervised dense information retrieval with contrastive learning.  _Trans. Mach. Learn. Res._ , 2022, 2022.  URL https://openreview.net/forum?id=jKN1pXi7b0.
  * Jiang et al. (2023) A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. de Las Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed.  Mistral 7B.  _CoRR_ , abs/2310.06825, 2023.  doi: 10.48550/ARXIV.2310.06825.  URL https://doi.org/10.48550/arXiv.2310.06825.
  * Kaplan et al. (2020) J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei.  Scaling laws for neural language models.  _CoRR_ , abs/2001.08361, 2020.  URL https://arxiv.org/abs/2001.08361.
  * Lester et al. (2021) B. Lester, R. Al-Rfou, and N. Constant.  The power of scale for parameter-efficient prompt tuning.  In M. Moens, X. Huang, L. Specia, and S. W. Yih, editors, _Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, EMNLP 2021, Virtual Event / Punta Cana, Dominican Republic, 7-11 November, 2021_ , pages 3045–3059. Association for Computational Linguistics, 2021.  doi: 10.18653/V1/2021.EMNLP-MAIN.243.  URL https://doi.org/10.18653/v1/2021.emnlp-main.243.
  * Li and Liang (2021) X. L. Li and P. Liang.  Prefix-tuning: Optimizing continuous prompts for generation.  _CoRR_ , abs/2101.00190, 2021.  URL https://arxiv.org/abs/2101.00190.
  * Li et al. (2023) Z. Li, X. Zhang, Y. Zhang, D. Long, P. Xie, and M. Zhang.  Towards general text embeddings with multi-stage contrastive learning.  _CoRR_ , abs/2308.03281, 2023.  doi: 10.48550/ARXIV.2308.03281.  URL https://doi.org/10.48550/arXiv.2308.03281.
  * Loshchilov and Hutter (2019) I. Loshchilov and F. Hutter.  Decoupled weight decay regularization.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net, 2019.  URL https://openreview.net/forum?id=Bkg6RiCqY7.
  * Mangrulkar et al. (2022) S. Mangrulkar, S. Gugger, L. Debut, Y. Belkada, S. Paul, and B. Bossan.  PEFT: State-of-the-art parameter-efficient fine-tuning methods.  https://github.com/huggingface/peft, 2022.
  * Meng et al. (2024) R. Meng, Y. Liu, S. R. Joty, C. Xiong, Y. Zhou, and S. Yavuz.  SFR-Embedding-Mistral: Enhance text retrieval with transfer learning.  Salesforce AI Research Blog, 2024.  URL https://blog.salesforceairesearch.com/sfr-embedded-mistral/.
  * Mikolov et al. (2013) T. Mikolov, K. Chen, G. Corrado, and J. Dean.  Efficient estimation of word representations in vector space.  In Y. Bengio and Y. LeCun, editors, _1st International Conference on Learning Representations, ICLR 2013, Scottsdale, Arizona, USA, May 2-4, 2013, Workshop Track Proceedings_ , 2013.  URL http://arxiv.org/abs/1301.3781.
  * Muennighoff (2022) N. Muennighoff.  SGPT: GPT sentence embeddings for semantic search.  _CoRR_ , abs/2202.08904, 2022.  URL https://arxiv.org/abs/2202.08904.
  * Muennighoff et al. (2023a) N. Muennighoff, A. M. Rush, B. Barak, T. L. Scao, A. Piktus, N. Tazi, S. Pyysalo, T. Wolf, and C. Raffel.  Scaling data-constrained language models.  _CoRR_ , abs/2305.16264, 2023a.  doi: 10.48550/ARXIV.2305.16264.  URL https://doi.org/10.48550/arXiv.2305.16264.
  * Muennighoff et al. (2023b) N. Muennighoff, N. Tazi, L. Magne, and N. Reimers.  MTEB: Massive text embedding benchmark.  In A. Vlachos and I. Augenstein, editors, _Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics, EACL 2023, Dubrovnik, Croatia, May 2-6, 2023_ , pages 2006–2029. Association for Computational Linguistics, 2023b.  doi: 10.18653/V1/2023.EACL-MAIN.148.  URL https://doi.org/10.18653/v1/2023.eacl-main.148.
  * Neelakantan et al. (2022) A. Neelakantan, T. Xu, R. Puri, A. Radford, J. M. Han, J. Tworek, Q. Yuan, N. Tezak, J. W. Kim, C. Hallacy, J. Heidecke, P. Shyam, B. Power, T. E. Nekoul, G. Sastry, G. Krueger, D. Schnurr, F. P. Such, K. Hsu, M. Thompson, T. Khan, T. Sherbakov, J. Jang, P. Welinder, and L. Weng.  Text and code embeddings by contrastive pre-training.  _CoRR_ , abs/2201.10005, 2022.  URL https://arxiv.org/abs/2201.10005.
  * Qi et al. (2022) W. Qi, Y. Ruan, Y. Zuo, and T. Li.  Parameter-efficient tuning on layer normalization for pre-trained language models.  _CoRR_ , abs/2211.08682, 2022.  doi: 10.48550/ARXIV.2211.08682.  URL https://doi.org/10.48550/arXiv.2211.08682.
  * Radford et al. (2019) A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever.  Language models are unsupervised multitask learners.  2019\.
  * Scao et al. (2022) T. L. Scao, A. Fan, C. Akiki, E. Pavlick, S. Ilic, D. Hesslow, R. Castagné, A. S. Luccioni, F. Yvon, M. Gallé, J. Tow, A. M. Rush, S. Biderman, A. Webson, P. S. Ammanamanchi, T. Wang, B. Sagot, N. Muennighoff, A. V. del Moral, O. Ruwase, R. Bawden, S. Bekman, A. McMillan-Major, I. Beltagy, H. Nguyen, L. Saulnier, S. Tan, P. O. Suarez, V. Sanh, H. Laurençon, Y. Jernite, J. Launay, M. Mitchell, C. Raffel, A. Gokaslan, A. Simhi, A. Soroa, A. F. Aji, A. Alfassy, A. Rogers, A. K. Nitzav, C. Xu, C. Mou, C. Emezue, C. Klamm, C. Leong, D. van Strien, D. I. Adelani, and et al.  BLOOM: A 176b-parameter open-access multilingual language model.  _CoRR_ , abs/2211.05100, 2022.  doi: 10.48550/ARXIV.2211.05100.  URL https://doi.org/10.48550/arXiv.2211.05100.
  * Thakur et al. (2021) N. Thakur, N. Reimers, A. Rücklé, A. Srivastava, and I. Gurevych.  BEIR: A heterogeneous benchmark for zero-shot evaluation of information retrieval models.  In J. Vanschoren and S. Yeung, editors, _Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks 2021, December 2021, virtual_ , 2021.  URL https://datasets-benchmarks-proceedings.neurips.cc/paper/2021/hash/65b9eea6e1cc6bb9f0cd2a47751a186f-Abstract-round2.html.
  * Touvron et al. (2023) H. Touvron, T. Lavril, G. Izacard, X. Martinet, M. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample.  Llama: Open and efficient foundation language models.  _CoRR_ , abs/2302.13971, 2023.  doi: 10.48550/ARXIV.2302.13971.  URL https://doi.org/10.48550/arXiv.2302.13971.
  * Wang et al. (2022) L. Wang, N. Yang, X. Huang, B. Jiao, L. Yang, D. Jiang, R. Majumder, and F. Wei.  Text embeddings by weakly-supervised contrastive pre-training.  _CoRR_ , abs/2212.03533, 2022.  doi: 10.48550/ARXIV.2212.03533.  URL https://doi.org/10.48550/arXiv.2212.03533.
  * Wang et al. (2023) L. Wang, N. Yang, X. Huang, L. Yang, R. Majumder, and F. Wei.  Improving text embeddings with large language models.  _CoRR_ , abs/2401.00368, 2023.  URL https://arxiv.org/abs/2401.00368.
  * Wolf et al. (2020) T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush.  Transformers: State-of-the-art natural language processing.  In _Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations_ , pages 38–45, Online, Oct. 2020. Association for Computational Linguistics.  URL https://www.aclweb.org/anthology/2020.emnlp-demos.6.
  * Xiao et al. (2023) S. Xiao, Z. Liu, P. Zhang, and N. Muennighof.  C-pack: Packaged resources to advance general chinese embedding.  _CoRR_ , abs/2309.07597, 2023.  doi: 10.48550/ARXIV.2309.07597.  URL https://doi.org/10.48550/arXiv.2309.07597.
  * Zaken et al. (2022) E. B. Zaken, Y. Goldberg, and S. Ravfogel.  Bitfit: Simple parameter-efficient fine-tuning for transformer-based masked language-models.  In S. Muresan, P. Nakov, and A. Villavicencio, editors, _Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 2: Short Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022_ , pages 1–9. Association for Computational Linguistics, 2022.  doi: 10.18653/V1/2022.ACL-SHORT.1.  URL https://doi.org/10.18653/v1/2022.acl-short.1.
  * Zhang et al. (2024) B. Zhang, Z. Liu, C. Cherry, and O. Firat.  When scaling meets LLM finetuning: The effect of data, model and finetuning method.  _CoRR_ , abs/2402.17193, 2024.  doi: 10.48550/ARXIV.2402.17193.  URL https://doi.org/10.48550/arXiv.2402.17193.
  * Zhang et al. (2022) S. Zhang, S. Roller, N. Goyal, M. Artetxe, M. Chen, S. Chen, C. Dewan, M. Diab, X. Li, X. V. Lin, et al.  OPT: Open pre-trained transformer language models.  _arXiv preprint arXiv:2205.01068_ , 2022.

##  Appendix A Additional results

In Figure 7 (for full fine-tuning), Figure 8 (for block freezing), Figure 9
(for LoRA), and Figure 10 (for bias tuning), we present the IsoFLOP profiles
where downstream task performance is shown. The downstream performance is
measured on a subset of MTEB tasks described in Appendix B.

Figure 7:  IsoFLOP profiles for _full fine-tuning_ , where performance on MTEB
tasks is observed instead of the loss. Both axes use log-scale. The
performance remains approximately inversely correlated with the loss. However,
interestingly, there are no corresponding plateaus on the IsoFLOP profiles.
Figure 8:  IsoFLOP profiles for _freezing_ , where performance on MTEB tasks
is observed instead of the loss. For a given FLOP level, the model with the
optimal active block fraction is depicted.  Figure 9:  IsoFLOP profiles for
_LoRA_ , where performance on MTEB tasks is observed instead of the loss. For
a given FLOP level, the model with the LoRA rank is depicted.  Figure 10:
IsoFLOP profiles for _bias tuning_ , where performance on MTEB tasks is
observed instead of the loss.  Figure 11:  The optimal downstream task
performance achieved using all four different fine-tuning methods at given
budgets. This is a performance analogue of Figure 1 (which shows final
training loss instead).

Using the location of optimal loss given the computational budget for freezing
and bias tuning methods – as illustrated in Figures 2(b) and 4(a) – we project
optimal model sizes for each computational budget in Figure 12. Similar
figures for full fine-tuning and LoRA are presented in the main text in Figure
6(a) and Figure 6(b), respectively.

Figure 12: Optimal choices of model size given the computational budget, with
respect to the _final training loss_ , for block freezing and bias tuning.

Using the location of _optimal downstream task performance_ given the
computational budget for all the four fine-tuning methods – as illustrated in
Figures 7, 8, 9, and 10 – we project optimal model sizes for each
computational budget in Figure 13.

Figure 13: Optimal choices of model size given the computational budget, with
respect to the _performance on the downstream task evaluation_.

##  Appendix B Downstream task evaluation details

For each category present in the MTEB benchmark [Muennighoff et al., 2023b],
we take the data from the models evaluated in the original paper and select
the task that has the highest correlation with the category average. This
results in 8 tasks from the MTEB benchmark used for the downstream task
evaluation:

  1. 1.

EmotionClassification (classification),

  2. 2.

TwentyNewsgroupsClustering (clustering),

  3. 3.

SprintDuplicateQuestions (pair classification),

  4. 4.

AskUbuntuDupQuestions (re-ranking),

  5. 5.

SciFact (retrieval),

  6. 6.

STS15 (semantic text similarity)

  7. 7.

SummEval (summarization),

  8. 8.

Tatoeba (bitext mining).

For each of the task above, we specify its category in brackets. Since Tatoeba
is not featured in the MTEB
leaderboard,333https://huggingface.co/spaces/mteb/leaderboard we skip it when
calculating our approximation of the MTEB performance of the embedding models.

In principle, one could evaluate all the checkpoints on all the MTEB tasks.
However, we decided to evaluate on a representative subset in order to speed
up evaluation (evaluating some MTEB task takes more than 24 hours, and we have
close to 1000 checkpoints, so the overall cost is significant).

##  Appendix C Correlation between loss and MTEB performance

In our work, we use loss as the main metric used to formulate the conclusions
in form of the compute-optimal recipe and the scaling laws. This is in line
with other works investigating scaling laws in various contexts, like
[Hoffmann et al., 2022, Muennighoff et al., 2023a, Zhang et al., 2024].
However, we additionally evaluated the trained models on a subset of MTEB
benchmark. We calculated the Spearman’s rank correlation between the loss and
the MTEB performance and it is equal to $-0.892$, which is significant.
Additionally, Figure 14 visually compares the loss and MTEB performance.

Figure 14:  All the final checkpoints involved in our suite of experiments,
evaluated bath in terms of the final training loss and the performance on a
representative subset of MTEB benchmark. The correlation between these two
metrics is significant.

##  Appendix D Hyperparameters for the training setup

Below, we list the hyperparameters that we used for training.

    batch_size = 1024
    context_length = 75
    AdamW.weight_decay = 0.1
    tau = 0.025 # temperature parameter

###  D.1 Learning rate for block freezing and full fine-tuning

For all the experiments with block freezing and full fine-tuning, we base our
peak learning rate on the peak learning rate from Pythia [Biderman et al.,
2023]. We calculate our learning rate by dividing the Pythia learning rate by
ten. We use the cosine learning rate schedule with warmup, warmup period being
fixed to 10% of the training steps. Learning rate starts at zero, goes up to
the peak learning rate, then follows the cosine schedule back to zero.

###  D.2 Temperature and weight decay for partial tuning

We found the temperature to affect the training performance rather
significantly, while the weight decay was of a very slight importance. We run
a grid search for partial tuning, varying temperature and weight decay and
arrived at a choice of $\tau=\frac{1}{40}$ and weight decay $0.1$. We compared
the runs by calculating the results on a subset of MTEB benchmark. Grid
results for temperature are detailed in Table 1, with best scores for each of
the temperature value listed.

Table 1: A comparison of different temperature values Temperature-1 | Avg. score
---|---
50 | 0.5
40 | 0.51
35 | 0.5
30 | 0.5
10 | 0.47
1 | 0.3

###  D.3 Readout: last vs average

We found that using averaging of embeddings as embedding results in a better
performance over selecting the last embedding. We compare the results of the
best run for mean and last pooling on full MTEB without MSMarco subset. The
scores are presented in Table 2.

Table 2: Mean pooling works better than last pooling. Pooling method | Avg. score
---|---
Mean pooling | 0.47
Last pooling | 0.36

###  D.4 LoRA hyperparameters

In order to be able to compare the loss from LoRA with partial block freezing,
we leave most of the hyperparameters unchanged. We change the learning rate,
we find the optimal learning rate by running a grid search over
hyperparameters. The result shows that for 3 different rank choices, the same
learning rate works best. The exact values are presented in Table 3.

Table 3: Learning rate for LoRA. Learning rate | LoRA rank | Loss
---|---|---
$1102$ |  | 1.96
$1103$ | 8 | 1.35
$1104$ |  | 2.22
$1102$ |  | 1.96
$1103$ | 16 | 1.34
$1104$ |  | 2.22
$1102$ |  | 2.09
$1103$ | 32 | 1.32
$1104$ |  | 2.19

###  D.5 Bias tuning hyperparameters

Similarly, as with LoRA, for bias tuning to be comparable with partial block
freezing, we leave most of the hyperparameters the same. We vary the learning
rate, we find the optimal learning rate by running a grid search over
hyperparameters. The exact values can be seen in Table 4.

Table 4: Learning rate for bias tuning. Learning rate | Loss
---|---
$1102$ | 1.24
$1103$ | 1.37
$1104$ | 2.97

##  Appendix E Scaling laws

We aim to express the formula describing the behaviour of the loss for our
models as a function of the number of parameters $N$ and training tokens $D$.
Following Hoffmann et al. [2022] we start with:

$L(N,D)=C+\frac{A}{N^{\alpha}}+\frac{B}{D^{\beta}},$

where $A,B,C,\alpha,\beta\in\mathbb{R}$ are parameters to be fitted based on
the data.

Our setting differs from that of Hoffmann et al. [2022] as we train using the
contrastive objective instead of the next-token prediction task, and moreover
we apply various compute-efficient methods making some fraction of the
parameters of the model not trainable. Despite that, we can still observe
certain linear trends that the loss follows. For all the methods, when the
FLOP budget is not overly restricted, if we fix the amount of both all the
parameters as well as the trainable parameters and vary the amount of FLOP,
the loss behaves approximately linearly. We observe that for full fine-tuning,
the points are roughly collinear, as in Figure 15. On the other hand, for LoRA
and partial block freezing, we discover that while the lines for one model are
parallel, they are translated, depending on the amount of trainable
parameters, which can be observed in Figures 18 and 19.

Following the Frantar et al. [2023]’s approach, we replace the term
$\frac{B}{D^{\beta}}$ by $\frac{(a_{s}(1-S)^{b_{s}}+c_{s}}{D^{\beta}}$, where
$S$ is an additional parameter indicating the fraction of the model’s
parameters that are trainable, and $a_{s},b_{s}$ are real-valued constants.
Similarly, if we keep the forward token amount fixed and vary the amount of
model parameters, the linear trend can be observed. Again, the lines are
parallel and differ by an intercept, that seems to be linearly dependent on
the logarithm of forward tokens, shown in Figure 17. Therefore, we replace the
term $\frac{A}{N^{\beta}}$ by $\frac{a_{d}\cdot\log(D)+b_{d}}{N^{\beta}}$,
where $a_{d},b_{d}$ are real-valued constants. The final formula describing
the loss in our setting has the form:

$L(S,N,D)=C+\frac{a_{d}\cdot\log(D)+b_{d}}{N^{\alpha}}+\frac{a_{s}\cdot(1-S)^{b_{s}}+c_{s}}{D^{\beta}}.$

We split our dataset into train and test set, with test set being the points
corresponding to Pythia 2.8B which is the biggest model that we consider and
train set the rest. Similarly to Hoffmann et al. [2022], we use L-BGFS
optimization algorithm and Huber loss. We find that $\delta=0.001$ prevents
over-fitting and decreasing it further does not bring more benefits. We then
compare the base formula being fitted formula from Chinchilla to our modified
formula and find that our works better, as demonstrated in Figure 16. Both
formulas are fitted after a grid search on the initial values of their
coefficients. It can be observed that we are able to accurately predict the
loss for the bigger model for all the methods. Therefore, we hypothetize that
our formula can be used for prediction the performance of models of an even
larger scale.

Figure 15:  Loss on the $y$ axis and FLOP on the $x$ axis. Lines of the same
colours correspond to the same amount of parameters in the model. Excluding
points corresponding to severely constrained FLOP budget, the loss behaviour
can be approximated by straight lines.

(a) Full fine-tuning, our formula

(b) Full fine-tuning, standard formula

(c) LoRA, our formula

(d) LoRA, standard formula

(e) Block freezing, our formula

(f) Block freezing, standard formula

(g) Bias tuning, our formula

(h) Bias tuning, standard formula

Figure 16:  The actual loss on the $y$ axis and predicted loss on the $x$
axis. Red points are the training set and blue the test set. Figures 16(a),
16(c), 16(e), and 16(g) are the results of fitting our formula for different
fine-tuning methods, while Figures 16(b), 16(d), 16(f), and 16(h) are the
result of fitting the standard formula from Hoffmann et al. [2022]. It can be
observed, that in contrast to the standard formula, our results in plots that
are roughly linear, without additional trends visible.  Figure 17:  Logarithm
of loss on the $y$ axis and logarithm of the parameter count on the $x$ axis.
The colors represent different quantiles of the amount of tokens used in the
model fine-tuning. Linear dependency can be observed, with the trends being
approximately parallel for different quantiles.

Figure 18: Plots of losses for block freezing. Each subplot contains
datapoints corresponding to one model size. Loss on the $y$ axis and FLOP on
the $x$ axis. Lines of the same colours correspond to the same amount of
trainable parameters in the model. The loss behaviour can be approximated by
straight lines.

Figure 19:  Plots of losses for LoRA. Each subplot contains datapoints
corresponding to one model size. Loss on the $y$ axis and FLOP on the $x$
axis. Lines of the same colours correspond to the same amount of trainable
parameters in the model. The loss behaviour can be approximated by straight
lines.

##  Appendix F Robustness to increase in batch size and context length

Constrained by the computational budget, we performed the biggest line of
experiments using batch size 1024 and context length 75. However, an increased
context length as well as batch size might be desirable in many cases. Here,
we demonstrate that the scaling laws established by us extend to a setting
with batch size 2048 and context length 512. For those experiments, we trained
five pre-trained decoder-only models from the Pythia suite [Biderman et al.,
2023], with sizes 14M, 31M, 70M, 160M, and 410M. We consider five
computational budgets: $1.51015$, $61015$, $2.41016$, and $9.61016$ FLOP. We
fit the formula following the method from Appendix E, with the change of using
the number of training steps as the measure of data seen by the model instead
of forward tokens. We then use it to predict the loss for the changed
hyperparameters values. The results are visible in Figure 20. We conclude that
our findings are relevant for setups with a larger batch size and context
length.

(a) Full fine-tuning

(b) LoRA

(c) Partial block freezing

(d) Bias tuning

Figure 20:  Plots of predicted vs actual losses for the 4 finetuning methods.
The formula is fitted on the data coming from trainings with batch size 1024
and context length 75 (red points). It is then used to predict the losses for
data points trained with batch size 2048 and context length 512 (green
points).

##  Appendix G Computational resources

For all the trainings, we used a cluster with A100 GPUs with 40GB VRAM or with
80GB VRAM. We used 4 CPUs and 128GB of RAM per GPU. We estimate that the
results contained in this article were obtained by expanding somewhere in
between 30k and 40k GPU hours.

##  Appendix H Dataset

We utilize the BAAI dataset available at https://data.baai.ac.cn/details/BAAI-
MTP.

It has a custom licence that allows for conducting academic research.


1.6 Tbps Classical Channel Coexistence With DV-QKD Over Hollow Core Nested Antiresonant Nodeless Fibre (HC-NANF)

O. Alia(1)*, R. S. Tessinari(1),
T. D. Bradley(2),
H. Sakr(2),
K. Harrington(2),
J. Hayes(2),
Y. Chen(2),

P. Petropoulos(2),
D. Richardson(2),
F. Poletti(2)
G. T. Kanellos(1),
R. Nejabati(1),
D. Simeonidou(1)

(1) High Performance Networks Group, University of Bristol, Woodland Road, Bristol, UK

(2) Optoelectronics Research Centre, University of Southampton, Southampton, SO17 1BJ, UK

*Corresponding author<EMAIL_ADDRESS>
978-1-6654-3868-1/21/$31.00 ©2021 IEEE

We demonstrate for the first time the coexistence of a quantum-channel and 8$\times$200 Gpbs 16-QAM optical channels with launching powers as high as -9dBm/channel in a 2 km HC-NANF. Comparative analysis with single-mode fibre reveals that the quantum-channel could not be sustained at such power-levels.

§ INTRODUCTION

Quantum Key Distribution (QKD) technology has been considered as the ultimate physical layer security due to its dependencies on the physical laws of physics to generate quantum keys <cit.>. However, in order for QKD to become functional for practical scenarios, it must be integrated with the classical optical networking infrastructure. Coexisting quantum and classical channels represent a challenge for QKD to multiple factors such as the high optical power used on classical channels (orders of magnitude higher than for quantum communication) and the additional noise generated from such channels due to optical non-linear effects <cit.> or insufficient isolation between the classical and quantum channels. This additional noise degrades the quantum channel performance and was studied extensively <cit.>.

Coping with optical nonlinearity represents a major challenge for QKD systems and sharp filters are required to partly isolate the noise from the classical channels and improve the quantum signal. However, non-linear processes in glass fibres also generate excessive photons at the same wavelength as the quantum channel that cannot be mitigated, dramatically affecting its performance. HCFs offer a radical solution as they provide several attractive advantages compared to glass core fibres, allowing for an ultra-low optical mode overlap with the glass (and hence reduced optical nonlinearity and Rayleigh scattering), a lower latency, and a very low total chromatic dispersion. These desirable qualities enable the transmission of classical channels at high optical powers <cit.> while coexisting with quantum channels over the same medium . One associated problem of HCF is the increased fibre losses and the high loss of interfacing between HCF and single mode fibre (SMF). A novel HCF design that could reach a total loss value lower than that of conventional solid fibres was proposed in 2014 and is called Hollow Core Nested Antiresonant Nodeless Fibre (HC-NANF) <cit.>. The losses of HC-NANF have improved immensely from 1.3 dB/km <cit.> to 0.65 dB/km <cit.> and recently to 0.28 dB/km <cit.>. Moreover, a record-low loss of interconnection between HC-NANF and SMF of 0.15 dB has been achieved which is slightly higher than the theoretically-expected minimum loss of 0.08 dB <cit.>. HC-NANF was also used in long-haul WDM transmission with a record transmission distance of 618 km and 201 km using PM-QPSK and 16-QAM respectively. <cit.>.

In this paper, we take the advantage of the ultra-low nonlinearity of HC-NANF to demonstrate the coexistence of a 16 dB power budget discrete variable QKD (DV-QKD) technology <cit.> and 8 $\times$ 200 Gbps 16-QAM carrier-grade classical optical channels at an extremely high coexistence power of 0 dBm over a 2 km HC-NANF, revealing minimal effects on the quantum channel performance. We also compare the QKD performance using SMF and HC-NANF in terms of Secret Key Rate (SKR) and Quantum Bit Error Rate (QBER) to prove the superiority of HC-NANF for the coexistence of quantum and classical channels.

§ TESTBED

Experimental testbed design.

<Ref> shows the experimental system setup used to demonstrate the coexistence of eight classical channels with the quantum channel. The testbed facilitates the experimental evaluation of the nonlinear effects on the performance of the quantum channel created by the presence of classical channels spectrally close to the quantum channel using both HC-NANF and SMF. For the classical channels, two bandwidth variable transponder (BVT) switches (Facebook Voyager) are deployed to transmit eight coherent classical channels. The frequencies of the classical channels range from 193.60 THz to 193.25 THz with 50 GHz spacing between each channel and 75 GHz spacing between the quantum channel (193.70 THz) and the nearest classical channel as shown in <Ref>. Each classical channel provides a 200 Gbps, 16-QAM signal with a soft-decision forward error correction (SD-FEC) of 15% to enable an error-free transmission. For the quantum channel, IDQuantique Clavis3 DV-QKD system <cit.> is used which implements Coherent One Way (COW) protocol and has a fixed frequency for the quantum channel at 193.70 THz. As shown in <Ref>, the eight coherent output ports of the two BVTs are multiplexed using a wavelength selective switch (WSS) with 6 dB of insertion loss for a total throughput of 1.6 Tbps. The WSS is used as a multiplexer and a band pass filter to couple the classical channels into a single fibre and provide a 30 dB isolation. The WSS combined output feeds the input of a tunable band pass filter (TBPF) with 60 dB of isolation to further suppress the noise generated by the classical channels.
The classical channels and quantum channel are then coupled through a WDM multiplexer in a coexisting configuration and travels through a 2 km of SMF or HC-NANF to a WDM demultiplexer that passes the quantum channel (193.70 THz) while rejecting all other channels (classical channels). The rejected classical channels are forwarded towards an optical isolator with insertion loss of 2 dB to prevent the tunable laser used by the BVT Rx as a local oscillator from returning to the Bob-QKD unit and interfering with the QKD measurements. It also prevents the Amplified Spontaneous Emission (ASE) noise generated by the Erbium-Doped Fibre Amplifier (EDFA) which is used to amplify the classical signal travelling back to the Bob-QKD unit. The amplified classical channels are directed to the Voyager BVT Rx for coherent detection after the EDFA via a 1x8 optical splitter.
Furthermore, an optical isolator is used after the Alice-QKD unit to prevent the noise generated from the classical channels returning to the Alice-QKD unit and interfering with the quantum signal. Moreover, the variable optical attenuator (VOA) is used before the Bob-QKD unit to adjust the losses of the quantum channel to the minimum operation level of 10 dB to prevent over saturating the single photon detector in the Bob-QKD unit. It also allows us to obtain a direct comparison of the quantum/classical coexistence between the SMF and HC-NANF in terms of losses. The loss of the quantum channel using the HC-NANF and SMF is $\approx$ 10.5 dB. Moreover, the 2 km HC-NANF used in this experiment has a loss of 1.3 dB/km <cit.> and a total loss of 7 dB when accounting for the  interconnection losses between the HCF and SMF.

Coexistence spectrum  of  the   quantum  channel  (Red)  and  8$\times$200  Gpbs  classical  channels  (Yellow).

§ RESULTS

a) Average SKR versus launch optical power using HC-NANF. b) Average SKR versus launch optical power using SMF. c) Average QBER versus launch optical power using HC-NANF. d) Average QBER versus launch optical power using SMF.

<Ref> shows the experimental evaluation of SKR and QBER of the quantum channel in the presence of a different number of classical channels at different launch powers for HC-NANF and SMF. <Ref>a) and <Ref>c) present the changes in the SKR and QBER of the quantum channel when coexisting with different number of classical channels at different launching power in a HC-NANF. As shown in <Ref>a) and <Ref>c), the SKR and QBER values without the presence of any classical channel (no coexistence) are similar to the values when coexisting 8 classical channels at a launch power of -9 dBm per channel which is equivalent to the highest recorded coexistence power of 0 dBm. This is due to the ultra-low optical nonlinearity in HC-NANF. Moreover, <Ref>b) and <Ref>d) show the SKR and QBER when coexisting quantum and classical channels in SMF. As shown in <Ref>b) and <Ref>d) the SKR values drops from $\approx$ 2400 bps to $\approx$ 304 bps while the QBER values increase from $\approx$ 3% to $\approx$ 4.5% when coexisting 8 classical channels at a launch power of -25 dBm per channel which is equivalent to a coexistence power of -16 dBm. This significant 87.5% drop in the SKR and 50% rise in the QBER is due to high optical nonlinearity in SMF causing a high noise leakage to the Bob-QKD unit. Moreover, increasing the coexistence power by 1 dBm would cause the SKR to plummet to zero bps. Although the losses of the quantum channel are similar when using both SMF and HC-NANF and the coexistence power (0 dBm) in the HC-NANF is 40 times higher than coexistence power (-16 dBm) in SMF, the ultra-low nonlinear effects of the HC-NANF due to its hollow core preserve the SKR and QBER of the quantum channel. This proves the suitability of HC-NANF as an excellent transmission medium for coexistence of quantum and classical channels.

§ CONCLUSIONS

The coexistence of a DV-QKD channel and 8 $\times$ 200 Gbps classical channels was successfully demonstrated over a 2 km long HC-NANF for a record-high transmission of 1.6 Tbps. The SKR was preserved without any noticeable changes when coexisting the quantum channel with eight classical channels at 0 dBm coexistence power compared to a significant drop of 87.5% when using SMF at -16 dBm coexistence power which is 40 times lower than the power used in HC-NANF. This significant difference in the QKD performance proves the advantage of using HC-NANF to provide a seamless coexistence of quantum and classical channels with a minimal effect on the quantum channel performance regardless of the number and power of the classical channels.

§ ACKNOWLEDGEMENTS

This work was funded by EU funded project UNIQORN (820474) and the EPSRC Airguide Photonics Partnership Fund Award (ref: 517129) (EP/P030181/1). Part of the research leading to this work has been supported by the Quantum Communication Hub funded by the EPSRC grant ref. EP/T001011/1 and the ERC LightPipe project (grant n 682724).

Missing 'biblatex' package
The bibliography requires the 'biblatex' package.

Optical Society of America

journaltitleAdvances in Optics and Photonics
titleAdvances in quantum cryptography

journaltitleJournal of Lightwave Technology
titleLimitations on lightwave communications imposed by optical-fiber nonlinearities

IOP Publishing

journaltitleNew Journal of Physics
titleQuantum key distribution and 1 Gbps data encryption over a single fibre

journaltitlePhysical Review X
titleCoexistence of high-bit-rate quantum key distribution and data on optical fiber

booktitle21st ICTON
titleCoexistence of Discrete-Variable QKD with WDM classical signals in the C-band for fiber access environments

journaltitleJournal of lightwave technology
titleImpact of Raman scattered noise from multiple telecom channels on fiber-optic quantum key distribution systems

American Institute of Physics

journaltitleApplied Physics Letters
titleBackscattering limitation for fiber-optic quantum key distribution systems

Nature Publishing Group

journaltitleNature Photonics
titleTowards high-capacity fibre-optic communications at the speed of light in vacuum

journaltitleJournal of Lightwave Technology
titleNonlinearity-free coherent transmission in hollow-core antiresonant fiber

Nature Publishing Group

journaltitleNature Photonics
titleHollow-core photonic crystal fibres for gas-based nonlinear optics

Optical Society of America

journaltitleOptics express
titleNested antiresonant nodeless hollow core fiber

booktitleEuropean Conference on Optical Communication (ECOC)
titleRecord low-loss 1.3 dB/km data transmitting antiresonant hollow core fibre

booktitle45th ECOC
titleAntiresonant hollow core fibre with 0.65 dB/km attenuation across the C and L telecommunication bands

Optical Society of America

booktitleOptical Fiber Communication Conference
titleHollow core NANF with 0.28 dB/km attenuation in the C and L bands

Nature Publishing Group

journaltitleScientific Reports
titleLow loss and high performance interconnection between standard single-mode fiber and antiresonant hollow-core fiber

journaltitleJournal of Lightwave Technology
titleTransmission of 61 C-band Channels over Record Distance of Hollow-Core-Fiber with L-band Interferers

titleClavis3 QKD Platform