Combine large LLM with small LLM for faster inference

[ggerganov]

So I was thinking about the following idea.
It is probably completely bogus, but I would definitely investigate it when and if I had the time to, so maybe someone else would be interested as well.

---

Large LLM takes a lot of time to perform token inference. Lets say it takes 500ms per token.

A small LLM (or some other approach) can infer a token very fast. Lets say < 5ms.

Lets assume that the small LLM is correct 80-90% of the time.

The idea is the following:

• Before I run the large LLM inference for the next token, I infer it using the small LLM
• I now want to somehow partially evaluate the large LLM (let's say the first 10% of the layers) and get an approximate estimate for the next token
• If this estimate indicates a high probability for that token (i.e. above some threshold) - we stop and directly say that this is the new token. At this point we would have consumed (5ms for the small LLM + ~50ms for the large LLM)
• Otherwise, we proceed to evaluate the rest of the layers of the large LLM

In the described process, I would reach step 4 only for 10-20% of the tokens, but for the rest - I will take the shortcut in step 3.
Hence, I will have an efficient inference with a Large LLM.

Obviously, the biggest question is if step 2 is possible at all.
I suppose the answer is "no", but who knows.

[blackle]

One idea, which would require adding extra model parameters and doing some training:

Because llama (and other transformer language models) are residual networks, i.e. each block adds to the output of the previous instead of replacing it. This can be thought of each block refining the model's current understanding of what token will occur next.

Maybe what might be possible is to add a linear layer after each block, one that produces its own estimate for the token distribution. The last block already has this; it is the output of the network.

I'm thinking if each block can estimate the token probability, maybe you could compute it at each block, and stop early if it's converged. That is, if the subsequent blocks are not appreciably changing the token probabilities we can just stop early.

I did an experiment where I limited llama.cpp to only use the first N blocks of the network. it is surprisingly coherent at 27/32 blocks.

Deepmind has a paper called "pondernet" that has a similar idea, but I believe doing this would require more changes to the model.

[KerfuffleV2]

If this estimate indicates a high probability for that token (i.e. above some threshold) - we stop and directly say that this is the new token. At this point we would have consumed (5ms for the small LLM + ~50ms for the large LLM)

If this approach worked reliably, couldn't you just run only the first few layers from a big model and just skip the rest?

I actually played with this kind of thing a while, randomly disabling chunks of layers, only evaluating every other layer, stuff like that. Unfortunately, it seemed to make a pretty big difference in quality and also doing something like only evaluating the first 10 layers of a 30 layer model produced pretty incoherent output.

It may depend on the model and which layers are skipped. Unscientifically, it seemed like the first layers were the most important.

[Andrey36652]

@ggerganov Relevant paper https://arxiv.org/pdf/2303.08112.pdf

[abetlen]

@ggerganov I think point 2 is possible, it sounds similar to Confident Adaptive Language Modeling (CALM) from Google AI. Instead of a smaller pre-trained LLM they trained an early-exit classifier on intermediate layers which allowed them achieve roughly 50% reduction in inference time.

[FNsi]

I think that might like the full output using huge layers be interrupted after the first some layers,

Like if output 65b is reasonable,

Output65b * layers7b have/ layers s65b have

May not creat a reasonable result, because the 65b output need the rest layers to be fully adjusted.

[PriNova]

I had the following idea:
You could determine in which layer the tensor results are moving away or converging by comparing each layer result from the 7B and e.g. the 65B model with the dot_product_similarity. In this way, one could at least roughly determine which layers of the small model and which layers of the large model are leading to a certain result.

And if I understand correctly, you don't need to load all layers for such experiments. You can use the mapped address of the mmap function to load and unload the layers you want to determine. This would also mean that a large model can be used with a small amount of RAM. Just load the layers in chunks, save the result in scratchpad and using the result as input to the next chunks of layers. Loading and unloading the chunks would of course affect performance.

I'm also not sure if this would work, as my understanding of the ggml library is that the calculations are done lazily instead of eagerly. Is this correct?

[elephantpanda]

Hello I had the same idea: meta-llama/llama#256
I'm using cerebras model.
So I use the 111M model to guess the next 2-3 tokens.
Then I check them in parralell batches with the 2.7B model. However many tokens are checked correct it can add to the list.
This method means the more GPUs you have the faster you can predict. Not limited by the speed of your GPUs.
I have tested it and it works. I think it gives a speed increase but I haven't timed it yet.

There are many ways to modify this algorithm including using the small LMM to predict also the second most likely word and so on.

[etiennebalit]

This idea of using a small model to predict multiple steps ahead and batch the scoring of generated sequence is called Speculative Sampling.

Here's a relevant paper: https://arxiv.org/abs/2302.01318

It's a completely different method than CALM that tries to train a model to output a prediction of the next token at each layer of the LLM and exit early if the confidence is high enough.

Both methods could be implemented separately and might even work together.

I think it's worth it to create another issue for optimization of inference through early exit (CALM-style).

[elephantpanda]

BTW I tried Speculative Sampling, and while the small model guessed correct about 70% of the time so it could jump ahead 2 tokens at a time I think it didn't quite compensate the extra time needed to infer a double batch. But the proof of concept worked and it would definitely give a speed boost of about 60% if you had enough GPUs.

(FUnny I thought I invented something new but that paper came out a month ago!)

[elephantpanda]

OK. I thought of a new idea. Called "chaining". Which is a modification of Speculative Sampling but chaining more models together.

You take models of size 1B, 4B, 16B
Then with the 1B model you generate one token
Next with the 4B model you check the previous token giving you two tokens
Next with the 16B model you check the previous token giving you three tokens.

The idea is that if a model half the size runs twice as quick you can chain together these models because 1+1/2+1/4+1/8+....=2

You could play around with the numbers of models, their sizes, how many steps each one does etc. to get the best speed. Even better you could use a genetic algorithm to get the best set up.

These speculative sampling techniques will work best with low temperatures.

[mikeggh]

I've been considering something similar but with other non LLM networks using the same technique, so thank you for posting. I'm happy to have read the other comments regarding papers which extend the concepts.

"""If this estimate indicates a high probability for that token (i.e. above some threshold) - we stop and directly say that this is the new token. At this point we would have consumed (5ms for the small LLM + ~50ms for the large LLM)"""

In this case would if you wanted to continue, then would you insert that chosen token as an input token to allow the large model to continue with proper state of logits? Would this also require more static topk/rand values, and verification of best values for temp across each model (the small, and large)? I'm guessing the newer API for setting/getting states can be used to rewind the 10% and reinsert the token you've chosen with comparison to the smaller LLM.

[elephantpanda]

IMO The paper https://arxiv.org/abs/2302.01318 overcomplicates things by calculating conditional probabilities etc.
My method is much simpler. You use the small model to generate 2-3 tokens using the greedy algorithm, choosing the most likely token each time.
Then using the large model you large model, to simultaneously test the strings with 0,1,2,3,.. tokens added on and predict the next token.
If the predicted tokens match the small model's greedy algorithm you can use these tokens otherwise you can use the token generated by the first token in the batch.

I can't see any advantage in calculating all these conditional probabilities. Unless someone can prove me wrong?

[aljungberg]

As mentioned by @etiennebalit, this idea is related to speculative sampling (and what pauldog suggested at the end here is exactly that I believe). Speculative sampling was recently landed in HF transformers (they called it assisted generation), so that might be useful as a reference: huggingface/transformers#22211

Their implementation uses a heuristic to dynamically figure out how many tokens to generate at a time.