Speculative decoding in autoregressive generative artificial intelligence models

Abstract

Certain aspects of the present disclosure provide techniques and apparatus for generating a response to an input query using a generative model. The method generally includes generating a first plurality of token sets based on an input query and a first generative model. The first plurality of token sets are output to the second generative model for verification. A second plurality of token sets is speculatively generated when waiting to receive an indication of a token set selected from the first plurality of token sets. An indication of a token set selected from the first plurality of token sets is received. Tokens associated with the selected set of tokens from the second plurality of sets of tokens are output to a second generative model for verification, and the selected set of tokens is output as a response to the input query.

Description

[0001] Cross-references to related applications

[0002] This application claims priority to U.S. Patent Application S / N.18 / 479,672, filed October 2, 2023, entitled "Speculative Decoding in an Autoregressive Generative Artificial Intelligence Model," which claims priority and benefit to U.S. Provisional Patent Application S / N.63 / 454,605, filed March 24, 2023, entitled "Speculative Decoding in an Autoregressive Generative Artificial Intelligence Model," both of which are assigned to the assignee of this application and are incorporated herein by reference in their entirety.

[0003] introduction

[0004] This disclosure relates to generative artificial intelligence models, and more specifically to speculative decoding in generative artificial intelligence models.

[0005] Generative AI models can be used in a variety of environments to generate responses to input queries. For example, generative AI models can be used in chatbot applications, where large language models (LLMs) are used to generate answers to input queries, or at least responses to those queries. Other examples of the use of generative AI models include models that generate stable diffusion of images based on input text descriptions of the content of desired images, and decision transformers that predict future actions based on sequences of previous actions in a given environment.

[0006] Generally, using generative AI models to generate responses to queries can be computationally expensive. For example, in a chatbot deployment that uses a large language model to generate responses to queries formatted as text, the response to the query may be generated for each token (e.g., a word or a portion of a word) generated as part of the response using a one-round pass through the large language model. The output of each round pass may be a probability distribution of the token set (e.g., individual words or portions of multiple words), which can be selected from the token set, for example, by sampling or based on maximum likelihood. Since a one-round pass through the large language model is used to generate each word (or tokens) in the query response, the computational overhead can be modeled as the product of the number of words included in the response and the computational resource overhead of performing a one-round pass through the large language model (e.g., in terms of processing power, memory bandwidth, and / or other computational resources used), which generally increases with the number of parameters within the large language model.

[0007] Brief Overview

[0008] Certain aspects of the present disclosure provide a method for generating a response to an input query using generative artificial intelligence models. The method generally includes generating a plurality of sets of tokens based on the input query and a first generative model, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input query. The plurality of sets of tokens are output to a second generative model for validation. An indication of a selected set of tokens from the plurality of sets of tokens is received from the second generative model based on the input query and the plurality of sets of tokens. The selected set of tokens is output as the response to the input query.

[0009] Certain aspects of the present disclosure provide a method for validating a response to an input query generated using generative artificial intelligence models. The method generally includes receiving an input query and a plurality of sets of tokens generated by a first generative model, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input query. A probability distribution associated with each respective set of tokens in the plurality of sets of tokens is compared to a corresponding probability distribution generated by a second generative model for the respective set of tokens. A set of tokens is selected based on the comparison, and the selected set of tokens is output to the first generative model.

[0010] Certain aspects of the present disclosure provide a method for parallel speculative generation of a response to an input query using generative artificial intelligence models. The method generally includes generating a first plurality of sets of tokens based on the input query and a first generative model, each set of tokens in the first plurality of sets of tokens corresponding to a first portion of a candidate response to the input query. The first plurality of sets of tokens are output to a second generative model for validation. A second plurality of sets of tokens is speculatively generated while waiting to receive an indication from the second generative model of a selected set of tokens from the first plurality of sets of tokens, each set of tokens in the second plurality of sets of tokens corresponding to a second portion of a candidate response to the input query. The indication of the selected set of tokens from the first plurality of sets of tokens is received from the second generative model. Tokens from the second plurality of sets of tokens associated with the selected set of tokens are output to the second generative model for validation, and the selected set of tokens is output as the response to the input query.

[0011] Other aspects provide a processing system configured to perform the aforementioned methods and those described herein, a non-transitory computer-readable medium comprising instructions that, when executed by one or more processors of a processing system, cause the processing system to perform the aforementioned methods and those described herein, a computer program product embodied on a computer readable storage medium comprising code for performing the aforementioned methods and those further described herein, and a processing system comprising means for performing the aforementioned methods and those further described herein.

[0012] The following detailed description and accompanying drawings provide a thorough explanation of certain illustrative features of one or more aspects. BRIEF DESCRIPTION OF DRAWINGS

[0014] The accompanying drawings are intended to convey a thorough understanding of one or more aspects of the disclosure. Descriptions of specific devices and techniques are provided to give a thorough understanding of these aspects. In no way should the descriptions be taken to limit the scope of the disclosure.

[0015] Figure 1 An example of group speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0016] Figure 2 An example relationship between a draft model and a target model used in group speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0017] Figure 3 An example of parallel speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0018] Figure 4 An example pipeline for hybrid speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0019] Figure 5 An example of hybrid speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0020] Figure 6 An example of hybrid parallel speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0021] Figure 7 An example of hybrid group speculation decoding in a generative model is illustrated in accordance with aspects of the disclosure.

[0022] Figure 8 An example operation for generating a response to an input query using a generative artificial intelligence model is illustrated in accordance with aspects of the disclosure.

[0023] Figure 9 An example operation for verifying a response to an input query using a generative artificial intelligence model is illustrated in accordance with aspects of the disclosure.

[0024] Figure 10 Example operations for speculatively generating responses to input queries using generative artificial intelligence models and verifying responses to input queries in accordance with aspects of the present disclosure are illustrated.

[0025] Figure 11 An example processing system configured to perform various aspects of the present disclosure is depicted.

[0026] For the purpose of promoting an understanding of the principles of the application, reference will now be made to the embodiments illustrated in the drawings and specific language will be used to describe the same. It will, nevertheless, be understood that no limitation of the scope of the application is thereby intended, any alterations and further modifications in the illustrated embodiments, and any further applications of the principles of the application as illustrated therein are contemplated as falling within the scope of the application.

[0027] DETAILED DESCRIPTION

[0028] Aspects of the present disclosure provide apparatuses, methods, processing systems, and computer readable media for efficiently generating responses to input queries using generative artificial intelligence models.

[0029] Generally, generative artificial intelligence models generate responses to queries input into the model. For example, a large language model (LLM) deployed within a chatbot can generate a response to a query using multiple passes of the large language model, where each successive pass is based on the query and tokens (or words or portions thereof) generated using previous passes of the large language model. Generally, these large language models can include millions, or even billions, of weights or parameters within the model. Due to the size of these models and the operations performed on each token to predict what the next token generated in response to the query and previously generated tokens should be, it can be impractical, or even impossible, to deploy large language models on various devices that can have limited memory, storage, and / or processing capabilities relative to cloud computing instances on which large language models are typically run. Moreover, the computational complexity involved in generating a response to a query provided as model input can involve significant energy consumption, processing time, memory utilization, and / or other resource utilization that can prevent computing resources from being used for other tasks.

[0030] To improve the efficiency and throughput of large language models, speculative decoding techniques allow smaller language models (sometimes called draft large language models, also known as draft models) to operate alongside larger language models (sometimes called target large language models, also known as target models). In this scenario, the draft model speculatively generates additional tokens and the probabilities used to sample these additional tokens based on the currently accepted set of tokens. The target model generates tokens based on those generated by the draft model. To produce the final result, the target model can perform rejection sampling on a per-token basis to accept or reject tokens generated by the draft model, ensuring that the draft and target models have similar probability distributions.

[0031] In some respects, the draft model can be a pruned version of the selected target model so that the draft model and the target model have similar probability distributions. In other respects, the draft model can be a smaller version of the target model (e.g., trained on millions of tokens rather than on hundreds of millions or even billions of tokens).

[0032] Various aspects of this disclosure provide techniques for generating responses to queries input into a large language model using speculative decoding on a group-by-group basis. Generally, a draft model can generate one or more sets of tokens as candidate responses to the query. The target model can then perform sample rejection on a per-set basis. Tokens within the group can be selected based on the conditional probabilities of tokens within the group. By performing rejection sampling on a per-set basis, various aspects of this disclosure preserve a close relationship between the probability distributions within the draft and target models, while increasing the throughput (e.g., the number of tokens generated per second) of the draft and target models compared to draft and target models configured to generate responses on a per-token basis.

[0033] This disclosure provides techniques for generating responses to queries input into a large language model using speculative decoding techniques that run in parallel with a draft model and a target model, also referred to herein as "hybrid speculative decoding." In hybrid speculative decoding, the draft model speculatively generates one or more tokens, which are then validated by the target model. The draft model can be executed on an edge device, and the target model can be executed on a server (e.g., locally accessible from the edge device, hosted in a cloud computing environment, etc.). By executing both the draft model and the target model in parallel, the rate of token generation can be maximized or at least increased. Furthermore, in scenarios where the edge device is disconnected from the server, the edge device can continue generating tokens to produce responses to received queries.

[0034] Speculative Decoding in Generative Artificial Intelligence Models

[0035] In general, autoregressive token generation (e.g., in large language models) can take as input a history of tokens to generate an output. That is, autoregressive token generation can be represented by the following expression:

[0036] x t ~ p(x | x0, x1,..., x t-1 ) → x t+1 ~ p(x | x0, x1,..., x t-1 , x t )

[0037] where x t represents a sequence of tokens generated at time t, whose conditional probability p is conditioned on the selection of tokens x0through x t-1 , and x t+1 represents a sequence of tokens generated at time t+1, whose conditional probability p is conditioned on the selection of tokens x0through x t . In general, a single token can be generated each time the autoregressive model is executed, which means that N inferences can be performed to generate a sequence of N tokens. As discussed above, speculative decoding techniques can be used to accelerate token generation by using a draft model that is smaller in size than a target model to speculatively generate tokens, where the target model is used to validate tokens generated (speculatively) by the draft model.

[0038] In a speculative decoding pipeline, a draft model can autoregressively speculatively generate n tokens according to the following expression:

[0039]

[0040] where t corresponds to a time point, and corresponds to a conditional probability distribution associated with a token x selected at time t, conditioned on the selection of tokens x0through x t-1 .

[0041] A target model takes the n tokens generated and processes the n tokens in parallel according to the following expression to generate a probability distribution for each of the n tokens:

[0042]

[0043] where k corresponds to a token index with respect to the n tokens generated.

[0044] The target model can then validate the tokens generated by the draft model by comparing the distributions from the draft model and the target model to determine whether the tokens are accepted or rejected.

[0045] For some function f and some threshold a (also referred to as an acceptance rate), a given token is accepted when . The token can be accepted. Otherwise, the token can be rejected. The final token can then be generated based on a function The final token is generated at the first rejected position or the last position n.

[0046] Compared to iteratively generating tokens using a single autoregressive model, speculative decoding with acceptance rate a can result in a reduction in cost. The inference cost savings compared to iterative token generation can be expressed by:

[0047]

[0048] Considering this example, for N = 1000, C 目标 = 10, C 草稿 = 1, n = 4, a = 3, where N corresponds to the number of tokens, C AR corresponds to the computational cost of generating a set of tokens using acceptance rate a, C 目标 corresponds to the computational cost of generating a set of tokens using the target model, C 草稿 corresponds to the computational cost of generating a set of tokens using the draft model, and n corresponds to the number of tokens generated speculatively by a single pass through the autoregressive model. In such an example, speculative decoding can result in a 35% reduction in computational overhead compared to individual autoregressive iterative token generation.

[0049] However, as discussed, speculative decoding on a per-token basis can impose a limit on the rate at which tokens are generated, as the first token can be individually sampled by the draft model and verified by the target model before the next token is sampled by the draft model and verified by the target model. That is, generating a response to an input query using a per-token speculative decoding technique can involve executing the draft model and the target model for each token generated as part of the response to the input query, which can use a substantial amount of computational resources (e.g., processor time, memory, memory bandwidth, etc.) to generate the response.

[0050] Example group speculative decoding in generative artificial intelligence models

[0051] Figure 1 An example 100 of group speculative decoding (GSD) in a generative artificial intelligence model is illustrated in accordance with aspects of the present disclosure.

[0052] As illustrated, the draft model and the target model can be used in conjunction with one another (or otherwise together) to perform group speculative decoding of tokens to generate a response to a query received for processing by one or more generative artificial intelligence models. As discussed in further detail below, group speculative decoding of tokens can allow for multiple sets of tokens to be speculatively generated by the draft model for verification by the target model. As multiple sets of tokens can be generated by the draft model for verification by the target model, group speculative decoding can increase a token generation rate of the generative artificial intelligence model by generating multiple sets of tokens that can be accepted as a correct response, as a larger number of sets of tokens can increase a probability that at least one set includes one or more tokens that will be accepted as a response.

[0053] Given an input of a received query, the draft model generally forms a sampled token group, including one or more high-probability node groups from a probability distribution output over a set of potential tokens. Each high-probability node group can include a group selected based on various techniques, such as optimal k-selection (e.g., selecting the k tokens with the highest probabilities within a probability distribution), kernel-based selection (e.g., selection based on a sum of probabilities satisfying a threshold probability), and the like. Tokens not included in one or more groups can be considered a separate set group (e.g., a group including a single member). By selecting candidate token groups, the draft model can sample tokens based on the probability distribution; however, each group can be treated as a single selection with an aggregated group probability (e.g., as a sum of individual token probabilities) such that the draft model matches or at least approximates the target model such that a resulting probability of each group is a sufficient response to the received query. Grouping can achieve this match or at least approximation between probabilities computed by the draft model and the target model, as while individual token probabilities can not match as well, the aggregation of possible tokens can have probabilities that can be similar in the draft model and the target model.

[0054] When the sampled token group is input into the draft model in a next iteration of executing the draft model, the tokens in the sampled token group are input at the sampled locations and processed independently. A result can be a tree data structure 110, where the prompt is a root node 111 of the tree data structure and subsequent layers within the tree data structure 110 represent different tokens (or groups of tokens) combined with each of the previously selected tokens. At some point in time (e.g., after generating a tree having a defined depth corresponding to a maximum length of a sequence generated by the draft model), the draft model can output the generated tree data structure 110 to the target model for further processing. In some aspects, the tree data structure 110 can be output to the target model along with the selection probabilities and groups of each generated by the draft model.

[0055] In some aspects, the number of nodes at each level of the tree data structure 110 and the depth of the tree data structure 110 can be defined a priori. The number of nodes at each level of the tree can be defined globally, on a per-level basis, or in some other manner. For example, the number of nodes at any level of the tree data structure 110 can be defined based on a branching factor at an immediately preceding level of the tree data structure 110. For example, a branching factor of 2 for nodes at the n-th level of the tree data structure 110 can result in 2 nodes (tokens) being generated at the n+1-th level of the tree data structure 110 for each node (token) at the n-th level of the tree. Meanwhile, the depth of the tree can be defined based on a maximum number of tokens (e.g., words) that can be generated using any round of pass of the draft model. For example, if the draft model is configured to generate sequences of a maximum length of 5 tokens during any instance of speculative generation, the depth of the tree can be 6 (to include a root node at the first level of the tree data structure 110 corresponding to input to the draft model).

[0056] The draft model can be configured to trigger the target model to generate tokens based on various complexity and / or performance criteria. For example, the draft model can trigger the target model to generate tokens based on a complexity or performance criterion associated with a size of a generated tree. In other examples, the draft model can be configured to trigger the target model to generate tokens based on a time criterion associated with an expected amount of time for the target model to generate a set of tokens that a generated tree can be compared against. Generally, these complexity and / or performance criteria can set an upper limit on a number of tokens generated by the draft model for verification by the target model. In some aspects, this upper limit can be based on a number of nodes in the tree data structure, and can be influenced by, for example, branching factors defined for different levels of a tree data structure that sampled tokens are organized into, a depth of the tree data structure, and the like. A worst-case computational load for a last round of speculative token generation can be configured to be bounded by a memory bandwidth at a device on which the draft model is executed.

[0057] Based on the generated tree data structure 110, the groups, and the selection probabilities received from the draft model, the target model generates an output (probability) distribution for each partial path of the generated tree data structure 110. Each partial path generally starts from the root node of the generated tree data structure 110 and can be a contiguous path through the generated tree data structure 110 starting from the root node and terminating at any node in the generated tree data structure 110. The target model can generate the output distribution for each partial path using a single pass through the target model, such as by including all tree nodes in the generated tree data structure 110 as token inputs and performing masked self-attention and positional encoding for each partial path within the tree data structure 110. Based on the output distribution for each node in the generated tree data structure 110, the target model can select a set of tokens 112 by traversing from the root node of the generated tree data structure 110. If the target model includes multiple tokens in the selected set of tokens 112, a token can be selected from within the group using normalized group token probabilities (e.g., token probabilities conditioned on selection of the group including the token) and the selected token can be redefined as the root node of the generated tree data structure 110. The process discussed herein can be repeated with the selected token defined as the root node of the tree data structure 110 to generate subsequent tokens to be included in the selected set of tokens 112.

[0058] If a token group is rejected based on the probability distribution generated by the target model, a final token group can be selected based on a group-level rejection sampling method. A final token can be selected from the selected final token group based on normalized group token probabilities. If all token groups are accepted by the target model, a final token can be sampled according to a final leaf distribution (e.g., a probability distribution generated on leaf nodes in the tree data structure 110 corresponding to the last accepted or validated token for any partial path through the tree data structure 110), which can allow the target model to generate additional tokens relative to the token(s) speculatively generated by the draft model. Using the techniques discussed herein, the selected set of tokens 112 can then be communicated back to the draft model to be used as input for a subsequent round of token speculation and validated by the target model.

[0059] Figure 2 An example 200 illustrating relationships between tokens generated by a draft model and a target model for use in group-speculation decoding in generative models in accordance with aspects of the disclosure is illustrated.

[0060] The tree generated by the draft model generally allows for diverse speculation, where many candidate token sets are provided by the draft model for verification by the target model, which can increase the acceptance rate of tokens generated by the target model by the target model. As illustrated, accepting token number 4 of the token group generated by the draft model at time t = 1 210 results in accepting token number 1 of the token group generated by the draft model at time t = 2 220. Subsequently, at time t = 3 230, there can be two options: token number 3 or token number 5. Token number 3 can be accepted as a sampled token, while token number 5 can be rejected as a sampled token, as illustrated. At time t = 4 240, token 4 from the token set generated speculatively based on selecting token 3 at time t = 3 230 can be accepted as a sampled token based on the calculated probabilities of tokens 3 and 4 conditioned on accepting token 3 at time t = 3 230. Meanwhile, both tokens 4 and 5 from the token set generated speculatively based on selecting token 5 at time t = 3 230 can both be rejected based on rejecting token 5 at time t = 3 230.

[0061] A token tree, rather than a token sequence, can be fed to the target model as input. For example, a tree can be fed to the target model that includes the option of token 3 or 5 as a candidate token at time t = 3 230, the option of token 3 or 4 as a candidate token at time t = 4 240 in the case of selecting token 3 at time t = 3 230, and the option of token 4 or 5 as a candidate token at time t = 4 240 in the case of selecting token 5 at time t = 3 230. As illustrated, inputting a token tree data structure (e.g., as illustrated by the token sequence shown at times t = 1 through t = 4) to the target model at time t = 5 250 can allow the target model to generate an output that includes five tokens— one more than the sequence(s) included in the tree generated by the draft model— which can provide an increased acceptance rate of the token tree, and allow for generation of additional tokens to be included in the response generated by the draft model and the target model. In the example illustrated in FIG. 2, the selected set of five tokens includes token 4 at time t = 1 210, token 1 at time t = 2 220, token 3 at time t = 3 230, token 4 at time t = 4 240 (conditioned on selecting token 3 at time t = 3 230), and token 2 at time t = 5 250 (conditioned on selecting token 4 at time t = 4 240). Figure 2 In the example illustrated in FIG. 2, the selected set of five tokens includes token 4 at time t = 1 210, token 1 at time t = 2 220, token 3 at time t = 3 230, token 4 at time t = 4 240 (conditioned on selecting token 3 at time t = 3 230), and token 2 at time t = 5 250 (conditioned on selecting token 4 at time t = 4 240).

[0062] In general, new selection of a group of tokens can be performed each time the draft model generates a new distribution of one or more subsequent tokens to be included in a response to a received query processed by the draft model and the target model. As discussed, use of group selection of tokens can allow for closer matching between the draft model probability distribution and the target model probability distribution to minimize or at least reduce the likelihood of group rejection. In some aspects, when there is a low level of uncertainty (e.g., when the next token selection has low entropy), the draft model can choose not to form a group to minimize or at least reduce computational resource consumption.

[0063] In general, while using group speculative decoding (to use a generative artificial intelligence model to generate a response to an input query) can have a high level of efficacy, forming groups of tokens can reduce computational complexity of both the draft model and the target model while maintaining a high level of throughput (e.g., number of tokens generated per second by the draft model and the target model). In general, triggering the target model can be influenced by various considerations, such as a maximum tree size to keep model computation performance metrics within bounds, an estimate of when longer speculations have diminishing returns (e.g., where the computational overhead of generating a large tree of speculated tokens is unlikely to be valuable, taking into account the likelihood that the target model will reach tokens later in the tree), and other performance considerations.

[0064] In some aspects, the draft model can match (e.g., in terms of probability distribution) the target model, but can have faster inference performance on the same hardware than the target model. In general, a smaller model can generate many speculative tokens, but the likelihood that the generated tokens are rejected by the target model can increase. As discussed above, group speculative decoding can address this increase in likelihood of token rejection at the cost of increasing computational overhead for longer sequences.

[0065] In some aspects, at the draft model, a temperature parameter— or a parameter that influences the likelihood that the draft model selects a token with a lower probability of being an accurate token for inclusion in a set of tokens— can be tuned or otherwise selected to improve performance of group speculative decoding.

[0066] In some aspects, the draft model can be fine-tuned to match (or at least approximate) the target model and maximize (or at least increase) the probability that a speculative generated token generated by the draft model will be accepted as a valid token by the target model.

[0067] Generally, the token generation performance of group speculative decoding can be improved as compared to per-token speculative decoding. That is, for a given Kullback-Leibler (KL) divergence between a draft model and a target model, measuring how the probability distribution of the draft model differs from the probability distribution of the target model (considering the target model as the reference distribution), the number of tokens generated for each run of the target model can be greater for group speculative decoding as compared to per-token speculative decoding. Different group strategies (e.g., group size, additional tokens, etc.) can have different computational complexity characteristics; as such, the selection of a group strategy can be based on a tradeoff between computational complexity and performance, given bounds on the read bandwidth of the draft model and the target model, and hardware performance.

[0068] Example parallel speculative decoding in generative artificial intelligence models

[0069] As discussed above, group speculative decoding can allow for the combination of draft model speculative token generation and target model token acceptance to generate a final set of tokens that match the distribution of the target model, while increasing the rate at which tokens are generated. Moreover, aspects of the present disclosure provide for parallel speculative decoding, in which the draft model speculatively generates multiple sets of tokens in parallel, which can provide for increased token generation rates, while improving computational efficiency involved in generating responses to queries using generative artificial intelligence.

[0070] Figure 3 An example 300 of parallel speculative decoding (PSD) in a generative artificial intelligence model is illustrated in accordance with aspects of the present disclosure.

[0071] As illustrated, the draft model can run multiple instances of the speculative decoding process to generate multiple candidate sequences of tokens. Each iteration can be performed using different parameters, or seeds, to perform token sampling. Performing N iterations using the S tokens generated by each iteration can result in a total of N*S tokens being generated, which can be provided to the target model for processing.

[0072] For example, as illustrated in example 300, the draft model can initially perform multiple instances of the speculative decoding process to generate a first candidate sequence 302 and a second candidate sequence 304 for a prompt input into the draft model. As illustrated, N=2, S=4; thus, the draft model can generate two sequences, with 4 tokens in each sequence, for a total of 8 tokens. It should be recognized that this is merely an example, and any values of N and S can be used to define the number of candidate sequences generated by the draft model and the number of tokens included in each sequence, respectively.

[0073] The target model uses N*S tokens as input and processes N candidate sequences of tokens at a time run through the target model. The result of the target model processing the N*S tokens can be a set of N potentially accepted sequences of different lengths depending on where the target model rejects tokens in each of the N candidate sequences of tokens. Various techniques can be used to select a final sequence based on the N potentially accepted sequences. As a non-limiting example, a "greedy" technique can result in selecting the sequence with the longest length, where constraints (e.g., among sequences of the same length) are arbitrarily broken, which can result in the highest number of tokens being selected each time the target model is run. In some cases, the longest length selection can bias toward an effective final sampling distribution, which can cause the probability distribution of the target model to converge at least to some extent with the probability distribution of the draft model.

[0074] For example, as illustrated, the target model accepts two of the four tokens of the first candidate sequence 302 and accepts four of the four tokens of the second candidate sequence 304. Because more tokens are accepted from the second candidate sequence 304 than from the first candidate sequence 302, a "greedy" technique can result in the target model selecting the second candidate sequence 304 as the sequence upon which the draft model generates subsequent candidate sequences 312 and 314. In this case, no speculative token generation need be performed based on the first candidate sequence 302 because the first candidate sequence 302 has fewer accepted tokens than the second candidate sequence 304. However, in other cases, other techniques can be implemented to select one of the candidate sequences, including sequences that can not have the most accepted tokens.

[0075] While Figure 3Parallel speculative decoding based on per-token speculative decoding is illustrated, but it should be recognized that the N candidate token sequences can additionally or alternatively be generated based on the group speculative decoding techniques discussed above. For example, using the group speculative decoding techniques discussed above, multiple token sets generated by multiple instances of a draft model can be generated for verification by a target model. In some aspects, the token sets generated by the instances of the draft model can be organized into tree data structures (as discussed above), resulting in the generation of multiple tree data structures, where each individual tree data structure corresponds to a token set generated by a particular instance executed by the draft model using instance-specific parameters. The target model can attempt to verify the tokens in each token set to determine a selected token set on which to base subsequent speculative decoding. The longest path through each tree data structure can be identified by the target model as a candidate sequence. Using a “greedy” technique, the tokens in the candidate sequence with the longest path can be selected as the token set on which to base subsequent speculative decoding, and can be output by the target model to the draft model for use by the draft model as a root node of a subsequent tree data structure (e.g., a tree data structure corresponding to subsequent tokens generated as part of a response to an input query).

[0076] Hybrid speculative decoding examples

[0077] In some cases, it can be possible to perform a generative artificial intelligence model using autoregressive inference in a hybrid system. Examples of these hybrid systems can include systems in which these generative models are executed on edge devices (e.g., smartphones, tablets, laptops, desktops, etc.) and cloud servers, on edge devices and home servers, on different home servers, etc. As such, in some aspects, a draft model can be executed on a local device, while a target model can be executed on a remote device. Executing a generative artificial intelligence model in a hybrid system can speed up the generation of query responses, and can reduce the computational load of a server (e.g., by offloading some of the computational processes to another device). Furthermore, when a server executing a target model is available, target model accuracy can be preserved, and when an edge device is out of communication with a server, a large language model (or other generative artificial intelligence model) can remain available at the edge device (albeit with lower accuracy).

[0078] However, in hybrid systems, when both a draft model and a target model are hosted in a cloud system, cost reduction can be limited. Furthermore, the sequential nature of the inference (e.g., involving the speculative generation of tokens using a draft model, followed by the verification of the tokens using a target model) can limit the rate at which tokens are generated.

[0079] To further reduce the computational cost involved in generating query responses using generative artificial intelligence models, aspects of the present disclosure provide various techniques for speculative decoding in a hybrid environment. In various aspects, speculative decoding in a hybrid environment can be performed such that candidate token generation using a draft model (e.g., at an edge device) and token validation using a target model (e.g., at a server) can be performed in parallel. Given a change in the number of accepted tokens, the draft model can continue to generate subsequent draft token sets with padding added to account for the difference in the number of accepted tokens from the previous round of speculative token generation. The draft model and the target model can be continuously performed in substantial parallel, which can accelerate the generation of query responses using generative artificial intelligence models.

[0080] Figure 4 An example pipeline 400 for hybrid speculative decoding (HSD) in a generative model is illustrated in accordance with aspects of the present disclosure.

[0081] As illustrated, to maximize or at least increase throughput (e.g., number of tokens generated per second), a draft model 402 executed on an edge device 403 (which can be referred to as a local device or a first device) can continuously generate batches of candidate tokens (or token sets, corresponding to different responses to an input prompt having different lengths, conditioned on any previously accepted tokens) and output these tokens to a target model 404 executed on a server 405 (which can be referred to as a remote device or a second device) (e.g., in a cloud computing environment). The edge device 403 can receive a set of accepted tokens from the target model 404 from the server 405. The draft model 402 can prune the sample tree to conform to the set of accepted tokens. In some aspects, when the set of accepted tokens is empty (e.g., when the target model accepts no tokens), the draft model 402 can backtrack to the last token accepted by the target model 404 and restart speculative token generation from the last accepted token. For example, to backtrack, the draft model 402 can prune the generated sample token tree to the last accepted token and restart speculative generation based on the pruned tree.

[0082] In the example pipeline 400, the draft model 402 executed on the edge device 403 can generate a first set of speculative tokens 410 and output the first set of speculative tokens to the target model 404 executed on the server 405 for validation. In some aspects, the first set of speculative tokens 410 can include multiple subsets of speculative tokens, where each subset corresponds to one or more sequences of tokens generated by the draft model 402 using different instances of execution of different operational parameters (e.g., according to group speculative decoding, parallel speculative decoding, or other techniques).

[0083] In some aspects, when the target model 404 validates the first set of speculative tokens 410, the draft model 402 can generate a second set of speculative tokens 414. The second set of speculative tokens 414 can be generated based on the assumption that the target model 404 accepts 1 to S tokens from the first set of speculative tokens 410, where S corresponds to the number of tokens per sequence generated by the draft model 402. The draft model 402 can then receive the selected set of tokens 412 from the target model 404. Based on the selected set of tokens 412, the draft model 402 can identify a subset of the second set of speculative tokens 414 for analysis (e.g., it generates tokens conditioned on accepting the correct number of tokens from the first set of speculative tokens 410) and provide the subset to the target model 404 for validation. This process can continue, with the draft model 402 generating a third set of speculative tokens 418 while the target model 404 validates the second set of speculative tokens 414 and returns a selection 416 from the second set of speculative tokens 414 to the draft model 402, with the draft model 402 generating a fourth set of speculative tokens 422 while the target model 404 validates the third set of speculative tokens 418 and returns a selection 420 from the third set of speculative tokens 418 to the draft model 402, and so on. In other aspects, the draft model 402 can wait to generate the second set of speculative tokens 414 until the draft model 402 receives the selected set of tokens 412 from the target model 404.

[0084] Some of the processes described herein to speculatively generate tokens in parallel or substantially in parallel using a draft model and validate previously generated sets of speculatively generated tokens using a target model can be performed according to a group selection and a target model trigger condition selected to enable continuous or at least near-continuous operation of the draft model 402 and the target model 404, while minimizing or at least reducing the likelihood that the draft model 402 will backtrack from the last accepted token and start over with speculative token generation (e.g., to minimize or at least reduce the likelihood that the target model 404 will reject each of the speculatively generated tokens generated by the draft model 402). Moreover, the group selection and the target model trigger condition can be selected to avoid or at least reduce the likelihood of overburdening the computational capabilities on the edge device or server (e.g., to enable the edge device to generate tokens and the server to validate tokens within defined performance metrics that minimize or at least reduce the likelihood of encountering processing system bottlenecks, such as memory thrashing where data is repeatedly exchanged between processor on memory and processor off memory, network bandwidth bottlenecks, etc.).

[0085] In some aspects, hybrid speculative decoding can be performed on a per-token basis. In such cases, the draft model 402 can operate continuously and generate various branch paths based on new speculative token generations. In some aspects, the edge device 403 can perform multiple original speculative decoding instances (e.g., using different parameters input into the draft model 402) such that the draft model 402 generates multiple paths within a tree representing options that the target model 404 can validate. Based on some defined triggering criteria, the edge device 403 can provide the generated tree to the target model 404 running on a server 405 (e.g., in a cloud computing environment), which, as discussed above, can select a token sequence corresponding to a possible valid response to the input query and return the selected sequence to the draft model 402 executing on the edge device. The draft model 402 can then prune the generated tree to conform to the selected sequence returned by the target model 404. In this example, the token generation rate can correspond to the rate at which the draft model 402 generates tokens, with tree pruning based on target selection. In this example, inference accuracy and speed can have an inverse relationship, as a smaller draft model can generate tokens at a higher rate, but the likelihood that the target model 404 will accept the tokens speculatively generated by the draft model 402 is lower; meanwhile, a larger draft model can generate tokens at a lower rate, but the likelihood that the target model 404 will accept the tokens speculatively generated by the draft model 402 is higher.

[0086] Figure 5 An example 500 of hybrid speculative decoding in a generative model is illustrated in accordance with aspects of the present disclosure. In hybrid speculative decoding, a first set of speculative generated tokens 502 (also referred to as “round 1 draft tokens”) generated by a draft model (e.g., the draft model 402 illustrated in Figure 4 Figure 4 In various aspects, the draft model can be executed on a first device, while the target model can be executed on a second device.

[0087] ​According to various aspects, when the first set of speculatively generated tokens 502 are processed by the target model, the draft model generates a second set of draft tokens (also referred to as “round 2 draft tokens”) in a second round of inference, with hypotheses made for different numbers of accepted tokens. For example, as illustrated, the second set of draft tokens can include (1) a first subset 504 that assumes acceptance of the first draft token from the first set and can include a set of tokens speculatively generated based on accepting the first token; (2) a second subset 506 that assumes acceptance of the first and second draft tokens from the first set and includes a set of tokens speculatively generated based on accepting the first and second tokens; (3) a third subset 508 that assumes acceptance of the first through third draft tokens from the first set and includes a set of tokens speculatively generated based on accepting the first through third tokens; (4) a fourth subset 510 that assumes acceptance of all four tokens from the first set and includes a set of tokens speculatively generated based on accepting all four tokens. For scenarios in which fewer than the number of tokens included in the first set are assumed to be accepted, padding (e.g., null values, pre-defined constants, etc.) can be added so that each hypothesis has the same length. As illustrated, the second set of draft tokens can be generated in a batch process so that the appropriate selection of draft tokens conditioned on a particular set of accepted tokens from the first set of speculatively generated tokens 502 can be provided to the target model for validation.

[0088] Figure 6 An example 600 of hybrid speculative decoding in a generative model according to aspects of the disclosure is illustrated. In this example 600 (also referred to as an example of hybrid parallel speculative decoding (HPSD)), the parallel speculative decoding aspects discussed above can be used to generate various sets of candidate tokens for validation by a target model. The target model can operate continuously, as if performing speculative decoding on a per-token basis. However, instead of generating a token on a per-token basis, the draft model can speculatively generate a group of tokens based on a probability distribution that approximates (but does not necessarily equal) the probability distribution associated with the target model and generate a different tree data structure based on the input query (or prompt) and the generated group of tokens, as illustrated in Figure 6 Based on some defined triggering criteria, the edge device (also referred to as a local device or first device) can provide the generated tree data structure to the target model (also referred to as a remote device or second device) running on a server (e.g., in a cloud computing environment), which selects a sequence of tokens corresponding to a likely valid response to the input query, as discussed above, and returns the selected sequence to the draft model executing on the edge device. In general, the selected sequence of tokens can include the speculatively generated tokens that the target model can validate, plus one additional token generated based on the speculatively generated tokens. The draft model can then prune the generated tree to conform to the tokens validated by the target model.

[0089] As illustrated, different instances of the draft model using different parameters can independently generate first and second sets of speculative tokens 602, 604 based on the input prompt. The draft model (e.g., executing on a local or first device) can provide the first and second sets of speculative tokens 602, 604 to the target model (e.g., executing on a remote or second device) for validation. In parallel or at least substantially in parallel with the target model validating the first and second sets of speculative tokens 602, 604, the draft model can generate a first set of subsequent tokens 612 based on an assumption that the target model has accepted the first set of speculative tokens 602 and a second set of subsequent tokens 614 based on an assumption that the target model has accepted the second set of speculative tokens 604. In this example, as illustrated by the ellipses, the target model has accepted the second set of speculative tokens 604 over the first set of speculative tokens 602 (e.g., in a similar manner as described above with respect to Figure 3 Because the target model has accepted the second set of speculative tokens 604, the draft model can discard the first set of subsequent tokens 612 and provide the second set of subsequent tokens 614 for validation by the target model. Further, although not shown, the draft model can generate additional sets of tokens based on the second set of subsequent tokens 614 while the target model validates the second set of subsequent tokens 614.

[0090] The process of speculative token generation, pruning (e.g., by discarding unselected sets of tokens), and target model validation repeats until a termination condition is reached for the input (e.g., until no further tokens need to be generated to generate a correct response to the input query, until a maximum number of tokens has been generated, until a maximum amount of computational resources has been used to generate a response, etc.).

[0091] In this example, the overall token generation rate can correspond to the rate at which the draft model generates tokens. However, when the target model determines that none of the speculatively generated tokens in the generated tree are candidates included in the response to the query, the previously generated tokens can be discarded as the draft model can restart the speculative token generation using the tokens (up to the last validated token) as input into the draft model. As with the per-token based hybrid speculative decoding, in this example, inference accuracy and speed can have an inverse relationship, as smaller draft models can generate tokens at a higher rate, but the likelihood that the target model will accept the tokens speculatively generated by the draft model is lower.

[0092] Figure 7An example 700 of hybrid group speculation decoding (HGSD) in a generative model in accordance with aspects of the present disclosure is illustrated. In this example 700, the group speculation decoding aspects discussed above can be used to generate various sets of candidate tokens for validation by a target model. The target model (e.g., executing on a remote or second device) can operate continuously, as if performing speculation decoding on a per-token basis. However, instead of generating tokens on a per-token basis, a draft model (e.g., executing on a local device or first device) can speculatively generate groups of tokens based on a probability distribution that is equivalent to a probability distribution associated with the target model, and generate a tree (e.g., tree data structure) based on the input query (or prompt) and the generated groups of tokens. For example, based on some defined triggering criteria, the edge device can provide the generated tree to the target model running on a server (e.g., in a cloud computing environment), which, as discussed above, can select a sequence of tokens corresponding to a likely valid response to the input query and return the selected sequence to the draft model executing on the edge device. Generally, the selected sequence of tokens can include speculatively generated tokens that can be validated by the target model, plus one additional token generated based on the speculatively generated tokens. The draft model can then prune the generated tree to conform to the tokens validated by the target model.

[0093] For example, as illustrated, a tree-shaped data structure can be generated in which an input prompt is used as a root node of the tree data structure. A first set of tokens 702, including multiple groups (represented as different paths through the tree), can be speculatively generated by the draft model and provided to the target model for validation, as discussed above. When the first set of tokens 702 is validated by the target model, the draft model can speculatively generate a second set of tokens 704, in which different subsets of tokens in the second set of tokens 704 are generated based on an assumption that different sets of tokens from the first set of tokens 702 are accepted by the target model. In the example 700, a selected set of tokens 706 can correspond to the set of tokens validated by the target model, while the draft model generates subsequent sets of tokens. Generally, as discussed above, the tree data structure can be pruned based on validation of various speculatively generated groups of tokens by the target model, such that computing resources are not wasted on speculatively generating further groups of tokens based on previous sets of tokens that are not accepted by the target model.

[0094] Similar to the example 600 illustrated in Figure 6 by the draft model 104, the target model 106 can be executed on a server (e.g., in a cloud computing environment) and can be configured to validate the groups of tokens generated by the draft model 104. For example, the target model 106 can be configured to validate the groups of tokens generated by the draft model 104 based on a probability distribution associated with the target model 106. In some examples, the target model 106 can be configured to validate the groups of tokens generated by the draft model 104 based on a probability distribution associated with the target model 106 and a probability distribution associated with the draft model 104. In some examples, the target model 106 can be configured to validate the groups of tokens generated by the draft model 104 based on a probability distribution associated with the target model 106 and a probability distribution associated with the target model 106 and the draft model 104. Figure 7The overall token generation rate in the hybrid group speculation decoding illustrated by example 700 in FIG. 7 can correspond to the rate at which the draft model generates tokens. However, when the target model determines that no speculation-generated token in the generated tree is a candidate included in response to the query, the previously generated tokens can be discarded because the draft model can restart speculation token generation using the tokens (up to the last validated token) as input into the draft model. As with the per-token based hybrid speculation decoding, in this example, there can be an inverse relationship between inference accuracy and speed because smaller draft models can generate tokens at a higher rate, but the target model is less likely to accept the tokens that were speculatively generated by the draft model.

[0095] Example operations for speculation decoding in generative artificial intelligence models

[0096] Figure 8 Example operations 800 that can be performed by a computing device to generate a response to an input query using a generative artificial intelligence model are illustrated in accordance with aspects of the present disclosure. The computing device to perform operations 800 can be a device on which at least a draft model can be deployed, such as a smartphone, tablet computer, laptop computer, desktop computer, server, cloud computing instance hosted in a distributed computing environment, and the like.

[0097] As illustrated, operations 800 begin at block 810, where the computing device generates a plurality of sets of tokens based on an input query and a first generative model. In general, each set of tokens in the plurality of sets of tokens can correspond to a candidate answer to the input query. For example, an input query can be received from a user of the computing device at a prompt at which a text query can be input, such as by converting a natural language utterance into a textual representation of the input query.

[0098] At block 820, operations 800 continue, where the computing device outputs the plurality of sets of tokens to a second generative model for validation (e.g., by another computing device).

[0099] At block 830, operations 800 continue, where the computing device receives an indication of a set of tokens selected from the plurality of sets of tokens from the second generative model based on the input query and the plurality of sets of tokens.

[0100] At block 840, operations 800 continue, where the computing device outputs the selected set of tokens as a response to the input query (e.g., to a user or application, to another model, and the like).

[0101] In some aspects, each set of tokens in the plurality of sets of tokens includes a group of tokens that have a highest probability within a probability distribution associated with the first generative model on a universe of tokens on which the first generative model is based.

[0102] In some aspects, each token set of the plurality of token sets includes a token group that is selected based on a sum of probabilities associated with tokens in the token group that exceeds a threshold probability.

[0103] In some aspects, the plurality of token sets is represented as a tree data structure. Within the tree data structure, a root node can generally correspond to the input query, and different paths through the tree data structure can generally correspond to different token sets of the plurality of token sets. In some aspects, a depth of the tree data structure can correspond to a maximum number of tokens generated by a single pass through the first generative model, and a maximum size of the tree data structure can be set based on a computational complexity metric associated with generating a target token set by the second generative model.

[0104] In some aspects, operations 800 can further include pruning the tree data structure based on the selected token set. A subsequent plurality of token sets is generated based on the pruned tree data structure and the input query and output to the second generative model for validation. An indication of a subsequent token set selected from the subsequent plurality of token sets is received from the second generative model based on the input query, the pruned tree data structure, and the subsequent plurality of token sets. Outputting the selected token set as a response to the input query can include outputting the selected token set and the selected subsequent token set as a response to the input query.

[0105] In some aspects, each respective token set of the plurality of token sets can be generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

[0106] In some aspects, operations 800 can further include generating a subsequent plurality of token sets based on the input query and the plurality of token sets when the second generative model validates the plurality of token sets. Based on the subsequent plurality of token sets and the selected token set, a refined subsequent token set can be generated and output to the second generative model for validation.

[0107] In some aspects, a token set of the subsequent plurality of token sets can include padding to account for a number of tokens in the selected token set that is less than a maximum number of tokens (or a preconfigured number of tokens).

[0108] In some aspects, operations 800 further include receiving a token generated by the second generative model based on the selected token set. The received token can be output as an additional token after the selected token set.

[0109] In some aspects, the first generative model can correspond to a draft model in a speculative decoding pipeline, and the second generative model can correspond to a target model in the speculative decoding pipeline.

[0110] In some aspects, the first generative model and the second generative model can have equivalent probability distributions. In some aspects, the first generative model can have a probability distribution that approximates a probability distribution associated with the second generative model. The approximation of the probability distribution can be a probability distribution that falls within a threshold difference of the probability distribution associated with the second generative model, such that the probability distribution associated with the first generative model need not be an exact match to the probability distribution associated with the second generative model. In some aspects, the first generative model can be executed locally (e.g., on a local device or a local system), and the second generative model can be a model that is hosted remotely (e.g., on a remote system or a remote device).

[0111] Figure 9 Example operations 900 for verifying responses to input queries using generative artificial intelligence models are illustrated in accordance with aspects of the present disclosure. Operations 900 can be performed by a device (or a system having multiple devices) on which at least a target model can be deployed, such as a server (e.g., server 405 of FIG. 4), a cloud computing instance hosted in a distributed computing environment, and / or the like. Figure 4

[0112] As illustrated, operations 900 begin, at block 910, by receiving, from a device on which a first generative model runs, an input query and a plurality of sets of tokens, each set of tokens in the plurality of sets of tokens corresponding to a candidate response to the input query.

[0113] At block 920, operations 900 continue by comparing a probability distribution associated with each respective set of tokens in the plurality of sets of tokens to a corresponding probability distribution generated by a second generative model for the respective set of tokens.

[0114] At block 930, operations 900 continue by selecting, based on the comparison, a set of tokens from the plurality of sets of tokens.

[0115] At block 940, operations 900 continue by outputting, to the first generative model, an indication of the selected set of tokens.

[0116] ​In some aspects, the plurality of token sets is represented as a tree data structure. Within the tree data structure, a root node can generally correspond to the input query, and different paths through the tree data structure can generally correspond to different token sets of the plurality of token sets. In some aspects, a depth of the tree data structure can correspond to a maximum number of tokens generated by a single pass through the first generative model, and a maximum size of the tree data structure can be set based on a computational complexity metric associated with generating a target token set by the second generative model.

[0117] In some aspects, comparing the probability distribution associated with each respective token set of the plurality of token sets to a corresponding probability distribution generated by the second generative model for the respective token set includes generating the probability distribution for each respective token set based on a single pass through the second generative model. The single pass through the second generative model can be performed based on masked self-attention and positional encoding in the tree data structure.

[0118] In some aspects, the operations 900 further include generating an additional token based on the selected token set using the second generative model. The additional token (or an indication thereof) can be output to the first generative model in conjunction with (or as part of an indication of) the selected token set.

[0119] In some aspects, the first generative model can correspond to a draft model in a speculative decoding pipeline, and the second generative model can correspond to a target model in the speculative decoding pipeline.

[0120] Figure 10 Example operations 1000 that can be performed by a computing device to speculatively generate a response to an input query using generative artificial intelligence models and validate the response to the input query in accordance with aspects of the present disclosure are illustrated. The computing device to perform operations 800 can be a device on which at least a draft model can be deployed, such as a smartphone, tablet computer, laptop computer, desktop computer, server, cloud computing instance hosted in a distributed computing environment, and the like.

[0121] As illustrated, operations 1000 begin at block 1010 with the computing device generating a first plurality of token sets based on an input query and a first generative model. Generally, each token set of the first plurality of token sets can correspond to a candidate answer to the input query. For example, the input query can be received from a user of the computing device at a prompt at which an input text query can be entered, e.g., by converting a natural language utterance into a textual representation of the input query.

[0122] At block 1020, operations 1000 continue with the computing device outputting the first plurality of token sets to a second generative model for validation (e.g., by another computing device).

[0123] At block 1030, the operations 1000 continue with the computing device speculatively generating a second plurality of sets of tokens while waiting to receive an indication from the second generative model of a set of tokens selected from the first plurality of sets of tokens. Each set of tokens in the second plurality of sets of tokens generally corresponds to a second portion of a candidate response to the input query.

[0124] At block 1040, the operations 1000 continue to receive an indication from the second generative model of a set of tokens selected from the plurality of sets of tokens.

[0125] At block 1050, the operations 1000 continue to output the tokens from the second plurality of sets of tokens associated with the selected set of tokens to the second generative model for verification.

[0126] At block 1060, the operations 1000 continue to output the selected set of tokens as a response to the input query (e.g., to a user or application, to another model, etc.).

[0127] In some aspects, the operations 1000 further include receiving an indication of a second set of tokens selected from the second plurality of sets of tokens that is associated with the selected set of tokens. The selected second set of tokens can be output (e.g., to a user or application, to another model, etc.) as another portion of a response to the input query.

[0128] In some aspects, each set of tokens in the first plurality of sets of tokens includes a group of tokens that has a highest probability within a probability distribution associated with the first generative model over a universe of tokens.

[0129] In some aspects, each set of tokens in the first plurality of sets of tokens includes a group of tokens that is selected based on a sum of probabilities associated with tokens in the group of tokens that exceeds a threshold probability.

[0130] In some aspects, the first plurality of sets of tokens can be represented as a tree data structure. Within the tree data structure, a root node can correspond to the input query. Each path through the tree data structure can correspond to a set of tokens from the first plurality of sets of tokens.

[0131] In some aspects, each respective set of tokens in the first plurality of sets of tokens is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

[0132] In some aspects, the operations 1000 further include generating a refined subsequent token set based on the selected token set and the second plurality of token sets. The refined subsequent token set for validation can be output to the second generative model. While waiting to receive an indication of the selected second token set from the refined subsequent token set, a third plurality of token sets can be speculatively generated. In some aspects, the token sets of the subsequent plurality of token sets include padding to account for a number of tokens in the selected token set that is less than the maximum number of tokens.

[0133] In some aspects, the first generative model can correspond to a draft model in a speculative decoding pipeline. The second generative model can correspond to a target model in the speculative decoding pipeline.

[0134] In some aspects, the first generative model can have a probability distribution that approximates a probability distribution associated with the second generative model. The approximation of the probability distribution can be a probability distribution that falls within a threshold difference of the probability distribution associated with the second generative model, such that the probability distribution associated with the first generative model need not be an exact match to the probability distribution associated with the second generative model.

[0135] In some aspects, the first generative model can be executed locally (e.g., on a local device or a local system), and the second generative model can be a model that is hosted remotely (e.g., on a remote system or a remote device).

[0136] Example processing system for speculative decoding in generative artificial intelligence models

[0137] Figure 11 An example processing system 1100 is depicted for generating responses to queries input into a generative artificial intelligence model based on group speculative decoding, such as described herein, e.g., with respect to Figures 8 to 10 as described herein.

[0138] The processing system 1100 includes a central processing unit (CPU) 1102, which in some examples can be a multi-core CPU. Instructions executed at the CPU 1102 can be loaded, e.g., from a program memory associated with the CPU 1102 or from a memory partition (e.g., of the memory 1124).

[0139] The processing system 1100 also includes additional processing components customized for particular functions, such as a graphics processing unit (GPU) 1104, a digital signal processor (DSP) 1106, a neural processing unit (NPU) 1108, and a connectivity component 1112.

[0140] An NPU, such as NPU 1108, is generally a specialized circuit configured to implement control and arithmetic logic for executing machine learning algorithms, such as algorithms for processing artificial neural networks (ANN), deep neural networks (DNN), random forests (RF), and the like. An NPU can sometimes alternatively be referred to as a neural signal processor (NSP), a tensor processing unit (TPU), a neural network processor (NNP), an intelligent processing unit (IPU), a vision processing unit (VPU), or a graphics processing unit.

[0141] An NPU, such as NPU 1108, is configured to accelerate the execution of common machine learning tasks, such as image classification, machine translation, object detection, and various other predictive models. In some examples, multiple NPUs can be instantiated on a single chip, such as a system on a chip (SoC), while in other examples such NPUs can be part of a specialized neural network accelerator.

[0142] An NPU can be optimized for training or inference, or in some cases configured to balance performance between the two. For NPUs capable of performing both training and inference, the two tasks can generally still be performed independently.

[0143] An NPU designed to accelerate training is generally configured to accelerate the optimization of a new model, which is a highly computation-intensive operation involving inputting an existing dataset (generally labeled or tagged), iterating over the dataset, and then adjusting model parameters such as weights and biases in order to improve model performance. Generally, optimization based on error prediction involves passing back through layers of the model and determining gradients to reduce prediction error.

[0144] An NPU designed to accelerate inference is generally configured to operate on a complete model. Such an NPU can thereby be configured to input a new piece of data and quickly process the piece through a model that has already been trained to generate a model output (e.g., an inference).

[0145] In some implementations, NPU 1108 is part of one or more of CPU 1102, GPU 1104, and / or DSP 1106. These can be located on a user equipment (UE) or another computing device in a wireless communication system.

[0146] In some examples, connectivity component 1112 can include subcomponents such as one or more of third generation (3G) connectivity 1112a, fourth generation (4G) connectivity 1112b, fifth generation (5G) connectivity 1112c, Wi-Fi connectivity 1112d, Bluetooth connectivity 1112e, and other wireless data

[0147] The processing system 1100 can also include one or more sensor processing units 1116 associated with any manner of sensors, one or more image signal processors (ISPs) 1118 associated with any manner of image sensors, and / or a navigation processor 1120, which can include satellite-based positioning system components (e.g., GPS or GLONASS) as well as inertial positioning system components.

[0148] The processing system 1100 can also include one or more input and / or output devices 1122, such as a screen, a touch-sensitive surface (including a touch-sensitive display), physical buttons, a speaker, a microphone, etc.

[0149] In some examples, one or more processors of the processing system 1100 can be based on an ARM or RISC-V instruction set.

[0150] The processing system 1100 also includes memory 1124, which represents one or more static and / or dynamic memories, such as dynamic random access memory, static memories based on flash, etc. In this example, the memory 1124 includes computer-executable components that can be executed by one or more of the aforementioned processors of the processing system 1100.

[0151] In particular, in this example, the memory 1124 includes a token set generation component 1124A, a token selection reception component 1124B, a token selection component 1026C, and an output generation component 1124D, a generative model 1124E, and can optionally include a comparison component 1124F. The depicted components, as well as other, non-depicted components, can be configured to perform various aspects of the methods described herein.

[0152] In general, the processing system 1100 and / or its components can be configured to perform the methods described herein.

[0153] Example Clauses

[0154] Implementation details of various aspects of the present disclosure are described in the following numbered clauses:

[0155] Clause 1: A processor-implemented method comprising: generating, based on an input query and a first generative model, a plurality of token sets, each token set of the plurality of token sets corresponding to a candidate response to the input query; outputting the plurality of token sets to a second generative model for validation; receiving, based on the input query and the plurality of token sets, an indication of a token set selected from the plurality of token sets from the second generative model; and outputting the selected token set as a response to the input query.

[0156] Clause 2: The method of clause 1, wherein each token set of the plurality of token sets includes a group of tokens having a highest probability over a probability distribution associated with the first generative model.

[0157] Clause 3: The method of clause 1 or 2, wherein each token set of the plurality of token sets includes a group of tokens selected based on a sum of probabilities associated with tokens in the group of tokens that exceeds a threshold probability.

[0158] Clause 4: The method of any of clauses 1-3, wherein: the plurality of token sets is represented as a tree data structure, a root node of the tree data structure corresponding to the input query, and each path through the tree data structure corresponding to a token set from the plurality of token sets.

[0159] Clause 5: The method of clause 4, wherein a depth of the tree data structure corresponds to a maximum number of tokens generated by a single pass through the first generative model.

[0160] Clause 6: The method of clause 4 or 5, wherein a maximum size of the tree data structure is set based on a computational complexity metric associated with generating a target token set by the second generative model.

[0161] Clause 7: The method of any of clauses 4-6, further comprising: pruning the tree data structure based on the selected token set; generating a subsequent plurality of token sets based on the pruned tree data structure and the input query; outputting the subsequent plurality of token sets to the second generative model for validation; and receiving, from the second generative model based on the input query, the pruned tree data structure, and the subsequent plurality of token sets, an indication of a subsequent token set selected from the subsequent plurality of token sets, wherein outputting the selected token set as a response to the input query comprises: outputting the selected token set and the selected subsequent token set as a response to the input query.

[0162] Clause 8: The method of any of clauses 1-7, wherein each respective token set of the plurality of token sets is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

[0163] Clause 9: The method of any of clauses 1-8, further comprising: generating, based on the input query and the plurality of token sets, a subsequent plurality of token sets when the second generative model validates the plurality of token sets; generating a refined subsequent token set based on the subsequent plurality of token sets and the selected token set; and outputting the refined subsequent token set to the second generative model for validation.

[0164] Clause 10: The method of clause 9, wherein the token set of the subsequent plurality of token sets includes padding to account for a number of tokens in the selected token set that is less than a maximum number of tokens.

[0165] Clause 11: The method of any of clauses 1-10, further comprising: receiving a token generated by the second generative model based on the selected token set; and outputting the received token as an additional token after the selected token set.

[0166] Clause 12: The method of any of clauses 1-11, wherein: the first generative model corresponds to a draft model in a speculative decoding pipeline, and the second generative model corresponds to a target model in the speculative decoding pipeline.

[0167] Clause 13: The method of clause 12, wherein the draft model comprises a model trained to have a probability distribution that approximates a corresponding probability distribution of the target model.

[0168] Clause 14: The method of any of clauses 1-13, wherein: the first generative model comprises a model executing on a local system, and the second generative model comprises a model executing on a remote system.

[0169] Clause 15: A processor-implemented method comprising: receiving an input query and a plurality of token sets generated by a first generative model, each respective token set of the plurality of token sets corresponding to a respective candidate response to the input query; comparing a probability distribution associated with each respective token set of the plurality of token sets to a corresponding probability distribution generated by a second generative model for the respective token set; selecting a token set from the plurality of token sets based on the comparison; and outputting, to the first generative model, an indication of the selected token set.

[0170] Clause 16: The method of clause 15, wherein: the plurality of token sets are represented as a tree data structure, a root node of the tree data structure corresponding to the input query, and each path through the tree data structure corresponding to a token set from the plurality of token sets.

[0171] Clause 17: The method of clause 15 or 16, wherein comparing a probability distribution associated with each respective token set of the plurality of token sets to a corresponding probability distribution generated by a second generative model for the respective token set comprises: generating the probability distribution for each respective token set based on a single pass through the second generative model.

[0172] Clause 18: The method of clause 17, wherein the single pass through the second generative model is performed based on masked self-attention and positional encoding in the tree data structure.

[0173] Clause 19: The method of any of clauses 15-18, further comprising: generating, using a second generative model, an additional token based on the selected set of tokens; and outputting the additional token to the first generative model.

[0174] Clause 20: The method of any of clauses 15-19, wherein: the first generative model corresponds to a draft model in a speculative decoding pipeline, and the second generative model corresponds to a target model in the speculative decoding pipeline.

[0175] Clause 21: A processor-implemented method comprising: generating, based on an input query and a first generative 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 query; outputting the plurality of sets of tokens to a second generative model for validation; speculatively generating, while waiting to receive, from the second generative model, an indication of a selected set of tokens from the first plurality of sets of tokens, 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 a candidate response to the input query; receiving, from the second generative model, the indication of the selected set of tokens from the first plurality of sets of tokens; outputting, to the second generative model for validation, a token from the second plurality of sets of tokens that is associated with the selected set of tokens; and outputting the selected set of tokens as a response to the input query.

[0176] Clause 22: The method of clause 21, further comprising: receiving an indication of a second set of tokens from the second plurality of sets of tokens that is selected to be associated with the selected set of tokens; and outputting the selected second set of tokens as another portion of the response to the input query.

[0177] Clause 23: The method of clause 21 or 22, wherein each set of tokens in the first plurality of sets of tokens comprises a group of tokens that has a highest probability, over a probability distribution associated with the first generative model, on a universe of tokens.

[0178] Clause 24: The method of any of clauses 21-23, wherein each set of tokens in the first plurality of sets of tokens comprises a group of tokens that is selected based on a sum of probabilities associated with tokens in the group of tokens, the sum exceeding a threshold probability.

[0179] Clause 25: The method of any of clauses 21-24, wherein: the first plurality of sets of tokens is represented as a tree data structure, a root node of the tree data structure corresponding to the input query, and each path through the tree data structure corresponding to a set of tokens from the first plurality of sets of tokens.

[0180] Clause 26: The method of any of clauses 21-25, wherein each respective token set of the first plurality of token sets is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

[0181] Clause 27: The method of any of clauses 21-26, further comprising: generating a refined subsequent token set based on the selected token set and the second plurality of token sets; and outputting the refined subsequent token set to the second generative model for validation; and speculatively generating a third plurality of token sets while waiting to receive from the second generative model an indication of a second token set selected from the refined subsequent token set.

[0182] Clause 28: The method of clause 27, wherein the token sets of the subsequent plurality of token sets include padding to account for a number of tokens in the selected token set that is less than a maximum number of tokens.

[0183] Clause 29: The method of any of clauses 21-28, wherein: the first generative model corresponds to a draft model in a speculative decoding pipeline, and the second generative model corresponds to a target model in the speculative decoding pipeline.

[0184] Clause 30: The method of clause 29, wherein the draft model comprises a model trained to have a probability distribution that approximates a corresponding probability distribution of the target model.

[0185] Clause 31: The method of any of clauses 21-30, wherein: the first generative model comprises a model executing on a local system, and the second generative model comprises a model executing on a remote system.

[0186] Clause 32: A processing system comprising: a memory having stored thereon executable instructions; and one or more processors configured to execute the instructions to cause the processing system to perform the method of any of clauses 1-31.

[0187] Clause 33: A processing system comprising means for performing the method of any of clauses 1-31.

[0188] Clause 34: A computer-readable medium having stored thereon instructions that, when executed by one or more processors, cause the one or more processors to perform the method of any of clauses 1-31.

[0189] Additional Considerations

[0190] The previous description is provided to enable any person skilled in the art to practice the various aspects described herein. The examples discussed herein are not meant to limit the scope of the claims encompassing the disclosure, which is defined solely by the claims. Various modifications to these aspects will be readily apparent to those skilled in the art, and the generic principles defined herein can be applied to other aspects. For example, features described in the context of one example can be combined or eliminated in other examples. Numerous other changes can be made hereunder without departing from the scope of the disclosure. For example, various sequences, steps, or stages can be performed in an order different from that described. Various steps can be combined or omitted, and other steps can be added. Also, features described as part of one example can be combined in or excluded from another example. For example, any of the aspects described herein can be implemented using any number of the aspects described herein. In addition, the scope of the disclosure is intended to cover any resulting device or method encompassed by the claims, whether or not it is now claimed. It is understood that any aspect of the disclosure disclosed herein can be implemented by one or more elements of a claim.

[0191] As used herein, the term "exemplary" means "serving as an example, instance, or illustration." Any aspect described herein as "exemplary" is not necessarily to be construed as preferred or advantageous over other aspects.

[0192] As used herein, the phrase "at least one of" followed by a listing of two or more items means any of the listed items individually, or any combination of two or more of the listed items. As an example, "at least one of a, b, or c" means "a" or "b" or "c" or "a-b" or "a-c" or "b-c" or "a-b-c" or any combination thereof.

[0193] As used herein, the term "determining" encompasses a wide variety of actions. For example, "determining" can include calculating, computing, processing, deriving, investigating, looking up (e.g., looking up in a table, a database or another data structure), ascertaining and the like. Also, "determining" can include receiving (e.g., receiving information), accessing (e.g., accessing data in a memory) and the like. Also, "determining" can include resolving, selecting, choosing, establishing and the like.

[0194] The methods disclosed herein include one or more steps or actions for achieving the methods. The method steps and / or actions can be interchanged with one another without departing from the scope of the claims. In other words, unless a specific order of steps or actions is specified, the order and / or use of specific steps and / or actions can be modified without departing from the scope of the claims. Further, the various operations of methods described above can be performed by any suitable means capable of performing the corresponding functions. These means can include various hardware and / or software component(s) and / or module(s), including, but not limited to a circuit, an application specific integrated circuit (ASIC), or processor. Generally, where there are operations illustrated in figures, those operations can have corresponding counterpart means-plus-function components with similar numbering.

[0195] The following claims are not intended to be limited to the aspects shown herein but are to be afforded the full scope consistent with the language of the claims. Within the claims, the reference to an element in the singular is not intended to mean "one and only one" unless specifically so stated, but rather "one or more." Unless specifically stated otherwise, the term "some" refers to one or more. Any element expressed by the phrase "means for" is intended to functionally equate to the instances disclosed enabling that specific function, either by itself or in combination with other elements. All structural and functional equivalents to the aspects described throughout this disclosure that are enabled by this description are expressly incorporated by reference and intended to be encompassed by the claims. Moreover, any combination of the aspects described throughout this disclosure are expressly intended to be within the scope of the claims. Furthermore, any publication or other disclosure material presented herein that is incorporated by reference is expressly intended to be only for purposes of the claims and not to limit the aspects of the disclosure in any way.

Claims

1. A processing system, comprising: A memory, on which executable instructions are stored; as well as One or more processors, the one or more processors being configured to execute the executable instructions to cause the processing system to: A first plurality of token sets are generated based on an input query and a first generative model, each of the first plurality of token sets corresponding to a first portion of a candidate response to the input query; The multiple token sets are output to the second generative model for verification; While waiting to receive an indication from the second generative model for a set of tokens selected from the first plurality of token sets, a second plurality of token sets are speculatively generated, each of the second plurality of token sets corresponding to a second portion of the candidate response to the input query; Receive the instruction from the second generative model for the set of tokens selected from the first plurality of token sets; The tokens associated with the selected token set from the second plurality of token sets are output to the second generative model for verification; as well as The selected set of tokens is output as a response to the input query.

2. The processing system of claim 1, wherein the one or more processors are further configured to cause the processing system to: Receive an instruction for a second token set associated with the selected token set from the second plurality of token sets; and The selected second set of tokens is output as another part of the response to the input query.

3. The processing system of claim 1, wherein each of the first plurality of token sets comprises a group of tokens with the highest probability in a probability distribution associated with the first generative model on a token universe.

4. The processing system of claim 1, wherein each of the first plurality of token sets comprises a token group selected based on the sum of probabilities associated with tokens in the token group, the sum exceeding a threshold probability.

5. The processing system as claimed in claim 1, wherein: The first set of tokens is represented as a tree data structure. The root node of the tree data structure corresponds to the input query, and Each path in the tree data structure corresponds to a set of tokens from the first plurality of token sets.

6. The processing system of claim 1, wherein each of the first plurality of token sets is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

7. The processing system of claim 1, wherein the one or more processors are further configured to cause the processing system to: A refined subsequent token set is generated based on the selected token set and the second or more token sets; The refined set of subsequent tokens is output to the second generative model for verification; as well as While waiting to receive instructions from the second generative model for a second set of tokens to be selected from a refined set of subsequent tokens, a third or more token sets are generated speculatively.

8. The processing system of claim 7, wherein the token set in the subsequent plurality of token sets includes a plurality of tokens padded to account for tokens in the selected token set that are less than a maximum number.

9. The processing system as claimed in claim 1, wherein: The first generative model corresponds to the draft model in the speculative decoding pipeline, and The second generative model corresponds to the target model in the speculative decoding pipeline.

10. The processing system of claim 9, wherein the draft model comprises a model trained to have a probability distribution that approximates the corresponding probability distribution of the target model.

11. The processing system of claim 1, wherein: The first generative model includes models executed on the local system, and The second generative model includes models executed on remote systems.

12. A processor-implemented method, comprising: A first plurality of token sets are generated based on an input query and a first generative model, each of the first plurality of token sets corresponding to a first portion of a candidate response to the input query; The multiple token sets are output to the second generative model for verification; While waiting to receive an indication from the second generative model for a set of tokens selected from the first plurality of token sets, a second plurality of token sets are speculatively generated, each of the second plurality of token sets corresponding to a second portion of the candidate response to the input query; Receive an instruction from the second generative model for the set of tokens selected from the first plurality of token sets; The tokens associated with the selected token set from the second plurality of token sets are output to the second generative model for verification; as well as The selected set of tokens is output as a response to the input query.

13. The method of claim 12, further comprising: Receive an instruction for a second token set associated with the selected token set from the second plurality of token sets; as well as The selected second set of tokens is output as another part of the response to the input query.

14. The method of claim 12, wherein each of the first plurality of token sets comprises a group of tokens that have the highest probability within a probability distribution associated with the first generative model in a token universe.

15. The method of claim 12, wherein each of the first plurality of token sets comprises a token group selected based on the sum of probabilities associated with tokens in the token group, the sum exceeding a threshold probability.

16. The method of claim 12, wherein: The first set of tokens is represented as a tree data structure. The root node of the tree data structure corresponds to the input query, and Each path in the tree data structure corresponds to a set of tokens from the first plurality of token sets.

17. The method of claim 12, wherein each of the first plurality of token sets is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

18. The method of claim 12, further comprising: A refined subsequent token set is generated based on the selected token set and the second or more token sets; The refined set of subsequent tokens is output to the second generative model for verification; as well as While waiting to receive instructions from the second generative model for a second set of tokens to be selected from a refined set of subsequent tokens, a third or more token sets are generated speculatively.

19. The method of claim 18, wherein the token set in the subsequent plurality of token sets includes a plurality of tokens padded to account for tokens in the selected token set that are less than a maximum number.

20. The method of claim 12, wherein: The first generative model corresponds to the draft model in the speculative decoding pipeline, and The second generative model corresponds to the target model in the speculative decoding pipeline.

21. The method of claim 20, wherein the draft model comprises a model trained to have a probability distribution that approximates the corresponding probability distribution of the target model.

22. The method of claim 12, wherein: The first generative model includes models executed on the local system, and The second generative model includes models executed on remote systems.

23. A processing system, comprising: A means for generating a first plurality of token sets based on an input query and a first generative model, each of the first plurality of token sets corresponding to a first portion of a candidate response to the input query; A means for outputting the plurality of token sets to a second generative model for verification; A means for speculatively generating a second plurality of token sets while waiting to receive an indication from the second generative model for a token set selected from the first plurality of token sets, each of the second plurality of token sets corresponding to a second portion of the candidate response to the input query; A means for receiving from the second generative model an indication of a set of tokens selected from the first plurality of token sets; A means for outputting tokens associated with the selected token set from the second plurality of token sets to the second generative model for verification; as well as A means for outputting the selected set of tokens as a response to the input query.

24. The processing system of claim 23, further comprising: A means for receiving an indication of a second set of tokens selected from the second plurality of token sets and associated with the selected set of tokens; as well as A means for outputting the selected second set of tokens as another part of the response to the input query.

25. The processing system of claim 23, wherein: The first set of tokens is represented as a tree data structure. The root node of the tree data structure corresponds to the input query, and Each path in the tree data structure corresponds to a set of tokens from the first plurality of token sets.

26. The processing system of claim 23, wherein each of the first plurality of token sets is generated using a unique instance of the first generative model and a unique parameter as input to the unique instance of the first generative model.

27. The processing system of claim 23, further comprising: A means for generating a refined subsequent set of tokens based on a selected set of tokens and a second plurality of token sets; A means for outputting a refined set of subsequent tokens to the second generative model for verification; as well as A means for speculatively generating a third plurality of token sets while waiting to receive an instruction from the second generative model for a second token set selected from a refined subsequent token set.

28. The processing system of claim 23, wherein: The first generative model corresponds to the draft model in the speculative decoding pipeline, and The second generative model corresponds to the target model in the speculative decoding pipeline.

29. The processing system of claim 23, wherein: The first generative model includes models executed on the local system, and The second generative model includes models executed on remote systems.

30. A computer-readable medium having executable instructions stored thereon, the instructions performing operations including the following when executed by one or more processors: A first plurality of token sets are generated based on an input query and a first generative model, each of the first plurality of token sets corresponding to a first portion of a candidate response to the input query; The multiple token sets are output to the second generative model for verification; While waiting to receive an indication from the second generative model for a set of tokens selected from the first plurality of token sets, a second plurality of token sets are speculatively generated, each of the second plurality of token sets corresponding to a second portion of the candidate response to the input query; Receive an instruction from the second generative model for the set of tokens selected from the first plurality of token sets; The tokens associated with the selected token set from the second plurality of token sets are output to the second generative model for verification; as well as The selected set of tokens is output as a response to the input query.

Images