Vcc: Scaling Transformers to 128K Tokens or More by Prioritizing Important Tokens

Summary

This paper aims at designing efficient Transformers for ultra long sequences. The design is motivated from the intuition that only a small subset of special tokens, which are referred to as VIP tokens, are most relevant to the final prediction. Based on the intuition, a multi-resolution compression is proposed to reduce the computational cost. Experiments are conducted to show the efficiency and accuracy of the proposed model.

Strengths

+ Nice presentation. + The observation that only a small subset of special tokens are most relevant to the final prediction is insightful. The proposed task-driven multi-resolution compression technique is well-motivated and technically sound. + The experiments are extensive and show reasonable performance gain.

Weaknesses

+ The major concern for the proposed approach is that it **only applies to Transformer encoders**. Many important LLMs today are decoder-only, e.g., GPT, PaLM, LLaMA, etc. The authors formulate their method in an encoder layer (Sections 2 & 3), and show their method can lead to decent performance gains on encoder-only models and encoder-decoder models. However it seems non-trivial to generalize the approach to decoder-only models. + The compression also seems to **impede the use of relative positional encoding**. Although there are experiments on T5-based models, no details are provided on how the relative positional information is handled. + As discussed at the beginning of Section 3.3, the multi-resolution compression technique is proposed to minimize the approximation error of Eq. (12). However, **neither theoretical guarantees nor numerical simulations on the approximation error are provided**. There seems to be some theoretical results in the appendix but I can't find any analysis on the approximation error (defined in Eq. (33)). Since the multi-resolution compression is the main technical contribution of this paper, I encourage the authors to provide some theoretical analysis and/or numerical experiments on the approximation error in the main body of the paper. + **Minor issues**. + Line 121: $n_p$ is the size of the VIP-tokens: "size" seems ambiguous here and can be misunderstood as "dimensionality". I suggest writing _$n_p$ is the number of the VIP-tokens_. + Eq. (9): $\mathbf I_{n_p\times n_p}$ can be simplified to $\mathbf I_{n_p}$. + Line 223, "$s\in\cdots$" should be "_where_ $s\in\cdots$".

Questions

+ Is the proposed approach generalizable to decoder-only models? + Is the proposed approach compatible with relative positional encoding? How is the T5's positional encoding handled in the experiments? + Can you provide more illustrations on the approximation error of the compression technique?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations are discussed in Section 4. However, I recommend the authors to explicitly discuss the applicability to decoder-only models, if the proposed approach only applies to Transformer encoders.

Summary

This paper proposes VIP-token centric compression (VCC) as a method to improve the efficiency of Transformers for ultra long sequences. Based on the impact on approximating the representation of the VIP-tokens, VCC can compress the long sequence into a much smaller representation. Thorough experiments validated the effectiveness of the proposed method.

Strengths

* Long sequences input is a practical problem for current LLM-based applications. And the efficient technology for long sequences that this paper focused on has a high application value, in my opinion. * The paper is well-written, the method is generally novel and the effectiveness is validated by thorough experiments. * The authors plan to release the code and checkpoints.

Weaknesses

* FLOPs is missing. Runtime is a hardware-dependent metric, and it would be better to see also hardware-independent metrics FLOPs. * Why only experiments on encoder-only and encoder-decoder architectures? It would be better to see experiments or discussions on recently more popular decoder-only architectures. * KV Cache is a useful technique for accelerating Transformer inference and has been already integrated into widely used packages such as huggingface transformer and timm. It would be better to see the discussion on whether KV Cache and VCC are orthogonal and whether they can be used simultaneously. * Missing reference. It would be better to have a discussion on the difference with the previous [1] and concurrent [2] work that scales Transformer to 262k and one million, respectively. 1. Memorizing Transformers. In ICLR 2022. 2. Scaling Transformer to 1M tokens and beyond with RMT. In arXiv.

Questions

See Weaknesses.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Adequately addressed.

Summary

This work introduces a method to scale up the supported number of tokens for transformer-based models. It prioritize the important/dominant tokens and compress those unimportant tokens during the processing. The conducted experiments demonstrate the effectiveness of the proposed method, and show that the proposed method can be directly incorporated into existing pretrained models with additional finetuning.

Strengths

I think this work is trying to target an important problem, i.e., how to break through the limitation on the number of tokens to be processed. The proposed VIP-token centric compression is overall reasonable in terms of its high-level idea. The coducted experiments provide evidences for its efficiency benefits and performance superority when compared to the baselines with the same sequence length.

Weaknesses

The targeted problem and the main story are quite attractive at the first glance. However, the proposed methods are not that solid and insightful as expected, making I feel some parts somewhat overclaimed. To name a few, for the high-level idea of this work, evaluating the importance of VIP-tokens should be one of major parts. However, the proposed method for evaluating the importance of difference tokens are based on some intuitive assumptions. These assumptions are task-specific or case-specific, affecting the generalization capacity of the proposed method in fact. Besides the proposed method for evaluating the importance of tokens, other parts of the proposed method only introduces high-level idea deas and illustrates them with examples, without providing a rigorous modeling. The experiments are conducted mainly on ultra long sequences. This matches with the major story, but in fact, prioritizing important tokens should be generally benefitial for a wide range of tasks using transformer-based models. The authors have not provides discussion or analysis on these.

Questions

1. Are the proposed methods, which include but are not limited to the method for evaluating the importance of different tokens, generally applicable (e.g., for other tasks, or for other modalities)? If yes, pls explain more. If not, please analyze the limitations in detail and in-depthly. 2. The proposed compression and decompression are in fact implemented by matrix multiplication. How can you ensure the tokens of different importances (not referring to splitting as VIP or non-VIP) are compressed or decomposed as expected? 3. Why can placing VIP tokens in the front seats by re-ordering play the role of prioritizing them? How about repeating the VIP tokens (this is a more commonly used way for highlighting important tokens)? Pls compare these two, if possible.

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The biggest problem is that this paper is quite attractive at the beginning, but the main parts, especially for the proposed methods, are somewhat disappointing since they are not that solid and insightful as expected. Pls see the weakness part and my questions for more detailed comments.

Summary

VIP-token centric compression (VCC) aims to solve the quadratic attention complexity bottleneck in the self-attention component and the transformer encoder block. Roughly, it denotes some tokens as very important (VIP) and compresses the remaining sequence using a low-rank approximation. They benchmark on natural language data and do experiments to verify the time and space complexity from their theoretical analysis.

Strengths

They achieve impressive accuracy and performance numbers. Strong results on question-answering tasks that require long context and small runtime. They scale to 128k tokens and are able to tackle NarrativeQA, although the results are not as strong. The method is intuitive and the analysis and description of the method is clear. They were able to get strong results with pretrained weights from existing transformer encoder models.

Weaknesses

The idea of VIP tokens is similar to the global memory component of efficient transformers like ETC, BigBird, and Longformer. It suffers a similar problem that it's not clear which tokens are VIP. Experiments are only done on encoders. They mention in passing that the method can be adapted to decoding, which seems to be more in the vogue these days. Some aspects of the experiments like doing a different computation in the first 4 transformer layers make me think this method may be finicky and hard to apply in general.

Questions

How does this method perform in other large context non-language tasks like images or audio or even synthetic tasks like Long Range Arena where the VIP token selection is less obvious? It would be great to see decoder experiments?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

Authors should be commended for addressing weaknesses of the method like the selection of VIP tokens and encoder-only aspect.

Summary

The paper presents a new method to reduce the computational cost of transformer models on long input sequences. It proposes to use certain tokens (VIP tokens) that will be responsible for compressing the rest of the sequence into fewer tokens that can then be processed efficiently. In detail, a tree is constructed from the input features to the transformer block such that at each node we have the mean of the children of that node. To construct the feature sequence that will be processed, the averaged attention of the nodes of the tree with the VIP tokens is used to walk the tree until the leaves are reached. The reached nodes form the input sequence to the transformer block. Using the tree and the stored differences between the nodes and their children the new sequence is constructed after the transformer block. Thorough experiments show that the proposed attention achieves results on par with dense attention while being significantly faster to compute.

Strengths

- Both using a set of tokens to generate the "clustering" of the input sequence as well as using the tree of differences to recreate the full sequence are very interesting and original methods. - The experimental analysis is thorough and clearly shows that the proposed method works as well as dense attention in the examined benchmarks

Weaknesses

- The main benefit of the method is the runtime efficiency, however the experiments do not describe adequately what this measured runtime is. 1. It is mentioned that it is the time for a single sequence? How does the method scale wrt batch size? 2. It is not clear whether it is the forward pass time or the forward/backward. How does this method scale for training? - There is a lot of theoretical analysis that seems slightly unnecessary as it doesn't provide much clarity but rather confuses given that the main approximation that allows for the method is taken from [6] and is not analyzed further in this work. - The proposed method cannot trivially be used in an autoregressive setting as it is really hard to efficiently compress N tokens while information flows only from the first to the N-th.

Questions

My main questions have to do with the calculation of the runtime of the presented method. Are the numbers for a single sequence? In that case how does the method scale with the batch size given that different samples will have different paths in the tree which will result in many out of order copies that may make the method slower. In addition, is the calculated runtime for a forward pass or for a full training step?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

The authors adequately discuss the limitations of their work.

Thanks for the rebuttal

Thanks for the rebuttal. My concerns are adequately addressed. Thus, I maintain my rating.