Long Context Fine-tuning for Repetition Task

Introduction

This cookbook is part of a technical deep dive blogpost on long context finetuning that you can read here.

If you ask an LLM to repeat a sequence back to you, surely this should be easy, right? The answer is not straight forward and might be surprising to many!

A lot of the capabilities that we know and trust our LLMs to have, fall short at longer contexts!

To solve this repetition task a LLM should be able to use a simple induction head - that just copies a specific part of the input back out.

However, for this task at longer contexts non-finetuned models fail quite miserably!

In this notebook we will:

1. Use a previously created dataset of long input sequences (upto 128k tokens)
2. We will setup the repitition task, where we ask the model to repeat the last k words of the sequences created in Step 1.
3. Demonstrate how even the best LLMs fail at this simple repitition task.
4. We will fine-tune the model on ~1975 examples of this long-context task and show a radical improvement.

Install Libraries

Dataset Preperation

In order to create a long context dataset for our repetition task we extracted 2000 samples of varying length as shown below:

Repetition Task Definition

We used the code above to previously curate a dataset of long sequences by processing the FineWeb and RedPajama datasets and retrieving 2000 English documents of 32k and 128k context length each.

For each of these documents we want to setup a task prompt that we can pass into an LLM.

This was done as follows:

For the task prompts outlined above we need the LLM to, given an input sequence of arbitrary length, repeat the last k words of the sequence back to us. Where K is an random number beteween 1 and 100.

We also extract the correct last k words directly from the document and store this to use for comparision with ground truth later.

We have provided a JSON file from which you can load the task prompts.

As we can see from the single example above, our LLM is not great at this task.

Ideally the LLM should be able to use an induction head to repeat a previously seen sequence back out. An induction head is a key component in transformer models that specializes in pattern recognition and prediction. Like a pattern-matching expert, it identifies repeated sequences in text and uses previous occurrences to predict what comes next. For example, if a phrase appeared before and was followed by specific text, the induction head remembers this pattern and applies it to similar future situations. This capability is fundamental to how transformers process language, enabling them to learn from repetition and make informed predictions based on previously seen patterns. Think of it as the model's memory mechanism for recognizing and utilizing recurring patterns in text.

Use Levenshtein Distance to Evaluate

For this repetition task we need an exact comparision between the correct sequence of words to the LLM output sequence of words.

Since this is an exact matching task we will use Levenshtein Distance.

Levenshtein Distance measures how different two strings are by counting the minimum number of single-character changes (including inserting a character, deleting a character, or replacing a character) needed to turn one string into another.

For example the levenshtein distance between kitten and sitting is 3 since we need 3 operations to for from one to the other.

kitten → sitten  (replace 'k' with 's')
sitten → sittin  (replace 'e' with 'i')
sittin → sitting (insert 'g')

Total Levenshtein Distance = 3 operations

Think of it like measuring the "editing effort" needed to transform one word into another. The lower the number, the more similar the strings are. A distance of 0 means the strings are identical, while larger numbers indicate more differences.

For our purpose we will use ratio = 1 - (leven_distance / (len1 + len2)) to obtain a score between 0 and 1.

• 0 implies that the two strings are very different
• 1 implies that that two strings are identical

For our repetition task higher is better!

Next we will loop over the first 25 task items and see how well our Llama 3.1 70B model performs at this task!

As we can see above, Llama3.1 70B performs suboptimally at this repetition task.

Fine-tune on Repetition Task

Below we will fine-tune a smaller Llama 3.1 8B model on 1975 examples of this repetition task and see if we can get it to outperform on this task.

Once the Model is finetuned we can assess how well it performs.

Deploy Model and Run Evals

Before we can run the evaluations we need to deploy our finetuned model as a Dedicated Endpoint(DE). After fine-tuning completes, access your model through the Together AI dashboard. Go to Models, select your fine-tuned model, and select Deploy. Choose from the available hardware options - we'll use a single A100-80GB GPU for this example.

We can see that even for the smaller 8B model the Score improves from 0.37 to 0.81 after finetuning when compared to the 70B untuned model.

We can also see that after finetuning the model more often outputs the correct number of words, with the length difference decreasing from 103.44 before finetuning to 15.08 afterwards.

To learn more about our fine-tuning API read the docs here!