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.