Paper: Chunk, Align, Select: A Simple Long-sequence Processing Method for Transformers (ICLR-2024)

Reviewer comment (Novelty):
The method could be counted as a form of dynamic pruning technology, and there are already several works in that area. For example, Transkimmer [1] already has dynamic token pruning, and there are other works in similar directions [2].

Author response:
As mentioned in the related works, token pruning is a common technique used to reduce the overhead of training and inference. However, most existing methods are only suitable for short sequence tasks and are often not applicable to the case of long sequence inputs. In contrast to previous work, we propose a selection strategy in the context of chunking to choose hidden representations, significantly improving model performance. Although the use of RL for token pruning has been explored, these methods typically only use the final output signal (such as cross-entropy loss in TR-BERT [1]) as a coarse-grained reward for training the selection strategy. In our method, we additionally utilize attention signals to design the more fine-grained reward for each action. Our experiment results demonstrate the effectiveness of our approach.

[1] Ye D, Lin Y, Huang Y, et al. Tr-bert: Dynamic token reduction for accelerating bert inference[J]. arXiv preprint arXiv:2105.11618, 2021.

Reviewer comment (Novelty):
The usability of RL for token pruning is not a particularly surprising method, and chunking + local attention is a very standard policy for efficiency gain.

Author response:
As mentioned in the related works, token pruning is a common technique used to reduce the overhead of training and inference. However, most existing methods are only suitable for short sequence tasks and are often not applicable to the case of long sequence inputs. In contrast to previous work, we propose a selection strategy in the context of chunking to choose hidden representations, significantly improving model performance. Although the use of RL for token pruning has been explored, these methods typically only use the final output signal (such as cross-entropy loss in TR-BERT [1]) as a coarse-grained reward for training the selection strategy. In our method, we additionally utilize attention signals to design the more fine-grained reward for each action. Our experiment results demonstrate the effectiveness of our approach.

[1] Ye D, Lin Y, Huang Y, et al. Tr-bert: Dynamic token reduction for accelerating bert inference[J]. arXiv preprint arXiv:2105.11618, 2021.

Reviewer comment (Evaluation):
The computational timecost brought by retrieval or selection based methods are always a significant issue. I did not find such discussions in main context and appendix about the latency part. The author should discuss and measure it to show the tradeoff between performance gain and latency increase.

Author response:
At the inference stage, the average time overhead of the selection process (Select) and the whole process (All), and their ratio (Ratio)  for a single sample are shown below. In addition, it should be noted that in our experiments, we used the V100 which has a relatively slow computation speed. If more advanced GPUs are used, the latency will be significantly reduced. 

|                | arXiv  | PubMed  | GovReport | SummScreen | Multi-News |  WCEP  | NarrativeQA |
| :------------: | :----: | :-----: | :-------: | :--------: | :--------: | :----: | :---------: |
| Selection (ms) | 157.5  |  70.4   |   293.9   |   103.4    |    74.2    |  9.8   |    112.4    |
|    All (ms)    | 9661.5 | 13799.1 | 104953.7  |  30424.0   |  17664.2   | 1356.2 |    689.6    |
|   Ratio (%)    |  1.63  |  0.51   |   0.28    |    0.34    |    0.42    |  0.72  |    16.30    |

Reviewer comment (Reproducibility):
As the selector FFN is randomly initialized but the backbone is pre-trained well, I think in the first 10k training iterations, the training might be unstable if following the alternative updates on a well-trained model and a newly-initialized model. The author provides limited details about this.

Author response:
Actually, we were surprised to find that the transformer’s MLE loss with label smoothing decreases quite normally during the training phase in most cases. We believe this may be because, although the selector is randomly initialized, it only participates in the selection of the transformer’s hidden representations and does not severely disrupt the representation information. Therefore, the training of the transformer is relatively stable. We have submitted the figures of MLE loss changes during the training process to the supplementary materials and included it in the appendix of the paper.

Reviewer comment (Experiments):
The ablation study on the chunk length might be interesting and important. A smaller chunk length brings better granularity and a larger chunk length accelerates the inference. Is the selected 512 length the best length for chunks?

Author response:
In our experiments, we found that with the same input length, chunk sizes 512 and 1024 have similar model performance, and the model performance significantly decreases when the chunk size is 256. Therefore, in order to process longer sequences while maintaining sound model performance, we set the chunk size to 512.

Reviewer comment (Reproducibility):
For footnote 1 of Page 3, how you deal with cases that a sentence is longer than a chunk as sometimes the sentence segmentation tool does not work well to split the sentence to desired length?

Author response:
In fact, although it is rare, we have indeed observed some sentences that are longer than the chunk size. In this case, we will directly discard the part of the sentence that exceeds the chunk size.

Reviewer comment (Experiments):
The selector is a small-size 8M FFN. Did you consider to scale up the parameters and also change the FFN to a RoBERTa model with binary classification head? This might be better for the action decision.

Author response:
In our method, it is feasible to improve the model performance by increasing the scale of the parameterized selector. However, considering the time overhead brought by the selection process, we hope that the increased overhead is small and the system still has good performance. In my opinion, changing FFN to RoBERTa can improve the model performance, but the time overhead of the selection process may be large.

Reviewer comment (Experiments):
If I understands the paper well, the selector is not pre-trained and is only randomly initialized FFN. Then you fine-tune the BART and selector together on specific task. Did you try to freeze BART and pre-trained the FFN selector on the same pre-training corpus of BART? Personally, I think if your selector FFN is only 8M parameters, the pre-training might not be helpful. But if the selector size is scaled up, this might be good to the model.

Author response:
Because we consider the compatibility of our method with existing pretrained language models as one of the contributions, we did not conduct additional pre-training. Theoretically, the current selector’s model size is relatively small, making it difficult to achieve good generalization through pre-training. However, pre-training a larger scale selector may improve the model effect, but the time latency will also increase.

Reviewer comment (Novelty):
Some missing related works: TRIME (Zhong et. al., 2023) and LongMem (Wang et. al., 2023) are related to this method in terms of chunking and memorizing; Landmark Attention (Mohtashami and Jaggi, 2023) is related to this method in terms of S and E tokens.

Author response:
Thank you for bringing these papers to our attention. We appreciate your suggestion and agree that these works are indeed relevant to our study. We will review these papers and discuss their relation to our method in the revised manuscript.

Reviewer comment (Novelty):
Missing related work. Long sequence Transformer is a hot topic, including two main directions: efficient computation or length extrapolation (train-short-test-long). In the realms of long-sequence Transformers, there appears to be a noticeable omission in the exploration of length extrapolation. Here are two related studies: 1) A Length-Extrapolatable Transformer 2) Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation.

Author response:
We thank the reviewer for these provided related studies and will modify the manuscript to include the content for the length extrapolation.

Reviewer comment (Experiments):
The proposed RL objectives lack solid motivation. It is well-known that RL objectives are hard to implement due to the variance of reward scores. There lacks necessary evidence to show how these RL objectives work and whether these RL objectives are necessary (compared to traditional MLE loss).

Author response:
In our method, we use a transformer as a simulation environment to optimize a reinforcement learning selection strategy for selecting crucial token hidden representations. We design rewards using the attention and output signals of the transformer, and obtain and update states using the transformer’s hidden representation. Due to the lack of ground truth to determine the importance of hidden representations, and the tendency of attention scores to shift as training progresses, we believe that reinforcement learning methods, especially these algorithms like PPO that can maintain the exploration characteristics, can handle this situation well. In addition, since the “Chunk” stage transforms long sequences into multiple blocks, selecting from one block at a time step can also be intuitively seen as a Markov Decision Process. In our paper, there are many empirical results that prove the effectiveness of our RL selection strategy. For example, in Table 4, we observed that adding the “Select” module can greatly improve the performance of the model. In fact, in our early research, we did not have an RL selection module, and the performance of the model at that time was worse than the current strong baselines. In addition, in Figure 3, from the left graph, we can see that with the progress of training, the number of selected hidden representations first drops rapidly, and then converges to around our preset value (2048). From the right graph, we can see that even though the input length at the inference stage is exponentially increased by adding pad tokens, the number of selected hidden representations still grows slowly. This indicates that our selector has the ability to filter out low-contribution representations. Based on these results, we believe that our RL objective is effective and necessary.

Reviewer comment (Experiments):
The authors only conduct experiments on BART models. It remains unclear whether the proposed method still works on more recent models, like LLAMA.

Author response:
Our method is applicable to most pre-trained encoder-decoder models, including BART, PEGASUS, T5, LED and PRIMERA. For decoder-only models like LLAMA, we believe that although our method cannot be applied directly, it still has high reference value.

Reviewer comment (Evaluation):
How about the PPL scores on long text modeling? It is a widely-used metric to evaluate the performance of long-text language modeling.

Author response:
Following recent long-sequence processing work such as Unlimiformer [1], we use commonly used metrics such as ROUGE and BERTScore in specific downstream tasks for more persuasive model performance comparisons. Thank you for your suggestion, we have added the PPL results on the test set.

|                        | arXiv | PubMed | GovReport | SummScreen | Multi-News | WCEP | NarrativeQA |
| ---------------------- | ----- | ------ | --------- | ---------- | ---------- | ---- | ----------- |
| BART$_{base}$          | 3.46  | 3.55   | 3.74      | 3.89       | 3.90       | 3.93 | 3.15        |
| BART$_{base}$ + SimCAS | 1.93  | 1.72   | 1.83      | 2.38       | 2.11       | 2.18 | 1.55        |

[1] Bertsch A, Alon U, Neubig G, et al. Unlimiformer: Long-range transformers with unlimited length input[J]. arXiv preprint arXiv:2305.01625, 2023.

Reviewer comment (Experiments):
The proposed method depends on cross-attention, and we could not introduce it to the encoder-only or decoder-only Transformer model. If it can outperform existing LLMs, this would not be a weakness.

Author response:
In our experiment, we utilized the encoder-decoder transformer BART as the backbone. This choice was made because, among pre-trained language models of a relatively smaller scale, this type of transformer architecture often exhibits superior performance in downstream tasks. Simultaneously, our method is well-suited to the characteristics of the encoder-decoder transformer, resulting in a low degree of coupling between our method and this type of transformer. For encoder-only and decoder-only transformers, our method may not be applied directly, but it still holds sound reference value for  long-sequence processing of encoder-only model and decoder-only model, and the environment construction of RL-based selection policy.

Reviewer comment (Experiments):
Experiments on other tasks are limited. This paper mainly experiments on summarization tasks, but does not experiment on CNN/DM or Xsum, which are the most compared data. For other tasks, this paper only does Narrative QA. More experiments on document translation or classification would be an advantage.

Author response:
The CNN/DM and XSum datasets are commonly used for text summarization tasks. However, the average input text length for the former is 766.1, while for the latter it is 430.2. Therefore, standard full-attention models can already handle them well on consumer-grade GPUs, eliminating the need to introduce long sequence processing techniques. Based on your suggestions, we have incorporated additional experiments on the IMDB document classification dataset. The results are presented as follows.

|   Method   | Accuracy |
| :--------: | :------: |
|  RoBERTa   |   95.0   |
|  BIGBIRD   |   95.2   |
| Longformer |   95.7   |
|   SimCAS   |   96.1   |

Reviewer comment (Presentation):
There is a typo in Figure 3, the blue line should be selected length.

Author response:
Thank you for pointing out this issue. We will correct this error in the new version.

Reviewer comment (Evaluation):
I would like to see a comparison regarding latency.

Author response:
Based on a V100 and the backone BART$_{base}$, we tested the average inference time (ms) of a single sample in each test set for SLED and SimCAS, and the results are as follows:

| Method |  arXiv  | PubMed  | GovReport | SummScreen | Multi-News |  WCEP  | NarrativeQA |
| :----: | :-----: | :-----: | :-------: | :--------: | :--------: | :----: | :---------: |
|  SLED  | 28375.3 | 25323.0 | 153560.5  |  59012.7   |  37253.0   | 3120.6 |   1711.5    |
| SimCAS | 9661.5  | 13799.1 | 104953.7  |  30424.0   |  17664.2   | 1356.2 |    689.6    |