Faster Language Models with Better Multi-Token Prediction Using Tensor Decomposition

## Abstract

Abstract We propose a new model for multi-token prediction in transformers, aiming to enhance sampling efficiency without compromising accuracy. Motivated by recent work that predicts the probabilities of subsequent tokens using multiple heads, we connect this approach to rank- 1 1 1 1 canonical tensor decomposition. By generalizing it to a rank- r 𝑟 r italic_r canonical probability decomposition, we develop an improved model that predicts multiple tokens simultaneously. This model can also be interpreted as a mixture of experts, allowing us to leverage successful techniques from that domain for efficient and robust training. Importantly, the overall overhead for training and sampling remains low. Our method demonstrates significant improvements in inference speed for both text and code generation tasks, proving particularly beneficial within the self-speculative decoding paradigm. It maintains its effectiveness across various model sizes and training epochs, highlighting its robustness and scalability.

## 1 Introduction

Autoregressive transformer models (Vaswani, ) have become a cornerstone in natural language processing tasks due to their ability to model complex sequential data.
However, one significant limitation of these models is the inefficiency in sampling during inference, as they generate tokens one at a time, leading to increased latency in practical applications (Fournier et al., ; Fields et al., ).
Accelerating the inference process without compromising the model’s performance is thus a critical challenge.

Recent efforts have explored multi-token prediction to address this inefficiency. A simple yet effective approach (Gloeckle et al., ) involves using multiple heads to predict the next $n$ tokens simultaneously.
This method approximates the joint probability of the next $n$ tokens by assuming conditional independence given the previous context.
Mathematically, given a sequence $(x_{1},x_{2},\ldots,x_{t})$ this approximation can be expressed as:

$$ $P_{\theta}(x_{t+n:t+1}|x_{t:1})\approx\prod_{s=1}^{n}P_{\theta}^{(s)}(x_{t+s}| x_{t:1}).$ (1) $$

This equation represents a rank-$1$ tensor approximation of the joint probability distribution, effectively treating future tokens as independent of each other given the past tokens.
While this assumption simplifies computation and can be combined with speculative decoding (Leviathan et al., ) to accept some of the predicted tokens, it remains a crude approximation that may limit token acceptance rates due to its disregard for token interdependencies.

To improve upon this, we propose a more accurate approximation of the joint distribution by introducing a sum over multiple rank-$1$ terms. Specifically, we generalize the approximation to a rank-$r$ canonical decomposition (Harshman, ; Kolda & Bader, ; Cichocki et al., ):

$$ $P_{\theta}(x_{t+n:t+1}|x_{t:1})\approx\sum_{\alpha=1}^{r}w_{\alpha}\prod_{s=1} ^{n}P_{\theta}^{(s)}(x_{t+s}|x_{t:1},\alpha),$ (2) $$

where $w_{\alpha}\geq 0$ are learnable weights satisfying $\sum_{\alpha=1}^{r}w_{\alpha}=1$.

The proposed formulation in equation 2 accounts for dependencies among future tokens by effectively considering a mixture of expert models, each capturing different aspects of the token distribution.
By leveraging this rank-$r$ decomposition, we aim to enhance the accuracy of multi-token predictions, thereby increasing token acceptance rates during speculative decoding and reducing overall inference time.
Thus, our main contributions are as follows:

- •
We identify the limitations of existing multi-token prediction methods that predict tokens independently.
- •
We introduce a novel model that employs a rank-$r$ canonical probability decomposition to better approximate the joint distribution of future tokens.
- •
We demonstrate that our approach can be integrated into existing transformer architectures with minimal overhead, resulting in more efficient sampling without significant increases in computational cost.

## 2 Method

### 2.1 Overall concept

Figure: Figure 1: Schematic representation of the proposed model that predicts several tokens at once for a given sequence $x_{1},x_{2},\ldots,x_{t}$. We present the case of $n=3$ predicted tokens $x_{t+1},x_{t+2},x_{t+3}$ and, accordingly, three heads which generate factor matrices $P_{\theta}^{(1)}$, $P_{\theta}^{(2)}$, and $P_{\theta}^{(3)}$ of the canonical decomposition and linear layer that generates weights $w$ are depicted.
Refer to caption: extracted/5948551/pictures/scheme.png

We propose a model architecture that differs from traditional transformer models by enabling simultaneous prediction of multiple tokens through a rank-$r$ Canonical Polyadic (CP) tensor decomposition (Harshman, ) of the joint probability distribution.
In Figure 1 we provide a corresponding schematic illustration, the content of which will be disclosed later in this section.

The joint probability of the next $n$ tokens given the input sequence $x_{t:1}$ can be represented as a $n$-dimensional tensor:

$$ $A\in\mathbb{R}^{V\times V\times\ldots\times V},\quad A[x_{t+1},\ldots,x_{t+n}] =P_{\theta}(x_{t+n:t+1}|x_{t:1}),$ (3) $$

where $V$ is the vocabulary size.
The tensor $A$ encapsulates the probabilities of all possible combinations of the next $n$ tokens.
In Gloeckle et al. () it was proposed to approximate this joint distribution by assuming that future tokens are conditionally independent given the past as shown in equation 1.
We draw special attention to the fact that this may be interpreted as a rank-$1$ CP approximation to the tensor $A$.
While computationally efficient, such approximation ignores dependencies among the future tokens.

To better capture these dependencies, we propose to approximate the joint distribution using a rank-$r$ CP tensor decomposition according to equation 2.
In order to ensure that $P_{\theta}$ from this equation is indeed a probability tensor, it is sufficient to undertake that

$$ $w_{\alpha}\geq 0,\quad\sum_{\alpha=1}^{r}w_{\alpha}=1.$ (4) $$

The difference between equation 2 and standard CP-decomposition is an additional constraint on the factors of decomposition, i.e., each factor, $P_{\theta}^{(s)}$ should be non-negative and sum up to $1$ along one mode:

$$ $\sum_{x_{t+s}=1}^{V}P_{\theta}^{(s)}(x_{t+s}|x_{t:1},\alpha)=1,\quad s=1,2, \ldots,n.$ (5) $$

This is easily achieved by taking *softmax* operation along the mode direction.

Thus, for the given input sequence $x_{t:1}$ we compute its embeddings $e_{t:1}$ using the encoder of the autoregressive transformer model.
Focusing on the last embedding $e_{t}$, we aim to predict the next $n$ tokens by parametrizing the factors of the decomposition as simple functions of $e_{t}$.
We introduce $n$ heads each corresponding to one of the next $n$ tokens.
For each position $s=1,2,\ldots,n$ the conditional probabilities are defined as:

$$ $P_{\theta}^{(s)}(x_{t+s}|x_{t:1},\alpha)=\text{softmax}\left(W^{(s)}_{\alpha}e _{t}\right)_{x_{t+s}},$ (6) $$

where $W^{(s)}_{\alpha}\in\mathbb{R}^{V\times E}$ are the weight matrices for each head and component, $V$ is the vocabulary size and $E$ is the embedding dimension.
The mixture weights $w_{\alpha}$ are computed in a similar way using an additional linear layer:

$$ $w=\text{softmax}\left(W_{h}e_{t}\right),$ (7) $$

where $W_{h}\in\mathbb{R}^{r\times E}$.

### 2.2 Training procedure

In training, we maximize the log-likelihood of the predicted $n$ tokens.
The computation of the log-likelihood is straightforward: first, the embeddings are calculated by the transformer backbone (it has the same cost as for the next token prediction).
We need to evaluate the logarithm of the likelihood, so using equation 6 directly is not numerically stable.
Instead, we compute everything using the logarithms of the probabilities.
For each pair of sequences $x_{t:1}$ and $x_{t+n:t+1}$, we evaluate the logarithm of the mixture weights $w$ (the computational cost corresponds to a matrix-by-matrix product and logsoftmax operation), then use equation 6 to compute $n$ matrices of the size $V\times r$

$$ $C^{(s)}_{\theta,\alpha}=\log P_{\theta}^{(s)}(x_{t+s}|x_{t:1},\alpha),$ (8) $$

to calculate logarithms of the conditional probabilities in a stable way with *logsumexp* operation:

$$ $\begin{split}L&=\log\left(P_{\theta}(x_{t+n:t+1}|x_{t:1})\right)\approx\log \left(\sum_{\alpha=1}^{r}w_{\alpha}\prod_{s=1}^{n}P_{\theta}^{(s)}(x_{t+s}|x_{ t:1},\alpha)\right)=\\ &=\log\left(\sum_{\alpha=1}^{r}w_{\alpha}\prod_{s=1}^{n}\exp(C^{(s)}_{\theta, \alpha})\right)=\log\left(\sum_{\alpha=1}^{r}\exp\left(\log{w_{\alpha}}+\sum_{ s=1}^{n}C^{(s)}_{\theta,\alpha}\right)\right).\end{split}$ (9) $$

### 2.3 Auxilary load balancing loss

Each term of the summation in equation 9 corresponds to a single expert, which predicts its own probabilities for each token.
We have found, that optimizing such loss directly leads to the effects, similar to the ones observed in Mixture Of Experts (MoE) framework (Masoudnia & Ebrahimpour, ; Cai et al., ): one expert (i.e., rank-1 term in our case) dominates the others, leading to worser likelihood even in the presence of larger number of parameters.
Note, that such interpretation and connection is not well-known in the low-rank approximation community, and can be investigated further on.
To obtain the balance between different experts, we utilize the achievements from the MoE communities and propose to use an auxiliary balancing loss on $w$.

It is well known that a critical challenge in training MoE models is ensuring equitable utilization of all experts (Zhou et al., ).
Without proper balancing, some experts may become dominant, handling a disproportionate share of the data, while others remain underutilized.
To address this, we incorporate an auxiliary balancing loss.
This auxiliary loss penalizes imbalances in the expert weights and encourages to distribute the workload evenly across all experts.

Formally, the auxiliary loss can be represented as:

$$ $\mathcal{L}_{\text{aux}}=\sum_{\alpha=1}^{r}\left(\frac{n_{\alpha}}{N}-\frac{1 }{r}\right)^{2},$ (10) $$

where: $r$ is the number of experts, $n_{\alpha}$ is the number of tokens with maximal weight on expert $\alpha$, and $N$ is the total number of tokens.
This formulation ensures that each expert $\alpha=1,2,\ldots,r$ is utilized approximately equally, mitigating the risk of certain experts becoming bottlenecks.

Empirical observations have demonstrated that training the model without the auxiliary loss or using the auxiliary loss values proposed in previous works leads to training instability and eventual failure.
The auxiliary loss is pivotal in maintaining a balanced distribution of token assignments among experts, which is essential for stable convergence and effective learning.
Therefore, careful tuning of the auxiliary loss coefficient is necessary to achieve optimal performance.
By ensuring balanced expert utilization through the auxiliary loss, the model enhances the accuracy of multi-token predictions, which increases token acceptance rates during speculative decoding, thereby reducing overall inference time.

### 2.4 Sampling method

Our sampling scheme is similar to the one proposed in Gloeckle et al. ().
We sample candidates from the proposal distribution (our approximation to the joint distribution of the next tokens) and then accept them or reject according to the recommendations of the draft model (which is the same model that predicts the next token).

For the rank-$1$ case the sampling is easy: probability distributions are computed for each token independently, and sampling is done from the computed distributions.
For our canonical rank-$r$ representation we need to use sequential sampling which is autoregressive, but only works with the factors of decompositions.
This makes sampling $dim$ tokens from our rank $r$ model just a bit slower, than $1$ token from the base model.

Note that the first marginal distribution
$P(x_{t+1}|x_{t:1})$ is given by the first head directly, and we just need to average among $\alpha$:

$$ $P_{\theta}(x_{t+1})=\sum_{\alpha=1}^{r}w_{\alpha}P_{\theta}^{(1)}(x_{t+1}|x_{t :1},\alpha),$ (11) $$

which can be also computed using logsumexp operation.
From this distribution, we sample the first token $x_{t+1}$.

Given $x_{t+1}$ we can now compute the marginal distribution:

$$ $P_{\theta}(x_{t+2}|x_{t+1})=\sum_{\alpha=1}^{r}w_{\alpha}P_{\theta}^{(1)}(x_{t +1}|x_{t:1},\alpha)P_{\theta}^{(2)}(x_{t+2}|x_{t:1},\alpha),$ (12) $$

which is also reduced to matrix-by-matrix products, *logsoftmax* and *logsumexp* operations, and can be implemented by updating the unnormalized logits of the experts with incorporation of $\log P_{\theta}^{(1)}(x_{t+1}|x_{t:1},\alpha)$ into them.

The sampling of the following tokens is also straightforward.
Given sampled $x_{t+1},\ldots,x_{t+s-1}$ we then compute the probability:

$$ $P_{\theta}(x_{t+s}|x_{t+1},\ldots x_{t+s-1})=\sum_{\alpha=1}^{r}w_{\alpha} \prod_{k=1}^{s-1}P_{\theta}^{(k)}(x_{t+k}|x_{t:1},\alpha).$ (13) $$

### 2.5 Speculative Decoding

Speculative decoding (Chen et al., ; Leviathan et al., ) is a technique designed to accelerate the inference process of autoregressive models by generating multiple tokens in parallel, thereby reducing the latency associated with sequential token generation.
In traditional autoregressive sampling, tokens are generated one at a time, with each new token conditioning on the previously generated tokens.
This sequential nature inherently limits the speed of generation, especially for lengthy outputs.

Our sampling method seamlessly integrates with the speculative decoding framework by enhancing its capacity to handle multi-token predictions, as can be seen from Algorithm 1.
The usual setup for a speculative decoding consists of a draft model and a base model.
In our case we implement a modification of a self-speculative decoding algorithm, as described in Zhang et al. (); Elhoushi et al. ().
So, as a base model we take a next-token prediction model and as a draft model – the prediction for $dim$ tokens forward obtained from a full “CP-head”.

Self speculative decoding with rank $r$ model inherits two nice small benefits from more simple rank-1 model: first generated sample from the draft model is always accepted and one additional token from the base model is generated.
This means, that in one pass of the draft model with the base model we will obtain at least 2 tokens.
Due to this fact it seems beneficial to use that type of models even with moderate quality of the draft model.

Figure: Algorithm 1 Self-speculative decoding with rank-r experts

## 3 Experiments

In this section, we present a comprehensive evaluation of our proposed multi-token prediction approach.
Experiments are designed to assess the efficacy of different ranks and auxiliary loss configurations, the capability to fine-tune only the prediction head, and the impact on inference speed for large-scale models.

### 3.1 Training different ranks and auxiliary loss models

Figure: Figure 2: Losses for the tiny transformer model with different CP-rank values trained on the TinyStories dataset.
Refer to caption: extracted/5948551/pictures/tiny_stories_joint_loss_rank_varination.png

Figure: Figure 3: Losses for the rank-8 tiny transformer model trained on the TinyStories dataset with different auxiliary loss penalties compared to the baseline (i.e., the rank-1 model).
Refer to caption: extracted/5948551/pictures/aux_variation_joint_prob_loss.png

For our experiments we’ve chosen the multi-head tiny transformer model with 56.3 M parameters based on the code in Karpathy ().
We consider the case of $4$ heads and added RoPe positional encodings as in Su et al. ().
Training was conducted on the Tiny Stories dataset (Eldan & Li, ) using various ranks for the CP-decomposition.
The objective was to observe how increasing the rank influences the joint loss and loss on the first token. Because the quality of our final generation depends only on the quality of the first head, we tracked both those metrics.
Additionally, we experimented with different sizes of the auxiliary loss penalty to ensure balanced expert utilization.

As illustrated in the left graph in Figure 2, increasing the rank from $1$ to higher values leads to a consistent decrease in joint loss, indicating a better approximation of the joint probability distribution.
This trend underscores the model’s enhanced capability to capture inter-token dependencies with higher ranks.

Contrary to the joint loss, right graph in Figure 3 shows that the loss for the first token remains largely unchanged across different ranks.
It is worth noting, that probability distribution for the first token as a function of last layer embeddings $m$ is given by $\sum_{\alpha=1}^{r}w_{\alpha}(m)C_{\alpha}(m)$ (in notation of equation 9) and both $C_{\alpha}$ and $w_{\alpha}$ are linear, which makes this function equivalent to a simple linear head.
So, after convergence we expected the same loss for all of the ranks.
As follows from the reported results, this is exactly what happened and this consistency confirms that our model maintains optimal training for the initial token prediction, ensuring that the foundational aspects of the sequence generation remain robust.
The loss on the first token is especially crucial, because with a speculative decoding we are improving sample for a big model, which is in our self-speculative case is a first head.
We also note that from Figure 2 it follows that all of the inference speedup will be obtained without compromising quality.

Figure 3 presents the effect of varying the auxiliary loss penalty size. We observed that with a very small penalty, the joint loss mirrors that of the rank-1 model, suggesting insufficient balancing among experts. Conversely, an excessively large penalty led to prolonged convergence times, as depicted in the figure. Then we identified an optimal penalty size balancing expert utilization without hampering training accuracy.

The efficiency of our model in speculative decoding was evaluated by measuring the acceptance rate of drafted tokens.
Figure 4 and Table 1 illustrates that on the Tiny Stories dataset, models with higher ranks achieved up to a around 30% increase in accepted drafts.
This allowed us to reduce inference time even for this tiny (“nanoGPT”) model for which the head is responsible for a significant percentage of computational time, which is not the case for larger models.

Figure: Figure 4: Speculative decoding performance for the tiny transformer model with different CP-rank values trained on the TinyStories dataset from scratch.
Refer to caption: extracted/5948551/pictures/tiny_stories_temp1_bars.png

**Table 1: Results with the speculative decoding for the tiny transformer model with different CP-rank values trained on the TinyStories dataset from scratch.**
| Rank | Loss | Avg. draft tokens accepted | Time per token (with speculative decoding) |
| --- | --- | --- | --- |
| 1 | 3.23 | 1.67 | 0.0336s |
| 3 | 2.88 | 2.01 | 0.0328s |
| 5 | 2.69 | 2.07 | 0.0303s |
| 8 | 2.66 | 2.15 | 0.0326s |

### 3.2 Head-only fine-tuning for PyCode model

To evaluate the flexibility of our approach, we fine-tuned only the prediction head of the PyCodeGPT (Zan et al., ) model across different ranks on the Github Code dataset by CodeParrot. (^1^11 See [https://huggingface.co/datasets/codeparrot/github-code](https://huggingface.co/datasets/codeparrot/github-code).)
This experiment aimed to determine whether partial model updates could yield performance improvements without the computational overhead of full model fine-tuning.

Figure 5 and Table 2 demonstrates speculative decoding performance for the experiments we conducted for different rank values.
From the reported results it follows that even when only the head is fine-tuned, increasing the rank leads to marginal improvements in joint loss.
Additionally, we can see that speculative decoding benefits from higher ranks, albeit to a lesser extent (approximately 9% increase in accepted drafts) compared to the full model training.

### 3.3 Inference time benchmarking

Figure: Figure 5: Speculative decoding performance for trained head of the PyCode transformer model with different CP-rank values.
Refer to caption: extracted/5948551/pictures/pycode_temp07_bar_headonly.png

**Table 2: Average number of accepted draft tokens for the PyCode model.**
| Rank | Loss | Average Draft Tokens Accepted |
| --- | --- | --- |
| 1 | 2.07 | 1.52 |
| 3 | 1.88 | 1.64 |
| 5 | 1.80 | 1.65 |

**Table 3: Inference time for one forward pass comparison for Llama and Rocket models.**
| Rank | Llama 8B Barebone | Llama 8B Head | Llama 8B Full | Rocket 3B Full |
| --- | --- | --- | --- | --- |
| Barebone | 0.1761 | - | 0.1761s | 0.0154 |
| Rank 1 | 0.1761 | 0.0132 | 0.1893 | 0.0160 |
| Rank 3 | 0.1825 | 0.0129 | 0.1954 | 0.0162 |
| Rank 5 | 0.1865 | 0.0330 | 0.2195 | 0.0166 |

To determine the impact of modified head on the inference time of bigger language models we benchmarked the time of one forward pass of our approach on large-scale models with 3 billion and 8 billion parameters.
As reported in Table 3, the inference overhead for integrating the proposed multi-head layer remains minimal, even as the rank increases.
Note that for the value of CP-rank of $5$ we observe a significant time increase for head execution then the Llama model is considerd which is probably caused by its huge vocabulary size. However, for moderate-sized networks inference time remains limited and increases only slightly with increasing CP-rank. The obtained results correspond to the theoretical algorithmic complexity of our new layer. During inference computational complexity of barebone grows linearly (given KV caches), but computational complexity of rank-$r$ head is always the same. Our measurements were made with seq length varying from 1024 to 4096, but for many practical applications sequence length is bigger, which further justifies usage of rank-$r$ head in the case of models with a large context window.

## 4 Related work

Training large language models (LLMs) to predict multiple tokens all at once and in parallel can drive these models toward better sample efficiency.
Various approaches for multi-token predictions have been proposed recently.
In Stern et al. () several feed-forward decoder layers from the last encoder state are added for prediction of the next several tokens simultaneously, and in Miao et al. () this idea was further improved within the framework of the so-called Medusa heads that use tree attention mechanism.
In a number of works Song et al. (); Santilli et al. (); Fu et al. () it is proposed to generate multiple draft tokens in parallel on the basis of the Jacobi iteration methods, i.e., via solving a non-linear system of equations while auto-regressive decoding in LLM.
In the work Bhendawade et al. () multiple tokens are predicted by adding streaming embeddings initialized from upper layers, with the token tree reduced by early exiting.

Thus, this direction of research is actively developing today, however, the approaches outlined above have several limitations, including the need for significant changes in the original architecture of the model and limited speedup.
Therefore, of particular interest is the recent work Gloeckle et al. (), where it was proposed to approximate the joint probability of the next several tokens using multiple heads but assuming conditional independence given the previous context.
As we have already noted above, this approach remains a crude approximation that may limit token acceptance rates in the speculative decoding approach due to its disregard for token interdependencies.
To improve upon this, we considered in this work a more accurate approximation of the joint distribution in the form of the CP-decomposition.

To effectively implement the proposed scheme, we paid attention to connection of the used weighted CP-decomposition with the Mixture of Experts (MoE) technique. MoE is a widespread approach to enhance capabilities of LLMs with the most popular one being Sparse-Gated MoE introduced in Shazeer et al. ().
MoE implementations can be either sparse or dense with sparse version being more popular, but there are many usages of both options, as in Dou et al. () and Pan et al. ().
While many parameters of MoE approach can be tweaked (Cai et al., ), the most common option is using MoE inside a transformer block, as in Zhou et al. ().
We also note that MoE usage is not limited to LLMs and, for example, in Oldfield et al. () it is applied to computer vision model.

In this work, as an application of the proposed model for multi-token prediction, we consider its use as part of the speculative decoding scheme, which was proposed in Leviathan et al. () and nowadays has become a common technique in the domain of inference acceleration.
While initial framework solves the problem of inference optimization of a model given a faster draft model, there are different methods to obtain this draft model.
Early works proposed blockwise decoding as in Stern et al. ().
This line of work is similar to ours, as the model, used for speculative decoding, is exactly the same, as base model.
Later more techniques for self-speculative decoding were developed, namely in Elhoushi et al. () it is proposed to use only particular layers of the base model to obtain draft model and in Hooper et al. () the base model consists of cycles, which also allows to skip layers to obtain a draft model.
Self speculative decoding and multi-token prediction naturally go well with each other.
This combination may require modification in model architecture as in Bhendawade et al. (), but it is possible to modify only heads as in Gloeckle et al. () to enable faster application of the approach to existing LLMs, and we use such approach in our work.

## 5 Conclusion

In this work, we propose a new model for multi-token prediction in transformers based on the Canonical Polyadic (CP) tensor decomposition of the joint probability distribution.
The results indicate that our model can be efficiently trained across a wide range of ranks, with higher ranks consistently yielding lower joint losses. This improvement underscores the model’s ability to better capture the dependencies among future tokens, leading to more accurate predictions.

We observed a direct correlation between lower joint losses and enhanced speculative decoding performance.
Specifically, our approach significantly increased the acceptance rates of predicted tokens, with notable improvements of up to 50 % of draft tokens accepted.
The factor matrices of our decomposition of the joint probability tensor are generated by several heads that use shared model trunk, which practically makes it possible to minimize extra costs during inference and convert higher draft token acceptance to faster inference times.

The ability to fine-tune only the prediction head of the model while maintaining competitive performance highlights the flexibility of our approach.
This capability allows for targeted improvements without the computational overhead associated with full model retraining.
Benchmarking inference speed for bigger models demonstrated that our method introduces negligible inference overhead, ensuring that in many practical cases the benefits of improved performance for draft model do not come at the cost of increased latency.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: Please address the issues above.
soundness: 3
strengths: -- The paper studies an interesting problem to speed updecoding by predicting multiple tokens in parallel at higher acceptance rates than typical speculative sampling approaches.

-- The proposed solution seems straightforward to implement.

-- The contribution to identifying issues with existing multi-token training approaches and proposing a higher rank alternative is novel.
summary: This paper borrows key idea from Gloecke et al. [1] to train multi-token predictors instead of single next word predictor. This work identifies a key flaw in [1] which is that the distributions for multiple $n$ future tokens are independent of each other thus ignoring the token  interdependency. This work interprets this as a rank-1 approximation to the full distribution tensor of $n$ next tokens and proposes to improve this to a higher rank estimate. This higher rank estimate is achieved by $r$ multiple heads defining $r$ different distributions and using their mixture for the $n$-future token prediction. Training and inference method for this is discussed followed by an observation that the multi-token predictor can be used in self-speculative sampling approach where the next word prediction is made faster by using proposal distribution that predicts multiple next tokens. The experiments are mainly performed on nano-GPT model architecture trained on TinyStories dataset and also finetuning the PyCodeGPT model.
weaknesses: -- The evaluation leaves a lot to be desired. Experiments are done on small datasets and small models but more concerningly, little else is provided aside from loss curves of training runs and token acceptance rates for the scheduled sampling approach. As an example, performance of these models on various benchmarks to estimate the quality of these trained models would aid in better assessment of the approach. Also, it is unclear if this approach empirically scaled to larger datasets and models effectively in terms of speed and performance.

-- Comparison to other speculative sampling approaches with various draft models will give abetter idea about the improvement on speed and resources with the proposed approach.

-- There is room for improvement in presentation. Figure 1 doesn't help with understanding the paper better and is confusing. Algorithm 1 can also be described more clearly. Currently, it hinges on the reader's prior understanding of speculative decoding.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: -How does the method combine with beam search?

-Does the speedup increase or decrease as a function of model size?
soundness: 3
strengths: -Tackles an interesting and important problem

-The method is written clearly. 

-I also find the connections to tensor decomposition interesting.
summary: After reading the author response, I thank the authors for providing clarifications to Table 1 and Table 3 and answering some of my other questions. However, I still think it is important to compare to other baselines (e.g. EAGLE, Medusa as another reviewer points out) in addition to Gloeckle et al. 2024. Therefore I am keeping my score the same. 

----

The paper studies multi token prediction in transformer language models. Vanilla autoregressive degressing is expensive for long outputs since it only decodes one output at a time.

The authors are inspired by the work of Gloeckle et al. 2024. In Gloeckle et al. 2024, given a context x_{t:1} the next n tokens are predicted independently (with multiple heads). As the authors point out, this amounts to a rank-1 tensor approximation of the joint probability distribution.

In this work, the authors explore higher ranks (r > 1) using CP decomposition. They draw a connection to mixture-of-experts and propose a auxiliary load balancing strategy so all the weight is not on one expert (component). 

They then perform experiments validating their work.
weaknesses: Some confusion on experimental results: I'm a bit confused as to how much of a speed up the author's approach gives over both the approach of Goeckle et al. 2024 (i.e. rank=1) and also vanilla non-autoregressive decoding for the same level of quality. 

For example in Table 1: I see that in Table 1 the  final column (time per token) is not much different across all the rows?

Moreover I don't quite understand Table 3.

Comparisons: I think the authors need additional baselines in addition to just ablations of their own approach from the related work. For example, as another reviewer suggested EAGLE and Medusa:
https://github.com/SafeAILab/EAGLE
https://arxiv.org/abs/2401.10774

Related work: The authors should also cite and discuss related work in non-autoregressive decoding (typically for neural machine translation) that has been developed for a while e.g. see below and citations therein. In particular it would be useful to discuss how the authors' approach compares and contrasts with these works.

https://arxiv.org/abs/1711.02281
https://arxiv.org/abs/2204.09269
https://arxiv.org/abs/2012.15833

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 2
questions: See above.
soundness: 2
strengths: (1) The method is well-motivated and explained clearly. The connection to MoE, which motivates a load-balancing auxiliary loss, is also interesting.

(2) The paper seeks to improve inference speed in large models, which is an important problem.
summary: One existing form of speculative decoding involves predicting k tokens at a time independently, which can be thought of as a rank-1 decomposition of the k-order joint probability tensor over those tokens. This paper instead proposes to predict the factors for a rank-r decomposition. They evaluate two instantiations of this idea: training a LM from scratch to predict this decomposition, and taking an existing LM and fine-tuning additional heads to predict this decomposition. Their experiments show that higher rank decompositions lead to higher acceptance rates in speculative decoding.
weaknesses: While the method seems interesting and promising, the paper's experiments seem disorganized and insufficient to fully demonstrate the effectiveness of the method.

(1) The majority of the results are for a 56.3M parameter trained on TinyStories, which is a very limited evaluation setting, both because the dataset is synthetic and because the setting involves retraining. There are also some experiments on head-only tuning for PyCodeGPT in Table 3, but the results in that setting are not very strong --- increasing the rank does not actually seem to actually improve inference speed for many of the models. The paper would benefit from more thorough evaluation and stronger results (especially on non-synthetic datasets, and on speeding up existing models rather than requiring retraining: for example, the evaluations done in https://arxiv.org/pdf/2211.17192 (Table 3) would improve this paper).

(2) The majority of the experiments section seems to involve analysis rather than results: only tables 1 and 3 report inference times, which are the main results. I would suggest moving other plots (token acceptance rate, first token vs joint loss, etc.) to a separate analysis section.

(3) There are a substantial number of issues with the experiment design that would be beneficial to address: (a) In Figure 3, it seems like hyperparameters are being selected using the test set; I would suggest using a dev set instead. (b) To make comparisons fair, I would suggest training each rank for the same amount of wall-clock time, rather than number of steps, in case higher ranks require more time per forward pass. (c) The self-speculative setup makes the results hard to interpret because each rank uses a different target model. I would suggest that each method be speculative with respect to the same target model. (d) The paper would be clearer if the experiments were described concretely: for example, the paper states that "Our measurements were made with seq length varying from 1024
to 4096" (lines 408-409), but it's not clear which experiments use which sequence lengths.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1.	In line 113, the authors denote the input sequence as x_{t:1} and the corresponding embeddings as e_{t:1}. According to the description, the embeddings are the representations of the final transformer layer, while in Figure 1, the same value is denoted as z_t. Do z_t and e_t means the same representation, or e_t means the “input” embeddings? This notation is somewhat confusing.

2.	Are there any results on the acceptance rate for Llama 8B, not just inference time?

__Typos__:

1.	In line 116, a comma is missing before "the conditional probabilities ...".
2.	In line 150, "Note, that" should be revised to "Note that".
soundness: 2
strengths: 1.	This work identifies the limitations of recent multi-token prediction standards and proposes a more generalized approach.

2.	The experimental results demonstrate the method's effectiveness, and the ablation study underscores the importance of the introduced components.
summary: This paper focuses on speculative decoding methods that incorporate additional prediction heads. The authors conceptualize current standard approaches as rank-1 canonical tensor decomposition and propose a generalized method that extends from rank-1 to rank-r canonical tensor decomposition to approximate the joint distribution of future tokens. To enhance model training, an auxiliary loss is introduced to address weight imbalances. Experimental results highlight several key findings:

1.	Increasing the ranks results in a decrease in joint loss.

2.	The first token appears to have no correlation with different ranks.

3.	The method is effective even when only the prediction heads are trained.

The proposed approach achieves notable speedups compared to autoregressive models and rank-1 baselines.
weaknesses: 1.	The work lacks comparison with existing state-of-the-art methods such as Medusa, Eagle, etc., which belong to the same research domain.

2.	In the code generation setting, the performance of averaging two accepted draft tokens is not promising.

3.	There are several typos in this version that need revision.

Meta-review: This paper aims to speed up the generation process of language models. Motivated by the observation that predicting multiple future tokens in parallel does not consider inter-token dependencies, and that it can be viewed as rank-1 tensor decomposition, this work proposes to extend that to a rank-r decomposition. Experiments show that the proposed method has higher acceptance rates compared to the baseline.

Strengths:
1. The proposed idea is well motivated and the extension to rank-r decomposition is neat.

Weaknesses:
1. Reviewers mentioned that experiments are only conducted on small models with small datasets, but this is addressed by authors' rebuttal.
2. Reviewers mentioned that baselines are lacking and this work should be compared to EAGLE and Medusa.
3. Some reviewers mentioned that the presentation of this paper can be improved.

Overall, the idea is very interesting and the proposed extension to rank-r decomposition is very neat. However, experiments can be further improved by adding more contemporary baselines, and the new larger experiments authors conducted during rebuttal can be added to the paper. I'm recommending rejection based on reviewers' opinions, but I wouldn't mind if the paper gets accepted.

Decision: Reject