Speculative Decoding with CTC-based Draft Model for LLM Inference Acceleration

Inference acceleration of large language models (LLMs) has been put forward in many application scenarios and speculative decoding has shown its advantage in addressing inference acceleration. Speculative decoding usually introduces a draft model to assist the base LLM where the draft model produces drafts and the base LLM verifies the draft for acceptance or rejection. In this framework, the final inference speed is decided by the decoding speed of the draft model and the acceptance rate of the draft provided by the draft model. Currently the widely used draft models usually generate draft tokens for the next several positions in a non-autoregressive way without considering the correlations between draft tokens. Therefore, it has a high decoding speed but an unsatisfactory acceptance rate. In this paper, we focus on how to improve the performance of the draft model and aim to accelerate inference via a high acceptance rate. To this end, we propose a CTC-based draft model which strengthens the correlations between draft tokens during the draft phase, thereby generating higher-quality draft candidate sequences. Experiment results show that compared to strong baselines, the proposed method can achieve a higher acceptance rate and hence a faster inference speed.

Paper

Similar papers

Peer review

Reviewer NRvU6/10 · confidence 3/52024-07-11

Summary

The paper proposes a novel architecture and training technique for LLM speculative decoding, aiming to improve the reliability and acceptance rate of generation candidates. Unlike Medusa, the proposed method replaces the draft module with a Transformer and utilizes CTC-based loss instead of CE loss. For CTC training, pseudo-labels generated from the base model, rather than ground-truth tokens, are used. During inference, the draft head generates probability candidates, and the CTC beam search process produces the final candidates for evaluation. Experimental results show that the CTC-drafter achieves higher speedup due to an increased acceptance rate and a greater number of accepted tokens.

Strengths

* Combining CTC loss with sequence prediction in text-only LLMs is an interesting approach that could inspire further research. * The results suggest that CTC-drafter is promising in terms of inference acceleration without significant overhead.

Weaknesses

* CTC is known to suffer from the conditional independence problem, i.e., each token's prediction does not depend on other tokens. Given this, it is unclear how CTC can address the low acceptance rate. The Hydra paper emphasized the importance of token-level sequential dependency, but CTC loss does not seem to support this due to its conditional independence. This is a critical point for the paper’s motivation, so a thorough justification of CTC loss is necessary. * Is the CTC output decoding during inference the same as the well-known CTC prefix beam search? If so, it should be clarified that this is known as prefix search decoding. * There is no ablation study on (Transformer Layer + CE loss) or (Linear layer + CTC loss). It is unclear which changes contribute more to performance improvements.

Questions

* How many beams are used for CTC decoding? Is the number of beams (or candidates) the same as in Medusa, ensuring a fair comparison? * The architecture of the Attention Draft Model is somewhat difficult to understand. Does it take a single hidden embedding vector as input? Is this vector expanded (duplicated? or repeated?) to create an input sequence for the Transformer-based draft model? What is the input of the Transformer-based draft model, and are positional encodings inserted at the beginning? * The explanation of the CTC-related part could be improved. For example, it would be helpful to mention that different candidates can be expanded to multiple "alignments" before the CTC blank-collapse.

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The paper adequately addresses the limitations.

Authorsrebuttal2024-08-12

Dear Reviewers, We are writing to kindly request your feedback. As the author-reviewer discussion phase is nearing its end, we are eagerly awaiting your response to address any questions or concerns you may have. We apologize for any inconvenience and appreciate your time and efforts.

Reviewer NRvU2024-08-12

Thank you

Thank you for further clarification. These detailed explanations make more sense than the information included in the original manuscript. In fact, this is quite "new" information to me (ex: 4 previous tokens, 2x copying, etc.) and possibly to other reviewers. I increased the soundness score from 2 to 3, and the overall score from 5 to 6. By the way, I wonder if that information was given in the original version. We might have asked different questions about architectural details. I note this to let the meta reviewer decide.

Authorsrebuttal2024-08-12

Many thanks for your time and valuable suggestions to help improve our paper. We will follow your suggestions to include necessary clarifications in the revised version.

Reviewer 6n7r6/10 · confidence 4/52024-07-14

Summary

The authors study the setup of speculative decoding where multiple tokens are being generated at parallel. In this setup the authors use the idea of connectionist temporal classification to train a draft model which generates multiple tokens in parallel.

Strengths

- The method shows significant improvement in token acceptance rates. - The speedups reported are quite impressive.

Weaknesses

- The speedups are unclear, it is calculated using the token acceptance rate rather than on a real system - Details are not clear about at what layers the authors are taking the intermediate representation - The details of training overhead are unclear - It will be great if authors can compare to - Draft & Verify: Lossless Large Language Model Acceleration via Self-Speculative Decoding

Questions

See weakness section

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have listed some weaknesses however it is not clear the cases where the idea will fail. Specifically how dependent is data drift between training the draft model and inference.

Authorsrebuttal2024-08-11

Supplementary Explanations by Authors

Thanks for your valuable suggestions. We reorganized the description of our method, hoping it can help understand our method better. Overall, the architecture of our method can be decomposed into three parts: the base model, the drafter which includes the Attention Draft module and LM head, and the CTC-related module which is used to calculate CTC-loss for training and performs CTC-style decoding at inference. The main structures of the three parts are as follows. - The base model in our method is based on the autoregressive generation framework like Llama which only uses Transformer decoder as the main structure. - For the drafter, the Attention Draft module employs one Transformer layer as its structure including masked multi-head self-attention sublayer and Feed Forward sublayer. Compared to Medusa which employs FFN as the structure of each Medusa head without interacting between Medusa heads, our method can introduce sequence correlations between the preceding generated words and next several words to generate, besides the CTC-related module. - The CTC-related module involves non-autoregressive generation of the next N tokens (N is a hyperparameter, which is set to 4 for current CTC-drafter) where blank character and repetitive tokens are introduced into the raw generated sequence which will be processed by CTC Transform module to produce the final draft for the current decoding timestep. Our method works in a different way during training and inference. During training, the base model will generate the whole sequence in an autoregressive manner which is used as the ground truth to distill the drafter. To generate the draft for the drafter, at each timestep, the Attention Draft module accepts as the input the 4 representations of the current position (assuming position i) and its preceding positions (i-3, i-2, i-1) generated by the last Transformer layer of the base model, and upsamples the input representations by 2 times via copying the input representations. Then through the Transformer layer of the Attention Draft module, the output representations are fed to LM head. LM head will generate 8 raw draft tokens in a non-autoregressive way and the final draft tokens will be produced by CTC Transform module. The CTC-loss is used to count the probabilities of all the draft sequences that can be transformed into the ground truth sequence via dynamic programming and the sum of all these sequences is maximized as the training loss. In this way, the probability distribution is drawn towards the draft sequences that can derive ground truth sequence being allocated bigger probability. This means the candidate draft sequence with more reasonable sequence correlations will be selected as the winner at a greater likelihood. At inference, the base model does not generate the whole sequence in advance. Instead, at each timestep, it is used to generate the representations fed to the Attention Draft module and verify the generated draft tokens according to the probability by performing teacher forcing decoding to generate the draft tokens. If the probability to generate the draft tokens is greater than the set threshold, the base model will accept the draft tokens and uses them as the following generated tokens. Then the base model continues to encode these generated draft tokens in an autoregressive way. If the probability is smaller than the threshold, the base model will reject the draft tokens and generate the next token on its own and meanwhile encode the generated token in an autoregressive way, too. Then the latest generated 4 representations by the last Transformer layer of the base model are used as the input of the Attention Draft module for the next timestep and the decoding goes to the next timestep. Here the Attention Draft module works in the same way as in training to generate input representations for LM head. Here LM head still generates N tokens for the next N positions in a non-autoregressive way with each position reserving top k tokens (k is a hyperparameter, which is set to 10 for current CTC-drafter), then the token sequence of the N positions with the highest probability will be selected as the raw draft via the token tree structure used in Medusa. The final draft generated via CTC transform will be fed to the base model to verify and the base model will decide to use the draft tokens as the next several output tokens or to generate the next token by itself where the former decision can generate several tokens at once and hence improves the decoding speed.

Authorsrebuttal2024-08-12

Dear Reviewers, We are writing to kindly request your feedback. As the author-reviewer discussion phase is nearing its end, we are eagerly awaiting your response to address any questions or concerns you may have. We apologize for any inconvenience and appreciate your time and efforts.

Reviewer 6n7r2024-08-12

Thank your for response

Appreciate the author response. Thank you for providing clarifications. I will bump up the score to 6.

Authorsrebuttal2024-08-12

Many thanks for your time and valuable suggestions to help improve our paper. We will follow your suggestions to include necessary clarifications and experiments in the revised version.

Reviewer EXv75/10 · confidence 3/52024-07-22

Summary

The paper proposes a novel framework, CTC-drafter, to accelerate speculative decoding in large language models (LLMs). The authors introduce the use of Connectionist Temporal Classification (CTC) as a training objective, replacing the traditional cross-entropy loss. This method aims to improve context modeling and generate adaptive candidate sequences, which purportedly enhances the speed and accuracy of the speculative decoding process. The paper demonstrates the effectiveness of CTC-drafter through experiments on various benchmarks and compares its performance with existing methods such as Medusa and Hydra.

Strengths

- Originality: Introducing CTC as a training objective for speculative decoding is a novel approach. - Quality: The theoretical framework is well-developed, with clear definitions and derivations. - Clarity: The paper is generally well-written and logically structured, making it accessible to a broad audience.

Weaknesses

- Experimental Validation: The experiments do not fully validate the claims. There is a need for more extensive testing across different datasets and model architectures to ensure the generality and robustness of the proposed method. - Comparative Analysis: While comparisons are made with Medusa and Hydra, the analysis lacks depth. More detailed insights into why - CTC-drafter performs better or worse in specific scenarios would be beneficial.

Questions

1. Can the authors provide more details on the hyperparameter settings and training configurations used in the experiments? 2. How does the performance of CTC-drafter vary with different model architectures and dataset sizes? 3. Can the authors elaborate on the computational overhead introduced by the CTC loss and how it compares to the benefits in inference speed?

Rating

5

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors have discussed several limitations, including the need for more training tricks to enhance the draft module and the uncertainty regarding the optimality of the current draft model structure.

Authorsrebuttal2024-08-11

Supplementary Explanations by Authors

Thanks for your valuable suggestions. We reorganized the description of our method, hoping it can help understand our method better. Overall, the architecture of our method can be decomposed into three parts: the base model, the drafter which includes the Attention Draft module and LM head, and the CTC-related module which is used to calculate CTC-loss for training and performs CTC-style decoding at inference. The main structures of the three parts are as follows. - The base model in our method is based on the autoregressive generation framework like Llama which only uses Transformer decoder as the main structure. - For the drafter, the Attention Draft module employs one Transformer layer as its structure including masked multi-head self-attention sublayer and Feed Forward sublayer. Compared to Medusa which employs FFN as the structure of each Medusa head without interacting between Medusa heads, our method can introduce sequence correlations between the preceding generated words and next several words to generate, besides the CTC-related module. - The CTC-related module involves non-autoregressive generation of the next N tokens (N is a hyperparameter, which is set to 4 for current CTC-drafter) where blank character and repetitive tokens are introduced into the raw generated sequence which will be processed by CTC Transform module to produce the final draft for the current decoding timestep. Our method works in a different way during training and inference. During training, the base model will generate the whole sequence in an autoregressive manner which is used as the ground truth to distill the drafter. To generate the draft for the drafter, at each timestep, the Attention Draft module accepts as the input the 4 representations of the current position (assuming position i) and its preceding positions (i-3, i-2, i-1) generated by the last Transformer layer of the base model, and upsamples the input representations by 2 times via copying the input representations. Then through the Transformer layer of the Attention Draft module, the output representations are fed to LM head. LM head will generate 8 raw draft tokens in a non-autoregressive way and the final draft tokens will be produced by CTC Transform module. The CTC-loss is used to count the probabilities of all the draft sequences that can be transformed into the ground truth sequence via dynamic programming and the sum of all these sequences is maximized as the training loss. In this way, the probability distribution is drawn towards the draft sequences that can derive ground truth sequence being allocated bigger probability. This means the candidate draft sequence with more reasonable sequence correlations will be selected as the winner at a greater likelihood. At inference, the base model does not generate the whole sequence in advance. Instead, at each timestep, it is used to generate the representations fed to the Attention Draft module and verify the generated draft tokens according to the probability by performing teacher forcing decoding to generate the draft tokens. If the probability to generate the draft tokens is greater than the set threshold, the base model will accept the draft tokens and uses them as the following generated tokens. Then the base model continues to encode these generated draft tokens in an autoregressive way. If the probability is smaller than the threshold, the base model will reject the draft tokens and generate the next token on its own and meanwhile encode the generated token in an autoregressive way, too. Then the latest generated 4 representations by the last Transformer layer of the base model are used as the input of the Attention Draft module for the next timestep and the decoding goes to the next timestep. Here the Attention Draft module works in the same way as in training to generate input representations for LM head. Here LM head still generates N tokens for the next N positions in a non-autoregressive way with each position reserving top k tokens (k is a hyperparameter, which is set to 10 for current CTC-drafter), then the token sequence of the N positions with the highest probability will be selected as the raw draft via the token tree structure used in Medusa. The final draft generated via CTC transform will be fed to the base model to verify and the base model will decide to use the draft tokens as the next several output tokens or to generate the next token by itself where the former decision can generate several tokens at once and hence improves the decoding speed.

Authorsrebuttal2024-08-12

Dear Reviewers, We are writing to kindly request your feedback. As the author-reviewer discussion phase is nearing its end, we are eagerly awaiting your response to address any questions or concerns you may have. We apologize for any inconvenience and appreciate your time and efforts.

Reviewer NRvU2024-08-09

Thank you for the response

Thank you for the response and additional experiments. I hope the discussion about CTC independence and beam search method to be included in the revised version. Regarding the Attention Draft Model's architecture, my confusion is this: Medusa uses different heads for each position (i+1, i+2, ...). When you are saying that the Transformer is used for Attention Draft, is it mean that each Medusa head is a Transformer or all heads are integrated into a single Transformer? * If Former (per-position): Query is only a single timestep input. Do you need a Transformer for this non-sequential input? * If Latter (unified): Is Attention Draft Module Auto-regressive? How can it generate all tokens at once?

Authorsrebuttal2024-08-11

Thank you for the further discussion

Many thanks for the further valuable suggestions. We will definitely include the discussion of CTC independence and beam search in the revised version. We would like to give reorganized description about our method to help understand the Attention Draft module’s architecture. Overall, the architecture of our method can be decomposed into three parts: the base model, the drafter which includes the Attention Draft module and LM head, and the CTC-related module which is used to calculate CTC-loss for training and performs CTC-style decoding at inference. The main structures of the three parts are as follows. - The base model in our method is based on the autoregressive generation framework like Llama which only uses Transformer decoder as the main structure. - For the drafter, the Attention Draft module employs one Transformer layer as its structure including masked multi-head self-attention sublayer and Feed Forward sublayer. Compared to Medusa which employs FFN as the structure of each Medusa head without interacting between Medusa heads, our method can introduce sequence correlations between the preceding generated words and next several words to generate, besides the CTC-related module. - The CTC-related module involves non-autoregressive generation of the next N tokens (N is a hyperparameter, which is set to 4 for current CTC-drafter) where blank character and repetitive tokens are introduced into the raw generated sequence which will be processed by CTC Transform module to produce the final draft for the current decoding timestep. Our method works in a different way during training and inference. During training, the base model will generate the whole sequence in an autoregressive manner which is used as the ground truth to distill the drafter. To generate the draft for the drafter, at each timestep, the Attention Draft module accepts as the input the 4 representations of the current position (assuming position i) and its preceding positions (i-3, i-2, i-1) generated by the last Transformer layer of the base model, and upsamples the input representations by 2 times via copying the input representations. Then through the Transformer layer of the Attention Draft module, the output representations are fed to LM head. LM head will generate 8 raw draft tokens in a non-autoregressive way and the final draft tokens will be produced by CTC Transform module. The CTC-loss is used to count the probabilities of all the draft sequences that can be transformed into the ground truth sequence via dynamic programming and the sum of all these sequences is maximized as the training loss. In this way, the probability distribution is drawn towards the draft sequences that can derive ground truth sequence being allocated bigger probability. This means the candidate draft sequence with more reasonable sequence correlations will be selected as the winner at a greater likelihood. At inference, the base model does not generate the whole sequence in advance. Instead, at each timestep, it is used to generate the representations fed to the Attention Draft module and verify the generated draft tokens according to the probability by performing teacher forcing decoding to generate the draft tokens. If the probability to generate the draft tokens is greater than the set threshold, the base model will accept the draft tokens and uses them as the following generated tokens. Then the base model continues to encode these generated draft tokens in an autoregressive way. If the probability is smaller than the threshold, the base model will reject the draft tokens and generate the next token on its own and meanwhile encode the generated token in an autoregressive way, too. Then the latest generated 4 representations by the last Transformer layer of the base model are used as the input of the Attention Draft module for the next timestep and the decoding goes to the next timestep. Here the Attention Draft module works in the same way as in training to generate input representations for LM head. Here LM head still generates N tokens for the next N positions in a non-autoregressive way with each position reserving top k tokens (k is a hyperparameter, which is set to 10 for current CTC-drafter), then the token sequence of the N positions with the highest probability will be selected as the raw draft via the token tree structure used in Medusa. The final draft generated via CTC transform will be fed to the base model to verify and the base model will decide to use the draft tokens as the next several output tokens or to generate the next token by itself where the former decision can generate several tokens at once and hence improves the decoding speed.

Reviewer EXv72024-08-09

Given that the experiments are indeed extensive, I will increase my score from 4 to 5.

Program Chairsdecision2024-09-25

Decision

Accept (poster)

© 2026 NYSGPT2525 LLC