Long context LLM: stretching RoPE with YaRN

Intro

0:00 Hello and welcome back to episode twenty-two, where we're going to give our model a context window four times as long, by stretching how its rotary embeddings turn. We do it first with no training at all, and then with fifteen minutes of it. Let's begin with a quick look at where we're going to be by the end of the episode.

0:19 By the end of today, our web model, which was trained on windows of one thousand and twenty-four tokens, reads pages of four thousand and ninety-six. Read as it was trained, its loss climbs past the window, to about five point nine eight at the end of the page. Stretched with YaRN and fine-tuned for fifteen minutes, it stays flat, and ends at about three point zero seven.

0:40 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on longer contexts today. And each episode has a page here as well, with its chapters and its code.

1:00 Now, before we get to today's work, let's quickly go over what we did last time. We fine-tuned our web model into an agent, which works with files and hands parts of a task to other roles. Each subagent is our model again, in a fresh context, and its answer alone goes back to the agent. On the long task, adding up four files, the agent's own context held about one hundred and fifty-nine tokens. Done in one context, the task needed up to six hundred and eighty-seven of the one thousand and twenty-four tokens our model can read. Subagents keep each context small, and today, we make the window larger.

1:36 Our web model was pretrained on windows of a thousand and twenty-four tokens. In episode ten, we saw that rotary embeddings let a model read past its window, and that its loss got worse when it did. Today, we're going to see why: past the window, some pairs of channels turn to angles that the model didn't see in training. We stretch the rotations four ways, from two papers: position interpolation, NTK-aware interpolation, NTK-by-parts, and YaRN, which adds a temperature to attention. We measure each on long web pages, position by position, and with a test that hides a number in a long text and asks for it at the end. Then we fine-tune with two of them, for fifteen minutes each, on pages of four thousand and ninety-six tokens. And we work out what a longer window costs in memory. All right, let's get started with what happens past the window.

Past the window

2:29 In episode ten, rotary embeddings turned each query and each key, a pair of channels at a time, by an angle that grows with the token's position. Each of a head's thirty-two pairs turns at its own speed, theta i, which is ten thousand to the minus two i over sixty-four. A query's score against a key depends on the difference of their angles, so on the distance between the two tokens, times each pair's speed. Pair zero turns a radian a token, about a hundred and sixty-three full turns in a window of a thousand and twenty-four, so in training, it went round many times. The last pair, thirty-one, turns about seven point eight degrees in a whole window. At four thousand and ninety-six tokens apart, it's about thirty-one degrees, a difference that no query and key in training had. So for those distances, attention's scores come from angles the model has no training for, and the loss shows it.

3:20 To see it, we're going to write long dot py, which measures the loss at each position of long pages. It imports argparse for its arguments, json for its results, numpy, torch and its functional module. From our own files, it takes load checkpoint and the data folder, then the results folder, which hellaswag dot py set up in episode nineteen, and pick device. Factor is four, as the pages are four windows long. We measure sixty-four pages, and the results average the loss in stretches of two hundred and fifty-six positions, for the charts. Long spans finds the pages longer than a length, in the held-out tokens or in the training ones. Memmap opens the tokens of the web as an array, without reading the whole file into memory. Above it, we say what the function gives back: where each web page of more than length tokens starts and ends, each one ending with end of text. So flat non zero gives the positions of all the end of text tokens, as an array of indices. A page starts one after the end of the page before it, and the first one at zero. Keep marks the pages of more than length tokens, and the function returns the tokens, with where each of those pages starts and ends. Long pages takes the first length plus one tokens of the first n long held-out pages, as the line over its body reads. Each page becomes a batch of one on the device, and the extra token is there so that the last position has a target too. Then we write two ways of measuring a page, with a loss at each of its positions. Full losses runs the whole page through the model in one pass, so each token is predicted from all the tokens before it. Cross entropy with reduction none keeps a loss for each position, where by default it averages them. Window losses gives the model no more than the window it was trained on. The first window runs as it is, and past it, windows a quarter of one apart each predict their last quarter. So each token past the first window sees from three quarters of a window to a whole one before it. Episode ten's context dot py slid the window a token at a time. For a window of a thousand and twenty-four tokens, that would take a pass for each token past the first window, about two hundred and forty times as many passes as our thirteen. Main's arguments are the checkpoint, the factor, the number of pages and the device, which defaults to pick device, the Mac's GPU here, through PyTorch's M P S backend. Argparse gives each one a type and a default, and a help text where its name isn't enough. We load the model, put it in evaluation mode, which turns dropout off, and name the results after the checkpoint's folder. Window is the block size the model was trained with, and the end of its line says it's the window it was trained on. Length is the window times the factor, four thousand and ninety-six tokens for us, and pages are the long pages, on the device. Results keeps what we measure, with the window, the length, the pages and the bin. Then we print two header lines: what's measured, and a column for each window's worth of positions, from zero to a thousand and twenty-three, and so on. There are two methods so far: the window, and none, the whole page at once. For the window, each page goes through window losses. For none, it goes through full losses, and the comment over that line says why we can: rotary embeddings turn by any position, but the model saw a window's worth. Torch stack puts the pages' losses in one tensor, a row for each page. View splits the positions into windows, a thousand and twenty-four at a time, and into bins. Mean over dimensions zero and two then averages over the pages and over the positions in each window, which leaves a number for each window. Each method keeps its bins, its windows, and its loss in the window and past it, and prints a row. At the end, the results go to a JSON file in results, long, named after the checkpoint's folder, pretrain web, in a folder that makedirs creates the first time.

7:20 Let's run it on the web model from episode nineteen. With the sliding window, the loss is about three point one in the first window, and it stays about there in the three after it, as each token sees most of a window either way. Reading the whole page at once gives three point one in the first window too, as up to there, the two are one computation. Then it climbs, to about four point six in the second window, five point seven in the third, and five point nine in the fourth.

7:46 On a chart, with a dot for each stretch of two hundred and fifty-six positions, the sliding window stays flat along the page. Reading the whole page follows it to the end of the window, and right after it, the loss jumps, then keeps rising the further we go, to about five point nine eight at the end. So the model can't use what's further back, and worse, reading it raises the loss well above the sliding window's. A model that sees more of the page should predict at least as well as one that sees a window of it, and that's what we're after today.

Position interpolation

8:18 The first fix we're going to try is from Extending Context Window of Large Language Models via Position Interpolation, by Shouyuan Chen and others at Meta, from 2023. Its abstract says position interpolation extends the context of Llama models up to thirty-two thousand, seven hundred and sixty-eight tokens, with fine-tuning of a thousand steps at most, and keeps their quality on texts that fit the original window. Figure one draws the idea for a model trained on two thousand and forty-eight positions. Reading four thousand and ninety-six positions as they are is extrapolation, out to positions the model hasn't seen, in red. Position interpolation scales the positions down, so the four thousand and ninety-six fit into the range the model was trained on. Figure two shows why extrapolating goes wrong: an attention score, as a function of the distance, stays between minus one and one where it was fitted. Past that range, it can reach thousands, which breaks attention. So in section two point three, the rotation for position m becomes the rotation for m times L over L prime, the old window over the new one, which for us is a quarter.

9:24 In our code, position interpolation divides each pair's speed by four, which turns each position by the angle of a quarter of it. So a token at position four thousand and ninety-six turns as far as one at a thousand and twenty-four did, and the whole page fits in the angles of one window. The slowest pair goes back to seven point eight degrees at most, which the model knows. The cost is at the other end, on the fast pairs. In training, two neighbouring tokens were one radian apart on pair zero, and now they're a quarter of a radian apart, as if they were a quarter of a token apart, which no two tokens were in training. So the fast pairs, which tell nearby tokens apart, now carry distances the model hasn't seen either, and the paper fine-tunes the model afterwards, so it can learn them.

NTK-aware, and NTK-by-parts

10:11 The other three come from YaRN: Efficient Context Window Extension of Large Language Models, by Bowen Peng and others, also from 2023. Its abstract says that models with rotary embeddings fail to generalise past the length they were trained on. YaRN extends their context with ten times fewer tokens and two and a half times fewer training steps than the methods before it. Section three point one makes the point we just saw: position interpolation stretches all the pairs equally, which removes the high frequencies, the fast pairs that tell nearby tokens apart. So NTK-aware interpolation, first posted on Reddit in 2023 by a user called bloc97, spreads the slowing across the pairs, scaling the high frequencies less and the low frequencies more. Section three point two looks at each pair's wavelength, the tokens it takes to go once round. Some pairs have a wavelength longer than the context in training, so they didn't make a full turn, and the paper suggests that those pairs give the model absolute positions, as each token sits at an angle of its own. Pairs with short wavelengths went round many times, so to them, what matters is the distance between tokens. So NTK-by-parts counts each pair's turns in the window, r, and sets a ramp, gamma, between two bounds, alpha and beta. It leaves the pairs that turn more than beta times as they are, and interpolates the ones that turn less than alpha times. With the paper's values for Llama, alpha is one and beta is thirty-two.

11:37 NTK-aware interpolation does it by changing the base of the speeds, the ten thousand. For a factor of four and heads of sixty-four channels, the base becomes ten thousand times four to the sixty-four over sixty-two, about forty-one thousand, eight hundred. Pair zero's speed is the base to the power of zero, so it stays at a radian a token. The power of sixty-four over sixty-two is chosen so that the last pair's speed is divided by four, which puts it within the angles of the window, like position interpolation does. The pairs in between are slowed more the slower they are. So some of the slow pairs are slowed by less than four, and far enough along the page, they turn past the angles they had in a window, and we're going to see where that shows up.

12:20 For our model, a window of a thousand and twenty-four tokens splits the thirty-two pairs into three groups. Pairs zero to five turn from thirty-nine to a hundred and sixty-three times in a window, which is more than thirty-two. Pairs six to seventeen turn from about one point two to twenty-nine times, and pairs eighteen to thirty-one turn less than once. So with NTK-by-parts, the first six keep their speeds, and nearby tokens are told apart as they were in training. The middle twelve are slowed by a factor between one and four, on the ramp, and the last fourteen are divided by four, so the whole page fits in the angles they had in a window.

YaRN's temperature

12:57 Section three point three is YaRN, which is NTK-by-parts with one more change. The paper observed that a temperature on attention's scores, before the softmax, lowers the loss evenly across the longer window, whatever the text and the position. In equation fourteen, the scores, q transposed times k, are divided by t, as well as by the square root of the head's size. Rotary embeddings make that cheap: multiplying both the query and the key by the square root of one over t does it, by scaling the rotations' cosines and sines. Equation fifteen is the temperature they recommend, for a factor s: the square root of one over t is zero point one times the natural log of s, plus one.

13:38 For us, the head's size is sixty-four, and t divides the scores before the softmax. With a factor of four, the square root of one over t is about one point one four. So the queries and keys are both made one point one four times longer, and the scores about one point three times larger. Larger scores make the softmax sharper, so attention puts more of its weight on the tokens that score highest. With four times as many tokens to spread its weight over, that keeps attention about as focused as it was in a window. That's my intuition for it, while the paper sets the formula from what it measured on Llama.

Stretching in code

14:14 In model dot py, we import replace from dataclasses, next to dataclass, as we're going to make a config from another one with a few fields changed. The config gets two new fields, right after rope, for the method and the factor. Rope scaling is the method, a string or None, and its comment lists them: rotations stretched past the window trained on, P I, N T K, parts or yarn. Rope factor is by how much, so the window trained on is block size over rope factor. The defaults are None and one, so the checkpoints we've saved, which have neither field, load with their rotations as they were. Scalings lists the four names, for the check and the scripts, and fast turns and slow turns are NTK-by-parts' two bounds, thirty-two and one. A pair of channels that turns more than thirty-two times in the window trained on keeps its speed, one that turns less than once is slowed by the factor, and those in between by some. Rope speeds gives each pair's speed, starting from the speeds we had, in radians per position, faster for the first pairs and slower for the last. With P I, each speed is divided by the factor: each pair slowed by the factor, so position four thousand and ninety-six turns as far as a thousand and twenty-four did. With N T K, the base becomes ten thousand times the factor to the head size over the head size minus two. A bigger base leaves the fastest pair as it is, and slows the slowest by the factor. Parts and yarn share the ramp, as YaRN is NTK-by-parts with a temperature. Turns is how many times each pair goes round in the window trained on, its speed times the window, over two pi. Keep is the ramp, gamma in the paper: zero at one turn or fewer, one at thirty-two or more, in a line between, and clamp holds it in that range. Each speed is then keep times its own speed, plus one minus keep times its speed over the factor, which is equation thirteen. Rotary takes the method, the factor and the window trained on, and gets the speeds from rope speeds. Over that line goes what the function returns: the angles for positions start to start plus T minus one, one per pair of channels, as cosines and sines. The angles are each position times each pair's speed, a row per position, turned into cosines and sines, as before. For YaRN, m is zero point one times the log of the factor, plus one, and the cosines and sines are multiplied by it. That's YaRN's temperature, and the comment explains it: queries and keys both lengthened by m, so attention's scores grow by m squared, a sharper softmax over the longer context. It's the length scaling trick from the paper's section three point three, and it adds no work to attention, as the cosines and sines are computed once, for all the blocks. Stretch makes the config of a model that reads further: replace copies the config, with the block size multiplied by the factor, along with the method and the factor in their fields. The GPT has no table of positions with rotary embeddings, so a stretched config loads the checkpoints we saved as they are, whatever its block size, and the embeddings and the blocks stay as they were trained. Above the return, a line on what it's for: the weights we have, with a window factor times as long, and the rotations stretched to fit it. In the model's constructor, a method needs rotary embeddings and one of the four names, or we raise a value error, as a mistyped method would otherwise give the rotations we had without a word. Its message lists the four names, and says that they need rope set to true. And forward passes the method, the factor and the window trained on to rotary. C is a short name for the config, as the call needs five of its fields, and the window trained on is block size over the factor, a thousand and twenty-four for the web model stretched by four. That's what the line over it now adds: a stretched model was trained on a window rope factor times shorter. Then the tests, in tests slash test long dot py. Tiny builds a small model with rotary embeddings, from a seed, as in our other tests. The first checks position interpolation: at four times a position, its angles match the ones we had at that position, for the first thousand and twenty-four positions. That's what its name says, position interpolation turns position four p as far as p. The cosines and sines are compared to within ten to the minus four, as float32 rounds the larger angles a little. The second checks that N T K keeps pair zero's speed, and divides the last pair's by four. Its last check gets a line of its own: the pairs in between are each slowed by some, and by no more than the factor. The third runs for parts and for yarn, and counts each pair's turns. Turns is the count we worked out for the thirty-two pairs, for a window of a thousand and twenty-four. The pairs with thirty-two or more keep their speeds, and torch equal checks that they didn't change at all. The ones with one or fewer are divided by four, and the ones between are slowed by more than one and less than four. The fourth checks YaRN's temperature: each pair's cosine squared plus sine squared, which is one for a plain rotation, is m squared. Over a hundred positions, that holds for each pair, so the temperature lengthens each query and key by m, and the turn stays a rotation. The fifth stretches the small model by a factor of one with each method. It loads the small model's state dict into the stretched one, checks that the block size stays thirty-two, and compares their logits on a batch of random tokens. They don't change, so stretching starts from the model we had. The sixth stretches it by four, with each method, and loads the small model's weights into it, as the fifth test did. Its block size is now a hundred and twenty-eight. It reads a hundred tokens, past the window of thirty-two it was built with, and checks the cache against the whole sequence, ninety tokens in one go and then one at a time. So the cache turns each new token by its own position, past the window trained on, with each method. And the last checks the two errors: a method without rotary embeddings, and a method we don't have. Pytest raises passes when the code inside it raises a value error, and fails the test when it doesn't. In long dot py, the import gets the scalings and stretch, and the factor's comment now says it's also how far a model trained on one window is stretched. The window it was trained on is the block size over the factor, so a stretched checkpoint reports the thousand and twenty-four it was trained on. Configs has one entry per method, none and the four, each the trained config stretched by the factor, where none keeps the rotations as they were. A checkpoint that's already stretched is measured as it is, with its own method. So it's one set of weights, read with another config each time. In the loop, model dot config is set before each method, as forward reads the config each time it runs, so one model in memory serves all five.

21:10 Let's run the tests with pytest, where dash v lists each case. All fourteen cases pass: the speeds of each method, the temperature, stretching by one with each of the four, a stretched model past its window with each of the four, and the two errors.

21:24 Now long dot py on the web model again, with no training. The window and none print the rows we had. Position interpolation is about four point four in the first window, well above the window, as the fast pairs no longer tell nearby tokens apart, and it stays about there along the page. N T K starts at about three point two, close to the window, and holds in the second window, then rises, to about four point eight in the last. NTK-by-parts is flat along the page, at about three point five past the window on average. And YaRN is flat too, at about three point one eight past the window, about zero point zero four above the sliding window, with no training at all.

22:05 On a chart, the five methods along the page, with the sliding window at the bottom. Without stretching, the loss rises right after the window. Position interpolation is high from the start, as it changes the angles of the nearby tokens too. N T K follows the window until about twenty-five hundred positions, then rises, as some of the slow pairs, which it slows by less than four, turn past the angles they had in training. NTK-by-parts stays flat, and YaRN stays flat and lower, so the temperature is worth about zero point three of loss here. In the first window, YaRN costs about zero point zero six against the model as it was trained. Past it, YaRN's model sees the whole page, and still predicts a little worse than a window does. So on web pages, with no training, the tokens far back don't help it yet.

The pass key

22:52 The loss averages over all the tokens of a page, and most of them are predicted from the words right before them. So to see whether the model can use a token thousands of positions back, we're going to write passkey dot py, with the test from the position interpolation paper's section three point three. It imports random for the keys, and from model dot py, the scalings, the KV cache and stretch. The prompt is the paper's, word for word. An introduction says there's an important info hidden inside a lot of irrelevant text. The filler is five short sentences about the grass, the sky and the sun, repeated, and the key is a five-digit number, said twice. The question at the end asks for it, with the start of the answer, the pass key is. That way, a pretrained model, which wasn't taught to answer questions, can carry on with the number. Depths are where the key goes in the filler, from the start, at zero, to right before the question, at one, and each length and depth gets ten keys. Prompt builds a prompt of about length tokens, with the key at a depth. It counts the tokens a filler takes, on ten repeats, and the tokens of the rest, and fits in as many repeats as there's room for, minus one, to stay under the length. Before is how many of them go ahead of the key. Answer takes the model and the prompt's ids. A KV cache, from episode four, keeps the keys and values, so the prompt runs through the model once, and then each new token on its own. Over the cache, we write how it answers: greedy, the likeliest token each time, through the KV cache. It writes six tokens with argmax, so running it again gives the answer it gave before. Six tokens are enough for the answer, as our tokenizer writes a five-digit number, with its space and full stop, in about four. Main reads the checkpoint, the factor, a whole number this time, as it counts the lengths, and the device. It loads the model in evaluation mode, and names the results after the checkpoint's folder, as long dot py does. The window is the one the model was trained on, the block size over the rope factor. The lengths are one to four windows: a thousand and twenty-four tokens, two thousand and forty-eight, three thousand and seventy-two, and four thousand and ninety-six. The configs come from long dot py too, each method stretching the trained config, or a stretched model as it is. Results keeps the window, the lengths, the depths and the trials, along with each method's grid, and we print a header, with a column for each length. Each method gets a random generator from a seed of zero, so, as the comment says, each method gets the keys the others get, and they're compared on one set of prompts. For each length and depth, ten keys: a random five-digit number, the prompt as a tensor on the device, and the answer, decoded back to text. A key is found when that text holds it, and found counts them, as Python adds a true as one. The grid keeps the keys found at each length and depth, a row per length, and by length is the share found over the five depths. Each method prints a row with a percentage for each length as soon as it's done, and the results go to a JSON file in results, pass key, named after the checkpoint.

26:07 Let's run it on the web model, with no training, and each method stretching it four times. For each method, it tries fifty keys at each length. As it was trained, with no stretching, it finds about forty percent of them in prompts of one window, and zero percent past it. Position interpolation finds zero percent at any length. N T K finds about ninety percent in a window, sixty-four in two and twenty-four in three, and zero in four. NTK-by-parts finds about eighty-eight percent in three windows and seventy in four, and YaRN about eighty and fifty-six.

26:40 As lines, by the length of the prompt. Without stretching, the model finds some keys in a window, and none past it, even with the key right before the question. It attends to the tokens more than a window back as well, and they come in at angles it has no training for. Position interpolation finds none at all, even in a window, which fits its loss: the nearby tokens are out of place, and copying a number needs them. N T K finds most keys in two windows, and fewer after, like its loss. NTK-by-parts and YaRN keep finding keys three thousand tokens back, which a sliding window of one thousand and twenty-four tokens can't do at all. NTK-by-parts finds a few more than YaRN here, though its loss was higher, so the temperature helps the loss more than it helps copying a number. Most of the keys they miss are at the two ends: at the start of the longest prompts, the furthest back, and right before the question. There, the model carries on with words instead of the number.

Fifteen minutes at 4,096

27:37 Both papers fine-tune after stretching, so the model can learn the new angles. We're going to write long train dot py, which fine-tunes a stretched model on long pages. It imports the summary writer for TensorBoard, our checkpoints, and from long dot py, the factor, full losses, long pages and long spans. Accum is four pages per step, one at a time, and the comment says why: two at once don't fit in a laptop's memory. With four thousand and ninety-six tokens, attention's scores take a lot of it, and we're going to work out how much at the end. The learning rate is five times ten to the minus five, after a warmup of ten steps, and then it stays where it is. The YaRN paper did that too, with twenty steps of warmup and a rate of two times ten to the minus five for Llama 2. We evaluate each twenty steps, on eight long held-out pages. Main seeds torch and numpy's random generator, so a second run picks the pages the first one did. It loads the web model, keeps its window, a thousand and twenty-four, and stretches its config with the method, by four, so length is four thousand and ninety-six. The training pages are long spans on the training tokens, and it prints how many there are, with their tokens in all. Valid is the first eight long held-out pages, from long pages. Page picks a random long page, and a random place in it, as its comment says. Integers gives a whole number from its first argument up to its second, so i is a page, and s is a start that leaves room for four thousand and ninety-seven tokens before the page ends. It returns four thousand and ninety-six inputs and their targets, the tokens one along, as a batch of one page. AdamW with no weight decay, as the paper's, and a TensorBoard run named after the method. Evaluate runs without gradients, as nothing in it trains, and in evaluation mode. It measures the held-out pages with full losses, in one pass each, and splits their loss into the window trained on and past it. Both go to TensorBoard and to the terminal, and the model goes back to training mode at the end. The loop runs until the minutes are up. Each step sets the learning rate for the warmup, as its comment says, so the first steps' updates don't knock the model off what it learned. The rate climbs by a tenth of five times ten to the minus five at each step, and from step ten on, it stays at five times ten to the minus five. Then four pages, each from a random place of its own, and each with its own forward and backward pass, so a step's memory is one page's. Divided by four, the gradients add up to the mean over the pages, as a batch of four would give, which is gradient accumulation, the trick from episode nineteen. Then the gradients are clipped, as before, and the optimiser takes a step. The step's loss, the four pages' average, goes to TensorBoard as the training loss. Each five steps, it prints the loss, the minutes and the tokens a second, and each twenty, it evaluates. At the end, a last evaluation, unless the last step just had one, and save checkpoint writes the weights, the merges and the config. The model goes to checkpoints, long, then the method, with its stretched config, so long dot py and pass key dot py measure it as it is. The method comes first on the command line, and the minutes second, fifteen by default.

30:59 First with position interpolation, which needs it most. There are about twenty-one thousand nine hundred training pages of more than four thousand and ninety-six tokens, about one hundred and eighty million tokens in all. At step zero, the held-out loss is about four point five four in the window and four point seven zero past it, on these eight pages. The laptop trains about twenty-three hundred tokens a second, one page of four thousand and ninety-six tokens at a time, with the gradients of four pages added up for each step. By the end, the loss is about three point three six in the window and three point three zero past it, after one hundred and twenty-seven steps in fifteen minutes.

31:38 Then we fine-tune with YaRN, for fifteen minutes too. It starts much lower, at about three point two nine in the window and three point two six past it. By step twenty, it's down to about three point one nine past the window, and the evaluations after it stay within a few thousandths of that. The model starts close to one that reads the longer window, so there's little left for it to learn, which is what the YaRN paper says about needing fewer tokens. By step one hundred and twenty, it's at about three point two five in the window and three point one nine past it. Each run reads about two million one hundred thousand tokens, against the three billion six hundred million tokens the model was pretrained on.

32:19 Past the window, on the eight held-out pages, position interpolation's loss falls fast in the first twenty steps, and then more slowly, still going down at the end. YaRN drops a little in the first twenty steps and stays there, as it started close to a model that reads the longer window. The YaRN paper reports that too, with fewer tokens and fewer steps than the methods before it, as its abstract said.

32:42 Now long dot py on the two fine-tuned models, each against a sliding window of its own weights, on the held-out pages we've measured all along, starting with the model fine-tuned with position interpolation. It's at about three point two three in the first window, down from four point four one with no training. In the last window, it's at about three point two zero, against three point two four with its own sliding window, so it uses the tokens far back. Then the model fine-tuned with YaRN, against a sliding window of its own weights too. It's at about three point one three in the first window. In the last, it's at about three point zero eight, against three point one four with its own sliding window, and three point one one for the model we started from.

33:26 Side by side, the loss past the window, for the models before and after fine-tuning. Fifteen minutes take position interpolation from the highest loss most of the way down, while it's still above the sliding window. YaRN, fine-tuned, is the lowest, below the sliding window of the model we started from. So with fifteen minutes on long pages, the stretched model predicts ordinary web pages better than a window can, as it uses tokens more than a window back.

33:54 And then the pass key, after fine-tuning, for both models. Position interpolation, which found none before, finds about fifty-six percent in a window and fifty-four percent in four. Then the model fine-tuned with YaRN, with as many keys at each length again. It finds about eighty percent in a window, and forty-four percent in four, below the fifty-six it found before fine-tuning. So fifteen minutes of web pages lowered its loss and didn't help it find the key, as nothing in those pages asks for a number from four thousand tokens back.

What a longer window costs

34:26 A longer window costs memory in two places. The KV cache from episode four keeps a key and a value for each token, in each of the web model's twelve blocks and twelve heads. That's sixty-four channels each, at four bytes a number, so about seventy-four thousand bytes a token. A full window of a thousand and twenty-four tokens takes about seventy-five megabytes, and four thousand and ninety-six take about three hundred, four times as much, as the cache grows with the tokens. That's what grouped-query attention, in episode eighteen, makes smaller. Attention's scores grow with the square of the length. PyTorch's attention on the Mac computes them all, a score for each pair of tokens in each head. So a block's scores for one page of four thousand and ninety-six tokens take about eight hundred megabytes, sixteen times the fifty they take for a window. Training keeps them for each of the twelve blocks, for the backward pass, and in my tests, one page took about twenty gigabytes on the Mac while fine-tuning. That's why long train dot py reads one page at a time. Flash attention, from episode thirteen, doesn't keep the whole matrix, which is how long contexts train on GPUs.

35:33 The docs get a page on longer contexts, with the two papers, the four methods and the code that stretches the rotations. It has the results too: the loss by position, the pass keys, the two fine-tuning runs and what the window costs in memory. The page goes in the list on the docs' home page, and at the end of the navigation. And the README has the commands, to measure a model past its window, to look for the pass key, to fine-tune with a method and to run the tests.

What's real, and what's scaled down

36:02 YaRN is in models that you can download today. Qwen3's model card, from Alibaba's Qwen team, says the model supports thirty-two thousand seven hundred and sixty-eight tokens natively, and that they validated it up to one hundred and thirty-one thousand and seventy-two with YaRN. To turn it on, the card adds rope scaling to the model's config, with the type yarn and a factor of four, from thirty-two thousand seven hundred and sixty-eight positions, which is our setup, with thirty-two times as many tokens. And it notes that the frameworks implement static YaRN, with a factor that doesn't change with the input's length, which can lower the quality on shorter texts. We saw that too: about zero point zero six of loss in the first window, with no training. So the card advises adding it when long contexts are needed, and leaving it out otherwise.

36:50 The mechanism is the one in large open models: rotations stretched past the window trained on, and then, in most cases, some training at the longer length. Qwen3 stretches from thirty-two thousand seven hundred and sixty-eight tokens to one hundred and thirty-one thousand and seventy-two with YaRN, by a factor of four, like ours. DeepSeek-V3's technical report uses YaRN with a factor of forty, to go from four thousand tokens to one hundred and twenty-eight thousand, with the alpha, beta and temperature we used. It trains for a thousand steps at thirty-two thousand tokens, then a thousand at one hundred and twenty-eight thousand, with batches of one thousand nine hundred and twenty sequences in the first stage. Ours goes from one thousand and twenty-four tokens to four thousand and ninety-six, with fifteen minutes on a laptop, about two million tokens. And we measured it on web pages and with the pass key. Their models are also measured on long documents and on harder tests, which ask questions about the whole of a long text.

Wrap-up

37:47 So, today we gave our web model a window four times as long. Past the window it was trained on, the slow pairs of channels turn to angles it didn't see, and its loss climbs. Position interpolation slows all the pairs by four, which fits the angles, and blurs the nearby tokens. NTK-by-parts slows the slow pairs and keeps the fast ones. YaRN adds a temperature, and with no training, it keeps the loss flat to four thousand and ninety-six tokens. It finds the pass key thousands of tokens back, which a window can't. And after fifteen minutes of fine-tuning on long pages, its loss past the window is about three point one one, below the sliding window's three point one four, so it uses the tokens far back on ordinary web pages.

38:30 Let me know in the comments how long a context you'd want from a model like ours, and what you'd have it read. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.