Better and Faster Large Language Models via Multi-Token Prediction

People waste too much time building out stuff that is really bad in AI right now.

Of course, everything is, but instead of taking on the task of patching that up, the better approach would be to pretend there will be something that is a lot better than GPT-4 in the near future (because there will be) and design a differentiated product under that premise.

I feel your pain and excitement in equal measure.

It can be hard to know where to start with some of these concepts, especially so given that a lot of recent developments (e.g. RAG) are developing so rapidly that there’s unlikely to be a reference book you could refer to anytime soon that would be current.

That said, I do find that documentation is getting better depending on where you look. The documentation for higher level tools like LlamaIndex is a good starting point for understanding the concepts (not so much in terms of explaining the concepts, but showing where they fit into the overall picture, then you can deep-dive elsewhere on the different parts).

YouTube has always been a mixed bag of very little solid information in a sea of non-experts trying to attract clicks for the latest trends, so it’s not a great starting point IMHO.

Most of the time I do not know where the terms fit in the big picture.

Nor do the majority of "AI" experts and consultants that I see on LinkedIn, Twitter or in podcasts.

The S/N ratio is very low in this field. Just pick some documentation from "industry leaders" like Langchain and see that not only is it already and always outdated, it sometimes simply contradicts itself.

In the "blockchain hype" this was similar, so I guess it's a trait of the hype train.

I recommend checking out Machine Learning Q and AI by Sebastian Raschka

Check out Lilian Weng's blog https://lilianweng.github.io/posts/2023-01-27-the-transforme...

Strongly recommend watching Andrej Karpathy’s “Lets build GPT-2” videos on YouTube which dives into an actual PyTorch implementation, then download the code and study it carefully. Then study “Spreadsheets is all you need” to see what the internal data structures look like.

Author here -- that's a very good point and as I understand work in progress in different teams. Training autoencoders for language is actually super easy given the small amount of information contained in text (compared to vision/video), the hard part is making the model focus on the semantic part if all signal we have comes from exact match in token space. Hence Yann LeCun's ideas on joint embedding predictive architectures. Note also that there is always a trade-off between auxiliary tasks giving more signal but shifting the focus. In our case, we noticed degradation if the number of predicted tokens is too high. So latent prediction methods need to sort out what is useful.

Not a stupid question in my opinion.

The problem is that once you have the vectors representing the answers you need something like another model that goes back to a word representation of said answers. Something like a diffusion model but for text. Additionally, the function that this diffusion model will approximate won't be injective, but at best surjective and at worst not even a function (in the mathematical meaning) since many textual represantions are possible given an embedding, and most of those won't be valid (not grammatically valid, no sense sentences, ...).

Finally remember that the embeddings are a "lossy" representation of some datum and so the inverse function will lose a lot of the nuances/context/... .

LLMs avoid the problems above by predicting the next (now next n tokens) in a way that is self consistent with the query and the previous n tokens, so the function they approximate should be mostly surjective.

You still need to convert between words and sentence vectors somehow. You could try using a faster model for that, but I suspect that the output quality will suffer.

My understanding is that tokenization is part of the bottleneck. When you break up a sentence to tokens, each token gets a vector representation. The dictionary of all tokens would be infinite if it was at the sentence level

the problem is then the total number of computation drops dramatically therefore leads to much less "thinking" power. i think the idea originated from an understanding that when we write/speak, we have an overall idea. my current hypothesis is it's probably an illusion.

you may want to search for "filler" papers to read.

Also a noob here, if we encoded, trained on and synthesized sentence vectors, wouldn’t that move the AIs ability to create novel things up from sentences to words?

I.e. we currently operate on words (roughly) so the AI can only use words it knows but can synthesize unique sentences from words. If the AI operates on sentences, wouldn’t it only be able to regurgitate sentences it has seen before? So it could synthesize novel paragraphs, but not sentences?

I’m not convinced that sentences are a useful abstraction for AI (in English, anyways). They’re barely useful to humans. Check out your average chat conversation, email, YouTube comment, etc. There’s a very good chance the sentences aren’t actually sentences, or that they haven’t even bothered to use punctuation.

I just don’t think sentences map to a semantic device. A sentence could be two words or half an English paper depending on the writer. It could traverse a half dozen ideas or a single one. Where a sentence ends generally is more about the writer than the semantics.

That's hierarchical prediction. On one level you predict the style of the paragraph, on another the tone and form of the sentence, and on the third the next word. However, this form of prediction is quite difficult since predictions from the layers affect each other.

1. Bananas are yellow due to a biochemical process known as carotenogenesis, which involves the synthesis and accumulation of carotenoid pigments.

2. Bananas are yellow due to a specific carotenoid called beta-cryptoxanthin, which gives the fruit its characteristic yellow hue.

3. Bananas are yellow due to a gradual increase in the concentration of carotenoid pigments as the fruit ripens and chlorophyll levels decrease.

4. Bananas are yellow due to a series of enzymatic reactions that convert starch into sugars and break down the green chloroplasts, revealing the underlying yellow carotenoids.

5. Bananas are yellow due to a change in the pH levels within the fruit cells during ripening, which triggers the production of yellow carotenoid pigments.

6. Bananas are yellow due to a genetic trait inherited from their wild ancestors, which enabled the development of carotenoid pigments as a way to attract seed dispersers.

7. Bananas are yellow due to a complex interplay between various plant hormones, such as ethylene and abscisic acid, which regulate the ripening process and pigment formation.

8. Bananas are yellow due to a metabolic shift from chlorophyll synthesis to carotenoid synthesis as the fruit reaches maturity.

9. Bananas are yellow due to a natural defense mechanism that involves the production of carotenoid pigments, which protect the fruit from oxidative stress during ripening.

10. Bananas are yellow due to a evolutionary adaptation that helps the fruit stand out against the green foliage, making it more visible to potential seed dispersers.

The output of an LLM is usually randomly sampled from the top few highest probability next token/word predictions, but the model itself has no idea which word the sampler will pick. It presumably has some conceptual plan of what could follow "a", or any of it's other suggestions, but any such plan (high level prediction) is then rethought from scratch once "a" is generated.

The model not only can, but has to, change it's mind after each word generated, so this "planning ahead" is very ephemeral - more like a freestyle rapper making it up on the fly than someone thinking deeply about how best to reply and how to express it.

To some degree, attention is already a mechanism to make computations from previous tokens useful later. (You can think of the KV cache as a representation of the text so far and all the models thoughts on it.) And since language models are trained on sequences end-to-end, I think this is likely to happen. Multi-token prediction encourages this behavior explicitly but only for the small n token window you define.

That said, there are many works attempting to increase the compute utilization of transformer language models (early exit, mixture of depths) and novel architectures (SSMs etc.).

I always presumed that "a pigment" is one token.

I'm a total amateur in this field though.

I wonder if this means that these types of models would perform better (or worse?) for languages that do not have this sort of forward-looking grammar.

The alternative theory is that any word starting with a vowel sound is exceedingly uncommon after ’a’ in its training set, so it doesn’t need to plan ahead, just predict the distribution of the most likely next words and choose.

Which is my understanding of how they work and the dynamic at play.

Maybe look at it another way: ask GPT to complete the following

  Bananas are yellow due to a

  Bananas are yellow due to an

  Bananas are yellow due to a pigment called bromelain.

  Bananas are yellow due to an organic compound called bromelain, which is a yellow pigment.

  The token following "due to" is "a" with 55% probability, "an" with 45% probability.

And if you didn't actually include any facts about bromelain in the pretraining data, LLMs absolutely could autocomplete this with something about "an optical illusion." GPT-3 made factual mistakes like that pretty routinely, but I recall it figured out the grammatical rules of "a" and "an."

I don't think the concept actually needs to be pre-activated as you said, though I agree with faabian that this "preactivation" probably does happen in some implicit/emergent sense.

This post is interesting: https://clementneo.com/posts/2023/02/11/we-found-an-neuron

The output of most LLMs is stochastic. The core LLM is given token, and outputs a set of ranked tokens, with a “confidence”, to go next. Then there’s normally a filtering and search stage, where those ranked token are either feed back into the LLM to get more ranked tokens and used to for a short probability tree. I.e. if we pick the top N-ranked tokens and put them back in, each of those tokens results in a new set of N-ranked tokens.

By looking at that tree some basic filtering is done. Such as picking the branch that has the highest summed confidence, or the branch that has the fewest repeated tokens, or the fewest tokens that match with input tokens, or more often some combination of the above, plus a random choice weighed by summed confidences.

That how you can give a LLM with complete fixed weights, which is all LLM, the same input multiple times, but get different outputs.

So to answer your specific question, it can “change its mind”. Every token produced creates a new opportunity for the stochastic output filters to pick a new path through all the possible outputs.

I think you're not looking at this from the right perspective. A LLM is designed to sample text, that follows the training distribution. It is not designed to tell you the "most likely" text that follows, and we don't actually want that. This would mean you have no diversity in your outputs.

In your example, sampling a 0 in 40% of cases and a 1 in 60% of cases does make sense for chat applications.

For applications where we do care about the most likely sentence (e.g. question answering), then beam search helps, as others have mentioned.

Another thing to consider is that the model can "look ahead" and precompute what the future tokens might be. And it can then use this to predict the current token. In fact, some work have been investigating this, such as [1].

And a final note, predicting one token at a time is what we are doing as humans when we speak, so clearly it is not a wrong approach. We are doing this "look ahead" in our mind before speaking.

[1] https://arxiv.org/abs/2404.00859

You're mixing training loss (cross-entropy/surprisal of next token) and post-training prediction decoding (done e.g. with beam search)

It's called the Markov assumption. It was basically the single most important piece of mathematics in the field for decades. It allowed us to solve otherwise intractable problems given the limited compute budgets of the time.

This is how they work and it's a real problem when doing prediction with low temperatures.

IIRC you see weird patterns in LLM outputs since "an" is often less likely than "a" so you end up with fewer nouns beginning with vowels than you would expect.

https://arxiv.org/pdf/2404.19737

Language models factor the joint probability p(y, x) as p(y, x) = p(y|x) p(x) which is exact. I.e. if you train a language model on your distribution and sample with temperature 1, you will get the exact same distribution out. If you sample at lower temperature or even greedily, evidently, you will get other distributions.

What you’ve described is basically the problem with greedy sampling in the decoder. Many other local optimization sampling strategies exist (e.g., beam search) and there’s been a lot of work on more global sampling (e.g., speculative decoding).

This is a fascinating point.

If I'm reading you right, you're saying that a simple way to do this would be to calculate logits for not just the next token, but also n+1 -- all at the same time. If one of the n+1 logits is chosen, then do an infill on the skipped token for the next step, then resume.

This could get us around the example that you gave for only a linear increase in the vocabulary size -- so looking an extra token ahead only increases vocab size by a factor of 2, and looking at a third token is a total factor of 3.

This seems really promising!

0: p=0.40 1: p=0.60 which suggests that 1 is the next bit and leads to a suboptimal starting point for predicting the bit after that. The error is even more prominent with longer sequences as the joint probability distribution becomes more unfactorizable into marginal distributions (as I would expect any minimal algorithmic description of real-world data to be).

Can someone explain this part a bit more? I'm not seeing the issue. From what I see, if the first token (t1) output is a zero, then the next token (t2) would have probabilities 0:p=.90 and 1:p=.10. (And t2 0/1:p= .50/.50 if t1=1)

Mathematically, those line up with the initial distribution, so what's the concern? That's how conditional probability works.

So it will not get worse in performance but only faster

A bit confused by this statement. Speculative decoding does not decrease the performance of the model in terms of "accuracy" or "quality" of output. Mathematically, the altered distribution being sampled from is identical to the original distribution if you had just used regular autoregressive decoding. The only reason you get variability between autoregressive vs speculative is simply due to randomness.

Unless you meant performance as in "speed", in which case it's possible that speculative decoding could degrade speed (but on most inputs, and with a good selection of the draft model, this shouldn't be the case).

The problem with speculative decoding is that there are hardly any models that support it and adding support takes extra GPU time. If speculative decoding also improves planning performance, then it will be more readily adopted.

They just throw the predictions for +1 and +2 away, and only generate them for more efficient training.

The abstract doesn't make that clear, but from the description of figure 1: "During inference, we employ only the next-token output head. Optionally, the other three heads may be used to speed-up inference time"

Maybe you can use all three heads if you take the top prediction from all of them, but that prevents you from doing any of the common sampling strategies. I'm not sure how many people actually run an LLM with temperature 0 outside of benchmarks, unless they do something even better than applying a temperature

The n+1-th token is discarded if it is unlikely given the n-th token.

Thanks for the feedback! Let me try to state it better:

In the end, we only use the next-token head for generating. So which parts of the 2-token target H(X) + H(Y) are "auxiliary" in the sense that they help learning and which are "wasted"? H(X | Y) and I(X; Y) are useful for next-token generation while, by definition, H(Y | X) is the information quantity not related to the next token X. So we could say: "multi-token prediction trades the useful information I(X; Y) from H(Y) for the wasted computations on H(Y | X)". However, note that H(Y | X) is a next-token entropy for predicting Y from the prefix (C, X). If the attention mechanism allows to transfer computations already made for predicting Y|X to the next step, these computations may actually not have been wasted -- it was just pre-computations.

Well, it feels like architecturally there's a lot of scope to do better for coding tasks specifically. Like, if you had FAIR level resources and wanted to train a really great Java coding model for example it would make sense to train the model to predict ASTs rather than tokens. You'd still need some kind of joint normal LLM for predicting comments, identifier names and so on, but you wouldn't model the program itself as a stream of tokens. Instead it would predict things like "add an if block", "add a method call block with 4 parameters" and so on.

You could also train the model to expect certain context window positions to be reserved for things like "type members at the current cursor" and then integrate the inferencing loop with IDE/LSP-style static analysis. This would allow the model to see more information than is actually contained in the text.

I think the reason we're not seeing models like this right now is the cost of doing such research combined with the fact that AI people are all Python-heads, and Python doesn't benefit from much IDEs.

If the LLM had an affordable model it would always generate enough tokens for the task at hand. The fact that this particular method would require more tokens would be irrelevant. If you don't have an affordable model, then you would always be at the mercy of the LLM being biased towards answering with an estimate instead of the actual answer.

Also, most speculative decoding strategies produce identical output compared to running the model sequentially. If the prediction is wrong, the token gets discarded and the speedup is lost.

Might be good to have some flexibility in where those particular tokens are placed, but yeah -- I could see value in creating a "pool" of tokens that should be used at some point in the future in the answer.