Title: Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers
Venue: NeurIPS 2023
Accepted: yes
Keywords: Transformers

Context-pruning

Efficient Transformer
Abstract: Autoregressive Transformers adopted in Large Language Models (LLMs) are hard to scale to long sequences. Despite several works trying to reduce their computational cost, most of LLMs still adopt attention layers between all pairs of tokens in the sequence, thus incurring a quadratic cost. In this study, we present a novel approach that dynamically prunes contextual information while preserving the model's expressiveness, resulting in reduced memory and computational requirements during inference. Our method employs a learnable mechanism that determines which uninformative tokens can be dropped from the context at any point across the generation process. By doing so, our approach not only addresses performance concerns but also enhances interpretability, providing valuable insight into the model's decision-making process. Our technique can be applied to existing pre-trained models through a straightforward fine-tuning process, and the pruning strength can be specified by a sparsity parameter. Notably, our empirical findings demonstrate that we can effectively prune up to 80\% of the context without significant performance degradation on downstream tasks, offering a valuable tool for mitigating inference costs. Our reference implementation achieves up to $2\times$ increase in inference throughput and even greater memory savings.

Dynamic Context Pruning for Efficient and Interpretable Autoregressive Transformers

Abstract

Autoregressive Transformers adopted in Large Language Models (LLMs) are hard to scale to long sequences. Despite several works trying to reduce their computational cost, most of LLMs still adopt attention layers between all pairs of tokens in the sequence, thus incurring a quadratic cost. In this study, we present a novel approach that dynamically prunes contextual information while preserving the model's expressiveness, resulting in reduced memory and computational requirements during inference. Our method employs a learnable mechanism that determines which uninformative tokens can be dropped from the context at any point across the generation process. By doing so, our approach not only addresses performance concerns but also enhances interpretability, providing valuable insight into the model's decision-making process. Our technique can be applied to existing pre-trained models through a straightforward fine-tuning process, and the pruning strength can be specified by a sparsity parameter. Notably, our empirical findings demonstrate that we can effectively prune up to 80% of the context without significant performance degradation on downstream tasks, offering a valuable tool for mitigating inference costs. Our reference implementation achieves up to 2 × increase in inference throughput and even greater memory savings.

1 Introduction

The introduction of Transformers [Vaswani et al., 2017] in Large Language Models (LLMs) has profoundly influenced the landscape of Natural Language Processing (NLP), due to their appealing scaling properties [Kaplan et al., 2020] and their ability to train efficiently on modern hardware architectures designed for extensive parallel computing. As LLMs grow larger and more complex, the challenges associated with training and deploying them become more prominent. Especially challenging is the quest for processing increasingly longer sequences, as pure self-attention layers scale quadratically in sequence length during train and inference.

To address this limitation, several efforts focus on efficient implementations of the attention mechanism on dedicated hardware [Dao et al., 2022, Touvron et al., 2023], or on algorithmic procedures to directly tackle the quadratic complexity. The latter direction has led to numerous variants sacrificing the generality of the standard attention mechanism in favor of more efficient alternatives [Tay et al.,

2020, Kitaev et al., 2020, Choromanski et al., 2020b, Katharopoulos et al., 2020, Zaheer et al., 2020, Shi et al., 2021, Lin et al., 2022, Zhu and Soricut, 2021, Dai et al., 2020], some of which are illustrated in Fig. 1. Specifically, a large number of these methods focus either on sparsifying the attention weights, reducing the size of the available context to each token, or compressing the number of tokens to reduce the size of the attention matrix.

These methods, however, are inherently static, in the sense that each token is either forced to attend to a fixed pre-specified context window, or the input context is compressed to a fixed dimensionality, regardless of the information content of the input sequence. Furthermore, a performance gap still exists with respect to pure self-attention in many applications, thus implying the existence of a non-trivial trade-off between the span of the attention context and the model's capabilities [Dao et al., 2022, Sun et al., 2021, Beltagy et al., 2020].

To address these challenges, and enhance inference efficiency, while staying faithful to pure selfattention, we pose the following question:

Can we dynamically prune past content based on the available context, while preserving as much as possible the expressivity of the model?

In response to this question, we introduce a novel method for context pruning in Transformer-based decoder architectures. Our approach adds a minimal amount of additional training parameters that enable individual tokens to dynamically remove portions of the input sequence in a layerwise fashion. Once part of the context is removed, it is disregarded for the remaining part of the autoregressive generation process, leading to reduced memory usage and computational requirements during inference. To this end, we also design a dynamic data structure that implements efficient insertion/removal of tokens from the context while supporting batched inference. In contrast to traditional methods relying on local or sparse attention, which may not capture the nuances and dynamic nature of the data over long contexts, ours leverages contextual cues to dynamically determine the relevance of the available information through a learned mechanism. This is achieved by making use of a sparse sigmoid function [Peters et al., 2019, Martins et al., 2020]. As demonstrated by our experimental evaluations, this allows us to extract and utilize essential details in a more adaptive and accurate manner. The degree of pruning can be effectively controlled through a hyperparameter that effectively accounts for the sparsity level.

Our technique serves as a modular building block for existing pre-trained models and can be easily integrated through a minimal fine-tuning stage. For our study, we focus on GPT-2 models [Radford et al., 2019] as they are publicly available and widely benchmarked, but due to the uniformity of modern architectures, our approach can be straightforwardly extended to any autoregressive Transformer. Moreover, since our method is based on context pruning, it can be seamlessly combined with other approaches aimed at improving inference efficiency, such as quantization, weight pruning, approximate attention, or other hardware optimizations.

We find that up to 80% of the context can be successfully pruned, with minimal deterioration in terms of perplexity and zero-shot performance, while requiring significantly fewer resources during inference. We showcase how these improvements can lead to measurable practical gains, by providing an efficient implementation that reduces memory usage for caching during token generation. More specifically, for larger context sizes we get up to 50% wall-time latency reduction for each generation step, while still decoding with up to 2 × larger batch sizes, leading thus to significant performance benefits. These findings highlight the potential of context pruning as a powerful technique to enhance the efficiency and interpretability of Transformers in NLP.

2 Related Work

Despite exhibiting human-level performance on a number of challenging tasks, LLMs are resourceintensive and inefficient. While the human brain consumes roughly the amount of energy equivalent to a dim light bulb, top-performing GPT models require multiple GPUs with ∼ 80GB of memory each for inference [Strubell et al., 2019, Frantar and Alistarh, 2023a]. Several research efforts have been focusing on improving their efficiency and memory requirements from several different angles.

Weight Pruning and Quantization. Modern LLMs have high memory and compute requirements for both training and testing. To address this limitation, a number of research efforts [Kwon et al.,

Attention

Local Attention

Sparse Attention

Adaptively Sparse Attention

Layer specific

Current Token

Attended Tokens

Masked Tokens

2022, Frantar et al., 2023, Frantar and Alistarh, 2023b] have resorted to the established practice of weight pruning [Hassibi et al., 1993] to efficiently compress the original model to a more manageable size. Remarkably, a large percentage of the original weights can be safely removed, resulting in only marginal perplexity growth [Bahl et al., 1983]. An alternative approach to reduce the memory and compute is quantization [Dettmers et al., 2022, Yao et al., 2022, Xiao et al., 2022, Frantar et al., 2022], which reduces the precision of the model's numerical representation. Quantization schemes [Dettmers et al., 2022] enable 8-bit matrix multiplication for both feed-forward and attention projection layers resulting in significantly improved memory allocation without incurring any performance degradation.

Efficient Transformers and context pruning. One primary constraint of Transformer-based models is their quadratic complexity with respect to the length of the input sequence. Extensive research explores alternatives that exhibit sub-quadratic scaling, resulting in three main strategies [Lin et al., 2022]. The first replaces the attention mechanism with an alternative operation that features more favourable scaling with the input sequence length [Peng et al., 2021, Katharopoulos et al., 2020, Choromanski et al., 2020a, Schlag et al., 2021]. While several recent methods in this category show promise, none have emerged as a definitive winner, and most state-of-the-art language models still rely on the standard attention mechanism [Touvron et al., 2023, Chowdhery et al., 2022]. The second approach proposed to compress the length of the input context, controlling the complexity of the attention operation but unavoidably sacrificing potentially relevant information from the original input [Lee et al., 2019, Wang et al., 2020, Jaegle et al., 2021]. The third approach involves pruning the attention matrix, preventing each token from attending to every other token within the context [Zaheer et al., 2020, Martins et al., 2020, Lee et al., 2023]. This line of research is motivated by the theoretical finding highlighting that sparse Transformers retain the expressivity of their dense counterparts [Yun et al., 2020]. Many methods in this category employ specially designed attention masks that aim to zero out as many entries as possible, often based on principles of locality, randomness, or a combination of both. The main drawback of these methods is their mostly static nature, meaning that every token is compelled to attend to a fixed context window and disregard the rest of the context regardless of its specific role within the input sequence. Our approach falls within this last category, and enables dynamic sparsification of the attention matrix for decoder models, without resorting to any potentially restricting inductive biases about its structure.

Implementation Speed-up Recently, hardware-optimized implementations [Dao et al., 2022, Touvron et al., 2023] have been proposed with the aim of optimizing computational resources during the training phase of Transformers [Hoffmann et al., 2022]. On the other hand, as recent breakthroughs have led to widespread adoption of these models [Ouyang et al., 2022, OpenAI, 2023, Köpf et al., 2023], performance during inference becomes more relevant by the day. In decoderbased autoregressive Transformers, the backbone architecture of most current state-of-the-art LLMs, inference involves evaluating and generating tokens one by one, using cached previous activations to avoid redundant computations. In contrast to training, the inference is memory bound [Shazeer, 2019, Ivanov et al., 2021, Pope et al., 2022]. Compute is under-utilized, especially when deploying larger models, as the time required to transfer model parameters and activations to hardware memory far exceeds the actual computational time. This is further exacerbated by recent trends to ever-increase the model size and enable longer context windows. As a result, batch decoding, a promising direction for more efficient utilization of hardware resources, is impeded.

A

transformer

is

a

deep

learning

model

.

It

Is

A

transformer

is

a

deep

learning

model

.

A

transformer

is

a

deep

learning

model

.

It

A

transformer

is

a

deep

learning

model

(Position) ID:

0

1

2

3

4

5

6

7

8

9

ID: 0

ID: 1

ID: 2

ID: 3

ID: 4

ID: 5

Empty

Empty

Memory buffer for layer

Stored data in the form (ID, cached activations for ID)

ID: 0

ID: 1

ID: 2

ID: 3

ID: 4

ID: 5

ID: 6

Empty

ID: 7

ID: 1

ID: 2

Empty

ID: 4

ID: 5

ID: 6

Empty

ID: 7

ID: 1

ID: 8

Empty

ID: 4

ID: 5

ID: 6

Empty

Adaptively Sparse Attention at layer

3 Methodology

Background. We operate on sequences of text tokens T ∈ { 0 , 1 , . . . , n vocab } n , where n is the length of the sequence and n vocab is the vocabulary size. Tokens are embedded into X 0 ∈ R n × d using an embedding layer, where d is the embedding dimension of the model. When necessary, we use the superscript ℓ ∈ { 1 , 2 , . . . , L } to denote the representations and weights at different layers. One layer of the Transformer-decoder architecture [Vaswani et al., 2017] is defined as

where MHA stands for Multi-head self-attention defined as

Here Q i = XW Q i , K i = XW K i , and V = XW V i are the queries, keys and values and SA denotes the single-head self-attention. The weight matrices W Q i , W Ki , W V i ∈ R d × p linearly project the input embedding into the head dimension p . Finally, W O ∈ R d × d is the output projection. The feed-forward part of the Transformer is defined as

where σ FF is a nonlinearity, and W F 1 , W F 2 are linear layers with typical dimensions W F 1 ∈ R d × 4 · d and W F 2 ∈ R 4 · d × d . A final projection layer W logits ∈ R d × n vocab is used to project back to the vocabulary space and predict the next token from the representations X L . We are focusing on Pre-LN [Xiong et al., 2020] decoder-only architectures, meaning that attention is causally masked, i.e. every input token i attends to the first i tokens in the input sequence. Conceptually, our method acts by predicting these attention masks using a learned mechanism in a layer-wise manner, with the introduction of additional constraints to make sure causality is preserved (i.e. if a token is dropped, it will remain dropped in the future). During inference, however, our method can efficiently be implemented by erasing tokens from the key-value cache commonly adopted in autoregressive attention models.

Background: key-value cache. In autoregressive Transformers, inference can be optimized by reusing pre-computed activations (keys and values) to accelerate the sequential generation of tokens [Ott et al., 2019, Vaswani et al., 2018, Wolf et al., 2020], bringing down the computational cost to generate a single token to O ( n ) from O ( n 2 ) (where n is the sentence length). Most existing sparse attention techniques ignore the specifics of this process and focus on sparsifying each attention operation separately. As non-attended tokens can still be attended to by subsequent tokens, memory benefits are limited. By contrast, our approach is compatible with this setting, allowing us to design an efficient batched data structure where dropped tokens are effectively removed from the computation.

3.1 Adaptively Sparse Attention

We allow the network to selectively drop parts of the context that are no longer required. An illustration of our proposed method can be seen in Fig. 2. At each layer, we introduce the parameters W ℓ Q int , W ℓ K int ∈ R d × r for dimension r ∈ R , that calculate the interaction queries and keys Q ℓ int , K ℓ int ∈ R n × r , as Q ℓ int = X ℓ W ℓ Q int and K ℓ int = X ℓ W ℓ K int . We then calculate the interaction of token k with token j at layer ℓ as:

where σ ( · ) denotes the sparse sigmoid function introduced in Section 3.2 and β ℓ ∈ R a scalar parameter per layer, that controls the initial sparsity as seen in Fig. 3 (right). Indices in Q ℓ int , K ℓ int ∈ R n × r refer to the rows of the matrices. We can then modify the self-attention

For j > k we set I ℓ k,j = 0 , which leads to masking entries in the self-attention, corresponding to the regular causal masking. We also impose that a token cannot drop itself, thus I ℓ k,k = 1 . We want to preserve information regarding the current token as its predictions are particularly important in determining the next token for the regular language modeling task that we are considering. Small values of I ℓ n,j impose partial masking of the corresponding token in the attention, and complete masking occurs when I ℓ n,j = 0 . The cumulative product over tokens j +1 → k in Eq. (6) imposes that dropping a token (when σ ( . ) → 0 ) has an irreversible effect, as it will remain dropped for all subsequent tokens, and hence for the remaining of the generation process. The complexity of the pruning logic is O ( n · d · r + n 2 · r ) , which is lower than the one of the self-attention operation for r < d .

Our mechanism allows layers to act independently, meaning that different sparsity patterns are encountered across layers. We also experimented with tying the model's dropping decisions with depth by imposing that a token dropped at a given layer cannot be attended to in subsequent layers. However, we observed worse results and hence did not pursue this further. This is perhaps expected, given numerous results and interpretability studies regarding sparsity patterns of attention heads at different layers [Ramsauer et al., 2020, Hao et al., 2021].

3.2 Sparse Sigmoid

In Eq. (6), we use σ ( · ) , as a sigmoid-like function to let the network decide when and what to drop. We favour binary decisions, leading to interaction values of either 0 or 1 . Inspired by the α -entmax function introduced in Peters et al. [2019], Martins et al. [2020], we define the α -sigmoid (based on the entropies proposed by Tsallis [1988]) as:

where

̸

By varying α during the training, we can control the sparsity in the network, i.e. regulate the softness of the pruning mechanism. In practice, we start from small values of α = 1 and increase it according to a cosine scheduler, as shown in Fig. 3. Small values of α allow meaningful gradient signals to pass through the dropping mechanism, which is crucial at the beginning of training. On the other hand, larger values of α lead to sparse results desired during inference. We thus increase α to values leading to very sparse solutions, as illustrated in Fig. 3. In practice, during inference, we replace σ ( · ) with the step function, that corresponds to α →∞ . We also initialize the biases parameters β ℓ in (6) to a positive value, ensuring that tokens at the beginning of training have a prior towards not being dropped. This strategy also facilitates fine-tuning existing pretrained models, as our module

Training Steps

1

3

5

7

Value of

3

2

1

0

1

2

3

0.0

0.2

0.4

0.6

0.8

1.0

Step function

0

100

200

300

400

500

600

700

800

900

1000

Context Length

0.0

2.0

4.0

6.0

8.0

10.0

0.0

0.2

0.4

0.6

0.8

1.0

will initially default close to the identity function. The α -sigmoid along with the training schedule on α allows for good signal propagation properties for the gradients [Noci et al., 2022]. We also explored using a regular sigmoid with a varying temperature [Kim et al., 2022], leading to suboptimal nonbinary predictions and instabilities during training. Training with our sparse sigmoid also directly eliminates the need to have any auxiliary network [Lee et al., 2023].

3.3 Regularized Objective

We augment the regular language modeling objective with a regularization that incentivizes the network f to drop parts of the sequence. We fine-tune pretrained models, with parameters θ , using the objective:

is the regular cross-entropy loss for the language modeling task based on the original and shifted input tokens T , and

is the sparsity loss, encouraging the model to prune the context. In total ( Ln ( n -1)) / 2 entries of I ℓ i,j are learned, as indicated in Eq. (6). We choose γ > 0 to enforce different levels of sparsity. In general, for a current position i in the context, we define as sparsity, the percentage of the previous tokens dropped, i.e. ( tokens ≤ i dropped ) /i .

4 Experiments

We fine-tune pretrained GPT-2 models 1 , that support a context size of up to 1024 tokens, on a subset of the English Wikipedia 20220301.en and English bookcorpus datasets. We keep a separate test set where we report perplexity after training. All models shown, for a fair comparison, were fine-tuned using the same lightweight training setup as described in Appendix A. When using our adaptive sparse attention, we use a cosine scheduler for the α parameter as displayed in Fig. 3 and specify r = 64 for the dimensions of W ℓ Q int , W ℓ K int . More ablations regarding optimization and variations of our dropping mechanism are provided in Appendix B. Unless otherwise stated, results refer to GPT-2small models. We use the term dense for the regular GPT-2 models, fine-tuned without any additional W Q int , W K int parameters.

where

Baselines. We compare against the baselines presented in Fig. 1. Local Attention refers to a causal attention mask, where each token attends to the previous k tokens in the sequence, including itself. This can also be interpreted as restricting the receptive field of the model. Sparse Attention refers to the baselines from Child et al. [2019], Lin et al. [2022], where each token i attends to the tokens satisfying (1) ⌊ i/k ⌋ = ⌊ j/k ⌋ and (2) the tokens k -1 , 2 · k -1 , . . . , ⌊ i/k ⌋ · k -1 (numbering starts from zero). We fine-tune these baselines using the same aforementioned fine-tuning procedure, for different choices of k , leading to different levels of sparsity, depending on the current context size.

Data structure. Real-world deployment of our approach exhibits numerous challenges due to the nature of batched generation. In particular, we highlight differences in prompt length (initial prefix), different final lengths (termination criteria), and uneven dropping of tokens across different sentences. Maximum performance is achieved when the key-value cache is represented as a contiguous block of memory, and any masking resulting from padding or removed tokens ('holes') will result in a decrease in efficiency. To this end, we devise an efficient batched data structure that allows for efficient insertion and deletion of tokens (leveraging the set nature of the self-attention operation), while (i) allowing the underlying storage to be processed as a contiguous block of memory and (ii) ensuring that the load factor of the data structure is high enough to guarantee a performance speed-up. More details are provided in the Appendix A.

4.1 Results

Perplexity vs sparsity. We first study how context-pruning changes for different levels of sparsity in Fig. 4. Depending on the current context size, our method allows for up to 80% of the context to be successfully pruned, i.e. removed, with no performance loss in terms of perplexity (-0.085 average gain in perplexity when context size is 1000 tokens for 80.35% of sparsity compared to the dense counterpart). Our method also adapts to the current context size, meaning a network trained with specific sparsity regularization exhibits different levels of sparsity depending on the current context size. Compared to the baselines, our method exhibits consistently lower perplexity results for the same level of sparsity.

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

10

15

20

25

30

35

Context size 100

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

10

15

20

25

30

35

Context size 500

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

10

15

20

25

30

35

Context size 1000

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

10

15

20

25

30

35

Perplexity

Context size 1-1024

Dense

=0.007

=0.01

=0.03

=0.1

=0.3

=1.0

=3.0

=10.0

Dense

Local Attention

Sparse Attention

Ours

Dense

Local Attention

Sparse Attention

Ours

Zero-Shot Performance. To test general model capabilities and complement perplexity evaluations, we provide results on several zero-shot tasks [Dettmers et al., 2022] in Fig. 5. Similar trends hold overall; our approach retains or even outperforms the performance of the dense baseline, even for cases with high sparsity. These tasks involve scenarios where the model is required to perform without any specific training or prior exposure to the target domain. The results obtained validate that the models' general capabilities can be retained, even under high levels of sparsity.

Computational Analysis. We analyse the gains in terms of FLOPs and required memory when generating new sequences due to caching in Fig. 6. Our dropping mechanism introduces additional computational costs for the calculation of Q ℓ int , K ℓ int and the logic behind dropping via Eq. (6). Due to the relatively small chosen parameter r , i.e. the output dimension of the interaction weights W ℓ Q int , W ℓ K int , these are nevertheless minimal. Although the raw FLOPs benefit when using sparse models does not seem very significant, as aforementioned, inference is predominately memory-bound. The attention thus takes a significant proportion of real-time inference [Dao et al., 2022]. On the contrary, dense matrix multiplications used for all linear projections are very efficient. Memory

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

0.36

0.38

0.40

0.42

Mean zeroshot accuracy

Local Attention

Sparse Attention

Dense

Ours

0.0

0.2

0.4

0.6

0.8

Sparsity

0.38

0.40

0.42

0.44

0.46

0.48

0.50

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

benefits, on the other hand, are substantial, as the memory required for caching is a linear function with respect to sparsity, with a negative slope. Sparser solutions will thus additionally allow us to generate more sequences in a batched fashion. This is particularly relevant for bigger models, also longer sequences, where batch decoding is a major challenge [Shazeer, 2019].

Increasing sparsity

Increasing sparsity

0.00

0.05

0.10

0.15

0.20

0.25

0.30

0.35

GFLOS

FLOPs for Context-size 1000

embedding-layer

logits-layer

feed-forward

qkvo-calculation

attention

drop-tokens

0.000

0.002

0.004

0.006

0.008

GB

Cached Memory for Context-size 1000

keys+values_attention

keys_int

Throughput. We demonstrate how reduced context and reduced memory requirements can lead to significant real-world time throughput in Fig. 7. Initially, our pruned networks are slower in terms of latency for small context lengths, because of the additional cost associated with the logic behind pruning. Nevertheless, they quickly surpass the dense baseline that struggles as the context size increases. This verifies the fact that although raw FLOPs benefits look unsubstantial, in fact, this leads to significant gains due to the specific memory profile of Transformers' inference. Crucially, our pruned networks can support a much bigger batch size, leading to significant throughput gains. More specifically, for long context sizes, our GPT-2small model offers an additional 98% margin in throughput for a loss in perplexity of only 0 . 316 , with respect to the dense counterpart. Similarly, our GPT-2medium model can yield 189% additional throughput for only 0 . 084 loss in perplexity for a context size of 1000 tokens. In particular, the same model (for γ = 1 . 0 ) provides a higher throughput than a GPT-2small model, while achieving 3 . 769 lower perplexity. As context windows become larger by the day in state-of-the-art models, we expect these gains to become even more relevant.

256

512

1024

2048

4096

8192

16384

10

12

14

16

18

20

22

24

26

Average Perplexity

GPT-2-small

7293 tokens/s

GPT-2-medium

2634 tokens/s

GPT-2-large

1241 tokens/s

GPT-2-xl

569 tokens/s

Average throughput

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

Dense

256

512

1024

2048

4096

8192

16384

Tokens/s

10

12

14

16

18

20

22

24

26

Perplexity when context size is 1000

GPT-2-small

4115 tokens/s

GPT-2-medium

1521 tokens/s

GPT-2-large

765 tokens/s

GPT-2-xl

399 tokens/s

Throughput when context size is 1000

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

Dense

0

200

400

600

800

1000

2000

4000

6000

8000

10000

Tokens/s

Throughput per generation step

0

200

400

600

800

1000

Context size

2×10 2

3×10 2

4×10 2

6×10 2

Seconds

Time per generation step

Optimal batch sizes

baseline

baseline × 0.96

baseline × 1.31

baseline × 1.92

baseline × 3.00

Dense

=0.007

=0.01

=0.03

=0.1

=0.3

=1.0

=3.0

=10.0

The

site

of

V

auc

ou

le

urs

(

North

)

was

selected

for

Air

Service

use

in

January

1918

.

The

general

region

was

chosen

as

the

center

for

the

future

American

Sector

and

the

site

was

one

of

those

available

after

three

years

of

use

by

the

French

.

It

was

located

75

kilometers

from

the

front

line

,

considered

safe

for

use

by

pursuit

(

fighter

)

aircraft

.

The

terrain

use

by

pursuit

(

fighter

)

aircraft

.

The

terrain

Generated tokens

Pruned

Context

NOUN

PUNCT

DET

ADP

PROPN

AUX

VERB

NUM

ADJ

CCONJ

PRON

PART

ADV

Part of speech

0.0

0.1

0.2

0.3

0.4

0.5

0.6

Alive probability

1

2

3

4

5

6

7

8

9

10

11

12

Layer

0

5

10

15

20

25

30

35

40

45

Number of dropped tokens

NOUN

PUNCT

DET

ADP

PROPN

AUX

VERB

NUM

ADJ

CCONJ

PRON

PART

ADV

0

200

400

600

800

1000

0

5

10

Tokens dropped layer 5

0

200

400

600

800

1000

Context size

0

5

10

Tokens dropped layer 6

Interpretability. Fig. 8 provides insights into the interpretability aspect of the model's decisionmaking process. It is observed that token removal predominantly occurs when encountering stop words (punctuation), which aligns with the intuition that local information within a sentence becomes less relevant after its completion. Furthermore, it is worth noting that layers at varying depths exhibit distinct behaviours, reinforcing our rationale for dissecting token removal decisions across depth. The variance in sparsity distribution across different depths indicates the necessity of conducting additional interpretability research to obtain valuable insights in the interactions of the tokens within the model. We provide more insights towards this direction in the Appendix C.

5 Discussion

We proposed Adaptively Sparse Attention, a novel approach to dynamically prune the context in decoder-only Transformer architectures. Our results indicate that our technique performs favourably compared to competitive baselines in terms of the ratio between perplexity and sparsity of the attention weights. Remarkably our approach also significantly reduces the computational and memory requirements without affecting its final performance. We practically showcase these benefits achieving more than double the throughput at cases. Adaptively sparse attention comes with two additional practical advantages: first, it can be seamlessly integrated into existing pre-trained models via a cheap fine-tuning step; second, it represents an orthogonal contribution to the burgeoning research line aimed at increasing the level of efficiency of modern LLMs. As such, we envision its combination with existing techniques like weight pruning and quantization to be a promising avenue for future research.

References

Lalit R Bahl, Frederick Jelinek, and Robert L Mercer. A maximum likelihood approach to continuous speech recognition. IEEE transactions on pattern analysis and machine intelligence , (2):179-190, 1983.

Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05150 , 2020.

Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence , volume 34, pages 7432-7439, 2020.

Daniel Bolya, Cheng-Yang Fu, Xiaoliang Dai, Peizhao Zhang, Christoph Feichtenhofer, and Judy Hoffman. Token merging: Your vit but faster. arXiv preprint arXiv:2210.09461 , 2022.

Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904.10509 , 2019.

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, David Belanger, Lucy Colwell, and Adrian Weller. Masked language modeling for proteins via linearly scalable long-context transformers, 2020a.

Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. arXiv preprint arXiv:2009.14794 , 2020b.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Zihang Dai, Guokun Lai, Yiming Yang, and Quoc Le. Funnel-transformer: Filtering out sequential redundancy for efficient language processing. Advances in neural information processing systems , 33:4271-4282, 2020.

Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memoryefficient exact attention with io-awareness. Advances in Neural Information Processing Systems , 35:16344-16359, 2022.

Tim Dettmers, Mike Lewis, Younes Belkada, and Luke Zettlemoyer. Llm. int8 (): 8-bit matrix multiplication for transformers at scale. arXiv preprint arXiv:2208.07339 , 2022.

Elias Frantar and Dan Alistarh. Massive language models can be accurately pruned in one-shot. arXiv preprint arXiv:2301.00774 , 2023a.

Elias Frantar and Dan Alistarh. Sparsegpt: Massive language models can be accurately pruned in one-shot, 2023b.

Elias Frantar, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. Gptq: Accurate post-training quantization for generative pre-trained transformers. arXiv preprint arXiv:2210.17323 , 2022.

Elias Frantar, Sidak Pal Singh, and Dan Alistarh. Optimal brain compression: A framework for accurate post-training quantization and pruning, 2023.

Yaru Hao, Li Dong, Furu Wei, and Ke Xu. Self-attention attribution: Interpreting information interactions inside transformer. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 12963-12971, 2021.

Babak Hassibi, David G. Stork, and Gregory J. Wolff. Optimal brain surgeon and general network pruning. IEEE International Conference on Neural Networks , pages 293-299 vol.1, 1993.

Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In Proceedings of the IEEE international conference on computer vision , pages 1026-1034, 2015.

Jordan Hoffmann, Sebastian Borgeaud, Arthur Mensch, Elena Buchatskaya, Trevor Cai, Eliza Rutherford, Diego de Las Casas, Lisa Anne Hendricks, Johannes Welbl, Aidan Clark, et al. Training compute-optimal large language models. arXiv preprint arXiv:2203.15556 , 2022.

Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint arXiv:2106.09685 , 2021.

Andrei Ivanov, Nikoli Dryden, Tal Ben-Nun, Shigang Li, and Torsten Hoefler. Data movement is all you need: A case study on optimizing transformers. Proceedings of Machine Learning and Systems , 3:711-732, 2021.

Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and Joao Carreira. Perceiver: General perception with iterative attention, 2021.

Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020.

Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020.

Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and Kurt Keutzer. Learned token pruning for transformers. In Proceedings of the 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining , pages 784-794, 2022.

Nikita Kitaev, Łukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. arXiv preprint arXiv:2001.04451 , 2020.

Andreas Köpf, Yannic Kilcher, Dimitri von Rütte, Sotiris Anagnostidis, Zhi-Rui Tam, Keith Stevens, Abdullah Barhoum, Nguyen Minh Duc, Oliver Stanley, Richárd Nagyfi, et al. Openassistant conversations-democratizing large language model alignment. arXiv preprint arXiv:2304.07327 , 2023.

Woosuk Kwon, Sehoon Kim, Michael W. Mahoney, Joseph Hassoun, Kurt Keutzer, and Amir Gholami. A fast post-training pruning framework for transformers, 2022.

Heejun Lee, Minki Kang, Youngwan Lee, and Sung Ju Hwang. Sparse token transformer with attention back tracking. In The Eleventh International Conference on Learning Representations , 2023.

Juho Lee, Yoonho Lee, Jungtaek Kim, Adam R. Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant neural networks, 2019.

Tianyang Lin, Yuxin Wang, Xiangyang Liu, and Xipeng Qiu. A survey of transformers. AI Open , 2022.

André Martins, António Farinhas, Marcos Treviso, Vlad Niculae, Pedro Aguiar, and Mario Figueiredo. Sparse and continuous attention mechanisms. Advances in Neural Information Processing Systems , 33:20989-21001, 2020.

Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation in transformers: Theoretical perspectives and the role of rank collapse. arXiv preprint arXiv:2206.03126 , 2022.

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. arXiv preprint arXiv:1904.01038 , 2019.

Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems , 35: 27730-27744, 2022.

Denis Paperno, Germán Kruszewski, Angeliki Lazaridou, Quan Ngoc Pham, Raffaella Bernardi, Sandro Pezzelle, Marco Baroni, Gemma Boleda, and Raquel Fernández. The lambada dataset: Word prediction requiring a broad discourse context. arXiv preprint arXiv:1606.06031 , 2016.

Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention, 2021.

Ben Peters, Vlad Niculae, and André FT Martins. Sparse sequence-to-sequence models. arXiv preprint arXiv:1905.05702 , 2019.

Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Levskaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling transformer inference. arXiv preprint arXiv:2211.05102 , 2022.

Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Hubert Ramsauer, Bernhard Schäfl, Johannes Lehner, Philipp Seidl, Michael Widrich, Thomas Adler, Lukas Gruber, Markus Holzleitner, Milena Pavlovi´ c, Geir Kjetil Sandve, et al. Hopfield networks is all you need. arXiv preprint arXiv:2008.02217 , 2020.

Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM , 64(9):99-106, 2021.

Imanol Schlag, Kazuki Irie, and Jürgen Schmidhuber. Linear transformers are secretly fast weight programmers, 2021.

Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150 , 2019.

Han Shi, Jiahui Gao, Xiaozhe Ren, Hang Xu, Xiaodan Liang, Zhenguo Li, and James Tin-Yau Kwok. Sparsebert: Rethinking the importance analysis in self-attention. In International Conference on Machine Learning , pages 9547-9557. PMLR, 2021.

Emma Strubell, Ananya Ganesh, and Andrew McCallum. Energy and policy considerations for deep learning in nlp. arXiv preprint arXiv:1906.02243 , 2019.

Simeng Sun, Kalpesh Krishna, Andrew Mattarella-Micke, and Mohit Iyyer. Do long-range language models actually use long-range context? arXiv preprint arXiv:2109.09115 , 2021.

Yi Tay, Mostafa Dehghani, Dara Bahri, and Donald Metzler. Efficient transformers: A survey.(2020). arXiv preprint cs.LG/2009.06732 , 2020.

Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

Constantino Tsallis. Possible generalization of boltzmann-gibbs statistics. Journal of statistical physics , 52:479-487, 1988.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, et al. Tensor2tensor for neural machine translation. arXiv preprint arXiv:1803.07416 , 2018.

Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.

Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowicz, et al. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 conference on empirical methods in natural language processing: system demonstrations , pages 38-45, 2020.

Guangxuan Xiao, Ji Lin, Mickael Seznec, Julien Demouth, and Song Han. Smoothquant: Accurate and efficient post-training quantization for large language models. arXiv preprint arXiv:2211.10438 , 2022.

Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference on Machine Learning , pages 10524-10533. PMLR, 2020.

Zhewei Yao, Reza Yazdani Aminabadi, Minjia Zhang, Xiaoxia Wu, Conglong Li, and Yuxiong He. Zeroquant: Efficient and affordable post-training quantization for large-scale transformers. Advances in Neural Information Processing Systems , 35:27168-27183, 2022.

Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. O (n) connections are expressive enough: Universal approximability of sparse transformers. Advances in Neural Information Processing Systems , 33:13783-13794, 2020.

Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems , 33:17283-17297, 2020.

Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can a machine really finish your sentence? arXiv preprint arXiv:1905.07830 , 2019.

Zhenhai Zhu and Radu Soricut. H-transformer-1d: Fast one-dimensional hierarchical attention for sequences. arXiv preprint arXiv:2107.11906 , 2021.

A Experimental Setup

We use the pretrained GPT-2 models from huggingface. Parameters of the architecture for these models are provided in Table 1.

Training. We fine-tune pretrained models on a subset of the English Wikipedia 20220301.en and English bookcorpus datasets, for a total of 25000 steps with a batch size of 6 . The datasets are provided by huggingface at https://huggingface.co/datasets/wikipedia and https: //huggingface.co/datasets/bookcorpus respectively. We use a learning rate of 1 e -4 for the small and medium models and 5 e -5 for the large and xl models with the Adam optimizer. We do not use any weight decay or any scheduler for the learning rate. For the self-attention operations we use flash-attention as provided by the scaled_dot_product_attention in pytorch-2.0 2 .

Evaluation. For the zero-shot accuracy experiments, we report accuracy on the WinoGrande [Sakaguchi et al., 2021], HellaSwag [Zellers et al., 2019], PIQA [Bisk et al., 2020] and LAMBADA [Paperno et al., 2016] datasets, similar to Dettmers et al. [2022]. As samples in these datasets have different lengths, we report as sparsity the mean sparsity over the samples for which predictions were made.

Efficient Memory Allocation. As we explained in Section 4, the computational cost of our method is greatly affected by the underlying data structure used for representing the key-value cache. Conceptually, the data structure should implement the following methods (all batched):

push() : inserts a new token ( K , V , and K int ).

get() : returns the keys/values added so far as a contiguous block of memory, as well as a binary mask used to represent padding and potential gaps due to removed tokens.

remove() : given a binary mask of the same shape as that returned by get() , removes the specified tokens from the data structure.

Ideally, the insertion and deletion operations should be as efficient as possible, while guaranteeing that the memory block returned by get() is as packed as possible (high load factor). Additionally, the data structure should support dynamic resizing as more tokens are inserted. A simple (yet inefficient) baseline consists in implementing the above interface as a dynamic array (i.e. a preallocated buffer that is dynamically resized once full) where erased tokens are simply masked out. Such an implementation, while correct, does not result in any memory and computation savings. Instead, motivated by the intuition that self-attention is a set operation - meaning that tokens do not need to be stored in sequential order - we recycle the memory slots of erased tokens. We insert new tokens at the leftmost available position in the data structure (Fig. 2), and ensure that the load factor (the ratio between the length of the longest sentence n and the capacity of the buffer) is always greater than a specified value. We choose a load factor of 0 . 9 and dynamically consolidate the data structure when the effective load factor falls below this threshold. We also mention that the asymptotic cost of these operations does not have a significant impact on the final performance, as the overall cost is still dominated by the O ( n ) cost of the self-attention mechanism (for a single generated token). In our implementation, both push() and remove() have a cost of O ( n ) , while get() has a cost of O (1) since it simply returns a view of the memory buffer. We also experimented with asymptotically faster implementations for push() and remove() (at O (log n ) using a priority queue), but found these to be slower in practice due to inefficient use of the GPU.

B Training Results and Ablations

Training Curves. We provide an example of a training curve in Fig. 9. Initially, the loss of the model is high, as at initialization many of the tokens are dropped, despite the introduced β ℓ parameters. Higher values of β ℓ can significantly mitigate this phenomenon. We noticed however that very high values of β ℓ lead to worse final solutions in terms of achieved complexity for a given level of sparsity. For our experiments we initialize β ℓ = 2 . 0 , ∀ ℓ ∈ { 1 , 2 , . . . , L } .

0

5000

10000 15000 20000 25000

Steps

2

4

6

8

10

Language Modeling Loss

0

5000

10000 15000 20000 25000

Steps

1

2

3

4

5

6

7

8

Sparse-sigmoid

0

5000

10000 15000 20000 25000

Steps

100

200

300

400

500

600

700

800

Non-pruned tokens when context-size is 1000

Setting α . As the value of α in the sparse-sigmoid functions rises, solutions become sparser, as indicated by the sparsity in Fig. 9 (right). During inference, we want to replace the sparse-sigmoid function with a simple step function, thus we should make the sparse sigmoid as similar as possible to the step function during training, to mitigate any inconsistencies caused by functional differences. In practice, we found no benefit from increasing α to values larger than 8 . 0 . The speed by which α is increased also plays a significant role. Increasing α too quickly does not allow for the new interaction parameters to be optimized correctly, leading to suboptimal solutions. We found that a cosine scheduler for 25000 steps was enough to lead to solutions of adequate sparsity. We present results when optimizing for a different number of steps in Fig. 10.

0

10000 20000 30000 40000 50000

Steps

2

4

6

8

10

12

Language Modeling Loss

0

10000 20000 30000 40000 50000

Steps

1

2

3

4

5

6

7

8

Sparse-sigmoid

0

10000 20000 30000 40000 50000

Steps

0

200

400

600

800

1000

Non-pruned tokens when context-size is 1000

Steps: 1000

Steps: 2500

Steps: 5000

Steps: 8000

Steps: 12000

Steps: 25000

Steps: 50000

Setting r . For our experiments, we used r = 64 for the embedding dimensions of the interaction weights. We experimented with different dimensions in Fig. 11. Larger dimensions lead to higher sparsity for the same perplexity but also require more memory to store the interaction keys. We found r = 64 to be a good compromise and tradeoff between extra memory and performance for the same sparsity.

1

2

4

8

16

32

64

128

256

512 1024

Dimension

r

0.0

0.5

1.0

1.5

2.0

2.5

3.0

Language Modeling Loss

0

5000

10000

15000

20000

25000

Steps

100

200

300

400

500

Non-pruned tokens when context-size is 1000

Propagation of Pruning with Depth. We also experimented with tying the dropping decisions with depth. In this case we modify Eq. (6) as:

Dropping a token in this case at a layer ℓ directly enforces that the token is dropped for all subsequent layers. Such a choice is inspired by similar pruning techniques in Transformer encoder models [Bolya et al., 2022] that usually, reduce the number of tokens with depth. We found, however, that this choice led to some disadvantages.

Firstly, we determined that sparsity is not a monotonic function of depth. Similar results have been shown before [Ramsauer et al., 2020]. Typically, the middle layers can be more sparse compared to the first and last layers. Secondly, for deeper models (GPT-2large and GPT-2xl ), we found that the large number of elements over which the cumulative product is taken over in Eq. (13), led to difficulties in learning. This can be perhaps expected given that no such training objective was optimized during the pre-training phase of these models. Finally, propagating the pruning decisions across depth significantly complicates the challenges for an efficient implementation.

Freezing Initial Parameters. We experimented with training just the interaction weights W ℓ Q int , W ℓ K int ∈ R d × r , β ℓ ∈ R , leaving the rest of the parameters of the network frozen. This led to an average increase in the perplexity of 9 . 285 for the same levels of sparsity, showing that modifications to the network's weights/logic are still necessary.

C Additional Results

Sparsity per Layer. In Fig. 12 we present the average level of sparsity per layer. In Fig. 13 we present similarly the number of tokens kept per layer for different sizes of the initial context.

Tokens that Cause Pruning. In Fig. 8 we presented which kind of tokens cause the pruning to take place. For the same example and settings, we present here exactly which tokens are dropped and when, in Fig. 14.

Context Switch. To better understand how well the dropping mechanism works, we create an artificial example, where we concatenate together text from three distinct, independent contexts. More specifically, we concatenate together the texts, as seen in Table 2. Inspecting the attention matrices in Fig. 15, reveals that the network has largely discovered the switches in the context and learned to ignore most of the preceding text if that comes from a different context.

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

Layer 1

Layer 2

Layer 3

Layer 4

Layer 5

Layer 6

0

200

400

600

800

1000

Context size

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

Layer 7

0

200

400

600

800

1000

Context size

Layer 8

0

200

400

600

800

1000

Context size

Layer 9

0

200

400

600

800

1000

Context size

Layer 10

0

200

400

600

800

1000

Context size

Layer 11

0

200

400

600

800

1000

Context size

Layer 12

dense

Reg. 0.007

Reg. 0.01

Reg. 0.03

Reg. 0.3

Reg. 1.0

Reg. 3.0

Reg. 10.0

0

200

400

600

800

1000

Number of tokens not dropped

Reg. 0.007

Reg. 0.01

Reg. 0.03

0

1

2

3

4

5

6

7

8

9

10 11

Depth

0

200

400

600

800

1000

Number of tokens not dropped

Reg. 0.3

0

1

2

3

4

5

6

7

8

9

10 11

Depth

Reg. 1.0

0

1

2

3

4

5

6

7

8

9

10 11

Depth

Reg. 3.0

0

1

2

3

4

5

6

7

8

9

10 11

Depth

Reg. 10.0

0

200

400

600

800

1000

Context size

Lebanon, officially the Republic of Lebanon or the Lebanese Republic, is a country in Western Asia. It is located between Syria to the north and east and Israel to the south, while Cyprus lies to its west across the Mediterranean Sea; its location at the crossroads of the Mediterranean Basin and the Arabian hinterland has contributed to its rich history and shaped a cultural identity of religious diversity.

A motorcycle (motorbike, bike, or trike (if three-wheeled)) is a two or three-wheeled motor vehicle steered by a handlebar from a saddle-style seat. Motorcycle design varies greatly to suit a range of different purposes: long-distance travel, commuting, cruising, sport (including racing), and off-road riding. Motorcycling is riding a motorcycle and being involved in other related social activity such as joining a motorcycle club and attending motorcycle rallies.

Zorba's Dance is an instrumental by Greek composer Mikis Theodorakis. The song featured for the dance, which has become known as sirtaki, in the 1964 film Zorba the Greek, for which Theodorakis wrote the soundtrack, and became renowned around the world. It is now commonly played and danced to in Greek tavernas. The film's track has since been recorded as a standalone song by many different musicians from around the world.

Ideal Generation Speed-up. In Fig. 7 we presented throughput in a realistic scenario where 'holes' formed in the memory buffer in an uneven fashion between samples in the batch. We expect that as we employ our method for larger models and contexts, the maximum feasible batch size will be reduced, mitigating partly this unwanted phenomenon. To evaluate the maximum possible speed-up,

The

site

of

V

auc

ou

le

urs

(

North

)

was

selected

for

Air

Service

use

in

January

1918

.

The

general

region

was

chosen

as

the

center

for

the

future

American

Sector

and

the

site

was

one

of

those

available

after

three

years

of

use

by

the

French

.

It

was

located

75

kilometers

from

the

front

line

,

considered

safe

for

use

by

pursuit

(

fi

ghter

)

aircraft

.

The

terrain

context-0

context-1

context-2

context-0

context-1

context-2

context-0

context-1

context-2

Layer 1

context-0

context-0

context-1

context-1

context-2

context-2

Layer 2

context-0

context-0

context-1

context-1

context-2

context-2

Layer 3

context-0

context-0

context-1

context-1

context-2

context-2

Layer 4

context-0

context-0

context-1

context-1

context-2

context-2

Layer 5

context-0

context-0

context-1

context-1

context-2

context-2

Layer 6

context-0

context-1

context-2

Layer 7

context-0

context-0

context-1

context-1

context-2

context-2

Layer 8

context-0

context-0

context-1

context-1

context-2

context-2

Layer 9

context-0

context-0

context-1

context-1

context-2

context-2

Layer 10

context-0

context-0

context-1

context-1

context-2

context-2

Layer 11

context-0

context-0

context-1

context-1

context-2

context-2

Layer 12

we generate samples using the same prefix and the same sampling strategy across samples. The holes generated in this case are the same across the batch. We demonstrate the throughput achieved in this case in Fig. 16. The benefits in this case are a lot more clear.

Investigating Dropped Tokens. In Transformer encoder models, it has been shown that similar tokens can be successfully pruned in deeper layers (e.g. Bolya et al. [2022]). Similarity, in this case, is measured by calculating the cosine similarity of the keys K across different tokens. To test how feasible such a pruning strategy is in our scenario, we calculate for the keys K of each token, the minimum cosine distance to other tokens in the sequence. We then group these distances based on whether the token was subsequently dropped or not, in Fig. 17. We can see that tokens that have other tokens very similar to them in the sequence are pruned. However, tokens with no other similar ones in the sequences can also be dropped.

We experimented with the framework from Bolya et al. [2022], which has exhibited very successful results for transformer encoder models out of the box, i.e. without any additional fine-tuning. We found, nonetheless, that the perplexity quickly diverged, i.e. increased, even for small levels of sparsity. This constitutes another indication that pruning decoder models requires additional effort. Compared to encoder models, decoders make a different prediction for each of the tokens in the input, and token similarity by itself is not good evidence for pruning. We finally highlight that even if the method from Bolya et al. [2022] achieved comparable results, computational benefits will be

256

512

1024

2048

4096

8192

16384

32768

10

12

14

16

18

20

22

24

26

Average Perplexity

GPT-2-small

7292 tokens/s

GPT-2-medium

2637 tokens/s

GPT-2-large

1241 tokens/s

GPT-2-xl

587 tokens/s

Average throughput

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

Dense

256

512

1024

2048

4096

8192

16384

32768

Tokens/s

10

12

14

16

18

20

22

24

26

Perplexity when context size is 1000

GPT-2-small

4129 tokens/s

GPT-2-medium

1521 tokens/s

GPT-2-large

764 tokens/s

GPT-2-xl

401 tokens/s

Throughput when context size is 1000

GPT-2-

small

GPT-2-

medium

GPT-2-

large

GPT-2-

xl

Dense

0

200

400

600

800

1000

2000

4000

6000

8000

10000

12000

14000

Tokens/s

Throughput per generation step

0

200

400

600

800

1000

Context size

2×10 2

3×10 2

4×10 2

6×10 2

Seconds

Time per generation step

Optimal batch sizes

baseline

baseline × 0.92

baseline × 2.15

baseline × 3.23

baseline × 6.15

Dense

=0.007

=0.01

=0.03

=0.1

=0.3

=1.0

=3.0

=10.0

0.0

0.2

0.4

0.6

0.8

1.0

Cosine distance

0

1

2

3

4

5

6

Density

Dropped

Not dropped

negligible, as no memory benefits can be achieved in this case and generation is performed one token at a time.

Attended Tokens across Layers. We provide additional visualization of the attended tokens in Fig. 18. Different layers exhibit different levels of sparsity and attend to different tokens.

How sparsity affects specific NLP tasks Evaluating GPT models pretrained by language modelling objective is in general challenging and different techniques have been developed based on the final intended use. This challenge motivated us to showcase zeroshot accuracy to downstream tasks, in addition to evaluating perplexity for upstream tasks. To better check if some tasks are affected more by sparsity compared to others, we provide per-task zeroshot performance, see Figure 19. We

Layer 1

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 2

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 3

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 4

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 5

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 6

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 7

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 8

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 9

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 10

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 11

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

Layer 12

Gard ening   is   the   practice

of   growing

and   cultivating

plants

as   part

of   h ort iculture

.  In   gardens

,  orn amental

plants

are   often

grown

for   their

flowers

,  foliage

,  or   overall

appearance

;  useful

plants

,  such   as   root   vegetables

,  leaf   vegetables

,  fruits

,  and   herbs

,  are   grown

for   consumption

,  for   use   as   d yes ,  or   for

medicinal

or   cosmetic

use .  Gard ening   ranges

in   scale

from

fruit   or chards ,  to   long   bou levard

plant

ings   with   one   or   more

different

types

of   shr ubs ,  trees

,  and   herb aceous

plants

,  to   residential

back   gardens

including

plant

ings ,  all   the   way   to   container

gardens

grown

inside

or   outside

.  Gard

ening   may   be   very   specialized

,  with   only   one   type

of

also include zeroshot results for additional tasks, that require long-range dependencies. It becomes clear that some tasks are affected more than others, according to our intuition.

Inference Compute An additional benefit of our pruning strategy is that we can accommodate longer initial contexts for the same hardware. Table 3 showcases the maximum potential context windows for different batch sizes. Our strategy accommodates inference with much longer contexts, for a fixed device.

Additional Models As the original context grows, pruning the context makes an increasingly significant difference, as more opportunities for pruning exist. Generally, enlarging the context length of Transformers is a fascinating area of research with novel techniques being currently proposed. In our experiments, we are limited by the maximum positional encodings of the base GPT models. To evaluate how our method performs in larger contexts, we additionally present preliminary results

0.2750

0.2775

0.2800

0.2825

0.2850

0.2875

0.2900

GPT2-small dense

GPT2-small with context pruning

0.5950

0.5975

0.6000

0.6025

0.6050

0.6075

0.6100

piqa

0.15

0.20

0.25

0.30

lambada_openai

0.5050

0.5075

0.5100

0.5125

0.5150

0.5175

winogrande

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

0.015

0.020

0.025

0.030

multirc

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

0.135

0.140

0.145

0.150

0.155

openbookqa

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

0.60

0.65

0.70

0.75

sciq

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

0.05

0.10

0.15

0.20

F1

on pythia language models. Pythia uses rotary position embedding, as many current state-of-the-art LLMs. We present results on 'Perplexity vs sparsity' in Figure 20. These concrete results demonstrate that our pruning technique works out of the box for different positional encodings as well. Longer contexts supported by our base model (for pythia models that is 2048) additionally allow for higher levels of sparsity for the same performance. In the future, we intend to also evaluate larger models, potentially with the use of LoRA [Hu et al., 2021]. We again highlight that our approach is not specific to a particular language model, as it can be applied to any autoregressive transformer architecture.

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

14

16

18

20

22

24

26

28

30

Perplexity

Pythia-70M

Dense

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

14

16

18

20

22

24

26

28

30

Perplexity

Pythia-160M

0.0

0.2

0.4

0.6

0.8

1.0

Sparsity

14

16

18

20

22

24

26

28

30

Perplexity

Pythia-410M

D Discussion

We presented a simple approach that can be applied to any decoder-based Transformer autoregressive model. In an era where big foundation models are taking the community by storm, we presented a relatively inexpensive technique that can be applied to these models and can lead to significant gains in terms of computational resources required to perform inference. We truly believe the balance between model performance and resource efficiency is crucial for the widespread adoption of large-scale language models, and our approach offers a promising avenue for achieving this goal. We hope that our technique offers more interpretable predictions and inspires future research.

Reproducibility. We have taken multiple steps to ensure the reproducibility of our experiments. We refer the reader to Section 4 and Appendix A for a complete description of the training protocol. We have also released the code as part of the supplementary material, including scripts on how to reproduce our results. Additionally, trained models will be released for further research.

Limitations. Our work focuses on autoregressive generation models based on Transformers. Although we argue that the same technique can be applied out-of-the-box to any such model, we focused on text generation models and specifically the GPT-2 family. Due to the uniformity of recent architectures, we expect that these results generalize to other popular models.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an efficient and interpretable context-based dynamic pruning method. They use additional query / key layers to generate dynamic attention masks and sparsify self-attention maps. They also introduce a sparse sigmoid and regularization term to control the sparsity. In experiments, they demonstrate the lowest performance degradation compared to previous static attention pruning methods at the same level of sparsity. The results of throughput and speed analysis show that the proposed method can achieve additional inference efficiency with minimal performance loss. They analyze the distribution of remaining contexts with respect to part of speech and the depth of the layer. They also demonstrate that the proposed method dynamically prune attention based on contexts through context switch experiments.\n\nSTRENGTHS:\n- The proposed method is well-motivated and easy to follow.\n- The proposed method efficiently increases interpretability and sparsity with minimal performance drop compared to previous research.\n- They provide an extensive analysis of the relationship between context length, throughput, and speed in various parameter sizes.\n- Based on the proposed pruning method, they devise efficient batched data structures for the optimized computation.\n\nWEAKNESSES:\n- Quantitative comparison of throughput and speed with local/sparse attention would contribute to a comprehensive understanding.\n- Further qualitative study of interpretability (Fig. 8) varying the sparsity level (gamma) would provide additional intuitive observations.\n\nQUESTIONS:\nDo you have expectations on how the zero-shot performance difference between the pruning strategies will change as the number of model parameters increases?\n\nLIMITATIONS:\nAs written in the Limitations section, scalability studies on larger language models (>7B) would provide further insights into the dynamic pruning method.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The paper demonstrates excellent writing with a clear flow of ideas.\n- The authors introduce a unique and innovative data structure that enables batch operations involving masked tokens.\n- 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.\n\nWEAKNESSES:\n- **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.\n- **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.\n- **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.\n\nQUESTIONS:\n- 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?\n- There seems to be a typo in Figure 5. The green curve should be labeled as \"Sparse Attention\".\n- 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a modified dynamic masking operation to the traditional multi-headed attention in transformers in order to allow models to learn to drop tokens at specific layers during training. In order to facilitate learning, they use a sparse sigmoid that is annealed to interpolate from a traditional sigmoid up to a step function in the infinite limit. They show that this dynamic sparse attention achieves lower perplexity at the same sparseness levels compared to other modified attention mechanisms, as well as better zero-shot accuracy and throughput. The proposed model shows promise for improving the efficiency and overcoming the quadratic time complexity of traditional dense transformers.\n\nSTRENGTHS:\nImproving the efficiency of vanilla dense transformers is an important problem, and being able to learn the sparsification operation rather than setting a static prior is an interesting direction. The proposed method is simple and computationally efficient, and the manuscript is well written and clear. The results demonstrate marked improvements over other sparse attention mechanisms at similar sparsity levels as well as increased throughput.\n\nWEAKNESSES:\nThe main concerns I have are with how well this adaptive sparsity can be used when training from scratch and on long range dependency benchmarks such as Long Range Arena [1]. Since the experiments are initialized from GPT-2, all dense information is present, the model simply has to learn which information it can safely ignore. However, when training from scratch this optimization problem becomes much more difficult since falling into a local optimum with respect to pruning early on can reduce the model’s effectiveness in the future as the $\\alpha$ annealing increases towards a step function. It would be interesting to see results (even on small models) on how much utilizing this mechanism during from scratch training affects perplexity, given the same compute budget. \n\n[1]  Long Range Arena: A Benchmark for Efficient Transformers. 2020. Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, Donald Metzler\n\nQUESTIONS:\n- The $\\alpha$-sigmoid operator is a bit confusing to me. How are argmax values calculated efficiently and how are gradients propagated through the argmax operator in equation 8?\n- Do the baseline GPT-2 models use FlashAttention in their implementation? If not, how does the throughput of the dynamically pruned model compare to one with FlashAttention?\n- How sensitive is the training to the correct annealing of the $\\alpha$ parameter?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nGiven the trend in large language models, it is a pretty important problem to search to efficient architectures. In this direction is the line of work to make the attention component efficient by introducing sparsity in the attention block and allowing every token to attend only a subset of the previous tokens.\n\nThe paper presents usage of Adaptively Sparse Attention whereby\n1) The network learns to drop parts of the context no longer required.\n2) The tokens dropped at every layer are independent, and different set of tokens might be chosen to be persisted at every layer.\n3) The network is trined using sparse sigmoid functions that introduce sparsity during training itself in contrast to some works that have tried to introduce sparsity only during inference.\n\nExperiments show that this methodology can lead to pretty strong performance (with minimal perplexity loss), even with 80% sparsity.\n\nSTRENGTHS:\nThe work presented in the paper is pretty innovative and impactful in this direction as finding efficient transformer architectures is key to sustain and grow the research and production usage of these large language models.\n\nThe model successfully exploits the fact that the attention matrix in these models are pretty sparse and additionally encourages that with sparse sigmoid-like function to mask out some tokens.\n\nExperiments show pretty strong performance even with 80% sparsity in the network.\n\nWEAKNESSES:\nGiven that the model is performing pretty well in terms of perplexity even with huge sparsity, it would be interesting to perform an analysis of whether there is a class of NLP tasks the is significantly affected (like knowledge intensive tasks?), or perhaps translation or similar tasks where the structure of the input sentence matters if the model is focusing to forget stopwords very early.\n\nQUESTIONS:\n- Was there any analysis done on which set of NLP tasks get affected the most with such sparsity?\n- For inference on a given device, do we have numbers on the increase in the max sequence length that can be supported?\n- Do these changes impact the training time throughput in any way?\n\nLIMITATIONS:\nNone discussed in paper, and nothing important that I can think of.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-07"}