Speculative decoding in autoregressive generative artificial intelligence models

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;

generate, using the generative artificial intelligence model, a second plurality of sets of tokens, based on the first plurality of sets of tokens;

while generating the second plurality of sets of tokens, select, by the generative artificial intelligence model, a set of tokens from the first plurality of sets of tokens; and

output, by the generative artificial intelligence model, the selected set of tokens from the first plurality of tokens and an associated set of tokens from the second plurality of sets 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 mobile device comprising the processing system of claim 1 .

14 . A processor-implemented method, comprising:

generating, based on an input prompt and a generative artificial intelligence model, a first plurality of sets of tokens;

generating, using the generative artificial intelligence model, a second plurality of sets of tokens;

while generating the second plurality of sets of tokens, selecting, by the generative artificial intelligence model, a set of tokens from the first plurality of sets of tokens; and

outputting, by the generative artificial intelligence model, the selected set of tokens from the first plurality of tokens and an associated set of tokens from the second plurality of sets of tokens as a response to the input prompt.

15 . The method of claim 14 , 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.

16 . The method of claim 14 , 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.

17 . The method of claim 14 , 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.

18 . The method of claim 17 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.

19 . The method of claim 17 , 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.

20 . The method of claim 17 , 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.

21 . The method of claim 14 , 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.

22 . The method of claim 14 , 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.

23 . The method of claim 14 , wherein the generative artificial intelligence model comprises a model including one or more non-autoregressive layers and one or more autoregressive layers.

24 . The method of claim 23 , 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.

25 . The method of claim 23 , 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.

26 . 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;

means for generating, using the generative artificial intelligence model, a second plurality of sets of tokens;

means for selecting by the generative artificial intelligence model, while generating the second plurality of sets of tokens, a set of tokens from the first plurality of sets of tokens; and

means for outputting, by the generative artificial intelligence model, the selected set of tokens from the first plurality of sets of tokens and an associated set of tokens from the second plurality of tokens as a response to the input prompt.

27 . 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;

generating, using the generative artificial intelligence model, a second plurality of sets of tokens;

while generating the second plurality of sets of tokens, selecting, by the generative artificial intelligence model, a set of tokens from the first plurality of sets of tokens; and

outputting, by the generative artificial intelligence model, the selected set of tokens from the first plurality of tokens and an associated set of tokens from the second plurality of tokens as a response to the input prompt.