search

hao-ai-lab / LookaheadDecoding

[ICML 2024] Break the Sequential Dependency of LLM Inference Using Lookahead Decoding
https://arxiv.org/abs/2402.02057
Apache License 2.0
1.11k stars 66 forks source link

Question on Initial guess tokens #8

Closed DaehanKim closed 10 months ago

DaehanKim commented 10 months ago

Hi! Thank you for sharing a good initial implementation of lookahead decoding!

I'm just curious on how to set initial W tokens when starting lookahead decoding. (I tried to find a code related to it but couldn't.)

It seems like this initial setup can impact decoding speed somehow since if those words are chosen randomly, no real match might be retrieved from the n-gram pools.

Viol2000 commented 10 months ago

Hello, the W tokens initialization is here: https://github.com/hao-ai-lab/LookaheadDecoding/blob/main/lade/decoding.py#L223

Actually, our method includes a multi-level lookahead window, and we only set one level at the beginning. It will run several warmup steps to fill the whole 2-D window.

The set_token function defines the token initialization method. You can change it by changing the code here: https://github.com/hao-ai-lab/LookaheadDecoding/blob/main/lade/decoding.py#L221C5-L221C26

Currently, we are copying the init tokens from the prompt. Actually, althrough tokens are randomly selected from the vocabulary, lookahead decoding can still converge and achieve speedup, thanks to the magic Jacobi iteration.