Decoding at the Speed of Thought:
Harnessing Parallel Decoding of Lexical Units for LLMs

Given a context, our model doesn’t restrict itself to predicting just the immediate next token. Instead, it ambitiously casts a look-ahead window of a fixed length $k$, aiming to predict the next $k$ tokens in one sweep. This block of tokens, denoted as $x_{t:t+k}$, is predicted as:

The probability distribution of the block $x_{t:t+k}$ is conditioned on the preceding context and the positional information of the tokens before the target token within this block. In practice, we append $k-1$ [PAD] tokens to the context input and extract the last $k$ logits to compute the probability distribution.

While the model attempts to predict $k$ tokens, not all predictions might be of high confidence. Thus, we introduce a mechanism to selectively accept tokens from this prediction. Specifically, only the first $l$ tokens that consistently maintain probabilities above a specified threshold $\beta$ are accepted as a span $S$. $l$ is estimated as follows:

where $m$ is the largest integer such that:

Parallel decoding method suffers more from the token repetition problem Gu et al. (2017) since the prefixing token is not ready for reference like the auto-regressive one. To address this issue, we implemented a straightforward method: for the decoded $k$ tokens, we check for consecutive token ids or instances where $x_{i-1}$ ends with $x_{i}$. If detected, the model immediately halts acceptance of new tokens.

A notable feature of our inference strategy is its adaptability. If the first token from the look-ahead window does not meet the confidence threshold, the model reverts to accepting only the first token. This ensures that in scenarios where the model isn’t confident about predicting a lexical unit, it equals to the conventional auto-regressive decoding strategy, ensuring robustness across diverse linguistic contexts. The inference process is illustrated in Figure 2.

Central to the methodology is the identification of “lexical units”. These are continuous sequences of tokens that captures semantically and linguistically coherent constructs. It has been observed that pre-trained language models, especially when fine-tuned on the target dataset, tend to predict tokens within these units with remarkable confidence.

To harness this observation, we establish an identification criterion. Specifically, a continuous span of tokens is deemed a lexical unit if the prediction probability of each token within this span surpasses a predefined threshold $\alpha$. This threshold serves as a confidence measure, ensuring that the identified spans actually represent coherent linguistic constructs.

Formally, given a span of tokens $x_{t:t+l}$, we define it as a lexical unit if the following constraints are satisfied:

1. 1.

   $P(x_{i})\geq\alpha$ for all $i\in[t,t+l]$, where $P(x_{i})$ is the prediction probability of the $i$-th token and $\alpha$ is the predefined threshold.

2. 2.

   $P(x_{t-1})<\alpha$ and $P(x_{t+l+1})<\alpha$, ensuring that the tokens immediately before and after the sequence have prediction probabilities lower than $\alpha$.

To facilitate the model to recognize these lexical units, we first fine-tune the original model $M$ on the target dataset $D$. This auto-regressive training refines the model into $M_{\mathrm{FT}}$. During the forward pass with $M_{\mathrm{FT}}$ on $D$, by comparing probability of the ground truth label against the threshold $\alpha$, we can effectively identify spans of tokens that the model predicts with high confidence. These high-confidence spans are then designated as lexical units. This process is elaborated in Algorithm 1.

Upon identifying the lexical units, we can proceed with data reconfiguration for continual training. First, as shown in Figure 3, given a single piece of data, multiple lexical units can be identified. Note that lexical units can contain multiple tokens or only a single token with coherent semantic meaning.

For every detected lexical unit, a new training instance is instantiated by replacing the tokens inside with trainable [PAD] tokens. This utilization of the trainable [PAD] token over other potential masking strategies ensures the integrity of positional information, which is critical for the model’s functioning. Further more, for the new training instance, the context tokens before it are left unchanged to provide complete context information, and the loss is only calculated for the tokens within the lexical unit. In this way, we can guarantee that the loss of each token is still calculated only once even though a single sequence is split into multiple ones.

The streamlined process of data generation, from lexical unit identification to data reconfiguration, is demonstrated in Figure 3. A more detailed illustration can be found in Algorithm 1. Note that the generated dataset $\bar{D}$ is mixed with the original dataset $D$ as the final reconfigured dataset $D^{\prime}$.

Tokens that are not encapsulated within any lexical units are trained in the traditional language modeling manner. In this scenario, the model predicts the next token in the sequence based on the preceding tokens, adhering to the standard auto-regressive nature of LLMs.

When it comes to Lexical Unit Tokens, the prediction mechanism deviates slightly. The model is trained to predict these tokens by considering not just the complete information from the tokens preceding the lexical unit, but also the positional information, which is actually embedded in the trainable [PAD] tokens, within the lexical unit. Note that unlike setting a look-ahead window with a fixed length $k$ during inference, the length of the identified lexical units can be varied during training. This is formulated in Equation (3).

where $x_{t}$ denotes the Lexical Unit Tokens to be predicted and $i$ denotes the initial position of the lexical unit to which $x_{t}$ pertains. Similar to inference, its prediction is conditioned on Context Tokens $x_{1},x_{2},...,x_{i}$ preceding the lexical unit, and the positional information $\mathrm{pos}(x_{i:t-1})$ of the Prefixing Tokens within the lexical unit. This formulation underscores that, although the training procedure remains unchanged, the underlying training dynamics differs based on a token’s association with a lexical unit.

As mentioned before, in practice, token ids of the Lexical Unit Tokens are replaced with pad_token_id, and labels for Context Tokens are set to $-100$ to exclude them from loss calculation. Note that for Lexical Unit Tokens, the model is trained to predict the original tokens rather than [PAD] token. [PAD] tokens are to provide the positional context.

The Forward Compression Ratio (FCR) quantifies the efficiency gains achieved by our Lexical Unit Decoding (LUD) method. In a conventional auto-regressive setting, each token generated necessitates a forward calculation, making the number of tokens generated equal to the number of forward calculations. However, LUD’s capability to decode multiple tokens in a single forward pass introduces a disparity between these two numbers. FCR is defined as:

A higher FCR value signifies greater computational efficiency improvement, as it indicates that fewer forward calculations are required to produce an equivalent number of tokens.

While the FCR offers a theoretical perspective on efficiency, the Wall-time Acceleration Ratio (WAR) provides a more pragmatic view. It measures the acceleration in terms of actual computation time, factoring in real-world considerations including hardware efficiency, potential parallelization, and other computational overheads. WAR is given by:

where, $t_{\mathrm{AT}}$ represents the average time taken to generate a single token using the auto-regressive approach, while $t_{\mathrm{LUD}}$ denotes the corresponding time for the LUD method. In practice, we measure the time consumption of solely the generation loop excluding data loading while setting batch size to $1$. The total time consumed is then divided by the number of generated tokens to get a precise estimate of the average time per token.

The efficiency of the Lexical Unit Decoding (LUD) approach is evident when examining both the FCR and WAR metrics. As we decrease the parameter $\beta$, we observe a consistent increase in both metrics. This indicates that the model is more aggressive in accepting multiple tokens at once when the confidence threshold is lower, leading to higher acceleration. However, this efficiency achieved by a smaller $\beta$ value can cause quality degeneration.

The quality of the generated content, both for text and code, shows a declining trend as we push for more acceleration by decreasing $\beta$. This degradation in quality underscores the trade-off between acceleration and quality. As we push the model to be more aggressive in its predictions, the chances of making errors increase, leading to a drop in the quality of the generated content. We do a deeper dive into the generation process in Section 5.

It’s apparent there’s some trade-off between the acceleration ratio and quality loss when varying $\beta$. However, we did notice that the generation quality can be maintained when $\beta$ is above a specific value, within which decreasing $\beta$ can lead to faster decoding. This means there is a “sweet pot” of $\beta$ that achieves fastest lossless decoding. While both text and code data reconfiguration is performed with $\alpha$ set to $0.85$, the best $\beta$ values are $0.9$ and $0.99987$ respectively, which doesn’t align with $\alpha$ strictly.

Another clear observation is with the same $\beta$ value, codes generation can be accelerated faster than text generation. The disparity between code and text generation is intriguing. One possible explanation is code sequences are viewed as a kind of low-entropy sequences while text sequences have much higher entropy Kirchenbauer et al. (2023). Codes often follow specific patterns and structures, making them more predictable. This predictability might allow the model to confidently generate larger chunks of code tokens at once, leading to faster decoding acceleration. This observation validates that our method can accelerate the decoding process adaptively. Note that models for natural language generation and code generation are fine-tuned only with in-domain dataset. However, we argue that using a more balanced dataset, a sweet pot can still be achieved for a general large language model.

However, the quality of code generation degrades much faster than text as $\beta$ decreases. This could be due to the fact that even minor errors in code can render it non-functional, whereas text might still be understandable even with minor inaccuracies. Noticeably, code decoding is also much more sensitive to $\beta$. We varied the beta value carefully and found that $0.99987$ can accelerate the decoding speed by 30% with 3% quality loss. We suspect that the code’s attribute of being more predictable makes the model confident on its decoding. But further experiments should be conducted to understand what’s happening inside the model. We leave it to our future work.