Better and Faster Large Language Models via Multi-Token Prediction
arxiv.org
arxiv.org
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 mean yes, this is what a rapidly expanding field looks like that's probing the boundaries of its problem space. Kind of like following physics in the early to mid 1900s. Different classes of problems have barely been tested against each other, much less fully explored themselves.
In some ways it reminds me of the earlier days of the internet when progress was still very rapid.
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.
As an outsider but avid reader of this stuff linked from HN, I would recommend the channel 3blue1brown. He's got several NN and AI related videos, and the couple I've seen were pretty good.
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 understand the prompts as a service is short term… but what is a long term product you see?
Why would you need an extra layer here?
This could mean: Instead of diving into langchain and trying to program your way out of a bad model, or trying to do weird prompts, just write a super clear set of instructions and wait for a model that is capable of understanding clear instructions, because that is an obvious goal of everyone working on models right now and they are going to solve this better than your custom workaround can.
This is not a rigid rule, just a matter of proportions. For example, you should probably be willing to try a few weird intermediary prompt hacks, if you want to get going with AI dev right now. But if most of what most people do will probably be solved by a somewhat better model, that's probably a cause for pause.
One wonders, if that's the case, how quickly an AI might improve if it has something close to Google's search site throughput. I mean fielding several billion queries a day, for a year — that would be some pretty stellar training right there I would think.
Yes, you can. Some of the big providers are fairly clear on where in their products this happens, and all offer a way out (mostly when paying for api access)
> One wonders, if that's the case, how quickly an AI might improve if it has something close to Google's search site throughput
Indeed. Another possibility is that user input will turn out to be increasingly less important for upcoming state of the art models.
They do take your feedback and presumably do something with it. Your actual queries are only indirectly useful since they might have private info in them.
Due to the multi-task training, it will however also get better. (This idea is already quite old, to predict multiple targets into the future as an auxiliary loss.)
Nice work.
The nice thing here is that you actually don't need another smaller model but the model itself already predicts the next N subwords.
Or maybe you mean it's not implemented in some of the common software? I'm not sure about that, but I thought it's a quite popular feature now.
https://github.com/ggerganov/llama.cpp/tree/master/examples/...
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).
If they don’t, I’m amazed they work as well as they do. Consider 2-bit sequence prediction with the following possible outcomes and associated probabilities:
00: p=0.36
01: p=0.04
10: p=0.30
11: p=0.30
0: p=0.40
1: p=0.60
Edit: now that I think about this a bit more, a cool research project that would be really simple to carry out might be to modify the cross-entropy loss function to consider only the nth future token in the text training data, and then plot LLM performance vs n, assuming that for all current LLM models we just have n=1.
My hypothesis is that you can mostly bypass all of the resource blow-up involved in predicting the joint probability distribution over the next 1 through n tokens (which scales as x^n) by just predicting the nth token directly, since doing so would implicitly require a better data model (at least for human-generated text; this wouldn’t be the case for all types of data).
I thought post-training prediction still only directly predicts the next token and beam search is sort of a meta-model applied over that (i.e., it is a model on top of the output of the model that performs next-token prediction—beam search considers at each iteration a subset of the current next-token predictions ranked by their probability to use as multiple starting points for predicting the next token, while keeping track of the joint probabilities to prune the set of candidate sequences at each step).
Seems like beam search would fail drastically in cases where the true (unknown) probability distribution over all sequences of tokens of length n has very low conditional probabilities for the first few tokens, each given the computed joint probability of the prior predicted tokens. That is, the true values of p(t2|t1), p(t3|t2,t1), p(t4|t3,t2,t1), ... as derived from the unknown p(t1,t2,...,tn) are very small, but very high when computed via a next-token prediction model.
I’m suggesting to modify both. Use cross-entropy of the nth token for training loss. Use cross-entropy of nth token for post-training prediction and then work backward from there to the beginning of your sequence prediction.
If you want to be better you need to switch to DDPMs for example (e.g. an encoder-only transformer to predict diffusion transition probabilities in parallel, then apply steps of denoising).
The problem is just that these don't work so well from auto regressive decoder transformers, and encoder-decoder architectures like e.g. Google's T5 have fallen out of favor since about LLAMA dropped.
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.
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!
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.
I wouldn't be surprised if we could predict token groups. When speaking off the cuff, people often rely on well-worn phrases and cliches.
Also, this is maybe a semantic point, but, I am not predicting any words I speak. Not in a statistical sense. I have intent behind my words, which means I have an abstraction of meaning that I want to convey and I assemble the correct words to do that. no part of that is "predictive"
Queues of words give cues to help memorize.
It is exactly designed to do that. A temperature of 0 this is what you are approximating. The crucial point though is that it is the most likely next word given the proceeding multi-token context, not just the previous token.
No, we specifically do want "most likely" to follow; the goal is to approximate Solomonoff induction as well as possible. See this recent paper by Hutter's team: https://arxiv.org/pdf/2401.14953
Quote from the paper:
"LLMs pretrained on long-range coherent documents can learn new tasks from a few examples by inferring a shared latent concept. They can do so because in-context learning does implicit Bayesian inference (in line with our CTW experiments) and builds world representations and algorithms (necessary to perform SI [Solomonoff Induction]). In fact, one could argue that the impressive in-context generalization capabilities of LLMs is a sign of a rough approximation of Solomonoff induction."
> In your example, sampling a 0 in 40% of cases and a 1 in 60% of cases does[n't] make sense for chat applications.
I didn't say anything about sampling. A sequence prediction model represents a mapping between an input sequence and a probability distribution over all possible output sequences up to a certain length.
My example uses a binary alphabet, but LLMs use an alphabet of tokens. Any chat application that expresses its output as a string of concatenated symbols from a given alphabet has a probability distribution defined over all possible output sequences. I'm simply comparing the fundamental limitations of any approach to inference that restricts its outcome space to sequences consisting of one symbol (and then layers on a meta-model to generate longer sequences by repeatedly calling the core inference capability) vs an approach that performs inference over an outcome space consisting of sequences longer than one symbol.
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.
It would be sort of like trying to make local optimization techniques less local by running the process multiple times from different starting points and choosing the minimal basin the process ended up in from out of the different runs. Quite a bit better, but in many cases not even close to the global optimum*.
For a slightly more apt analogy, it would be like:
1) Choose multiple starting points
2) For each starting point, perform local optimization until you get stuck in a basin (local minimum)
3) Keep the n overall lowest points, create m different perturbations for each point that pushes the point out of its basin, and go to 1), using these m*n points as your new set of starting points for the next round.
Note this process is totally agnostic to whatever the local optimization algorithm is. That’s why I called the beam search part of an LLM’s post-training prediction a “meta-model”, because it doesn’t matter if the core inference is performed by a transformer architecture or something entirely different.
*I say “in many cases”. But I am extremely curious for this particular case of inferring sequences of human-generated text the degree to which we fail to capture the true joint probability distribution via single-token prediction + beam search.
It’s quite possible we are very far off—perhaps this is the missing “ingredient” for generalized reasoning. On the other hand, as I said in my original post, I never would have guessed current SOTA LLMs only use single-token prediction, and I’m astonished they work as well as they do based on that, so maybe we’re not actually too far off. Without further research, it’s just speculation either way though.
*This might be just a detail/semantics, but for example, for a 3 word context, your starting points may look like "[some token][empty][empty]". Here your procedure simply reduces to a single run of beam search, not multiple, since beam search optimises locally for every turn, producing n different "perturbations" every turn. Let me know if I misunderstood you on this.
But inference combinatorial optimisation as a continous surrogate(which sounds like what you are conveying, and will naturally result in starting points as sentences and not just incomplete words/sequences) is something I never considered. There must be some literature around on this....lets see.
Lets say you ask an LLM
What makes bananas yellow?
Bananas are yellow due to a pigment called bromelain.
Would it be possible to capitalize on the work that has already been done when the LLM outputs "a"? Could the state of the neural net be somehow preserved for the next answer?
Which is my understanding of how they work and the dynamic at play.
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 think I’ve read something about this but I wonder if you could abstract attention to sentence/page levels and then only recalculate the parts that are relevant.
I.E. the KV cache is 'just' a time saving measure because an LLM goes back and calculates those values anyway. (Which is why per-token compute increases exponentially otherwise)
You're not wrong that you could make an LLM more stateful. There are plenty of ideas for that but it would
a) be far more compute intensive to train and run (especially train)
B)be susceptible to all of the issues that RNNs have.
C) most importantly, it would almost certainly just converge at scale with transformers. Labs run small scale, internal tests of architectures all the time and most of them basically come to this conclusion and abandon it
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'm a total amateur in this field though.
If you want to just see the tokens for yourself, though, just enter some text here:
my gut said "pig" and "ment" on this one, which happens to be right, but my gut would also say "para" and "graph" but no, "paragraph" is a single token which falls way outside the "normal" length I see of 3-4 characters
In either case, I do consistently see spaces between characters included as part of the token following the space.
" paragraph" (10 characters) is the longest token I've seen- and now I wonder what the longest token is
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.
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.
Did an LLM really make this basic grammatical error?
That's what makes it surprising.
Models do struggle with these tests, for sure, but from an analytical standpoint, a "next token predictor" should not be able to ever correctly land on the right token 100 tokens in the future.
Edit: Thinking about it, I suppose it is possible that the model can encode a "destination" in the first token. Like a pool shot that is artfully bounced off many bumpers to hit a ball, perhaps the LLM can encode a "path" to a destination token in the first token generated. Which might be even crazier as it suggests that the model is playing a meta-game with being able to precisely manipulate the individual layers of output, even though those layers are disparate from token to token.
One way to think of what all the intermediate layers are doing is to consider them as levels of a linguistic parse tree with the leaves (words) at the bottom and trunk ("sentence") at the top, except in the transformer the evolving embeddings at each level contain semantic as well as syntactic information. This largely hierarchical view of language was the motivation for the transformer design.
It seems we should really think of each layer of the transformer as an independent predictor, with increasingly abstract and more semantically complete information available as we ascend the transformer layers towards the output. Predict next token is only what the transformer is being trained to do at the output layer. At the inner layers (i.e. the bulk of what the transformer is doing), it will be predicting at these higher levels of representation held at those layers.
These models do struggle (although getting better) at ending on a given word rather than starting on it, and understandably so since random sampling and continual resetting after each output token (= new input sequence) means that planning ahead at level of word specificity is simply not an option. They have to continually adapt to next sampled token, and take it from there.
I'm guessing that ability to end on a chosen word is due to continued salience of that word during generation, prediction of sentence fragments using/ending with that word, and opportunistic stopping when it has been emitted and the sentence is complete. Kind of the same way you might do it yourself if you just started talking immediately without planning, while trying to end on a given word.
That is to break out the idea that characters are formed into words, and words into sentences, and a sentence is a sequence of "concepts" for the lack of a better description.
So have one NN which takes a sequence of tokens and predicts an moderately-dimensional "word vector", which is fed into another which predicts a high-dimensional "concept vector".
Then the "thinking layer" would map a sequence of "concept vectors" to "concept vectors", and then you'd have some layers which does the reverse of the input layers to output tokens which can be printed.
Thought being that by splitting it up like this you could swap out the decode and encode layers independently to translate, for example, and so on.
Just a shower thought.
If you wanted to create an embedding algorithm for phrases, you could and you could throw a transformer at it.
I don’t know how you get the output of higher levels to diffuse to phrase and word levels.
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.
Could it be just a smaller llm that takes as input both the semantic vector and the prompt, and is trained to predict the output tokens based on those? A model with high linguistic abilities and very little reasoning skills.
Sorry, that's where the limit of my knowledge is. I work on ML stuff, but mostly on "traditional" deep learning and so I am not up to speed with the genAI field (also, the sheer amount of papers coming out makes it basically impossible stay up to date of you're not in the field).
This actually sounds amazingly useful.
People are working on it though: https://arxiv.org/pdf/2305.09515
This isn't obviously the case, compare this "intelligent designer" view with evolution: there was no prior plan for rabbits. it's sufficient to create the appearance of design that sequential steps are simply probabilistically modulated by prior ones.
Consider a continuation of "the cat..." merely a distribution over all possible words suffices to create the illusion of a plan, suppose: "the cat sat..." then, "on.., the..." etc. follow from the training data.
I think there's a strong argument against trying to model entire sentences exactly because the system isn't modelling semantics: one should expect accuracy to drop off a cliff if there is no actual plan. ie., predicting "sat on the mat" from "cat" shouldnt be a valid prediction, because of the infinite number of possible continuations that as a whole is terrible (eg., what about "chased the mouse" etc.). The space of all possible sentences to continue from "the cat" is infinite, which much of that space actually useful; whereas the number of words is very small, very fininte, and many of them not useful.
The only reason that "the cat sat..", "the cat sat on..." is reasonable is because each sequential word can be modulated by the prompt to seem as if planned.
If the KQV doesn't encode information about likely future token sequences then a transformer empirically couldn't outperform Markov text generators.
Though, more simply, you can just take any LLM and rephrase it as a markov model. All algorithms which model conditional probability are equivalent; you can even unpack a NN as a kNN model or a decision tree.
They all model 'planning' in the same way: P(C|A, B) is a 'plan' for C following A, B. There is no model of P("A B C" | "A B"). Literally, at inference time, no computation whatsoever is performed to anticipate any future prediction -- this follows both trivially form the mathematical formalism (which no one seems to want to understand); or you can also see this empirically: inference time is constant regardless of prompt/continuation.
The reason 'the cat sat...' is completed by 'on the mat' is that it's maximal that P(on|the cat sat...), P(the|the cat sat on...), P(mat|the cat sat on the...)
Why its maximal is not in the model at all, nor in the data. It's in the data generating process, ie., us. It is we who arranged text by these frequencies and we did so because the phrase is a popular one for academic demonstrations (and so on).
As ever, people attribute "to the data" or worse, "to the LLM" no properties it has.. rather it replays the data to us and we suppose the LLM must have the property that generates this data originally. Nope.
Why did the tape recorder say, "the cat sat on the mat"? What, on the tape or in the recorder made "mat" the right word? Surely, the tape must have planned the word...
>It replays the data to us and we suppose the LLM must have the property that generates this data originally.
So to clarify, what you're saying is that under the hood, an LLM is essentially just performing a search for similar strings in its training data and regurgitating the most commonly found one?
Because that is demonstrably not what's happening. If this were 2019 and we were talking about GPT-2 it would be more understandable but SoTA LLMs can in-context learn and translate entire languages which aren't in their dataset.
Also RE inference time, when you give transformers more compute for an individual token, they perform better https://openreview.net/forum?id=ph04CRkPdC
You can make things up as you go, and the constraints emerge from the flow.
you may want to search for "filler" papers to read.
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.
H(X) + H(Y) = H(X | Y) + 2I(X ; Y) + H(Y | X)
By discarding H(Y | X) - which appears again when predicting at the following position - we observe that 2-token prediction increases the importance of I(X ; Y) by a factor of 2.
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.
Maybe this sort of multi more fiction takes their view into 1.1 dimensions? In any gas, there us s real argument for expanding that window, somehow, into two or more dimensions.
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.
Another possibility would be some sort of fixed knowledge base, which could be program language documentation or "common sense" like CYC wants to provide.
There’s multihead attention and multiple output heads as a concept in the paper.
Multihead attention is about focusing on different areas of the input in transformer architectures, and the biological analogy here is head as a central processing unit.
An output head refers to the final layer of a neural network, of which you could have more than one producing different outputs based on the same previous layers. This is also a loose biological analogy, but instead of head as cpu, think more along the lines of head being on one end of the body.
In neither case is there any analogy to a tape head that reads data.
- restating what it thinks is being asked of it
- expressing a high level strategy over what sort of information it might need in order to answer that question
- stating the information it knows
- describing how that information might inform its initial reasoning
etc...
I'd be concerned that going about this by having the model predict the next multiple tokens at any given time would essentially have the opposite effect.
Chain of thought prompting appears to indicate that a model is "smarter" when it has n + m tokens than when it just has n tokens as input. As such, getting the next 5 tokens for a given n might net worse results than getting the next 1 token at n, then the next 1 token at n + 1, and so on.
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.
If token/word +1 and +2 are predicted independently then surely often it won’t ?
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