Adaptive Rank Allocation: Speeding Up Modern Transformers with RaNA Adapters
Authors: Roberto Garcia
Jerry Weihong Liu
Daniel Sorvisto
Sabri Eyuboglu
Venue: ICLR 2025 Poster (2025)
Keywords: Large Language Models
Adaptive compute
Rank adapters
Neuron adapters
Primary area: foundation or frontier models, including LLMs

Abstract: Large Language Models (LLMs) are computationally intensive, particularly during inference. Neuron-adaptive techniques, which selectively activate neurons in Multi-Layer Perceptron (MLP) layers, offer some speedups but suffer from limitations in modern Transformers. These include reliance on sparse activations, incompatibility with attention layers, and the use of costly neuron masking techniques. To address these issues, we propose the Adaptive Rank Allocation framework and introduce the Rank and Neuron Allocator (RaNA) adapter. RaNA adapters leverage rank adapters, which operate on linear layers by applying both low-rank matrix decompositions and adaptive masking to efficiently allocate compute without depending on activation sparsity. This enables RaNA to be generally applied to MLPs and linear components of attention modules, while eliminating the need for expensive maskers found in neuron-adaptive methods. Notably, when compared to neuron adapters, RaNA improves perplexity by up to 7 points and increases accuracy by up to 8 percentage-points when reducing FLOPs by $\sim$44\% in state-of-the-art Transformer architectures. These results position RaNA as a robust solution for improving inference efficiency in  modern Transformer architectures.

Published as a conference paper at ICLR 2025
ADAPTIVE RANK ALLOCATION: SPEEDING UP MOD- ERN TRANSFORMERS WITH RANA ADAPTERS
Roberto Garcia Jerry Liu Daniel Sorvisto and Sabri Eyuboglu Institute of Computational and Mathematical Engineering, Stanford University *Department of Computer Science, Stanford University
ABSTRACT
Large Language Models (LLMs) are computationally intensive, particularly dur- ing inference. Neuron-adaptive techniques, which selectively activate neurons in Multi-Layer Perceptron (MLP) layers, offer some speedups but suffer from limitations in modern Transformers. These include reliance on sparse activa- tions, incompatibility with attention layers, and the use of costly neuron masking techniques. To address these issues, we propose the Adaptive Rank Allocation framework and introduce the Rank and Neuron Allocator (RaNA) adapter. RaNA adapters leverage rank adapters, which operate on linear layers by applying both low-rank matrix decompositions and adaptive masking to efficiently allocate com- pute without depending on activation sparsity. This enables RaNA to be generally applied to MLPs and linear components of attention modules, while eliminating the need for expensive maskers found in neuron-adie adaapive methods. Notably, when compared to neuron adapters, RaNA improves perplexity by up to 7 points and in- creases accuracy by up to 8 percentage-points when reducing FLOPs by ~444 in state-of-the-art Transformer architectures. These results position RaNA a robust solution for improving inference efficiency in modern Transformer archi- tectures.
INTRODUCTION
As Large Language Models (LLMs) have grown in popularity and size, they have begun consuming non-trivial amount of compute and time for training and inference (Kim et al. (2023). Pope et al. (2022)) ). Adaptive compute methods seek to speed up the inference stage of Transformers Vaswanni al. (2023)) the de facto LLM architecture, by identifying and avoiding redundant computations to save 1/O and floating-poin point operations (FLOPs). Commonly, these methods apply neuron adapters to the Multi Layer Perceptron (MLP) layers of the Transformer architecture, which dynamically ignore neurons depending on the input of the layer (Lee et al. (2024), Mirzadeh et al.] (2023), Liu et al. (2023 Zhang et al. (2024)). ). Utilizing these adapters leads inference speedups, as the amortized time complexity of an MLP layer reduces from ((dan-dhidddn)) to O(dyn "dadapted) where di is the input dimension of the MLP, dhidden is the hidden dimension, and dadapted is the average number of active neurons. For modern Transformer architectures, this results in practical speedups the expense of negligible model quality decrease for sufficient amount of average active neurons.
Unfortunately, adaptive compute methods using neuron adapters suffer from rapid performance degradation in modern Transformer architectures (Figs. Ia Tcp. We observe this issue arises from the limitations of the neuron-adaptive framework, which enforces dynamic neuron allocation based on activation values. Concretely, we believe this decline is attributed to this framework's lack of generalization across different layer types, reliance on sparse activation functions, and costly neuron masking techniques. First, neuron adapters, often tailored to MLPs, can not be directly applied to non-ML.P layers, such as the Query, Key, and Value modules of attention layers, which lack neurons. In addition, these methods frequently depend on indivaatoo-induuce sparsity, like that attained by ReLU (Agarapl (2019)), making them ineffective with non-sparse activation functions like SwiGLU (hazzer (2020), Zhang et al. (2024)). Further, while recent approaches like CATS (Lee et al. (0024) or ReLUfication (Mirzadeh et al. (2023) attempt to handle non-sparse activations, they inefficiently
"Corresponding author: robgarct @st edu

---

Published as a conference paper at ICLR 2025
0.0
0.11
0.4
8 821
039
Adtaar Toreel
bane made KAbk
and Briodoped
CXT radegeideddees
52% www.eaaaaaaieciooooo
wurtheeappe model
Ieteey Jre0
(a) Llama2-7b FLOPs.
Accuracy
v.s. (b) Llama2-7b Accuracy v.s. lency.
La- (c)
Pythia FLOPs.
Suite
Accuracy
V.S.
Figure I: RaNA improves accuracy-comput tradeoff over neuron adapters. y-axis shows accu- racy averaged over multiple downstream tasks (Sect. 5.T): for Figs. and Tc x-axis shows average FLOPs for forward pass with sequence length 512; for Fig. [b z-axis shows average per-token de- coding latency over a sequence of 492 tokens with initial context lengths ranging from to 1000. We compare RaNA-adapter models to neuron-adapted versions at various compression rates for (left) Llama2-7b and (right) Pythia models. Notably, RaNA accuracies decay slower as compression rates increase compared neuron-adapters.
compute neuron- activvatois before determining which neurons to exclude. This is concerning as Transformers architectures like Llama (Touvron et al. (2023), Gemma (Team et al. (2024) or Mis- tral (Jiang et al. (2023)) rely on non-sparse activation functions, which makes sparsity- adapters ineffective for them, pushing us to rely more computationally intensive neuron adapters.
neuron
To address these gaps the neuron-adaptive setup, we propose the Adaptive Rank Allocation frame- work, together with the Rank and Neuron Allocator (RaNA) adapter. We note that this aun-aadaptive framework is a more powerful generalization of the couron-adaptive one. Concretely, we devise this framework from the observation that any linear layer be decomposed into the product of two low-rank matrices and an adaptive router/masker, namely Linear(z) = Wx z A(r(x) O Br). Moreover, develop RaNA in this framework, leveraging specific A and B decompositions and masker m(x) for Transformer layers, like QKV or MLPs. Notably, unlike neuron adapters, RaNA be directly applied any linear layer without relying on activation function sparsity. Further, when used with modern MLP layers with SwiGLU activations, it enables better distribution of com- pute across the MLP's linear layers.
We empirically validate the rank-adaptive framework and the RaNA adapter. We demonstrate that, similar neuron-adaptive setups, the ranks of the AB matrix decomposition in the proposed RaNA adapters have sparse importances, depending on the input (Figs. 2a bb). allowing cally prune them. We also show that RaNA adapters attain the lowest error MLP layers when recovering the full MLP outputs, outperforming neuron-adaptive methods by 6.7, 18.1 and 7.4 percentage-points on Llama2-7b, Gemma-2b and Pythia- 160M respectively. Further, we show the effectiveness of RaNA in modern Transformer architectures by applying it to Llama (Touvron et al. (2023) and Gemma (Team et al. (2024)). Notably, RaNA closes the gap between full-model and adapted-model performance, on multiple compression rates (Tabs. D Concretely, at a 444% FLOP compression rate, RaNA achieves an average improvement of perplexity-points and percentage-points in benchmark curacy for Llama2-7b, while for Gemma-2b, it improves perplexity by 7 points and accuracy by percentage-points, compared prior adapters. Finally, assess the effectiveness of RaNA's ap- plicability to different types of layers and activations, we apply it to the Pythia suite (Biderman al.) (2023), a set of varied-sized GPTNeox (Black al.) (2022)) models. Here we also observe that, when comparing their performance to conventional neuron-adaptive methods, it consistently attains better perplexity and downstream task performance (Figs. Tc 4). Further, code is available at httpps//ggihu..co////e/// cot..b/oo//oeeeee///aaaaaa
to dynami-
outperforms the state-of-the-art method CATS (Lee et al. (0024)
RELATED WORK AND PRELIMINARIES
Here we discuss key work in neuron neuron-adaptive adaptive methods and examine neuron adapters for Transform- with sparse activations and with modern non-sparse activations.
2

---

popular ap- Published as a conference paper at ICLR 2025
Neuron-adaptive framework: Adaptive compute methods (Han al. (021
proach to speed up inference in Transformer (Vaswani et al. (2023) architectures. A concrete instance of them are neuron adapters, which dynamically allocate neurons of the MLP layers to different inputs using a masker. These neuron-adaptive methods for Transformers stem from the observation that neuron importance is sparse, and while Transformers allocate equal resources to all tokens, some tokens require less compute based on context (Li et al. (2023b)). Empirically, redun- dant computations are common due model parameterization (Frankle & Carbin (2019)).
Nevertheless, the neuron-adaptive framework is constraining, as it only allows us to dynamically allocate neurons, which are only present on MLP layers of Transformers. Additionally, neuron adapters do not generalize well to different activation functions in the MLP layers (Zhang et al. (2024). ) Hence, previous work has focused on creating adapters specific to certain activation func- tions, namely sparse activation functions and non-sparse activation functions.
Neuron adapters for sparse activations: In the case of Transformers leveraging sparse activation functions like ReLU (Agarap) (2019)) neuron adapters rely on the abundance of O-valued neuron activations (Li et al. ())3a))) These adapters commonly work by using a small MLP masker that predicts whether neuron going to be active for given input, then only computing that neuron if
its prediction is positive. Concretely, consider a conventional MLP layer with a ReLU activation:
MLPReLU
W down (ReLU(Ww
(1)
where Wap E R.d. and Wownn& E Rh.dd are the Up-Projection and Down-Projection matrices of the MLP layer. Then, the neuron adapted version of this MLP follows:
RL (3) = Wdown (m(x) ) @ ReLU((
(2)
where m(x): : R + 10 is the binary masker, which often is parameterized by small MLP. While effective in these type of models, given their strong reliance on the sparsity of neuron activa- tions, this adapters have unfortunately shown poor performance when applied to Transformers with non-sparse activations (Zhang et al. (2024))
Neuron adapters for non-sparse activations: Modern Transformer architectures, like those used in popular models like Llama (Touvron et al. (2023). Gemma (Team et al. (2024)), Mistral (Jiang eet al. )2033) or GPTNeox (Black et al. all.(222)) (2022)), leverage non sparse activation functions like GeLU
HHendrycks & Gimpel (2023)) for GPTNeox and the more popular SwiGLU (Shazeer (0200) for
the others. Concretely, an MLP layer with a SwiGLU activation follows:
MLP SWIGLu (x) = Wdown (SiLU(W, gate @ W Wapi))
(3)
where Wup E R.d.h, Waowm E RR.dd and W Wgate E RRd are the Up-Projection, Down-Projection and CateeProjection matrices of the MI.P layer.
Recent work has sought to improve the performance of conventional neuron adapters on non-sparse activations by developing methods specifically tailored towards them, such as CATS (Lee et al. (2024)) and ReLUfication (Mirzadeh et al. (202 While these approaches have demonstrated good performance in SwiGLU-based Transformers, they suffer from inefficiencies. The inefficiency stems from their reliance calculating exact activation values prior applying thresholding. For instance, CATS adapters compute the entire output of the Gate- Proje-Prooetion layer before selecting which neurons to retain. At large compression rates, this enforces sub-optimal FLOP allocation imbalance across the Up. Down and Gate projection layers of the MIP, where the Gate-Projection layer receives most of the FLOPs.
3 ADAPTIVE RANK ALLOCATION FRAMEWORK
We begin by presenting the Adaptive Rank Allocation Framework, which addresses the shortcom- ings of the neuron- adurtovv approach, specifically its limitation to neuron adaptation and its frequent reliance on neuron sparsity. Adaptive compute frameworks generally involve two key parts: dynam- ically allocated components and a function to determine their use based on the input. We make the observation that such adaptive compute setup can be applied at the granularity of linear layers, where the dynamically allocated components are the ranks of a weight matrix, while the function that determines which ranks are leveraged by given input is router masker. This is the main ob- servation that motivates our Adaptive Rank Allocation Framework, which can be applied to arbitrary linear layers without relying on actívation function sparsity.
3

---

Published as a conference paper at ICLR 2025
Concretely, consider any linear layer Linear(x) Wæ, where W R Under the Adaptive Rank Allocation framework, we replace this linear layer by:
Linear' (a)) A(r(x) @ Br)
(4)
This simple setup has two parts. First, we have a set of static matrices A and B, where A RRod and B E Rd.. Second, we have the adaptive component r(x), where r R' - IR.. Ideally we want r(x:) be sparse, as that allows us to save I/O and floating-poii operations. rank adapted layer from Eqn.Abecome becomes essentially low-rank matrix multiplication, where the the low-rank matrix (A diag(r(x)) B) has rank = [ Notably, the FLOPs consumed by such a matrix multiplication are proportional to the rank of such matrix, hence why sparsity is desirable.
r(xx) is sparse, our
Generalization of Neuron-Adaptive Methods: We note that the Adaptive Rank Allocation frame- work is strict generalization of the neuron-adaption one, where neuron adapters can be viewed as specific instance rank adapters with appropriate choices of A, B and r(z). In Prop. [ we illustrate this for ReLU-basee MLPs, but note the proof easily extends to other activation functions:
Proposition 1. Consider an MLP(x) ReLU layer (Eqn. 7 and its neuron adapted version MLP' (x) ReLU (Eqn. 2). Then, there exists a rank adapted ReLL (i.e. an MLP whose linear layers have been rank adapted) s.t. MLP MLPRREL(())
MLP' (I) (E)R for all x E IR.
Proof. Let M (x) = Adowaa (Tdown (r) ① Bdown (ReLU (Aup( (rapp(pppll (x) O Bagx))))-
Then, by setting Adown Wdowns Bdown I. Aup I, Bup WuP, Idown == Mdowns and
Tup = 1, follows that:
MLP'ReL Wdown (m²(x) @ I(ReLU(I(1 @ Wapz)]))
doomn (m2 (2) @ ReI.U(Wupar)) MLPP'
In summary, this result shows that rank adaptation generalizes neuron adaptation, highlighting the increased versatility and applicability of the Adaptive Rank Allocation framework.
4 RANK AND NEURON ALLOCATOR ADAPTERS
Here, we introduce Rank and Neuron Allocator (RaNA) adapters as instances of the Adaptive Rank Allocation framework. RaNA adapters leverage Linear Layer Rank adapters (Sect. 4.I at their core, which operate at the granularity of linear layers. Such Linear Layer Rank adapters employ approximately optimal A and B matrices in tandem with dynamic rank allocation mechanism based on the input distribution. Notably, while Linear Layer Rank adapters are effective in QKV, Ap-Projection, and Gate-Prooection layers, they are inefficient in Down- coonn-rroection layers when used as stand-alone adapters. To overcome this limitation, we propose RaNA adapters.
RaNA adapters operate at the level of QKV and MLP layers, combining Linear Layer Rank adapters with neuron-three phtrronttreeholddng techniques, a setup that has demonstrated the best empirical performance when adaptively compressing layers (Figs. 3a 3b
3c 3d). Notably, these adapters address the inefficiencies of previous non-sparse neuron adapters by avoiding the costly exact computation of specific linear layer outputs prior to applying neuron thresholding. We empirically demonstrate their effectiveness in model accuracy and perplexity compared to other adapters (Tabs. 2
4.1 LINEAR LAYER RANK ADAPTERS
The Adaptive Rank Allocation framework, unlike the neuron-adaptive one, does not impose limi- tations the A B matrices the routing function r(x) when adapting linear layer through Linear' (I) A(m(x) @ Br) s Wx. We leverage these properties to introduce the Linear Layer Rank Adapters. Here, we want to find the optimal A. B and r(x) that, in expectation over the input distribution, best recover the original output of the linear layer. Formally, we want:
argmin A.B. Fx A(r(x) @ Bx)
(5)
Evidently, optimally finding A, B and r is non-trivial. Hence, we choose to address this problem by breaking into two steps. First, we propose A and B matrices that solve a similar problem to Eqn. [5] which ignores the routing function. Then, with our election of A and B matrices, we propose an optimal input-dependent masker, which serves router, for Linear Layer Rank Adapter.
4

---

Published as a conference paper at ICLR 2025
A and B matrices: First, devise set of A and B matrices for our rank-adapted layer Linear' (z) = A(r(x) @ Bz). We start off by relaxing the optimization problem from Eqn. [5]
Concretely, we remove the router from the picture and frame the problem finding fixed low-rank matrices A, and B, of rank r that minimize our objective. Namely, the relaxed objective becomes:
argmin A,,, R, Ey (1IWa (,
(6)
Note that this relaxed problem is not only convenient, but also a decent choice, as it is equivalent to solving the original problem from Eqn. 5 with fixed router r(æ)) that always outputs I's for the first r elements the output vector and O's for the rest. In practice, we do not have access to the actual distribution of the inputs z, hence reframe the problem its empirical version:
argmin I[ A,
where each column of X E R+* namely 24. is hidden-state imput that the linear layer from our pretrained model observes practice. Notably, we used 32,000 samples our experiments.
We can then analytically find the optimal Ar and B, matrices that solve this problem.
Theorem 1. Let U, E Roy be the first r singular vectors W' then by the Eckart-Youe Young theorem (Ekkatt& & Young 1936 ) the Ar and B, that optimize the objective from Eqn. 7 are:
A, = U, B, U.W.
The proof for Theorem I is provided in Appendix A.1
Therefore, we pick our A.B matrices as A : U. B == U W. Next, we demonstrate how to pick a masker/router to effectively leverage these matrices in an input-adaptive manner.
B-Masker: With our A = U and B = UTW matrices at hand, we devise a sparse router r(x) function that minimizes the adapter's error (Eqn. 5). For simplicity, we opt for binary r(z), i.e. masker r(x) m(x) I' + {0, 1 that allows us to allocate different ranks (and hence different amount of FLOPs) to different input hidden states. Formally, we want an m that optimizes:
argmin (x)
[W Um((z))) @ U Wx;)lla,
s.t. (1)m(())) llll =r
(8)
argmax
[m(x)) U W Wx;ll?
s.t. E Eellm(z) =r
The constraint Eqn. 8 enforces that expected rank of the matrix (A diag(m(x)) B) is r. This constraint directly controls the FLOPs that the rank-adaptive model will consume average.
Observe that, for any hidden state v, we can compute the contribution of each rank our adapter to the output of the original linear layer. To do this, we examine the contribution of each column vector in A = U (equivalently, the contribution of each rank of A B) to the output o = Linner(x) R A(m(z) O Br). Specifically, i-th column vector of A, u, is given by (Bx)}, due to the orthogonality of the columns in A = U. This allows us identify the important ranks of our A, B decomposition for given input æ and create a sparse masker m())) ) for it (Figs. 2a 2b). which just keeps the most descriptive ranks. We call this the B-Masker, as it uses the B matrix to select the most important ranks for x. Formally:
note that the contribution to the Frobenius norm ||o| from the
m(x): = B B-masker(z); = ((xx) 2 t
where 1 1{-} is the indicator function, and the threshold t is picked so that, on average, desired amount of FLOPS is consumed by our adapted layer. We note that the B-masker is efficient to use when the matrix from our linear layer W E R(o.)) has an output dimension that is bigger than its inner dimension i, as this makes computing (Bx)² cheap. Hence, it is convenient to use this masker in practice for the Up-Projection,
Projection and QKV layers of our models.
MLP-Sigmoid Masker: An alternative option to the B-masker, which has the potential to be more efficient, and that works for any linear layer, is a small MLP masker with a Sigmoid activation function, just like that used commonly in the neuron-adaptive literature (Liu et al. (2023), Zhang et al. 22244). Concretely, we use the parametrization m(x=) = (CDDr where C E R D R i is the dimension of the hidden states r' is the inner dimension of the predictive masker, and r is the rank of the matrix W from the linear layer the masker corresponds to. In our experiments, train this masker on a binary cross-entropy loss to match the output of the B-masker.
5

---

Published as a conference paper at ICLR 2025
4.2 FUSING RANK ADAPTERS AND NEURON THRESHOLDING FOR RANA
Building on the insights from the previous section, we observe that Linear Layer Rank Adapters independently perform well for QKV, Up-Projection, and Gate- PraeePProection layers due to their use of tall and narrow matrices. However, they may encounter difficulties with Down-Projection layers, which involve short and wide matrices. To overcome this limitation, RaNA combines Linear Layer Rank Adapters with neuron- trtrronntthrsolddigg adapters. specifically for the Down-Projectes ther, RaNA implements a FLOP allocation procedure across the various components of the adapter, addressing the imbalance found in previous neuron-adaptive methods, where FLOP distribution is heavily skewed toward specific components.
layers. Fur-
RaNA in QKV layers: For the QKV layers, just replace the linear QKV layers with Linear Layer Rank Adapters, namely:
QKV(x) = Wa (<) = A(m(x) @ Bx)
(10)
where m(x) is B- B-masker, find it outperforms the MLP-Sigmoid masker (Fig. 3db.
RaNA in MLP layers: For MLP layers, we demonstrate how RaNA is applied to SwiGLU-based MLPs and note that its application MLPs with other activation functions follows the same ap- proach, excluding the GrojeProjection adapter. Namely, we use Linear Layer Rank Adapters for the Up-Projection and Gate-Projed Gate-Projection matrices and use neuron-thresholding for the Down-Projection matrices. Formally, our RaNA adapted MLP' (x) is described by the following:
MLP'(z) (x) = Down' (SWL' (x)) + Up'(x)))
where Up' (x)) Aup aap (2) 2 B Bupt),
(11)
Gate' (x)
(mgate (I) 2 B B)ate Down't (x:) Wdown (M (z) ② x)
where (xp(()) and Mgaae (x) are B-maskers and M.dow (2)) is simple neuron thresholding masker:
m(x); neuron- thresholding(2)) 1[ ||W!! P 2 t}
(12)
RaNA FLOP Allocation: Prior heuron-adapters for non-sparse activation functions are forced to ineffectively allocate FLOPs across their components (Sect. 2 On the other hand, RaNA's flexi- bility permits us to freely distribute FLOPS across the adapter components, therefore we propose FLOP allocation strategy specifically for RaNA. At the Linear Layer Rank Adapter level, we per- form a simple line search to balance FLOPs between the B-Masker and the target sparsity, selecting the configuration that minimizes the output error with respect to the original linear layer. Similarly, the MLP level, we conduct grid search to distribute FLOPs among the Up', Gate', and Down' layers, with each component further balancing FLOPs between its masker and target sparsity. We then retain the configuration that achieves the greatest error reduction in the MLP output. Please refer to Appendix A.4 a pseudocode implementation of RaNA.
EXPERIMENTS AND RESULTS
In this section, we present experiments and results for the following claims:
1. The contribution of ranks in Linear Layer Rank Adapters is sparse. Just like this is desirable property for neuron-ades addptive approaches to work well, it is also one for the Linear Layer Rank Adapters, which we use in RaNA. Notably, the distribution of the contributions of each rank in the AB decomposition (to the recovery of the original output of the linear layer) of our rank adapters is concentrated at and is heavy tailed (Figs. 2a 2b). This validates Linear Layer Rank adapters and practically means that can effectively mask out many of the ranks with near-O contributions.
2. RaNA adapters attain lowest errors on Transformer layers when reconstructing orig- inal layer outputs. Concretely, we conduct a study of the reconstruction error in the layer outputs of the Llama2-7b, Gemma-2b and Pythia-160M models using multiple compres- sion methods, targeting a ~50% reduction in the FLOPs of the compressed layers. When compared to other adapters, RaNA attains the lowest mean squared reconstruction error of the outputs of the original layers (Figs. 3a 3b 3cl 3db. Notably, RaNA achieves aver- age errors of 10.5% and 281, surpassing the 11.22% and 20.31% attained by CATS on MI.P layers for I. -lmma277b and Gemma-2h respectively. Likewise, in Pythia- 160M, RaNA

---

Published as a conference paper at ICLR 2025
achieves average errors of 7.87% and 0.36% in MLP and QKV layers respectively, out- performing the 15.23% and 1.29% errors obtained with neuron adapters and SVD-based adapters.
3. RaNA outperforms neuron adapters in perplexity and downstream tasks. Using RaNA, we attain lower perplexity and higher average accuracy on multiple NLP bench- marks in modern Transformers compared neuron-adapters. Notably, RaNA outperforms previous neuron adapters across multiple FLOP compression levels (Tabs. for SwiGLU based architectures. In the case of Llama2-7b, we achieve an improvement of 8 percentage-points in accuracy and 4 perplexity-point on a 42% model compression rate. Similarly, for Gemma-2b. 2b, we attain an improvement of percentage- -ooitttggg-pisssssseeeee in accuracy and 7 perplexity-poine on 46% model compression rate. We also show that RaNA not only preserves a theoretical advantage in the accuracy-FLoe FLOP trade-off but also practically in the accuracy-latency trade-off (Fig. [b). Further, RaNA outperforms conventional neuron- adapters in GeI.U-based Pythia models across multiple model sizes and compression levels (Figs. lc
Fig. Tap
5.1 EXPERIMENTAL. SETUP
Here we describe the settings under which we run our experiments.
Models and Adapters: We conduct experiments using Llama2-7b, Gemma-2b, and multiple GPT- Neox models from the Pythia suite. Notably, we fine-tuned Gemma-2b on the RedPajama dataset for ~422M tokens before applying adapters to it, so as to align it with the data used in subsequent adaptation and fine-tuning steps. In Sect. 5.3 we evaluate the following adapters, which are mainly implemented in Py PyToc..
RaNA: Our proposed adapter from Sect.
CATS: A state-of-the-as art adapter for MLP layers leveraging SwiGLU activations. We refer to their work (Lee et al. (2024)) for more details.
SliceGPT: A recent structured pruning method, which compresses linear layers by rotating and slicing them, including MLP and QKV layers. We refer to their work (Ashkboos et al. (2024) for more details.
Neuron-Adaptive: A standard neuron adapter for MLP layers with a small MLP masker, such as that leveraged by Zhang et al. (2024) or Liu et al. (2023). using 6% of MLP FLOPs for the masker, as done by Zhang et al. (2024).
with MLP maskers (LLRA): An adapter that applies Linear Layer Rank adapters leveraging an MLP-based masker (Sect. 4.1 to all linear layers in QKV and MLPs.
Lntear-Layer-aanke
For output error assessments in Sect. 5.3 we additionally use a fixed low-rank singular value de- composition (SVD) for comparison.
Datasets: We use the RedPajama (Computer 2023 dataset for Llama2-7b and Gemma-2b, and the Pile (Gao al. (2020) dataset for Pythia models when evaluating rank contribution sparsity (Sect. (.2). output errors (Sect. 5.3), and perplexity (Sect. [.33), and for devising any data-dependent adapter component (e.g. and matrices in RaNA, the activation threshold in CATS and the slicing and rotating procedure of SliceGPT). Downstream-tase performance is assessed using HellaSwag (Zellers et al. (2019), PIQA (Bisk et al. (2019), WinoGrande (Sakaguchi et al. (2019)), Arc-Easy Clark et al. (2018), Arc-Challenge (Clark et al. (2018)) and RACE (Lai et al. (2017)).
Fine-tuning: To assess accuracy and perplexity (Sect. 5.3 we fine-tune adapted models using the Huggingface library (Wolf al. (2020)) and LoRA adapters (Hu et al. (202I)) for ~31M tokens on Llama2-7b and Gemma-2b, with an AdamW optimizer, where learning rates were determined from various options, depending the model's performance following ~6M tokens of training. In a similar setup, Pythia models are fine-tuned for ~61M tokens with the exception of not leveraging LoRA adapters.
Performance Evaluations: We use LM-evaluation harnesses (Gao et al. (203 for downstream- task performance evaluation in zero-shot setting for Llama2-7b and Pythia models, and five-shot
7

---

Published as a conference paper at ICLR 2025
setting for Gemma-2b. Perplexity is measured held-out subset of each model's fine-tuning dataset. Further, FLOP compression assessed by measuring the average FLOPs required to decode 512-token sequences.
Latency Evaluations: For our latency cvaluations, we leverage 100 sequences from the RedPajama dataset, where adapted models are timed in the task of decoding sequence of 492 tokens with an initial context ranging from 1 to 1000 tokens. Evaluations are performed on an NVIDIA L40S GPU.
5.2 RANK CONTRIBUTION SPARSITY
key property of adaptive compute methods having sparse contributions from the pruned com- ponents to the module's ' output. In RaNA, adaptively prune the ranks of the AB matrix in the decomposition Wr RJ A(m(z) @ Bz), which effectively means pruning the column vectors of the A matrix (Sect. 4 Concretely, we aim for sparsity in the contribution of each column vector of the A matrix to the output of the linear layer, which we obtain by measuring (Bz)?- To measure this, we study the histograms of these contributions in Llama2-7b, Gemma-2b. and Pythia- Pythia-160M models, shown in Figs. 2a 2b
The distributions exhibit heavy tails with concentrations near zero, allowing us mask out irrelevant values and retain the most impactful ones.
Sayer
Toker Aug Layer
(a) Llama2-77 -b
(b) Gemma 2b.
Figure 2: The contribution of ranks in Linear Layer Rank Adapters is sparse for multiple layer types (Sect. (.2). Histograms outline the contribution of different column-vectors from the A matrix in the Linear Layer Rank Adapter decomposition Wa R A(m((x @ Bx) to the original layers for Llama2-76 7b (left) and Gemma-2b (right). Red dashed line indicates a 50% sparsity threshold.
5.3 RANA EVALUATIONS
Output Errors: To assess the compression capabilities of RaNA, we apply it to the MLP and QKV layers of Llama2-7b, Gemma-2b and Pythia-160M, Further, we measure the error that RaNA and other adapters induce in these layers when adapting them to consume ~500 of their FLOPs,
is a common compression ratio in the pruning literature (Ma et al. (2023) Ashkboos et al. (2024). )) For
the MLPs, we measure the normalized error [MLP(z) MLP (2)) where MLP' is the adapted MLP; MLP(=)]5
we do the analogous measurement for the QKV layers too. Intuitively, effective adapter produces small errors, as that allows it to recover the model's original behavior better. Notably, from Figs. we can observe that RaNA attains the lowest eиror across all layers when compared to secronnaddaten and other adapter types. Concretely, in the case of Llama2-78 -bb, RaNA attains an average error of 10.5%, while CATS attains an average error of 17.22% in MLP layers. Further, for Gemma-2b's s MLPs, RaNA achieves an error of 2818, while CATS attains an error of 20.31 %. Similarly, for Pythia-160M, RaNA achieves average errors of 7.87% and 0.36% across MLP and QKV layers respectively, while other approaches attain average errors of 15.23% and 0.36%. This demonstrates RaNA's capacity effectively compress modern Transformer layers, which we further show translates into practical downstream task performance and perplexity improvements.
3a
3c
3d
To examine the performance of RaNA adapters beyond their compression capabilities, we measure the perplexities and downstream task perfor- mance when applied Llama2-7b, Gemma-2b, and a set of varied sized Pythia models. For these evaluations, we first apply the given adapter to the MIP and/or QKV layers of the model at hand, targeting specific FL.OP reduction ratio. Concretely, for Llama and Pythia models we apply RaNA to MLP and QKV layers, while for Gemma for not adapting QKV layers in Gemma for simplicity, as they constitute just a small proportion of FLOPs (~5%) relative the MLP layers. Further, fine-tune the adapted models. Finally, measure their perplexity on a held-out subset of the dataset used for fine-tuning and measure their accuracies for multiple NLP benchmarks.
Downstream Task Performance and Perplexity:
only apply it to MLP layers. Notably,
opted
From Tabs. 2 we observe that RaNA outperforms the of CATS adapters across a
varied set of compression rates in SwiGLU based Transformers. Notably, not only does it outper-
8

---

Published as a conference paper at ICLR 2025
ParAA
Neuron
LLBA
Emror
Emrer
20
SVD
Layer
(a) Llama2-7b MLPP layers.
(b) Llama2-7b
(c) Gemma-2b MLP layers.
(d) Pythia-160m 16MM
QKV layers.
MLP layers.
Figure 3: RaNA adapters attain lowest errors on Transformer layers when reconstructing orig- inal layer outputs. y-axis shows the error percentage; z-aaxi shows layer number. Errors induced by different adapters are compared when compressing layers Llama2-7b, Gemma-2b and Pythia- 160M to 50% FLOPs. RaNA attains the lowest error consistently across model layers (Sect. 5.3
Method
FLOP Compression
Arc
Are Challenge
HellaSwag
PIQA
RACE
WinoGrande
Avg Acc
PPL
Rate
Easy
Llama2-7b
76.26%
43.52%
57.08%
78.07%
39.62%
69.14%
60.61%
6.39
RaNA
42%
68.81%
36.86%
51.96%
74.76%
39.71%
66.14%
56.37
8.04
CATS
42%
56.31%
27.13%
41.01%
68.28%
35.12%
57.22%
47.51
12.36
SliceGPT
42%
44.91%
21.84%
37.92%
61.64%
37.22%
55.72%
43.21%
13.69
RaNA
30%
72.85%
39.76%
54.63%
77.42%
40.48%
66.85%
58.67 %
7.29
CATS
32%
69.49%
36.01%
53.16%
76.28%
38.76%
65.59%
56.55%
7.55
SliceGPT
31%
52.36%
26.62%
42.83%
67.74%
38.28%
60.30%
48.02%
11.23
RaNA
17%
73.86%
41.55%
56.09%
77. 77.80%
39.52%
69.38%
59,70%
6.63
CATS
15%
75.17%
41.55%
56.96%
77.48%
39.52%
68.27%
59.82%
6.37
SliceGPT
17%
58.96%
31.40%
48.40%
72.52%
39.62%
63.38%
$2.38%
9.22
Table 1: RaNA outperforms neuron-adapters in perplexity and accuracy in Llama2-7b. Per- plexity is measured on ~300K tokens of RedPajama. Average accuracy is aggregated over the listed benchmarks. The compression rate outlines the average FLOP compression rate for decoding a 512- token long sequence.
forms when applied to MI.P and QKV layers, as shown for Llama2-7b (Tab. T.. but it also does when only applied the MLP layers shown for Gemma-2b (Tab. [ Particularly, for Llama2- 7b, RaNA improves over neuron-adaptive methods by attaining 4 less perplexity-poine and 8 more percentage- prrcentage-points in accuracy when reducing FLOPs by 42% on the overall model. Moreover, both RaNA and CATS outperform SliceGPT in both perplexity and downstream task performance (Tab. Fig. 5) for all FLOP compression rates, highlighting the benefits of adaptive compression meth- ods compared to static ones. Similarly, at 46% compression, for Gemma-22 2b, RaNA achieves 7 less perplexity-ponee points and more percentage-poinng points in accuracy than prior neuron adapters.
We attribute RaNA's strong performance to its notable compression capacity, its ability to more evenly distribute FLOPs across the Up-Project, Down-Prote Down-Project, and PaaeeProject layers, and its direct applicability to QKV and MLP layers, all of which posed challenges for previous neuron adapters. In addition, from Figs. Ic and 4 we observe that RaNA applied to MLP and QKV layers outper- forms neuron adapters across the varied sized set of GeLU based Pythia models in both average accuracy and perplexity. This highlights RaNA's general applicability to modern Transformers with non-sparse activations, such SwiGLU in Llama and Gemma GeLU in Pythia. Further, these results show RaNA's effectiveness beyond individual layer compression, demonstrating practical im- provements over neuron adaptive techniques in perplexity and downstream task evaluations. More- over, we attribute RaNA's and CATS' strong performance over SliceGPT, static structured pruning method, to their adaptive nature. Different from static approaches, RaNA and CATS dynamically adjust weight matrices based on the input, enabling a more efficient performance-compute trade-off practice.

---

Published as a conference paper at ICLR 2025
Method
FLOP Compression
Are Easy
Arc Challenge
HellaSwag
PIQA
RACE
WinoGrande
Avg
PPL
Rate
Acc
Gemman22
0%
59.47%
29.9 39.95%
46, 66.90%
70.24%
36.17%
56.83%
49.93%
11.05
RaNA
44%
58.84%
29.18%
42.34%
69.31%
34.83%
54.46%
48.16%
13.83
CATS
47%
51.43%
24.57%
35.07%
65.94%
28.42%
52.64%
43.01% %k
21.02
RaNA
32%
61.74%
31.66%
45.38%
70.40%
35.89%
51.62%
49,45%
11.74
CATS
34%
59.55%
29.61%
44.51%
70.95%
34.55%
54.38%
48.92%
12,60
RaNA
19%
59.26%
29.78%
46.42%
69.31 %
35.98%
55.01% %
49.29%
11.18
CATS
20%
62.79%%
32.94%
47.84%
72.25%
35.50%
55.41%
51.12%
11.41
Table 2: RaNA outperforms neuron-adapters in perplexity and accuracy in Gemma-2b. Per- plexity is measured on ~300K tokens of RedPajama. Average accuracy is aggregated over the listed benchmarks. The compression rate outlines the average FLOP compression rate for decoding a 512- token long sequence.
Telly rededs. divedagted
Latency Evaluations: From Fig la RaNA presents an im- proved theoretical accuracy- FLOPs trade-off over previous adaptive methods. To validate its practical benefits, we evalu- ate latency by constructing a accuracy- -latenc curve. Since a naive Py Torcch implementation lacks practical speedup, we in- tegrated a custom masked GEMV kernel to RaNA for both its linear-layer rank adapter and neuron adapter. Concretely, leverage the masked GEMV kernel from CATS, implemented in Triton (Tillet al. )2019))
untuted
6.0
0.2 330
Figure 4: Perplexities are mea- sured across adapted Pythia mod- cls, complementary measure- ment to accuracies outlined in Fig [c] y-axis shows perplexity mea- sured over ~300K tokens of the Pile dataset (Sect. 5.1); x-axis shows average FLOPs for for- ward pass with sequence length 512.
Leveraging a masked GEMV kernel in RaNA is possible as, in the linear-layer rank adapter, A(m(x) Bx) reduces a masked matrix-vector multiplication, with A as the matrix, m(x) as the mask, and Br as the vector. Similarly, in the neuron adapter, W (m((r) @ z) performs a masked multiplica- tion with W the matrix, m(x) as the mask, and as the vector. [b highlights RaNA's capacity to realize practi- cal speedups in various compression rates for Llama2-7b, po- sitioning RaNA as practical adapter for accelerating Trans- former inference.
6 CONCLUSION AND FUTURE WORK
We present the Adaptive Rank Allocation framework and RaNA adapters, designed address the limitations of ncuron-adaptive methods in modern Transformer architectures, By moving beyond neuron-basee basend adaptation, RaNA is effectively applicable both MLP and QKV layers, leverag- ing low-rank matrix decompositions and adaptive routers, Empirical results demonstrate that RaNA achieves greater performance over existing neuron- neuron-adaptive methods like CATS, providing improve- ments in perplexity and accuracy across benchmarks when reducing FLOPs. For instance, RaNA yields fewer cerplexity-point and improvement of percentage-pon pointss--pstsssseee in accuracy for Llama2- 7b when compressing FLOPs by 42%.
Future work seems promising. First, exploring the applicability of RaNA to other architectures, such as vision transformers or those leveraging different activation functions than the ones studied in this work could extend its impact. Additionally, investigating alternative matrix decomposition tech- niques and router configurations within Linear Layer Rank adapters could further improve RaNA's effectiveness. Finally, exploring a FLOP allocation strategy at the model level, rather than focusing solely individual layers, presents a promising opportunity improve RaNA's overall capacity.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
JL is supported by the Department of Energy Computational Science Graduate Fellowship under Award Number DE-SC00231
REFERENCES
Abien Fred Agarap. Deep learning using rectified linear units (relu), 2019. URL https:// arxiv. 0tt..///////03.008333555
Saleh Ashkboos, Maximilian L. Croci, Marcelo Gennari do Nascimento, Torsten Hoefler, and James Hensman. Slicegpt: Compress large language models by deleting rows and columns, 2024. URL
fttts:://axxi.oor//aeee
15024
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference on Machine Learning, pp. 2397-2430. PMLR, 2023.
Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi. PIQA: reasoning about physical commonsense in natural language. CoRR, abs/1911.11641, 2019. URL http:
/
/1tbs////9.... 1911. 11641
Sid Black, Stella Biderman, Eric Hallahan, Quentin Anthony, Leo Gao, Laurence Golding, Ho- race He, Connor Leahy, Kyle McDonell, Jason Phang, Michael Pieler, USVSN Sai Prashanth, Shivanshu Purohit, Laria Reynolds, Jonathan Tow, Ben Wang, and Samuel Weinbach. GPT- NeoX-20B: An open-source autoregressive language model. In Proceedings of the ACL Work- shop on Challenges & Perspectives in Creating Large Language Models, 2022. URL https: 7/ //arxiv.orge abs 2t.bb///22..000 06745
Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. Think you have solved question answering? try arc, the ai2 reasoning challenge, 2018. URL wtl.tttps:
abs/ 0.0..05557
Together Computer. Redpajama: an open dataset for training large language models, 2023. URL -
htttps://iithu: cow http:///ww..oootttteeeooooeeer/.
Carl Eckart and Gale Young. The approximation of one matrix by another of lower rank, Psychome- trika, 1(3):211-218, 1936. doi: 10.11007/BF02288337. /
BF BFC2288367
Jonathan Frankle and Michael Carbin. The lottery ticket ticke hypothesiss: Finding sparse, trainable neural networks, 2019.
/ttp:i//////vvv//// http:///wxxl...o//a//////// 03635
Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, Shawn Presser, and Connor Leahy, The pile: An 800gb dataset of diverse text for language modeling, 2020. URL htttss:/arxxi... abs Moss/1011 2211.1111111 00027
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang, Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL http:/ttt://eeeetee 10256836
Yizeng Han, Gao Huang, Shiji Song, Le Yang, Honghui Wang, and Yulin Wang. Dynamic neural networks: A survey, 2021. URL
2bb9 04906
Dan Hendrycks and Kevin Gimpel. Gaussian enror linear units (gelus), 2023. URL https:// wwtp://gg/bbbb/./.// 08415
11

---

Published as a conference paper at ICLR 2025
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, 2021.
Albert Q. Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chap- lot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, Lélio Renard Lavaud, Marie-Anne Lachaux, Pierre Stock, Teven Le Scao, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mistral 7b, 2023. URL https:
www.i..ogg.as/////ooo....at.taap /330000022
Sehoon Kim, Coleman Hooper, Thanakul Wattanawong, Minwoo Kang, Ruohan Yan, Hasan Genc, Grace Dinh, Qijing Huang, Kurt Keutzer, Michael W. Mahoney, Yakun Sophia Shao, and Amir Gholami. Full stack optimization of transformer inference: survey, 2023. URL https://
arxiv. wwwpv/orgggb/b//aeeee/o /2002 14017
Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy, Race: Large-scale reading comprehension dataset from examinations, 2017. URL httpps://axxiv..orr/abs////
04683
Je- Yong Lee, Donghyun Lee, Genghan Zhang, Mo Tiwari, and Azalia Mirhoseini. Cats:
Contextually-aware thresholding for sparsity in large language models, 2024. URL https:
//arxiiv.rrr/aas//400....
Zonglin Li, Chong You, Srinadh Bhojanapalli. Daliang Li, Ankit Singh Rawat, Sashank J. Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. The lazy neuron phenomenon: On emergence of activation sparsity in transformers, 2023a.
2210.06313.
Zonglin Li, Chong You, Srinadh Bhojanapalli,, Daliang Li, Ankit Singh Rawat, Sashank J. Reddi, Ke Ye, Felix Chern, Felix Yu, Ruiqi Guo, and Sanjiv Kumar. The lazy neuron phenomenon: On emergence of activation sparsity in transformers, 2023b. URL
2210.06315
Zichang Liu, Jue Wang, Tri Dao, Tianyi Zhou, Binhang Yuan, Zhao Song, Anshumali Shrivastava, Ce Zhang, Yuandong Tian, Christopher Re, and Beidi Chen. Deja vu: Contextual sparsity for efficient ilms inference time, 2023. URL /ttttp:://axii......t
www.ii..rr/aab//33332 17157
Xinyin Ma. Gongfan Fang, and Xinchao Wang. LIm-pruner: On the structural pruning of large language models, 2023.
Iman Mirzadeh, Keivan Alizadeh, Sachin Mehta, Carlo C Del Mundo, Oncel Tuzel, Golnoosh Samei, Mohammad Rastegari, and Mehrdad Farajtabar. Relu strikes back: Exploiting activation sparsity in large language models, 2023. /https:///artiti wrg. abtpgg/bbs//22 040564
Reiner Pope, Sholto Douglas, Aakanksha Chowdhery, Jacob Devlin, James Bradbury, Anselm Lev- skaya, Jonathan Heek, Kefan Xiao, Shivani Agrawal, and Jeff Dean. Efficiently scaling trans- former inference, 2022. URL httLhhtps:/// http://xxivv.. www. ott.bs//.rgb///////ooooo/../ss abs 2211.11111111 05102
Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. WINOGRANDE: an adversarial winograd schema challenge at scale. CoRR, abs/1907.1166m 2019. URL http:
//arxiv.org/abs/////
10641
Noam Shazeer. Glu variants improve transformer, 2020. URL ntttps//azxxv..og//
2002.05203
Gemma Team, Thomas Mesnard, Cassidy Hardin, Robert Dadashi, Surya Bhupatiraju, Shreya Pathak, Laurent Sifre, Morgane Rivière, Mihir Sanjay Kale, Juliette Love, Pouya Tafti, Léonard Hussenot, Pier Giuseppe Sessa, Aakanksha Chowdhery, Adam Roberts, Aditya Barua, Alex Botev, Alex Castro-Ros, Ambrose Slone, Amélie Héliou, Andrea Tacchetti, Anna Bulanova, An- tonia Paterson, Beth Tsai, Bobak Shahriari, Charline Le Lan, Christopher A. Choquette-Choo, Clément Crepy, Daniel Cer, Daphne Ippolito. David Reid, Elena Buchatskaya, Eric Ni, Eric Noland, Geng Yan, George Tucker, George-Christee Muraru, Grigory Rozhdestvenskly. Hen- ryk Michalewski, Ian Tenney, Ivan Grishchenko, Jacob Austin, James Keeling, Jane Labanowski, Jean-Baptiss Lespiau, Jeff Stanway, Jenny Brennan, Jeremy Chen, Johan Ferret, Justin Chiu,
12

---

Published as a conference paper at ICLR 2025
Justin Mao-Jones, Katherine Lee, Kathy Yu, Katie Millican, Lars Lowe Sjoesund, Lisa Lee, Lucas Dixon, Machel Reid, Maciej Mikula, Mateo Wirth, Michael Sharman, Nikolai Chinaev, Nithum Thain, Olivier Bachem, Oscar Chang, Oscar Wahltinez, Paige Bailey, Paul Michel, Petko Yotov, Rahma Chaabouni, Ramona Comanescu, Reena Jana, Rohan Anil, Ross Mcllroy, Ruibo Liu, Ryan Mullins, Samuel I. Smith, Sehastian Borgeaud, Sertan Girgin, Sholto Douglas, Shree Pandya, Siamak Shakeri, Soham De, Ted Klimenko, Tom Hennigan, Vlad Feinberg, Wojeiech Stokowiee, Yu hui Chen, Zafarali Ahmed, Zhitao Gong, Tris Warkentin, Ludovic Peran, Minh Giang, Clément Farabet, Oriol Vinyals, Jeff Dean, Koray Kavukcuoglu, Demis Hassabis. Zoubin Ghahramani, Douglas Eck, Joelle Barral, Fernando Pereira, Eli Collins, Armand Joulin, Noah Fiedel, Evan Senter, Alek Andreev, and Kathleen Kenealy. Gemma: Open models based on gemini research and technology, 2024. URL https: http:///axxvv.o.eeeeee 2203 08295
Philippe Tillet, Duane Kung. David Johnson, Nicolas Charrier, and Zachary Vitioo. Triton: intermediate language and compiler for tiled neural network computations. arXiv preprint arXiv:1905.00366, 2019.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet. Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023. URL[httpa / /wt.xvivvv./..////////aa//aaaav. wwgg abs/ 2307 09288
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need, 2023. URL hottss:/aarx:
org/abs/1706.00 03762
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, Joe Davison, Sam Shleifer, Patrick von Platen, Clara Ma, Yacine Jernite, Julien Plu, Canwen Xu, Teven Le Scao, Sylvain Gugger, Mariama Drame, Quentin Lhoest, and Alexander M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38 45. Association for Computational Lin- guistics, 2020. URL http:://aclatthoooryyy ort.
5
Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. Hellaswag: Can ma- chine really finish your sentence? CoRR, httbs/995.078300 2019. URL http::ttt////x..orggggg
abs/ 1905. 07830
Zhengyan Zhang. Yixin Song, Guanghui Yu, Xu Han, Yankai Lin, Chaojun Xiao, Chenyang Song, Zhiyuan Liu, Zeyu Mi, and Maosong Sun. Relu? wins: Discovering efficient activation functions for sparse 1lms, 2024. URL https: /tts:///axxiiii wwwppxwwwwo//b//////..33/
03804
13

---

Published as a conference paper at ICLR 2025
A APPENDIX
A.1 PROOF OF THEOREM 1
Here we prove Theorem
Proof. Let Mr be a rank matrix. The matrix M, that obtains the lowest possible error ([ M, is the rank-r singular value decomposition WX, namely M, U,E,V" Now, solving for Ar, Br:
A,B,X = M, U,E,V? = A,
where X is the pseudo-inverse of X. Hence, we have found the optimal A, for eqn. namely:
= = W..
A.2 ABLATIONS AND ADDITIONAL RESULTS
Compressing ML.ss only vs. MIPs + QKVs:
As an ablation study, we look at how compressing both MLPs and QKVs compare only com- pressing MLPs to the same FLOP compression ratio in Llama2-7b. In addition, we study how our FLOP-allocation algorithm for MLP layers impacts performance by evaluating adapted model that leverages the FLOP allocation procedure in the MLP layers and evaluating the same adapted model with the exception that uniformly allocates FLOPs across the components of each individual MLP layer.
Concretely, we compress three Llama2-7b models by ~31% % of their FL.OPs and evaluate their per- plexity on ~300K tokens of the RedPajama dataset, without fine-tuning. First, we look at vanilla RaNA adapted model, which leverages the FLOP allocation algorithm at the MLP level and that also compresses QKV layers. Second, we look at model with the same setup, except without compress- ing QKV layers, but still compressing ~-3 % of the overall total FLOPs. Third, we look at a model that compresses both QKV and MLP layers but does not leverage the FLOP allocation procedure for MLP layers.
Model Version
FLOP Compression Rate
PPL
Llama2- 7b MLP + QKV + FLOP Allocation
31%
8.40
Llama2-7b- - FLOP Allocation
31 %
8.79
Llama2-7b -MLP QKV (No FLOP Allocation)
32%
9.10
Table 3: Perplexity Evaluation for Different RaNA Settings.
As we can observe in Tab. 3 the best model is the one combining both MLP + QKV compression and the FLOP allocation procedure in the MLPs, while the two other versions fall behind in perplexity.
Accuracy-FLOPs trade-off including SliceGPT:
In Fig. we include the SliceGPT curve for the accuracy-FLOPs trade-off in Llama2-7b. There, we can observe that RaNA and CATS achieve a better accuracy-FLOPP trade-off than SliceGPT across various FLOP compression rates. Notably, RaNA and CATS adaptive compression meth- ods designed to reduce FLOPs rather than memory usage; hence, we focused our analysis on FLOP compression rates. On the other hand, SliceGPT is static compression method that targets both FLOP and memory reductions, with their work principally discussing their slicing compression rate, which tracks memory savings. We believe the observed performance gap attributed to this fun- damental trade-off between static and adaptive compression methods: static compression methods seek to attain both memory and FLOPs savings, at the cost a less favorable accuracy-FLOPs trade-off compared to adaptive approaches.
14

---

Published as a conference paper at ICLR 2025
0.550
0.475
acaetted models
cars adagted model
Stie-CPF adapted models
677
1.0
Baps
1610
Figure 5: Llama2-7b Accuracy v.s. FL.OSs.
A.3 FLOP COMPRESSION BREAKDOWN
Here we leave in Tab. the FLOP compression breakdown across MLP and QKV layers of the main adapted models.
Model
Total FLOP Compression
MLP FLOP Compression
QKV FLOP Compression
Gemma-2b-RaNA
44%
61%
Gemma-2b-CAn
47%
65%
0%
Gemma-2b 2b-RaNA
32%
45%
0%
Gemma-2b-CATS
34%
48%
0%
Gemma-2b-RaNa 2b-RaNA
19%
27%
0%
Gemma-2b-CATS
20%
28%
0%
Llama2-7b-RaNA
42%
47%
46%
Llama2-7b-CATS
42%
65%
0%
Llama2-7b-RaNA
30%
34%
33%
Llama2-7b-CATS
32%
50%
0%
Llama2-7b-RaNA
17%
19%
18%
Llama2-76-CATS 15%
23%
0%
Table 4: FLOP Compression Comparison for Different Models. Total, MLP, and QKV FLOP compression rates are reported for the Gemma-2b and Llama2-7b models with RaNA and CATS adapters.
15

---

Published as a conference paper at ICLR 2025
A.4 ALGORITHM
We provide a pseudocode implementation of RaNA below:
Algorithm RANA Layer Compression
procedure CoMPREsSLAVER(lave,
decomposition, prune ratio)
2:
// Find optimal rank and thresholds for compression
3: rank, threshold t- Compuuteppiimalllaanntttt (ccrmootetccooomoottttttnttnt
4:
Il Create compressed layer using low-rank approximation
5: compressed layer <- DDecomposeToRankY (layenN runk)
6: II Apply adaptive thresholding for dynamic sparsity
7: maskedJayer +- Thpphhohreshl tones threshold)
8: return masked layer
9: end procedure
Algorithm RANA MLP Transformation
1: procedure TRANSFORMMLLP(ni) input_data, prune_ ratios)
2:
// Transffomm each component with different pruning ratios
3: up masked + CompressLayer mmlp.up proj, prune ratios.up)
4: gate masked CompresslLayee Layer Layer(mlp-ggtee proj. prune. ratios.gate)
5:
Compute activation threshold for downstream pruning
6:
act threshold Compute AetivAationirrehshhhheeo
data, prune.
7:
II Construct efficient MLP with dynamic sparsity
8:
ef ficient _aalp-mppp--p CreateDynamicMLP (apkedd gate_ mate_masked, malp.down_pre act ach.thre)hlld) return ef ficient mlp
9:
11: end procedure
Algorithm 3 RANA Forward Pass
1: procedure FORWARDPASS(npuu, compressed layer)
2:
Dynamic pruning based activation magnitudes 3: activations t ComputeAActivationss (input)
4: pruned activs + ApplyDynamic (activations)
5:
I/ Efficient forward computation using compressed weights
6: output +
7: return output
8: end procedure
16


Reviews:
[{"review_id": "jssGggKLWo", "reviewer": "Reviewer_RtUt", "summary": "Authors propose RaNA (Rank and Neuron Allocator), aiming to timprove the efficiency of LLMs by compressing linear layers. Authors propose the adaptive rank allocation, and decompose all linear layer into a product of low-rank matrices and an adaptive router.", "questions": "See weaknesses", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is clearly written.\n2. The method is applicable to all linear layers. \n3. Error achieved by RaNA is far smaller than previous methods and SVD.", "weaknesses": "1. no wall-clock latency measurement. Measuring only FLOPs is not acceptable for a compression method on transformers.\n2. The paper lacks baselines from the pruning literature (WANDA, SliceGPT). These are baselines that must be compared for this method since RaNA achieved compression by low rank and sparsity.\n3. The proposed method takes individual layers into account while deciding sparsity. Some literature in LoRA (AdaLoRA) shows that different layers may need different ranks in the adaptation context. What is the gap between considering only one layer’s error and considering multiple layers’ errors while determining the rank to mask out?\n============ post rebuttal ==============\nAuthors resolved my concerns. I will raise my score to 6.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gxni5IoiWB", "reviewer": "Reviewer_LZp1", "summary": "This paper presents a novel model compression framework that leverages adaptive rank allocation via rank and neuron allocation (RaNA) adapters. Specifically, weights in linear layers are decomposed into two low-rank matrices A and B, with a learnable router in between that selects which ranks to activate. RaNA also supports adaptive FLOP allocation across adapters. The paper evaluates RaNA on LLMs such as Llama2-7B and Gemma-2b, and compares its performance to other recent neuron adaptation methods such as CATS. The evaluation compares local (layer-wise) reconstruction errors for RaNa vs. baselines, FLOPs, and accuracy on downstream tasks.", "questions": "* Please address the wall clock speedup comment in the weaknesses section.\n* On the topic of inference speedups, RaNA replaces a single GEMM in linear layers with two GEMMs (please correct me if I’m wrong) - I’d really like to understand how this might affect real-world inference latency. \n* How well does RaNA perform on networks with sparse activation functions?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "* The paper is reasonably well-written and easy to follow.\n* Model compression is a promising way of reducing LLM size and making it fit various deployment constraints such as target latency, parameter and memory-footprint. As such, this paper targets a relevant and important problem.", "weaknesses": "* I’m not a big fan of using FLOPs as an optimization metric - it is well-known at this point that FLOP reductions may or may not correspond to reductions in inference latency or wall-clock time. This is likely also the reason why the main baseline method used in the paper (CATS) introduces a custom GPU kernel to realize runtime speedups. Have the authors done any study on how the proposed method improves the inference latency or throughput of the decomposed network? A theoretical discussion may also suffice given the limited time.\n* I believe that RaNA needs to be compared to other related work on structured pruning to better understand its performance and specific scenarios where it does well. Examples might include Sheared Llama [1], ShortGPT [2] , SliceGPT [3], or Minitron [4].\n* The size (< about 7B parameters) and age (Llama2 was released nearly 2 years ago at this point) of the models being used for evaluation is a weakness of the paper. While I understand that scaling up to larger models is not always feasible given limited resources, adding results on newer models (eg: Llama3.1 8b) would help compare with other recent pruning/compression efforts.\n\nReferences:\n\n1. Xia, Mengzhou, et al. \"Sheared llama: Accelerating language model pre-training via structured pruning.\" arXiv preprint arXiv:2310.06694 (2023).\n2. Men, Xin, et al. \"Shortgpt: Layers in large language models are more redundant than you expect.\" arXiv preprint arXiv:2403.03853 (2024).\n3. Ashkboos, Saleh, et al. \"Slicegpt: Compress large language models by deleting rows and columns.\" arXiv preprint arXiv:2401.15024 (2024).\n4. Muralidharan, Saurav, et al. \"Compact language models via pruning and knowledge distillation.\" arXiv preprint arXiv:2407.14679 (2024).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "emD7Jk2DDG", "reviewer": "Reviewer_yoCw", "summary": "The work proposes a new method for input-dependent compute allocation in transformer layers that is not dependent on activation sparsity and can be applied to any linear layer. The idea is to decompose $Wx$ as $A(r(x) \\odot Bx)$ and find $A$, $B$ and $r$ such that the approximation error of $Wx$ is minimal. The derived 'RaNA'-adapters are  tested on MLP and QKV layers in a variety of Transformer models, demonstrating reduced approximation errors as well as improved performance on downstream tasks when compared to conventional neuron adapter baselines such as CATS (Lee+2024).", "questions": "When describing the downstream and perplexity results in L480, you state that QKV is only RaNA-fied for Llama and Pythia, but not Gemma. Could you please explain why RaNA was not applied to QKV in Gemma?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The key idea of this work to apply adaptive compute at the rank level of linear layers makes for a straightforward and generally applicable sparsification strategy that unlike other prior methods does not require specific architectural features or activation functions.\n\nThe formulation as an error minimization problem over the layer input distribution allows for a principled derivation of suitable decomposition and routing functions.  \n\nThe experimental results show improvements over existing neuron-adaptive methods.", "weaknesses": "While the experiments show that RaNA outperforms CATS in perplexity and accuracy, the setup offers little insight into where this improvement comes from. Table 1 compares RaNA with sparse MLP/QKV against CATS with sparse MLP only. Given that they are compared at an equal compression rate, I assume that the MLPs in the CATS model are more aggressively sparsified than the MLPs in RaNA. This makes it hard to conclude what causes the performance difference. Table 2 directly compares MLP-only sparsification but is for a different, smaller model Gemma-2b. It would be helpful to add RaNA with MLP-only sparsification for Llama2-7b to differentiate the influence of MLP vs. QKV sparsification. Furthermore, it would be useful to add an additional baseline like LLRA that like RaNA can sparsify MLP+QKV. To aid the comparison, could you also clarify the sparsification levels for each component (MLP, QKV) in both RaNA and baseline methods?\n\nL509 offers different hypotheses for the causes of RaNA's performance (compression capacity, FLOP distribution, QKV applicability) but an ablation study to investigate the factors is missing. Adding results that can shed light on RaNA's performance could greatly strengthen the analysis, for instance, by comparing RaNA with and without QKV sparsification, or with different FLOP distribution strategies.\n\nI suggest moving the proof for Proposition 1 into the main text and describing it as a concrete example to help the reader understand how the neuron-adaptive method is a special case of rank adaption. Specifically, I found it helpful to see that $B_{down}, A_{up}, r_{up} := I$ and $A_{down}, B_{up}$ adjustment are all it takes to get to the neuron-adapted version of the MLP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cwMMX2IrTx", "reviewer": "Reviewer_5oY9", "summary": "This paper introduces a dynamic rank allocation method for improving the computational complexity of a transformer model.\n\nThe paper describes a related body of work on neural adaptors, and recognizes that those that rely on sparse activation functions (ReLU) are hardly suitable for transformers, the majority of which now use non-sparse activations functions (e.g. SwiGLU). Neural adaptors that do not rely on sparse activations are better suited to transformers but suffer from inefficiencies, primarily because of the need to densely evaluate the output of the gating layer. Besides, these methods leave the question of optimizing attention layers (QKV projection) unanswered.\n\nThe paper further sets context with adaptive rank allocation, which consists in decomposing a linear into two lower-rank layers parameterized by A and B matrices. Importantly, the rank is a function of the input x, and is determined by a learnable router r(x). The paper notes that adaptive rank allocation is a generalization of neural adaptors, which can be reformulated as adaptive rank allocation.\nLinear layer rank adaptors can be simplified if the router is not input dependent, in which case the determination of an optimial decomposition into and B can be performed using the Eckart-Young theorem.\nFurther to this, a binary masker may be derived to output an input-dependent mask which, in expectation, leads to the target average rank. Specifically the binary masker requires calculating (Bx)^2. This operation is inexpensive when the output dimension is bigger than the input dimension and is thus suitable for QKV, up-projection and gate-projections.\nThe paper then tackles the case of down-projections: for these layers, a simple neural thresholding masker is used.\n\nExperimental results are shown to empirically prove that a small number of A column vectors capture most of the information, making low-rank approximation efficient.\nFurther experimental results show superior accuracy and perplexity vs other methods at similar compression ratios.", "questions": "Can you provide a pseudo-code implementation of the method? \n\nHow does the method compare against FlexTron (https://arxiv.org/pdf/2406.10260), Minitron (https://arxiv.org/pdf/2408.11796).\n\nWhat is the relationship between the neural thresholding employed for down-projection layers and CATS?\n\nAre the compression ratios mentioned in experimental results taking the contribution of the router FLOPs into account? Can you show a breakdown of total FLOPs between the main computation and the computation of the routers.\n\nHow is the FLOPs reduction achieved in practice? Are tensor sliced and how does this affect memory caching/coalescing? Which DL framework was used to implement the method? \n\nDoes batched inference work and achieve the desired speed-ups if masks are inconsistent between samples?\n\nCould you share end-to-end inference latency measurements and compare the achieved speed-up with the FLOPs reduction?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is very well written and is grounded with elaborate theory on the subject.\nThe topic of the paper is important for the field of transformers.\nThe method could be applied to a broad range of model architectures, including vision transformers.", "weaknesses": "The paper compares the method against a rather narrow range of baselines. It would be interesting to compare against pruning baselines, as these are related methods for optimizing the runtime performance of models.  Other baselines could include LLMPruner, SliceGPT, ShearedLLaMA, FlexTron, Minitron, etc.\n\nThe paper does not give much details on the practical implementation of the method. It would be useful to see more information on the training process. It would help readers to have an informative diagram to summarize the process. It would also be interesting to know how the method interfaces with DL frameworks (PyTorch, FlexAttention, etc.). An open-source implementation would be best.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)