How to generate the n-grams - which to keep, which to discard?

[bobqianic]

This is actually not my question, llama.cpp wants to implement this but encountered some problems.

ggml-org/llama.cpp#4207

[Viol2000]

Hi, thanks for your interest!

Sorry for not clearly presenting it in our blog and code. We will refactor the code for better readability.

We maintain a lookahead branch (a 2D window) with size (W , (N - 1)). It is in variable past_tokens and initiated here https://github.com/hao-ai-lab/LookaheadDecoding/blob/main/lade/decoding.py#L223C5-L223C17
Currently, we initialize the first level with the tokens randomly selected from the prompt. Initialized by the function set_token. Then we will fill out the #(N - 2) levels in the #(N - 2) warmup steps. Find discussion here: #8 .

An example of the lookahead branch is in the figure below (N = 4 and W = 5)

There are two points:

1. The first level (orange) is one token less.
2. We need #(N - 2) warmup steps to gradually fill the lookahead branch. It is not very important. Currently, we gradually fill the lookahead branch, but you can also initialize the whole lookahead branch with random tokens for simplicity. It will still finally converge and give speedup.

For each step. We will generate one token from each window slot, as yellow ones in the figure below. Then we do two things:

1. collect N-grams as in the figure above (4-grams).
2. Drop the oldest level and use the newly generated tokens to maintain the size of the lookahead branch. We will drop orange tokens and green token with number 1. And we form a new lookahead branch with green tokens with number 2-5, all red tokens and all yellow tokens.

About your question, the n-grams collected is stored in a hash map with its starting token as the hash map key and a set of all following (n-1)-grams as value (variable token_map in our code). For each key, we maintain a set of sizes GUESS_SET_SIZE to avoid a too-large verification branch size. The logit of maintaining its size can be found here: https://github.com/hao-ai-lab/LookaheadDecoding/blob/main/lade/decoding.py#L357C1-L385C1 . Currently, when its size is larger than GUESS_SET_SIZE, we will discard the n-gram that was added first. You can also explore using other discard policies.

Please feel free to ask if you have other questions.

[ggerganov]

For a given starting token and a pool of GUESS_SET_SIZE(n-1)-grams, how do you select the G n-grams to use for verification? AFAICT, GUESS_SET_SIZE can be large - let's say 60. And G is recommended to be equal to W so, let's say 5. How do you choose 5 from the set of 60?

[Viol2000]

@ggerganov Hi, thanks for your interest.

Sorry for my confusing expression. In my implementation, GUESS_SET_SIZE is actually G in the blog. By doing this, I can maintain the pool size starting with a given token when I add new n-grams into the pool. I just dropped the earliest token in that pool when its size was larger than G. Making the GUESS_SET_SIZE and G not the same is also OK; it can offer more selection policies. But for me, I just discard the earliest token in that pool and set GUESS_SET_SIZE = G.

It is OK to set G equal to W as a small number like 5 or 10. Or it can be slightly larger than W. I think It will be good enough to achieve speedups. We will face diminishing margins when you set a large G like 60. It is not recommended on person GPU.

[tdene]

@Viol2000 I remain confused about the outputs of the lookahead branch.

You say

And we form a new lookahead branch with green tokens with number 2-5, all red tokens and all yellow tokens.

What happens to the output tokens that correspond to input tokens of green 1-5 and orange 1-4? If they are not used in the next iteration of the lookahead algorithm, what are they used for? Are they used as speculation drafts alongside the cache or are they just not computed?

[Viol2000]

@tdene Thanks for your interest! You raised some excellent points!

These tokens' outputs are discarded. However, I do not think it is a waste to compute them. Although these tokens' outputs are discarded, these tokens themselves are used to build yellow tokens with stronger local information. For example, yellow 7 takes orange1-4, green 5 and red 6 in attention and give output.

Another way to do this is to put these tokens in kv-cache. So we can remove row orange 1-4 and green 1-5 and save flops. Moreover, it seems llama.cpp's implementation of lookahead decoding uses this method(correct me if I was wrong). However, if you put them in kv-cache, they actually missed the information on newly coming tokens. For example, if you put orange 1 in kv-cache, it should have been built 2 steps before and remained unchanged since then. At that time, blue 0 is not obtained, so the kv-cache embedding of orange 1 have an incorrect token's information rather than blue 0. And I think it will reduce the local information and reduce the acceptance rate of your guess tokens.

[ggerganov]

Another way to do this is to put these tokens in kv-cache. So we can remove row orange 1-4 and green 1-5 and save flops. Moreover, it seems llama.cpp's implementation of lookahead decoding uses this method(correct me if I was wrong).

It is not the case in llama.cpp. In the case from the figure we submit a full batch containing:

• blue token
• 4 orange
• 5 green
• 5 red
• 6 light blue (verification)

None of the lookahead tokens are stored in the KV cache so they will "see" the blue token from the current iteration.

[Viol2000]

Hi @ggerganov, thank you for providing clarification! I've noticed the following graph, and it appears to not include several rows. My initial assumption was that there might be an innovation in kv-caching past tokens. Nevertheless, I appreciate your substantial effort in implementing Lookahead Decoding in llama.cpp.

        // Example for W = 5, N = 4, G = 2:
        // (I = input, L = lookahead, V = verification)
        //
        // Batch:  0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20
        // T:        -2 -2 -2 -2 -1 -1 -1 -1 -1  0  0  0  0  0  0
        // Info:   I  L  L  L  L  L  L  L  L  L  L  L  L  L  L  V  V  V  V  V  V
        // Pos:    0  1  2  3  4  1  2  3  4  5  2  3  4  5  6  1  2  3  1  2  3   (+ n_past)
        // Logits: 1  0  0  0  0  0  0  0  0  0  1  1  1  1  1  1  1  1  1  1  1
        // ---------------------------------------------------------------------
        // Seq:    0
        //         1              1              1
        //         2  2              2              2
        //         3  3  3              3              3
        //         4  4  4  4              4              4
        //         5  5  5  5  5              5              5
        //         6                                            6  6  6
        //         7                                                     7  7  7
        // ---------------------------------------------------------------------
        //                                       |  |  |  |  |  |  |  |  |  |  |
        //                                       V  V  V  V  V  |  |  |  |  |  |
        //                                         j_tokens     |  |  |  |  |  |
        //                                                      V  V  V  V  V  V
        //                                                             id

[ggerganov]

This diagram represents the input batch containing 21 tokens. Later in the code, based on the sequence ids of the tokens in the batch, we construct the attention mask which is 21 x 21.

[tdene]

And we form a new lookahead branch with green tokens with number 2-5, all red tokens and all yellow tokens.

these tokens themselves are used to build yellow tokens with stronger local information. For example, yellow 7 takes orange1-4, green 5 and red 6 in attention and give output.

So the output of the red 5 line in the diagram (which pays attention to blue 0, orange 1, orange 2, orange 3, green 4, red 5), which you name "yellow 6", becomes the new red 5 in the next time-step's calculation?

And the output of the green 3 line in the diagram (which pays attention to blue 0, orange 1, orange 2, green 3), which you name "red 4", becomes the new green 3 in the next time-step's calculation?

And then the output of the orange 2 line in the diagram (which pays attention to blue 0, orange 1, orange 2), which you name "green 3", becomes the next orange 2 in the next time-step's calculation?

I was confused because it sounded, from that first quote, that you use the same "green tokens with number 2-5, all red tokens and all yellow tokens" that are shown in the diagrams as inputs. But "these tokens themselves are used to build yellow tokens with stronger local information", I'm now understanding that you don't use green tokens 2-5 from the diagram, you use the outputs of their lines?

[hashming]

@tdene

looking into the code, I think:

So the output of the red 5 line in the diagram (which pays attention to blue 0, orange 1, orange 2, orange 3, green 4, red 5), which you name "yellow 6", becomes the new red 5 in the next time-step's calculation?

yes

And the output of the green 3 line in the diagram (which pays attention to blue 0, orange 1, orange 2, green 3), which you name "red 4", becomes the new green 3 in the next time-step's calculation?

The red in the diagram is generated by green, not in the current decode iteration (when green and red already exists), but in the previous iteration when red did not exist.
If red does not exist, decode with (input_ids, orange, green) to get red;
if red already exists, decode with (input_ids, orange, green, red) to get yellow, and use green to update orange, use red to update green, and use yellow to update red.
and the output of green and orange in current decode iteration is discarded, as the author previously mentioned.

[learning-chip]

These tokens' outputs are discarded. However, I do not think it is a waste to compute them. Although these tokens' outputs are discarded, these tokens themselves are used to build yellow tokens with stronger local information. For example, yellow 7 takes orange1-4, green 5 and red 6 in attention and give output.

Hi @Viol2000 -- regarding this comment, if the green and orange rows of output are not used, why not just trim them out in the attention mask? This won't have any effect on the output yellow-3,4,5,6,7 tokens, and also won't affect the collected 4-grams.

Are the green and orange outputs somehow used to refine past states (y(t-2) and y(t-3))? If not, then the above mask should totally work fine. The information still flows from input {orange1-4, green 5, red 6} to the output yellow-7, using such remaining mask.

[learning-chip]

why not just trim them out in the attention mask?

I checked the related code:

LookaheadDecoding/lade/models/llama.py

Lines 445 to 451 in b756db3

	ret.out_logits = ret.logits[:,prefill_size - 1,:].to(input_ids.device)
	assert fill_level != -1
	if lguess > 0:
	ret.inp_logits = ret.logits[:,-len(past_tokens[fill_level])-lguess:-lguess,:].to(input_ids.device)
	ret.guess_logits = ret.logits[:,-lguess:,:].to(input_ids.device)
	else:
	ret.inp_logits = ret.logits[:,-len(past_tokens[fill_level]):,:].to(input_ids.device)

LookaheadDecoding/lade/decoding.py

Lines 274 to 275 in b756db3

	outputs = self.jforward_multilevel(
	**model_inputs,

In higher-level call, the full outputs.logits are never directly used. Only its segments out_logits, inp_logits, guess_logits are used.

• out_logits is at position prefill_size - 1 -- is the one "must-be-correct" token?
• guess_logits is at position -lguess: -- is the verification result? (still kept in the trimmed mask)
• inp_logits is at position -len(past_tokens[fill_level])-lguess:-lguess -- correspond to the yellow-3,4,5,6,7 tokens in your figure? (still kept in the trimmed mask)
  (also what does inp stand for?)

[Viol2000]

These tokens' outputs are discarded. However, I do not think it is a waste to compute them. Although these tokens' outputs are discarded, these tokens themselves are used to build yellow tokens with stronger local information. For example, yellow 7 takes orange1-4, green 5 and red 6 in attention and give output.

Hi @Viol2000 -- regarding this comment, if the green and orange rows of output are not used, why not just trim them out in the attention mask? This won't have any effect on the output yellow-3,4,5,6,7 tokens, and also won't affect the collected 4-grams.

Are the green and orange outputs somehow used to refine past states (`y(t-2)` and `y(t-3)`)? If not, then the above mask should totally work fine. The information still flows from input {orange1-4, green 5, red 6} to the output yellow-7, using such remaining mask.

@learning-chip Sorry for the late reply. I did not see your message. Even if the green and orange rows of output are not used, they are useful for generating yellow 3,4,5,6,7 tokens. In the current case, if the current step is t, the red tokens come from the t-1 step, green tokens come from the t-2 steps, and yellow tokens come from the t-3 steps. If you do not input these tokens, you are actually using the kv-cache of green and orange tokens, and the information of current input- deep blue 0- can not be obtained in the green and orange tokens. (they are in kv-cache and static) So they are actually out of date. In short, if you do not input green and orange rows, these tokens can not 'see' deep blue 0, and only red tokens can 'see' deep blue 0. If you input these tokens, in shallow layers, these tokens can obtain the information of deep blue 0 and thus affect the output of red tokens in deep layers.

I have some initial implementations that show that without green and orange rows as inputs, the generated n-grams have a lower quality to be accepted. I believe there are more trade-offs to do -- with lower computational costs and lower acceptance ratios or higher computational costs and higher acceptance ratios.

[Viol2000]

why not just trim them out in the attention mask?

I checked the related code:

LookaheadDecoding/lade/models/llama.py

Lines 445 to 451 in b756db3

	ret.out_logits = ret.logits[:,prefill_size - 1,:].to(input_ids.device)
	assert fill_level != -1
	if lguess > 0:
	ret.inp_logits = ret.logits[:,-len(past_tokens[fill_level])-lguess:-lguess,:].to(input_ids.device)
	ret.guess_logits = ret.logits[:,-lguess:,:].to(input_ids.device)
	else:
	ret.inp_logits = ret.logits[:,-len(past_tokens[fill_level]):,:].to(input_ids.device)

LookaheadDecoding/lade/decoding.py

Lines 274 to 275 in b756db3

	outputs = self.jforward_multilevel(
	**model_inputs,

In higher-level call, the full outputs.logits are never directly used. Only its segments out_logits, inp_logits, guess_logits are used.

• out_logits is at position prefill_size - 1 -- is the one "must-be-correct" token?
• guess_logits is at position -lguess: -- is the verification result? (still kept in the trimmed mask)
• inp_logits is at position -len(past_tokens[fill_level])-lguess:-lguess -- correspond to the yellow-3,4,5,6,7 tokens in your figure? (still kept in the trimmed mask)
  (also what does inp stand for?)

My naming is bad here; I will keep refactoring the code in the following weeks.

[learning-chip]

I have some initial implementations that show that without green and orange rows as inputs, the generated n-grams have a lower quality to be accepted.

Hi @Viol2000 thanks for the explanation. For this part I totally understand. By "trim the mask" I mean shortening the y-axis (output dimension, or q dimension), not shortening the x-axis (input dimension, or kv dimension). The output red tokens can still see the full inputs.