Hybrid Dynamic Pruning: A Pathway to Efficient Transformer Inference

Transformer models, including BERT[1], GPT [2], T5 [3] and others [4] [5] , have transformed Natural Language Processing (NLP) with their attention mechanism, achieving top performance in tasks such as question-answering [6], text classification [3], and machine translation [7]. The transformer architecture uses self-attention mechanism [8] and it is highly parallelizable on modern Graphical Processing Units (GPUs), providing major benefits over older models like Long Short Term Memories (LSTMs) and Recurrent Neural Networks (RNNs). This has led to fast progress in NLP, with models like BERT exceeding human performance in difficult tasks [9] and expanding their application to computer vision, including object recognition and detection [10], image classification [11], and segmentation [12].

Deploying large transformer models on devices with limited resources is challenging due to their high computational and memory requirements. For example, BERT-Base Transformer needs $440$ MB of memory and over $176$ Giga FLoating Point Operations (GFLOPs) [13] . The computations are particularly difficult because of the complex attention operations and the quadratic computational complexity related to the length of input sequences [14]. Attention operations in transformer models become increasingly dominant as the input sequence length grows. For BERT-Base Transformer deployed on edge platforms, with a sequence length of $512$, the attention operations account for about half of the total execution time, and this figure rises to $70\%$ when the sequence length extends to $768$ [15]. Therefore, finding efficient ways to handle attention operations is crucial for speeding up transformers.

Many studies utilized sparsity to mitigate the quadratic time and space complexity issue. Some techniques save computational effort by using fixed or static sparse attention patterns [16] [14] [17], but their performance is limited [18], since the sparse pattern in attention is naturally dynamic, and depends only on the input. Other techniques focus on dynamic sparsity, meaning there’s no fixed pattern for which parts are sparse (zero). For example, $A^{3}$ [19] uses various approximation methods to skip calculations of near-zero values, aiming to decrease computational demands, but it requires loading all data onto the chip, which doesn’t decrease off-chip DRAM access. SpAtten [20] introduces a cascaded token pruning mechanism that gradually eliminates less important tokens to simplify the workload using a Top-K strategy. Despite being tailored for dynamic decision-making in hardware, Top-K requires significant computational cost. Energon [15] uses a mixed-precision, multi-stage filtering technique to mimic the Top-K pruning approach, but it relies on a special unit to handle sparsity. AccelTran [21] prunes values in all transformer matrix multiplications if they fall below certain predetermined threshold. $A^{3}$, Energon, and AccelTran leverage unstructured sparsity to realize efficient deployment for Transformers which leads to non uniform data access and lower efficiency. It’s hard to predict these sparsity patterns, which can slow down performance.

Other research efforts have been directed towards the removal of attention heads in transformer models, based on the understanding that not all heads contribute significantly to the transformer’s performance. Researchers in [22] add trainable gating parameters attached to each head and regularized with $L0$ loss. In [23], the importance of each attention head is gauged based on its sensitivity to the overall loss. This sensitivity is utilized as an indirect measure to determine the significance of each head. In [24], a novel approach termed ’single-shot meta-pruner’ is presented. This method involves training a compact convolutional neural network with the specific purpose of identifying and selecting the attention heads that are crucial for preserving the distribution of attention within the model. All of these studies perform pruning at compile time not run time, require retraining to recover the accuracy drop. SpAtten [20] also performs a cascaded head pruning at run time using the Top-K approach, wherein the importance of each attention head is derived by aggregating the absolute values of it’s attention outputs of the head throughout the layers, resulting in an aggregate score of importance for each head. SpAtten uses a separate unit to perform head Top-K strategy, which also requires significant computational cost. $A^{3}$, Energon and AccelTran do not support head pruning.

To overcome the above challenges, we propose integer-based Hybrid Dynamic Pruning (HDP), an algorithm-architecture co-design to enable efficient attention inference. This method operates on multiple levels within the attention matrix, reducing complexity by concentrating on blocks, and heads. This method employs fine-grained block pruning for removing smaller, less critical blocks and incorporates head pruning to selectively eliminate less important heads, all based on the integer parts of inputs for decision-making. The pruning is done dynamically; applied during inference without relying on fixed patterns for pruning, and without fine-tuning or retraining. Our main contributions are as follows.

1. 1.

   We propose integer-based, fine-grained block pruning, which prunes unimportant small sized blocks, with the observation that self-attention primarily relies on a few critical query-key pairs, highlighting significant redundancy in the self attention mechanism. HDP utilizes the integer part of the input to identify and prune less important blocks, focusing subsequent operations only on the unpruned blocks for enhanced efficiency.

2. 2.

   We introduce an early head pruning strategy that identifies and eliminates less important heads based on the integer parts of the input at the initial stages of computation, unlike the method in [20] which performs pruning after all computations.

3. 3.

   We approximate attention calculation by breaking down the multiplication of Q and K into three components: $integerQ\times integerK$, $integerQ\times fractionalK$, and $fractionalQ\times integerK$. By summing these components together, we not only approximate the attention outcome but also achieve near-zero pruning, as the $fractionalQ-fractionalK$ multiplication is omitted. This method effectively reduces computational complexity while maintaining model accuracy during inference.

4. 4.

   We design and implement an ASIC-based architecture to efficiently execute HDP, utilizing encoder-only models to reduce the critical path to the half and enhance throughput and hardware utilization. Our architecture functions as a co-processor, compatible with existing neural network accelerators. Through carefully designed pipelines and architectural optimizations, we aim to significantly boost performance and reduce energy consumption.

The article is structured as follows: Section II provides background information on transformer acceleration. Section III details the methodology of the HDP framework. Section IV describes the hardware architecture of HDP co-processor. Section V outlines the experimental setup, the baselines used for comparison and discusses the results. Finally Section VI concludes the article.

1. A.

   Transformer Algorithm
   

   

   Transformers have shown state-of-the-art performance in natural language processing field since their initial introduction in 2017 [8]. They are often regarded as alternatives to classic CNNs and RNNs in various situations in real-world due to their exceptional efficiency and generality. Transformers consist of two essential components: the encoder and the decoder, which are formed by layering many transformer blocks. As depicted in Fig. 1, a block comprises three primary elements: linear layers, multi-head self-attention, and the normalization layer. The linear layers include the fully connected Layer (FC), Feed Forward (FFN), and the projection layer. The processes starts with transforming the input vectors into embedding vectors, which are then sequentially passed through a series of processing blocks. Within each block, the input is projected through the projection layer to $Query~(Q)$, $Key~(K)$ and $Value~(V)$ features using the weights $W^{Q},W^{K}$ and $W^{V}$. Following this, attention mechanisms are utilized on these features to comprehend the long-term relational dependencies present in the input sequence.
   The attention mechanism, as shown in Algorithm 1, splits the $Q,~K~$and $V$ matrices into multiple smaller sets $Q_{h}$, $K_{h}$, and $V_{h}$ equivalent to the number of heads $H$. Each of these sets forms an ”attention head”, see Fig. 1 right side. Inside each head the output is computed as follows:

   $$\text{Attention}(Q_{h},K_{h},V_{h})=\text{softmax}\left(\frac{Q_{h}{K_{h}}^{T}}{\sqrt{d_{k}}}\right)V_{h}.$$

   The process starts with the computation of the dot product between $Q_{h}$ and $K_{h}$, followed by scaling. This step computes the alignment or similarity between each pair of tokens, the fundamental components of a sequence. The resulting matrix represents each token’s relationship with every other token. Then a row-wise softmax is applied to obtain the attention probability. When processing a particular token, the attention probability weights tell the model how much attention to give to each token in the sequence. After that, the attention probability is multiplied with the Value vector $V_{h}$, resulting in a weighted sum that forms the output of the single head attention layer, the resulting output is a combination of information from all tokens, with each token’s importance and relevance determined and weighted by the attention mechanism. Each head independently calculates attention result. The results are concatenated to get the end result. The concept entails that each head possesses the ability to concentrate on distinct segments of the input sequence, or capture different types of relations within the data.
   The attention mechanism fulfills various functions: it enables the model to focus on the most significant tokens of the input sequence, captures long-term relationships within the sequence, and greatly enhances the model’s awareness of the context. This selective attention and context-aware processing are critical for complicated tasks such as language translation [25], text summarization [26], and other NLP tasks.

   0Input: $\{K,Q,V\}\in\mathbb{R}^{l\times d}$

   1 $l$: sequence length, $d$: input feature dimension

   2 Number of heads: H;

   3 $Q_{h},K_{h},V_{h}\leftarrow$ Split $Q,K,V$ into H chunks;

   4 $d_{h}=d/H$;

   5 for  $head_{id}$ $\leftarrow$ $0$ to H do

   6    $attention\_score={Q_{h}{K_{h}}^{T}}$ ;

   7    $attention\_score\in\mathbb{R}^{l\times l}$ ;

   8    $attention\_score=\frac{attention\_score}{\sqrt{d_{h}}}$ ;

   9    for $row_{id}$ $\leftarrow$ $0$ to $l$ do

   10       $attention\_prob=\text{softmax}(attention\_score[row_{id}])$ ;

   11      

   12    end for

   13   $result[head_{id}]=attention\_prob\cdot V_{h}$ ;

   14   

   15 end for

   16$attention\_out=\text{concat}(result[head_{id}])$ ;

   17

   17Output: $attention\_out\in\mathbb{R}^{l\times d}$ ;

   In addition to the multi-head attention component, transformer models also employ other processes such as normalization and FFN. The FFN in a transformer is composed of two FC layers. The function of the FFN can be mathematically expressed as:

   $$\text{FFN}(X)=\text{GELU}(XW_{1}+b_{1})W_{2}+b_{2},$$

   where $X,W_{1},b_{1},W_{2},$ and $b_{2}$ represent the input to the FFN, the weight matrices, and the bias vectors, respectively. The Gaussian Error Linear Unit (GELU), serves as a special activation function, providing non-linear transformation to the output of the first layer before it is passed on to the second layer [27], where it is widely used in Transformers. Our focus in this work is the attention layer as it is the bottleneck for Transformers models.
   

2. B.

   Complexity of Attention Layers
   The primary computational requirements of attention layers in a neural network model are pairwise dot-products performed on a set of $l$ vectors. As a consequence, the complexity is $O(l^{2}d)$. The computational cost associated with attention operations escalates as the length of the input sequence $l$ increases, while $d$ is usually fixed. By measuring the time consumed in attention layer for Bert-base model [1], on both embedded GPU and CPU platforms, it’s observed that attention operations consume half of the total computational time for input sequences longer than $512$. Furthermore, these operations consume almost $70\%$ of computational time for input sequences reach $768$ [15]. Therefore, in scenarios where processing long input sequences is necessary, attention operations emerge as the primary computational bottleneck. Also because of the advancements in linear layers through weight pruning [28], quantization [29], and specialized accelerators [30] [31], there is a need to optimize the attention mechanism to ensure balanced computational efficiency in transformer models.
   

3. C.

   Dynamic sparsity in Multi Head Self Attention
   Sparse attention methods are driven by the recognition that not all attention probabilities are significant. After applying a row-wise $Softmax$ function to get the $attention\_prob$, most likely a small subset of the scores in each row will impact the $attention\_out$ results. Similarly, sparse multi-head attention strategies are typically motivated by the recognition that not all heads in the attention mechanism have an equal impact on the final output. This suggests that some heads in the attention model might have a minimal impact on the overall outcome. These observations point to the existence of redundant heads and insignificant attention probabilities in the attention architecture, indicating that the model’s performance could be enhanced by focusing on the most crucial heads and removing unimportant $attention\_scores$.
   By examining the attention weights matrix ( $attention~probability$) generated from different heads across various layers and inputs, we can say that the significance of individual heads in a transformer’s multi-head attention mechanism depends only on the input data. Varying not only with the input data but also with their position within the model’s layered architecture. Fig. 2 presents the attention weights of various heads across different layers and inputs in a BERT-base model fine tuned on the SST-2 dataset [32]. As illustrated in Fig. 2(a), there is notable variation in the attention weights for the same head across different layers. For example, the attention values for the eleventh head, highlighted by the red box, show considerable fluctuation across Layers 9, 10, and 11. Additionally, the attention patterns of the same head within the same layer also exhibit significant differences when compared across various inputs. This is clearly depicted in Fig. 2(a) and Fig. 2(b), where the green boxes depict the differences in attention values for the same heads and layers when subjected to different inputs. Specifically, Head 0 and Head 1 in Layer 11 exhibit notably low attention weights for Input 1. In contrast, for Input 2, Head 1 and Head 2 in Layer 11 display significantly higher attention values, highlighting the data-dependent nature of attention mechanisms in the model. Furthermore, it is clear for most of the heads that it is only a subset of the attention weights which have high magnitude, where there is no fixed pattern for the important weights.
   

   

   

   This variability highlights the complex, dynamic nature of how transformers process information. For a given input, certain heads in specific layers may become more active, focusing on particular aspects of the data. This activation can differ from one input to another, reflecting the adaptive response of each head to the unique characteristics of the data it encounters. Similarly, the role of a head may differ across layers. Furthermore, the high magnitude attention weights subset is dynamic depending on the input sequence , and different heads have different subsets. This data-dependent behavior of attention heads is a critical consideration in model analysis and optimization and motivates us to explore effective methods to eliminate unimportant heads and query-key relations and save computations.
   

In this section, we discuss the algorithmic optimizations that enhance the transformer model’s efficiency and performance, focusing on block, head pruning and approximation techniques. These optimizations are key to reducing computational complexity and memory access.

1. A.

   Block Pruning
   For the attention score matrix, most of the query-key relations are not important and can be pruned safely, many methods have been used to prune these relations. Top-$K$ pruning method [20] is used to prune the attention weights where a whole row can be pruned, but this requires a retraining to recover the accuracy, also it requires a specialized hardware to get the $k$ most significant attention weights. Energon [15] avoided the Top-K selection and used the mean filtering as a practical approximation instead, but it still has a separate unit to perform this operation, also faces data duplication overhead due to the multi-round filtering method employed. In Energon and AccelTran [21] the pruning is done in an element-wise pattern which results in an irregular sparse matrix, where zeros are randomly spread over the matrix and this results in an irregular memory access and stalls in the hardware.
   To address these challenges we propose integer-based block pruning, where pruning decision is exclusively determined by the integer parts of the numbers. We employ a small block size for pruning to eliminate the necessity for retraining and to guarantee a more organized and hardware-compatible sparsity pattern.
   Initially, multiplication is conducted only on the integer parts of $Q$ and $K$ to obtain $Integer\_atten$. For each $2\times 2$ block, we calculate its importance, $\theta$, as the absolute sum of the values within the block. For each row of blocks, we determine the block pruning ratio, $\Theta$, using a method similar to that in Energon, which involves calculating the minimum, maximum, and mean importance values, along with a predefined pruning ratio, $\rho_{B}$, as shown in Algorithm 2 line 18. Blocks with importance, $\theta$, falls below the row-specific threshold, $\Theta$, are pruned and the mask value for the block is assigned to $0$. When a block is pruned, subsequent computations for that block are omitted. Conversely, if a block is retained (mask is 1), the final attention result is approximated, with the approximation technique detailed in B. The block pruning mechanism is expressed in details in Algorithm 2 lines 2 to 2 and in Fig. 3.
   

   

2. B.

   Approximation
   For blocks with mask value equals $1$, $\theta$ exceeds $\Theta$, we proceed to approximate the final attention outcome after obtaining the $Integer\_att$ result. This involves calculating two fractions: the product of $Q$’s fractional component with $K$’s integer component, and vice versa, yielding $Frac1\_atten$ and $Frac2\_atten$, respectively. The ultimate attention score is obtained by summing these two fractions with the $Integer\_att$. This method not only approximates the attention outcome but also enables near-zero pruning, which has minimal impact on model accuracy during inference [33]. Specifically, when two numbers are close to zero, their integer parts are zero, leading all three components to be zero. Consequently, the multiplication of fractional parts is omitted, resulting in effective near-zero pruning. This approximation process is highlighted within the black box in Fig. 3 and in Algorithm 2 lines 2 to 2.
   

3. C.

   Head Sparsity
   Not all heads are essential, and many can be pruned without affecting the overall performance. Unlike the method in [20], where head importance is assessed after completing all computations for the head, we introduce an early head pruning approach. To evaluate head importance, $\theta_{Head}$, we compute the absolute summation of all values in $Integer\_att$. Heads with $\theta_{Head}$ below a predefined threshold $\tau_{H}$ are completely pruned and the rest of computations of this head are skipped. The threshold $\tau_{H}$ is a parameter that will be profiled. The head pruning process is visually indicated by the red box in Fig 3 and described in Algorithm 2 at line 2.
   

Current attention accelerators and traditional CPUs/GPUs lack the capability to execute the proposed hybrid dynamic sparse attention technique efficiently. To address this gap, we are introducing a novel HDP accelerator. HDP is developed to function as a co-processor and is compatible with a variety of neural network accelerators for easy integration.

1. A.

   Architecture Overview
   The HDP architecture, depicted in Fig. 4, compromises multiple cores, the architecture of individual core is shown in Fig. 4 middle part. Each core is composed of an array of processing elements (PE Array), a Sparsity Engine (SE), an adder, and a softmax unit. The PE Array handles matrix multiplication tasks such as $Q\times K^{T}$ and $attention\_prob\times V$, also calculating importance values. The SE is tasked with identifying which blocks to prune and deciding whether a head should be pruned or not.
   

   

   Workflow: Once $Q,K,$ and $V$ are generated and quantized by another processor in fixed point 16 bit format and stored in memory, HDP processes each attention head sequentially. It employs tiled matrix multiplication for these operations. Initially, the integer components of $Q$ and $K$ are retrieved from off-chip memory into on-chip memory for the computation of $Integer\_Q\times Integer\_K$ using tiling. The SE then uses the computed importance values for each block to create a mask indicating which blocks are not pruned. This approach prevents unnecessary data fetching for pruned blocks, reducing memory access and computational overhead. Once $Integer\_Q\times Integer\_K$ computation is complete, and the head importance is assessed, the SE decides whether a head will be pruned. If so, the remaining computations for that head will be skipped and proceeds to the next head. For heads that remain unpruned, a Fetch Upon Mask (FUM) strategy is utilized. If the mask value is 0, indicating a pruned block, the corresponding K values will not be fetched, and the computation for that block is skipped. If the mask value is 1, the corresponding Q and K values are fetched, and the processing element (PE) array calculates the two remaining fractions ($Integer\_Q\times Frac\_K$ and $Frac\_Q\times Integer\_K$) simultaneously. These results, along with the integer results from the previous step, will be added together using an ADDER module to obtain the total $attention\_score$. After performing all $Q\times K$, a row-wise softmax will be applied to the $attention\_score$. Then, the PEs will be utilized to perform the calculation $attention\_score\times v$. Specifically, the first and second PEs in the first row calculate $Integer\_attention\_score\times Integer\_v$, while the third and fourth PEs in the first row compute $Integer\_attention\_score\times Frac\_v$. The first and second PEs in the second row calculate $Frac\_attention\_score\times Integer\_v$, and the third and fourth PEs in the second row calculate $Frac\_attention\_score\times Frac\_v$. Finally, all these results will be summed using the ADDER module to obtain the final output. The attention results for each tile will be immediately stored in DRAM memory upon completion. Consequently, host DNN accelerators can access these results to perform subsequent computations. In the sections that follow, we provide a more in-depth exploration of each module and the optimizations we have proposed to enhance their functionality.
   

2. B.

   Tiling and Dataflow
   A significant portion of the computational workload in transformer models is attributed to matrix multiplication, necessitating optimization to boost accelerator performance. We propose the implementation of tiled matrix multiplication, a technique primarily employed in GPUs [34], to enhance the efficiency of these operations in our accelerator design. Tiling enhances resource efficiency and enables parallel computation, as depicted in Fig. 5. The first $4\times 4$ tile from matrix $A$ is multiplied by the first $4\times 8$ tile from matrix $B$, with the partial results stored in a $4\times 8$ tile in matrix $C$, denoted by in all matrices. Subsequently, the process advances along in the tiles of $A$ and $B$ to accumulate additional partial sums in $C$. During this stage, an output stationary dataflow approach is employed, facilitating the reuse of partial sum outputs in the accumulator. Additionally, a local $A$ stationary strategy is implemented, meaning that while outputs are reused in the outer loop, inputs from matrix $A$ are retained and reused in the inner loop [35]. Next, the process proceeds along in all matrices, followed by a moving along .
   

3. C.

   Processing Elements
   The processing element, shown in Fig. 4 (right), serves as a fundamental computational unit within the accelerator handling all matrix multiplication operations. Operating in an output stationary mode, and behaving similar to systolic array PE; it receives rows from tiles of the first matrix and columns from tiles of the second matrix as inputs, one input at a time. It multiplies these values and stores the intermediate sums in accumulators until the entire row from the first matrix has been multiplied by the corresponding column from the second matrix. At this point, the accumulators hold the final results for the tile of the result matrix. In the case of $Integer~Q\times Integer~K$ multiplication, these results are also utilized to determine the block’s importance, as the output from a processing element corresponds to a block in the result matrix. The importance of the block, as illustrated in Fig. 4, is equal to the absolute sum of the accumulators.
   

4. D.

   Sparsity Engine
   The Sparsity Engine is responsible for determining the sparsity pattern at both block and head levels. Illustrated in Fig. 6, the Sparsity Engine’s internal architecture takes in importance scores from the PE and stores them in its internal memory. Additionally, it keeps track of the minimum, maximum, and total sum of these importance values for every row of blocks. Upon receiving the $END\_R$, which signals the completion of a full row in the result matrix or, equivalently, the multiplication of a row from the first matrix by all columns in the second matrix in $Integer~Q\times Integer~K$ multiplication, the engine calculates the block pruning threshold $\Theta$ for that specific row. This calculation is based on the equation provided in line 18 of Algorithm 2. Additionally, the engine generates the $Mask$ for the row by subtracting $\Theta$ from the importance values of the blocks. If the result is negative, the block falls below $\Theta$ and is marked for pruning.
   Furthermore, when the $END\_H$ is received, signaling the completion of the $Integer~Q\times Integer~K$ multiplication, the engine utilizes the $\theta_{Head}$ value it has computed representing the total sum of all importance values across the entire head. Upon receiving this flag, the engine compares it with $\tau_{H}$, an input parameter denoting the head pruning threshold. If $\theta_{Head}$ falls below $\tau_{H}$, the head is assumed redundant and is thus excluded, allowing the accelerator to bypass any further calculations for this head and proceed to processing the subsequent head.
   

5. E.

   Softmax Module
   Once the attention scores are obtained, a row-wise softmax function is applied defined as $e^{s_{i}}/\sum_{j}e^{s_{j}}$. For every input received to the module, the exponent is approximated using $2^{nd}$ order polynomial. The exponent results are stored in internal memory, and the sum of these exponent results is calculated. By the end of every row, the reciprocal of the sum is computed using a linear approximation. Then the exponent values are multiplied by the reciprocal to generate the softmax result.
   

In this work we presented HDP, a novel algorithm-architecture co-design to efficiently run dynamic sparse attention models. We first proposed a novel integer-based row-balanced block pruning to prune unimportant blocks in attention matrix and integer based head pruning to prune unimportant heads. Moreover, we propose an approximation method that reduces the computations and performs near-zero pruning. We also implemented this method in 2 co-processors architecture, HDP-Edge and HDP-Server, to accelerate algorithm on mobile and sever platforms.

This work was supported by the Khalifa University of Science and Technology under SOCL.