Speculative decoding in autoregressive generative artificial intelligence models

Abstract

Some aspects of this disclosure provide techniques and apparatus for generating responses to query inputs in a generative artificial intelligence model. An exemplary method generally includes receiving an input prompt and a plurality of sets of tokens generated based on a first generative artificial intelligence model, where each set of tokens in the plurality of sets of tokens corresponds to a candidate response to the input prompt; selecting one set of tokens from the plurality of sets of tokens using a second generative artificial intelligence model and recursive adjustment of a target distribution associated with the received plurality of sets of tokens; and outputting the selected set of tokens as a response to the input prompt.

Claims

1. A processing system, At least one memory location where executable instructions are stored, One or more processors, wherein the processing system includes: Based on an input prompt and a generating artificial intelligence model, the system generates a first plurality of sets of tokens, wherein each set of tokens in the first plurality of sets of tokens corresponds to a first portion of a candidate response to the input prompt. Using the generative artificial intelligence model, a second set of tokens is speculatively generated, wherein each set of tokens in the second set of tokens corresponds to a second portion of the candidate response to the input prompt based on the first set of tokens. While speculatively generating a second set of the aforementioned tokens, one set of tokens is selected from the first set of the aforementioned tokens. In response to the input prompt, the system outputs a selected set of tokens from the first set of tokens and an associated set of tokens within the second set of tokens. One or more processors configured to execute the aforementioned executable instructions, A processing system equipped with the following features.

2. The processing system according to claim 1, wherein, in order to select a set of tokens from a first set of multiple sets of tokens, one or more processors are configured to cause the processing system to select the longest sequence of accepted tokens from the first set of multiple sets of tokens.

3. The processing system according to claim 1, wherein a set of tokens among a second plurality of sets of tokens includes padding that takes into account the number of tokens in a selected set of tokens from the first plurality of sets of tokens.

4. The processing system according to claim 1, wherein a first plurality of sets of the tokens are represented as a tree data structure, and the root node of the tree data structure corresponds to the input prompt.

5. The processing system according to claim 4, wherein the depth of the tree data structure corresponds to the maximum number of tokens generated by a single pass through the generative artificial intelligence model.

6. The processing system according to claim 4, wherein the maximum size of the tree data structure is set based on a computational complexity metric associated with generating one set of tokens by the generative artificial intelligence model.

7. To select a set of tokens from a first set of multiple sets of tokens, one or more processors provide the processing system with The first token is rejected at the first level of the tree data structure representing the first set of tokens, Based on the rejection of the first token, a modified probability distribution is generated. From the tree data structure, discard or ignore the child tokens of the first token. Based on the adjusted probability distribution, the first level of the tree data structure is used to determine whether to accept or reject the second token. The processing system according to claim 4, configured as described above.

8. To select a set of tokens from a first set of multiple sets of tokens, one or more processors provide the processing system with The artificial intelligence model generates each set of tokens in the first set of tokens, and rejects each set of tokens within that set. Using the generative artificial intelligence model, tokens are sampled based on a target distribution that excludes the probability associated with each set of tokens in a first plurality of sets of tokens, and the selected set of tokens from the first plurality of sets of tokens includes the sampled tokens. The processing system according to claim 1, configured as described above.

9. The processing system according to claim 1, wherein the generative artificial intelligence model includes a generative artificial intelligence model trained to generate a plurality of tokens in response to an input prompt based on predictive prompt embedding.

10. The processing system according to claim 1, wherein the generated artificial intelligence model includes one or more non-autoregressive layers and a model including one or more autoregressive layers.

11. The processing system according to claim 10, wherein the one or more autoregressive layers comprise one or more layers on top of a stack of layers representing the generative artificial intelligence model.

12. The processing system according to claim 10, wherein the one or more autoregressive layers comprise one or more layers at the bottom of a stack of layers representing the generative artificial intelligence model.

13. A method implemented by a processor, Based on an input prompt and a generating artificial intelligence model, a first plurality of sets of tokens are generated, wherein each set of tokens in the first plurality of sets of tokens corresponds to a first portion of a candidate response to the input prompt. Using the generative artificial intelligence model, speculatively generate a second plurality of sets of tokens, wherein each set of tokens in the second plurality of sets of tokens corresponds to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens. The process involves speculatively generating a second set of the aforementioned tokens, and selecting one set of tokens from the first set of the aforementioned tokens. In response to the input prompt, the system outputs a selected set of tokens from the first set of tokens and an associated set of tokens within the second set of tokens. Methods that include...

14. The method according to claim 13, wherein selecting a set of tokens from a first plurality of sets of tokens includes selecting the longest sequence of accepted tokens from the first plurality of sets of tokens.

15. The method according to claim 13, wherein a set of tokens from a second plurality of sets of tokens includes padding that takes into account the number of tokens in a selected set of tokens from the first plurality of sets of tokens.

16. The method according to claim 13, wherein a first plurality of sets of the tokens are represented as a tree data structure, and the root node of the tree data structure corresponds to the input prompt.

17. The method according to claim 16, wherein the depth of the tree data structure corresponds to the maximum number of tokens generated by a single pass through the generative artificial intelligence model.

18. The method according to claim 16, wherein the maximum size of the tree data structure is set based on a computational complexity metric associated with generating one set of tokens by the generative artificial intelligence model.

19. Selecting a set of tokens from a first set of multiple sets of tokens is Rejecting a first token at the first level of the tree data structure representing a first set of the aforementioned tokens, Based on the rejection of the first token, a modified probability distribution is generated, From the aforementioned tree data structure, discard or ignore the child tokens of the first token, Based on the adjusted probability distribution, determine whether to accept or reject the second token at the first level of the tree data structure. The method according to claim 16, including the method described in claim 16.

20. Selecting a set of tokens from a first set of multiple sets of tokens is Rejecting each set of tokens in the first set of multiple sets of tokens generated by the generative artificial intelligence model, Using the generative artificial intelligence model, sampling tokens based on a target distribution that excludes the probability associated with each set of tokens in a first plurality of sets of tokens, wherein the selected set of tokens from the first plurality of sets of tokens includes the sampled tokens. The method according to claim 13, including the method described in claim 13.

21. The method according to claim 13, wherein the generative artificial intelligence model includes a generative artificial intelligence model trained to generate a plurality of tokens in response to an input prompt based on predictive prompt embedding.

22. The method according to claim 13, wherein the generated artificial intelligence model includes one or more non-autoregressive layers and a model including one or more autoregressive layers.

23. The method according to claim 22, wherein the one or more autoregressive layers comprise one or more layers on top of a stack of layers representing the generative artificial intelligence model.

24. The method according to claim 22, wherein the one or more autoregressive layers comprise one or more layers at the bottom of a stack of layers representing the generative artificial intelligence model.

25. A processing system, Means for generating a first plurality of sets of tokens, wherein each set of tokens in the first plurality of sets of tokens corresponds to a first portion of a candidate response to the input prompt, based on an input prompt and a generating artificial intelligence model. Means for speculatively generating a second plurality of sets of tokens using the generative artificial intelligence model, wherein each set of tokens in the second plurality of sets of tokens corresponds to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens; Means for selecting one set of tokens from the first set of tokens while speculatively generating a second set of the tokens, Means for outputting, in response to the input prompt, a selected set of tokens from the first plurality of tokens and an associated set of tokens within a second plurality of tokens, A processing system equipped with the following features.

26. The processing system according to claim 25, wherein the means for selecting a set of tokens from a first plurality of sets of tokens includes means for selecting the longest sequence of accepted tokens from the first plurality of sets of tokens.

27. The processing system according to claim 25, wherein a set of tokens from a second plurality of sets of tokens includes padding that takes into account the number of tokens in a selected set of tokens from the first plurality of sets of tokens.

28. The processing system according to claim 25, wherein a first plurality of sets of the tokens are represented as a tree data structure, and the root node of the tree data structure corresponds to the input prompt.

29. The processing system according to claim 28, wherein the depth of the tree data structure corresponds to the maximum number of tokens generated by a single pass through the generative artificial intelligence model.

30. The processing system according to claim 28, wherein the maximum size of the tree data structure is set based on a computational complexity metric associated with generating one set of tokens by the generative artificial intelligence model.

31. The means for selecting a set of tokens from a first set of multiple sets of tokens is Means for rejecting a first token at the first level of the tree data structure representing a first set of the tokens, Means for generating an adjusted probability distribution based on the rejection of the first token, Means for discarding or ignoring the child tokens of the first token from the tree data structure, 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, The processing system according to claim 28, including the following:

32. The means for selecting a set of tokens from a first set of multiple sets of tokens is Means for rejecting each set of tokens in a first set of multiple sets of tokens generated by the generative artificial intelligence model, Means for sampling tokens based on a target distribution that excludes the probability associated with each set of tokens in a first plurality of sets of tokens, using the generative artificial intelligence model, wherein a selected set of tokens from the first plurality of sets of tokens includes the sampled tokens, The processing system according to claim 25, including the following:

33. The processing system according to claim 25, wherein the generative artificial intelligence model includes a generative artificial intelligence model trained to generate a plurality of tokens in response to an input prompt based on predictive prompt embedding.

34. The processing system according to claim 25, wherein the generated artificial intelligence model includes one or more non-autoregressive layers and a model including one or more autoregressive layers.

35. The processing system according to claim 34, wherein the one or more autoregressive layers comprise one or more layers on top of a stack of layers representing the generative artificial intelligence model.

36. The processing system according to claim 34, wherein the one or more autoregressive layers comprise one or more layers at the bottom of a stack of layers representing the generative artificial intelligence model.

37. A non-temporary computer-readable medium storing executable instructions, wherein the executable instructions are executed by one or more processors. Based on an input prompt and a generating artificial intelligence model, a first plurality of sets of tokens are generated, wherein each set of tokens in the first plurality of sets of tokens corresponds to a first portion of a candidate response to the input prompt. Using the generative artificial intelligence model, speculatively generate a second plurality of sets of tokens, wherein each set of tokens in the second plurality of sets of tokens corresponds to a second portion of the candidate response to the input prompt based on the first plurality of sets of tokens. The process involves speculatively generating a second set of the aforementioned tokens, and selecting one set of tokens from the first set of the aforementioned tokens. In response to the input prompt, the system outputs a selected set of tokens from the first set of tokens and an associated set of tokens within the second set of tokens. A non-temporary computer-readable medium that performs an action, including [the specified function].

Images