TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer

Q1. How does attention handle autoregressive decoding?

Inference can be divided into two stages. The first stage involves calculating the key-value (kv) based on the input prompt. This is done using the following code:

# b, h, n, e, d
kv_outproduct = torch.einsum("... n e, ... n d -> ... n e d", k, v)
# 1, 1, n, 1, 1
index = torch.arange(n - 1, -1, -1).reshape(1, 1, -1, 1, 1).to(x)
decay = ratio.unsqueeze(0).unsqueeze(-1) ** index

kv_outproduct_with_decay = kv_outproduct * decay
kv = torch.sum(kv_outproduct_with_decay, dim=-3)

The second stage involves using the calculated kv and computing the output. This is done using the following code:

kv = ratio * kv + torch.einsum(
    "... n d, ... n e -> ... d e",
    k,
    v,
)
output = torch.einsum(
    "... n e, ... e d -> ... n d", q, kv
)

Please note that Stage 1 is executed only once, and the cost of each inference is mostly determined by Stage 2. The time and space complexity for Stage 2 are both $O(bhd)$, where b represents the batch size, h represents the number of heads, and d represents the head dimension, which is usually either 64 or 128.

We provide a speed comparison in the table below for different sequence lengths, measured in tokens/second/gpu. The larger the value, the better:

Q2. The contribution of this work.

Scaling up a linear attention-based language model and making it faster and better than transformer-based models is a non-trivial task. It requires extensive optimization of the whole system rather than a small part of the model. As shown in this paper, we make crucial enhancements to nearly every part of the model, such as the architecture, the attention mechanism, the inference algorithm, and the model parallel implementation.

Regarding the difference between lightning attention and flash attention, our lightning attention is IO-friendly linear attention, while the flash attention is IO-friendly softmax attention. We agree that Flash Attention sheds light on our lightning attention. By incorporating this idea, we have successfully addressed the issue of slow speed in the causal scenario of Linear Attention mentioned in "Flash: Transformer Quality in Linear Time". In addition, by leveraging the characteristics of Linear Attention, our lightning attention can achieve significant efficiency improvements compared to Flash 1/2. We provide a comparison of the efficiency between our method and Flash Attention 1/2 in the following section.

Q3. The computing method of lighting attention.

In our preliminary version, we use eq 10 implementations to avoid the slow looping issue mentioned in "Flash: Transformer Quality in Linear Time". However, we have further optimized our approach by leveraging the mechanisms of linear attention, resulting in enhanced efficiency. Through these optimizations, our lightning attention demonstrates significant improvements in efficiency compared to Flash 1/2. Below, we provide an efficiency comparison between our method and Flash Attention 1/2. The speed table is measured in seconds, while the memory table is measured in MB. Smaller values indicate superior performance in both speed and memory.

Speed table(seconds):

Forward:

Backward:

Memory Table(MB):

Forward:

Backward:

Q1. The contribution of TransnormerLLM.

The technical contribution of this paper is to build a linear attention-based large language model with faster speed and better performance than transformer-based ones. Prior to our work, no previously developed language models had achieved the same level of effectiveness as our model. Our approach involved extensive ablations to validate the network architecture, and significant engineering efforts to accelerate linear attention and scale up the model using advanced model parallel implementation. As a result, our model achieved unprecedented levels of performance, surpassing that of any prior efficient language models. We have pushed the boundaries of what linear attention can achieve in terms of model size, which we believe represents a significant advancement in the field.

Q2. Lightning attention algorithm, training efficiency, and baseline.

We understand your concerns about the fairness of comparison. However, it's worth noting that our comparison object is an optimized Flash attention Transformer, not a plain PyTorch Transformer. The comparison aims to demonstrate the improvements achieved by our approach over advanced models like Flash Attention Transformer in terms of speed and memory management. Detailed comparisons can be found in Tables 12 and 13 in our paper. We further optimized the algorithm and compared its speed and memory usage with Flash2 and Flash1. The units in the speed table are seconds, while the units in the memory table are in MB. Smaller values indicate better performance in both speed and memory.

Forward:

Backward:

Memory Table(MB):

Forward:

Backward:

Q3. The robust inference algorithm is trivial.

It is noteworthy that the issue of robust inference in linear transformers has not been explored previously in this field. The body of literature on linear attention with exponential decay is somewhat limited, and to date, no studies have been conducted on its corresponding inference. In this paper, we first reveal the existence of the problem and then provide a mathematically rigorous solution for it. We believe that our contribution to the growth of large language models is valuable.

Q4. Experiments result about RetNet.

It is noteworthy that our work and the RetNet paper were conducted simultaneously, which made it impossible for us to make a direct comparison during the research process. We also noticed that RetNet has been submitted to this ICLR conference, as seen in the link. Nevertheless, we evaluated one of their models to draw a comparison, and the results are presented in the following table. Our model performed significantly better than RetNet.

Q1. The contribution of this work.

Our TransNormerLLM is built on TransNormer and Flash Attention, with crucial enhancements, including architecture optimization, algorithm adaptation and optimization, and model parallel implementation. The novelty of this work lies in how to build a linear attention-based language model with faster speed and better performance. Prior to our work, none of the previously developed efficient language models, including those based on linear attention, sparse attention, RNNs, and LongConv, had achieved the same level of effectiveness as our model. By combining algorithmic developments with engineering efforts, we have pushed the boundaries of what linear attention is achievable regarding model size, which we believe constitutes a substantial advancement in the field.

Q2. The motivation of this work and the advantages compared to TransNormer.

Our work focuses on improving efficiency rather than accuracy compared to the standard TransNormer model to enable successful scaling. As we have claimed in the abstract: ”TransNormerLLM evolves from the previous linear attention architecture TransNormer by making advanced modifications”. The key contribution of this work is scaling linear attention to large language models with significant efficiency improvement and comparable performances to transformer-based LLM models. Note that TransNormerLLM also requires less memory compared to TransNormer-T1/T2:

Besides, the modeling differences between TransnormerLLM and Transnormer are extensive including position encoding, token mixing methods, etc. In addition, the optimization such as model parallelism, and lightning attention are none trivial and they require significant engineering efforts. These modifications made up our key contributions.

Q3. The motivation and significance of Lightning Attention.

We agree that Flash Attention sheds light on our lightning attention. By incorporating this idea, we have successfully addressed the issue of slow speed in the causal scenario of Linear Attention mentioned in "Flash: Transformer Quality in Linear Time". Additionally, we have further optimized our approach by leveraging the mechanisms of linear attention, resulting in enhanced efficiency. Through these optimizations, our lightning attention demonstrates significant improvements in efficiency compared to Flash 1/2. Below, we provide an efficiency comparison between our method and Flash Attention 1/2. The speed table is measured in seconds, while the memory table is measured in MB. Smaller values indicate superior performance in both speed and memory.

Speed table(seconds):

Forward:

Backward:

Memory Table(MB):

Forward:

Backward:

Q4. The robust inference algorithm.

In our work, the robust inference modification is indeed mathematically motivated, and they are completely equivalent in terms of theory. We have tested the numerical errors of Origin Inference and Robust Inference, as shown in the following table. The parameters used for the test were b=2, h=8, n=1024, and d=64, where b represents the batch size, h represents the number of heads, n represents the sequence length, and d represents the feature dimension. It can be observed that the robust inference algorithm does not have numerical precision issues.

Q5. The pre-training dataset.

We understand your concern about the specificity of the pre-training data, which can indeed affect a fair comparison with other models. However, it is common practice for LLM not to open-source pre-training data due to various reasons including privacy and security concerns. For instance, Llama1/2 do not disclose their pre-training data, thus it is impossible for us to obtain their training data information. Our evaluation protocol for LLMs adheres to standard practices, where we perform assessments via downstream tasks. To ensure fairness in model design comparison, all the ablation studies mentioned in our paper were conducted using the same dataset, ensuring internal consistency and validity of the conclusions drawn.

Q6. The details on the pre-training data characteristics and compute resources used.

For more detailed information about the pre-training data characteristics and compute resources used, please refer to Appendix D of our paper where we have provided a comprehensive description. To address your concern about potential data leakage, we've taken strict measures during data selection to ensure there's no occurrence of such issues. We understand the importance of maintaining the integrity of our experiments and ensuring that our results are accurately reflective of the model's performance.

Q7. The learning curve.

It is important to mention that TransNormer has surpassed Transformer in language modeling, as shown in table 4 of TransNormer paper. While TransNormerLLM primarily aims to enhance efficiency, it also yields a positive impact on performance.

As for the training curves, we present the training loss in the form of a table rather than a curve due to space constraints. However, we believe that the table still provides a clear visualization of the training process and performance improvements over time.

Q1. The meaning of exponential decay position encoding is mentioned in section 3.1.1.

Exponential decay position encoding indicates the term $\lambda^{s-t}$, where $\lambda$ $\in$(0,1). In this case, the QK attention scores are biased towards distanced tokens, with a proportional penalty for distance. In the first layer, it is combined with LRPE (Linearized Relative Positional Encoding) and is in the form of $a_{st}=\mathbf q_s^{\top} \mathbf k_t \lambda^{s-t}\exp^{i\theta(s-t)}$ where $\theta$ is the learnable parameter. While in the rest layers, it is in the form of $a_{st}=\mathbf q_s^{\top} \mathbf k_t \lambda^{s-t}$. In both cases, $\lambda$ is a non-learnable parameter to control the distance penalty in different layers as discussed in section 3.1.1 and equation(2). Such a form is fully compatible with Linear Attention, as it can be decomposed with respect to s and t separately. Thank you for the suggestion and we will clarify it in the revised version.

Q2. Although the training efficiency for TransNormerLLM models larger than 7B is presented, the performance on benchmarks is not included. It would be advantageous to include those results if feasible.

Thanks for your comments. Due to time constraints, we only include benchmark results for TransNormerLLM-385M, TransNormerLLM-1B, and TransNormerLLM-7B. As an initial exploration of scaling efficient LLM, we believe that 7B is a commonly used LLM size and is large enough to demonstrate that linear attention models can achieve better performance than transformers while being significantly faster. However, larger models, such as the 13B and 65B, are currently being trained. Once the training is completed, we will benchmark them. Thank you for your understanding and patience.

Q3. Although the pre-training performance is on par with that of Transformer models, have any experiments on fine-tuning the models been conducted?

The main objective of this paper is to scale up linear attention-based language models and confirm their effectiveness and efficiency. As most of our competitors use base models, we will only benchmark our base model to ensure fairness. However, we recognize that fine-tuning for specific applications is also crucial, and we plan to release our chat model soon.

Q4. In Algorithm 3, why is an explicit mask $\mathbf M \in \mathbb R^{N \times N}$ necessary? Is it for general attention?

Yes, it is necessary in causal attention. In Algorithm 3, the $\mathbf Q\mathbf K^{\top}$are calculated first as general attention, and the mask is used to achieve causal attention. However, in our efficient implementation, the explicit mask is no longer needed, and the model can achieve linear training and inference computation complexity. We provide an efficiency comparison of our method and Flash Attention 1/2 below, the units in the speed table are seconds, while the units in the memory table are in MB. Smaller values indicate better performance in both speed and memory:

Speed table(seconds):

Forward:

Backward:

Memory Table(MB):

Forward:

Backward: