Adaptive Rank Allocation: Speeding Up Modern Transformers with RaNA Adapters

Q6. Does batched inference work and achieve speed-ups if masks are inconsistent?

This is a good question, thanks for asking it! Our method and experiments focused on single-sequence decoding (batch size = 1). We found the batch size = 1 decoding scenario relevant due to its importance in systems where latency matters more than throughput (e.g. on-device workloads). Notably, due to their input dependent nature, neuron-adaptive adapters like CATS also mainly look at this setting. Therefore, we decided to leave the larger batch size scenario outside the current scope of our work. However, leveraging a masked GEMV for larger batch sizes and measuring latency would be an interesting direction of future work.

Q7: Lack of wall-clock latency measurements for RaNA.

Thank you for raising this point! In the revision, we leveraged an efficient implementation of RaNA using a masked GEMV Triton kernel and included a wall-clock comparison of RaNA and CATS [1] (Figure 1b). Please refer to the general response for more details.

References:

1. Lee, Je-Yong, et al. “Contextually-aware thresholding for sparsity in large language models”, 2024. URL https://arxiv.org/abs/2404.08763.
2. Ashkboos, Saleh et al. “Slicegpt: Compress large language models by deleting rows and columns”, 2024. URL https://arxiv.org/abs/2401.15024.
3. Zhang, Zhengyan et al. “Relu2 wins: Discovering efficient activation functions for sparse llms, 2024”. URL https://arxiv.org/abs/2402.03804.

We thank reviewer 5oY9 for their positive comments and helpful constructive feedback! We note that the wall-clock latency measurements and comparisons to other related work suggestions are addressed in the general response above. Further, we reply to specific suggestions below.

W1/Q2: Lack of baselines from the structured pruning literature and comparison to other compression methods.

Thanks for bringing up this point! In Section 5 of our updated manuscript, we include additional comparisons to SliceGPT [2], a cutting-edge open source pruning method, which was the most suggested by reviewers (Table 1, Figure 5). Please refer to the general response for more details.

W2/Q1. Adding practical implementation details

Thanks for requesting this! In our revised manuscript, we now incorporate implementation pseudocode in the appendix (sect. 7.2).

Q3. What is the relationship between the neural thresholding for down-projection layers and CATS?

This is a good question, thanks for bringing it up! The neuron masker method is indeed slightly different from the one CATS used. To be concise:

• The CATS neuron mask is $m_{CATS}(x) _i= |SiLU(GateProj(x))_i| >= t$.
• The RaNA neuron mask is $m_{RaNA}(x)_i = |SiLU(GateProj(x))_i \cdot UpProj(x)_i| \cdot ||W^{down}_i||_F >= t$. (Eqn. 12)

We argue this masker is more natural. Notably our proposed masker looks at $SiLU(GateProj(x))_i \cdot UpProj(x)_i$ since it is the input to the down-project matrix. Masking out based only on $SiLU(GateProj(x)) _i$ may be troublesome if $SiLU(GateProj(x))_i \cdot UpProj(x)_i$ results in a large value for the masked out entry. Moreover incorporating the $||W^{down}_i||_F$ is more natural if we are considering reducing the overall output error of the MLP.

Q4. Are compression ratios including router FLOPs? Can you show a breakdown of FLOPs distribution between main computation and maskers?

Thanks for this detailed question! Yes, the masker FLOPs are considered into the total FLOPs. Nevertheless, we note that element-wise operations over vectors were generally not considered in the FLOP computation.

We’d like to note that it's unclear how to devise a breakdown between main computation and masker FLOPs for CATS and RaNA, since the masker can be thought of as adding an elementwise operation on top of an intermediate variable that will be used in computing the output anyways. For example, for the linear-layer rank adapter masker (eqn. 9), for the masker we look at the output of $Bx$, element-wise square it and compare each element to a threshold. However, $Bx$ is still used in the overall computation anyway (eqn. 4). For this reason, in the case of RaNA and CATS, it is unclear how to determine what FLOPs belong to the main computation and which belong to the masker, since they are heavily dependent.

For the Pythia suite models, the breakdown is possible for our neuron-adaptive baseline. For the neuron-adaptive baseline, the masker is described as the MLP sigmoid masker from Sect. 4.1. As outlined in Section 5.1, we allocated 6% of the MLP FLOPs to the masker and the rest to the main computation, as done in [3].

Q5.1. How is the FLOP reduction achieved in practice?

Thanks for this question! FLOP reduction is achieved by applying a sparse mask to each linear layer in the MLP and QKV projection layers. For example, from Eq. 11, the Up-Project, Gate-Project and QKV layers are approximated using a linear-layer rank adapter, $\text{Linear}(x) \approx \text{Linear'}(x) = A (m(x) \odot B x )$, where $m(x) \in \\{0,1\\}^r$ is a sparse binary mask. After computing $v=m(x) \odot B(x)$, $v$ becomes sparse, allowing us to skip the product involving the columns of $A$ corresponding to the 0 entries in $m(x)$. In a similar fashion, the neuron-adapter leveraged in the Down-Project layers constitutes a matrix-vector product, where the vector is sparse, allowing us to skip computations. Moreover, we realize practical speedups associated with this theoretical FLOP reduction by leveraging a masked GEMV kernel, as we detail in the general response.

Q5.2. Which DL framework was used to implement the method?

Thanks for bringing this up! The method is implemented in PyTorch, as clarified in section 5.1 from the revised manuscript. Moreover, the efficient version of RaNA leverages a masked GEMV kernel.

We thank reviewer yoCw for their positive comments and helpful constructive feedback! We note that the suggestion on comparisons to other related work (specifically a method that also compresses MLP + QKV layers for Llama2) is addressed in the general response above. Further, we reply to specific suggestions below.

W1. Insights on MLP compression v.s. MLP + QKV compression and using different FLOP allocation strategies in RaNA.

This is a great question, thank you for raising it! In Section A.3, we provide an ablation study where we compare the perplexity of RaNA, before fine-tuning, on 3 different settings to assess the performance of the adapter when compressing different layers or using different FLOP allocation strategies. Concretely, we observed the RedPajama perplexity on a RaNA-adapted Llama2-7b model under the following settings, all targeting the same total FLOP compression ratio of ~31%:

• Vanilla RaNA adapter compressing MLPs + QKVs, leveraging the FLOP allocation strategy
• RaNA adapter only compressing MLPs, leveraging the FLOP allocation strategy
• RaNA adapter compressing MLPs + QKVs, uniformly allocating FLOPs across all the MLP components (Up-Project, Gate-Project and Down-Project layers).

Notably, the best adapted model was the one leveraging the vanilla RaNA adapter (8.40 ppl), with the other two settings falling behind in perplexity (8.79 ppl and 9.10 ppl). Moreover, the FLOP allocation strategy showed to be of high importance as the adapted model took a notable performance hit when ablating this component (0.7 ppl points).

W2. Comparison against baseline that sparsifies MLP and QKV layers like RaNA

Thanks for this suggestion! In Section 5 of our updated manuscript, we include additional comparisons to SliceGPT [2], a cutting-edge pruning method (Table 1, Figure 5). Please refer to the general response for more details.

W3. Clarify sparsification levels for each component (MLP, QKV) in RaNA and baseline methods.

This is a good point, thanks for bringing it up! Definitely, we added a breakdown of sparsification levels for MLP and QKV layers in RaNA and baseline methods in Table 4.

W4. Move the proof for Proposition 1 into the main text

Thanks for pointing this out! We are glad you found the proof helpful. We have moved the proof to Section 3 in the revised manuscript.

Q1. Why was RaNA not applied to QKV in Gemma?

Thanks for this question, this is a great observation! RaNA was not applied to QKV in Gemma since we did not expect much speed up from compressing those layers, as they consume much less flops than the MLPs. Concretely, the Query, Key and Value projection together consume only ~5% of the FLOPs relative to those consumed by MLPs in Gemma-2b, while they consume ~37% in the case of Llama2-7b. Therefore, we focused on MLP sparsification for Gemma. In our revision, we have added this clarification in Section 5.3 of the paper.

References:

1. Lee, Je-Yong, et al. “Contextually-aware thresholding for sparsity in large language models”, 2024. URL https://arxiv.org/abs/2404.08763.
2. Ashkboos, Saleh et al. “Slicegpt: Compress large language models by deleting rows and columns”, 2024. URL https://arxiv.org/abs/2401.15024.

We thank reviewer LZp1 for their positive comments and helpful constructive feedback! We note that the wall-clock latency measurements and comparisons to other related work suggestions are addressed in the general response above. Further, we reply to specific suggestions below.

Q1/W1: Lack of wall-clock latency measurements for RaNA.

Thank you for raising this point! In the revision, we leveraged an efficient implementation of RaNA using a masked GEMV Triton kernel and included a wall-clock comparison of RaNA and CATS [1] (Figure 1b). Please refer to the general response for more details.

W2: Lack of baselines from the structured pruning literature.

Thank you for this important point! In Section 5 of our updated manuscript, we include additional comparisons to SliceGPT [2], a cutting-edge pruning method (Table 1, Figure 5). Please refer to the general response for more details.

W3. Results on newer models (eg: Llama3.1 8b) would help compare with other recent pruning/compression efforts.

Thank you for bringing this point up! Unfortunately, we were limited in resources, which did not permit us to test RaNA in bigger models. Nevertheless, CATS and SliceGPT were both applied to Llama2-7b models, which motivated us to use it as our largest model for testing. Therefore, given the limited timing, we decided to keep our current analysis for Llama2 models only. However, we hope to look into assessing the performance of RaNA in Llama3 for the camera ready version.

Q2. How does replacing a single GEMM with two GEMMs in linear layers affect latency?

Thank you for this question! That’s correct. RaNA replaces a single general matrix vector multiply (GEMV) with two for Up-Project, Gate-Project and QKV linear layers. Concretely, the linear layers adapted using the linear-layer rank adapter of RaNA replace a single GEMV with one vanilla GEMV for $v=Bx$ and one masked GEMV for $A(m(x) \odot v)$.

Our experiments show that RaNA achieves practical inference speedups across various compression ratios and hint that decomposing a single GEMV with two does not impact performance drastically. Concretely, from Figure 1.b we can see that, as latency increases, RaNA’s downstream task performance converges to that of the full model, with only a ~0.85 percentage-point gap at ~40ms/token compared to the full model's ~46ms/token latency. These results show that despite decomposing a matrix multiplication into two operations, practical speedups are attained. Furthermore, the convergence of the downstream-task performance towards the full model performance as latency increases hints that the decomposition does not drastically impact speed.

Q3. How well does RaNA perform on networks with sparse activation functions?

Thanks for bringing this up! We would like to emphasize that the main objective of the RaNA adapters, as outlined in our introduction, is to improve adaptive compression for modern transformer architectures leveraging non-sparse activations. Notably, previous work has found conventional neuron adapters to work well for sparse activation functions, however, a lack of performance has been noticed when applied to non-sparse activation functions [3]. This has led to research on tackling this specific problem, for example ReLUfication [4], and CATS [1]. Our goal with RaNA aligns with previous work, as we seek to improve adaptive compression for modern Transformers with non-sparse activation functions. Therefore, for this work, we decided to not assess our method in sparse activation functions, as we determined it out of scope for the problem we are trying to address.

References:

1. Lee, Je-Yong, et al. “Contextually-aware thresholding for sparsity in large language models”, 2024. URL https://arxiv.org/abs/2404.08763.
2. Ashkboos, Saleh et al. “Slicegpt: Compress large language models by deleting rows and columns”, 2024. URL https://arxiv.org/abs/2401.15024.
3. Zhang, Zhengyan et al. “Relu2 wins: Discovering efficient activation functions for sparse llms, 2024”. URL https://arxiv.org/abs/2402.03804.
4. Mirzadeh, Iman et al. “Relu strikes back: Exploiting activation sparsity in large language models”, 2023. URL https://arxiv.org/abs/2310.04564.

We thank reviewer RtUt for their positive comments and helpful constructive feedback! We note that the wall-clock latency measurements and baseline suggestions are addressed in the general response above. Further, we reply to specific suggestions below.

W1: Lack of wall-clock latency measurements for RaNA.

Thank you for raising this point! In the revision, we leveraged an efficient implementation of RaNA using a masked GEMV Triton kernel and included a wall-clock comparison of RaNA and CATS [1] (Figure 1b). Please refer to the general response for more details.

W2: Lack of baselines from the structured pruning literature.

Thank you for this important point! In Section 5 of our updated manuscript, we include additional comparisons to SliceGPT [2], a cutting-edge pruning method (Table 1, Figure 5). Please refer to the general response for more details.

W3: What is the gap between considering only one layer’s error and considering multiple layers’ errors while determining the rank to mask out?

This is a great question, thank you for bringing it up! We acknowledge that considering multiple layers’ errors could potentially yield better results, as prior work suggests that some layers contribute more significantly to overall performance [3]. In this work, we are focused on the problem of extending neuron-adaptive methods to non-sparse activations. As such, we chose to focus on the simpler setting of layer-wise optimization rather than holistic optimization across layers, which introduces additional complexity. We recognize that a holistic FLOP allocation is an interesting direction to explore on improving RaNA, which we hope to investigate in future work.

References:

1. Lee, Je-Yong, et al. “Contextually-aware thresholding for sparsity in large language models”, 2024. URL https://arxiv.org/abs/2404.08763.
2. Ashkboos, Saleh et al. “Slicegpt: Compress large language models by deleting rows and columns”, 2024. URL https://arxiv.org/abs/2401.15024
3. LLMP: Ma, Xinyin, et al. “Llm-pruner: On the structural pruning of large language models”, 2023. URL https://arxiv.org/pdf/2305.11627.