Summary
The paper introduced a novel inference strategy for transformer models that focuses on inference efficiency. Instead of retaining all context tokens throughout the entire inference process, they gradually eliminate tokens as they move deeper into the layers. To determine which tokens to drop, they trained some small linear layers to predict the remaining relevant context. Their experiments revealed that approximately 80% of the tokens could be safely discarded without much adverse impact on downstream task performance or perplexity. As a result, this approach significantly reduces the computational resources required for inference when the context length exceeds 500 tokens.
Strengths
- The paper demonstrates excellent writing with a clear flow of ideas.
- The authors introduce a unique and innovative data structure that enables batch operations involving masked tokens.
- The experimental results indicate that in scenarios where long context (> 500 tokens) is involved in the inference process, the models can safely drop up to 80% of the tokens without any impact on perplexity.
Weaknesses
- **The choice of downstream tasks raises questions:** The downstream tasks evaluated in Figure 5 primarily involve small context sizes, and Figure 7 indicates that a smaller context leads to reduced throughput compared to the standard dense model. It would be more persuasive to demonstrate that task performance remains intact when longer contexts are required, while simultaneously achieving gains in inference efficiency.
- **Insufficient experiments with stronger base models:** It appears that the proposed method inevitably leads to performance degradation for larger and more capable base models, as evident from both Figure 5 and 7. This is likely because stronger base models are better at utilizing contexts, and additional contextual information enhances performance. To strengthen the paper's argument, the authors should include more results using larger and stronger base models (>1.5B parameters), such as LLaMA, Pythia, or even OPT.
- **Evaluating generation quality:** Though the authors perform evaluations on language modeling with perplexity, it does not necessarily align with generation quality. If would be helpful if the authors could further provide evidence that dropping context tokens does not affect generation quality.
Questions
- The paper introduces the use of sparse sigmoid to gradually enforce sparsity on the context tokens. It is important to understand why sparse sigmoid techniques were chosen and if they offer any specific advantages. Could other sparsity techniques, such as l0-regularization with hard-concrete distributions (Louizos et al., 2017), or techniques used in movement pruning (Sahn et al., 2020), achieve similar results?
- There seems to be a typo in Figure 5. The green curve should be labeled as "Sparse Attention".
- Figure 2 appears to be a bit confusing without proper legends. While it is clear that "X" denotes dropped tokens, the meaning of the red blocks is unclear. Could you provide an explanation or add appropriate legends?
Rating
7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed 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.
Limitations
The authors mentioned the limitation of the working being exclusively tested on autoregressive language models, and specifically GPT2 model family. The paper would be stronger if the author can show positive results on stronger base models.