Speculative decoding in autoregressive generative artificial intelligence models
Certain aspects of the present disclosure provide techniques and apparatus for generating a response to a query input in a generative artificial intelligence model. An example method generally includes receiving a plurality of sets of tokens generated based on an input prompt and a first generative artificial intelligence model, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input prompt; selecting, using a second generative artificial intelligence model and recursive adjustment of a target distribution associated with the received plurality of sets of tokens, a set of tokens from the plurality of sets of tokens; and outputting the selected set of tokens as a response to the input prompt.
1. A processing system, comprising:
at least one memory having executable instructions stored thereon; and
one or more processors configured to execute the executable instructions in order to cause the processing system to:
generate, based on an input prompt and a generative artificial intelligence model, a first plurality of sets of tokens, each set of tokens in the first plurality of sets of tokens corresponding to a first portion of a candidate response to the input prompt;
speculatively generate, using the generative artificial intelligence model, a second plurality of sets of tokens, each set of tokens in the second plurality of sets of tokens corresponding to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens;
while speculatively generating the second plurality of sets of tokens, select a set of tokens from the first plurality of sets of tokens; and
output the selected set of tokens from the first plurality of tokens and an associated set of tokens in the second plurality of tokens as a response to the input prompt.
2. The processing system of claim 1 , wherein to select the set of tokens from the first plurality of sets of tokens, the one or more processors are configured to cause the processing system to select a longest sequence of accepted tokens from the first plurality of sets of tokens.
3. The processing system of claim 1 , wherein sets of tokens in the second plurality of sets of tokens include padding accounting for a number of tokens in the selected set of tokens from the first plurality of sets of tokens.
4. The processing system of claim 1 , wherein the first plurality of sets of tokens is represented as a tree data structure and wherein a root node of the tree data structure corresponds to the input prompt.
5. The processing system of claim 4 , wherein a depth of the tree data structure corresponds to a maximum number of tokens generated by a single pass through the generative artificial intelligence model.
6. The processing system of claim 4 , wherein a maximum size of the tree data structure is set based on a computational complexity metric associated with generating a set of tokens by the generative artificial intelligence model.
7. The processing system of claim 4 , wherein to select the set of tokens from the first plurality of sets of tokens, the one or more processors are configured to cause the processing system to:
reject a first token at a first level of the tree data structure representing the first plurality of sets of tokens;
generate an adjusted probability distribution based on the rejection of the first token;
discard or ignoring children tokens of the first token from the tree data structure; and
determine whether to accept or reject a second token at the first level of the tree data structure based on the adjusted probability distribution.
8. The processing system of claim 1 , wherein to select the set of tokens from the first plurality of sets of tokens, the one or more processors are configured to cause the processing system to:
reject each set of tokens in the first plurality of sets of tokens generated by the generative artificial intelligence model; and
sample, using the generative artificial intelligence model, a token based on a target distribution that excludes probabilities associated with each set of tokens in the first plurality of sets of tokens, wherein the selected set of tokens from the first plurality of sets of tokens comprises the sampled token.
9. The processing system of claim 1 , wherein the generative artificial intelligence model comprises a generative artificial intelligence model trained to generate multiple tokens in response to the input prompt based on forecast prompt embeddings.
10. The processing system of claim 1 , wherein the generative artificial intelligence model comprises a model including one or more non-autoregressive layers and one or more autoregressive layers.
11. The processing system of claim 10 , wherein the one or more autoregressive layers comprise one or more layers at a top of a stack of layers representing the generative artificial intelligence model.
12. The processing system of claim 10 , wherein the one or more autoregressive layers comprise one or more layers at a bottom of a stack of layers representing the generative artificial intelligence model.
13. A processor-implemented method, comprising:
generating, based on an input prompt and a generative artificial intelligence model, a first plurality of sets of tokens, each set of tokens in the first plurality of sets of tokens corresponding to a first portion of a candidate response to the input prompt;
speculatively generating, using the generative artificial intelligence model, a second plurality of sets of tokens, each set of tokens in the second plurality of sets of tokens corresponding to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens;
while speculatively generating the second plurality of sets of tokens, selecting a set of tokens from the first plurality of sets of tokens; and
outputting the selected set of tokens from the first plurality of tokens and an associated set of tokens in the second plurality of tokens as a response to the input prompt.
14. The method of claim 13 , wherein selecting the set of tokens from the first plurality of sets of tokens comprises selecting a longest sequence of accepted tokens from the first plurality of sets of tokens.
15. The method of claim 13 , wherein sets of tokens in the second plurality of sets of tokens include padding accounting for a number of tokens in the selected set of tokens from the first plurality of sets of tokens.
16. The method of claim 13 , wherein the first plurality of sets of tokens is represented as a tree data structure and wherein a root node of the tree data structure corresponds to the input prompt.
17. The method of claim 16 wherein a depth of the tree data structure corresponds to a maximum number of tokens generated by a single pass through the generative artificial intelligence model.
18. The method of claim 16 , wherein a maximum size of the tree data structure is set based on a computational complexity metric associated with generating a set of tokens by the generative artificial intelligence model.
19. The method of claim 16 , wherein selecting the set of tokens from the first plurality of sets of tokens comprises:
rejecting a first token at a first level of the tree data structure representing the first plurality of sets of tokens;
generating an adjusted probability distribution based on the rejection of the first token;
discarding or ignoring children tokens of the first token from the tree data structure; and
determining whether to accept or reject a second token at the first level of the tree data structure based on the adjusted probability distribution.
20. The method of claim 13 , wherein selecting the set of tokens from the first plurality of sets of tokens comprises:
rejecting each set of tokens in the first plurality of sets of tokens generated by the generative artificial intelligence model; and
sampling, using the generative artificial intelligence model, a token based on a target distribution that excludes probabilities associated with each set of tokens in the first plurality of sets of tokens, wherein the selected set of tokens from the first plurality of sets of tokens comprises the sampled token.
21. The method of claim 13 , wherein the generative artificial intelligence model comprises a generative artificial intelligence model trained to generate multiple tokens in response to the input prompt based on forecast prompt embeddings.
22. The method of claim 13 , wherein the generative artificial intelligence model comprises a model including one or more non-autoregressive layers and one or more autoregressive layers.
23. The method of claim 22 , wherein the one or more autoregressive layers comprise one or more layers at a top of a stack of layers representing the generative artificial intelligence model.
24. The method of claim 22 , wherein the one or more autoregressive layers comprise one or more layers at a bottom of a stack of layers representing the generative artificial intelligence model.
25. A processing system, comprising:
means for generating, based on an input prompt and a generative artificial intelligence model, a first plurality of sets of tokens, each set of tokens in the first plurality of sets of tokens corresponding to a first portion of a candidate response to the input prompt;
means for speculatively generating, using the generative artificial intelligence model, a second plurality of sets of tokens, each set of tokens in the second plurality of sets of tokens corresponding to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens;
means for selecting, while speculatively generating the second plurality of sets of tokens, a set of tokens from the first plurality of sets of tokens; and
means for outputting the selected set of tokens from the first plurality of tokens and an associated set of tokens in the second plurality of tokens as a response to the input prompt.
26. The processing system of claim 25 , wherein the means for selecting the set of tokens from the first plurality of sets of tokens comprises means for selecting a longest sequence of accepted tokens from the first plurality of sets of tokens.
27. The processing system of claim 25 , wherein sets of tokens in the second plurality of sets of tokens include padding accounting for a number of tokens in the selected set of tokens from the first plurality of sets of tokens.
28. The processing system of claim 25 , wherein the first plurality of sets of tokens is represented as a tree data structure and wherein a root node of the tree data structure corresponds to the input prompt.
29. The processing system of claim 28 wherein a depth of the tree data structure corresponds to a maximum number of tokens generated by a single pass through the generative artificial intelligence model.
30. The processing system of claim 28 , wherein a maximum size of the tree data structure is set based on a computational complexity metric associated with generating a set of tokens by the generative artificial intelligence model.
31. The processing system of claim 28 , wherein the means for selecting the set of tokens from the first plurality of sets of tokens comprises:
means for rejecting a first token at a first level of the tree data structure representing the first plurality of sets of tokens;
means for generating an adjusted probability distribution based on the rejection of the first token;
means for discarding or ignoring children tokens of the first token from the tree data structure; and
means for determining whether to accept or reject a second token at the first level of the tree data structure based on the adjusted probability distribution.
32. The processing system of claim 25 , wherein the means for selecting the set of tokens from the first plurality of sets of tokens comprises:
means for rejecting each set of tokens in the first plurality of sets of tokens generated by the generative artificial intelligence model; and
means for sampling, using the generative artificial intelligence model, a token based on a target distribution that excludes probabilities associated with each set of tokens in the first plurality of sets of tokens, wherein the selected set of tokens from the first plurality of sets of tokens comprises the sampled token.
33. The processing system of claim 25 , wherein the generative artificial intelligence model comprises a generative artificial intelligence model trained to generate multiple tokens in response to the input prompt based on forecast prompt embeddings.
34. The processing system of claim 25 , wherein the generative artificial intelligence model comprises a model including one or more non-autoregressive layers and one or more autoregressive layers.
35. The processing system of claim 34 , wherein the one or more autoregressive layers comprise one or more layers at a top of a stack of layers representing the generative artificial intelligence model.
36. The processing system of claim 34 , wherein the one or more autoregressive layers comprise one or more layers at a bottom of a stack of layers representing the generative artificial intelligence model.
37. A non-transitory computer-readable medium having executable instructions stored thereon which, when executed by one or more processors, perform an operation comprising:
generating, based on an input prompt and a generative artificial intelligence model, a first plurality of sets of tokens, each set of tokens in the first plurality of sets of tokens corresponding to a first portion of a candidate response to the input prompt;
speculatively generating, using the generative artificial intelligence model, a second plurality of sets of tokens, each set of tokens in the second plurality of sets of tokens corresponding to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens;
while speculatively generating the second plurality of sets of tokens, selecting a set of tokens from the first plurality of sets of tokens; and
outputting the selected set of tokens from the first plurality of tokens and an associated set of tokens in the second plurality of tokens as a response to the input prompt.