LLM Compression with Jointly Optimizing Architectural and Quantization choices

# LLM Compression with Jointly Optimizing Architectural and Quantization choices
Source: [https://arxiv.org/html/2606.04063](https://arxiv.org/html/2606.04063)
11institutetext:UiT The Arctic University of Norway22institutetext:University of Oslo, Norway
22email:\{hoang\.l\.la,phuong\.ha\.hoai\}@uit\.no
\{truongl,amirhost\}@ifi\.uio\.no###### Abstract

Deploying large language models \(LLMs\) is challenging due to their significant memory and computational requirements\. While some methods address this by developing small or tiny language models from scratch, these approaches demand extensive GPU training\. Compressing pre\-trained LLMs for edge devices offers a compelling alternative\. Beyond pruning and quantization, Neural Architecture Search \(NAS\) enables effective compression, yet prior NAS approaches often limit the search space and decouple architecture from quantization\. We introduce a differentiable NAS framework that explores the entire space and jointly optimizes architectural configurations alongside mixed\-precision quantization for linear layers of LLMs\. Experiments demonstrate superior accuracy\-latency trade\-offs: our models achieve up to1\.4××1\.4\\texttimes\\timesfaster inference than sequential NAS\-then\-quantization baselines at comparable accuracy, or up to 6% higher average accuracy across seven reasoning tasks at equivalent latency\.

## 1Introduction

In recent years, Large Language Models \(LLMs\) have gained widespread attention, but their high computational and memory demands make deployment difficult on resource\-constrained devices like laptops and smartphones\. Rising privacy concerns with cloud\-based LLMs have fueled interest in on\-device inference, yet memory requirements remain a key barrier\.

Two main approaches address the deployment challenges of LLMs: developing novel lightweight language models and compressing existing pre\-trained LLMs\. The first method involves training small language models from scratch, such as TinyLlama\[[31](https://arxiv.org/html/2606.04063#bib.bib26)\]with 1 billion parameters, which requires 90 days on 16 A100\-40GB GPUs\. Similarly, training Phi\-2\[[16](https://arxiv.org/html/2606.04063#bib.bib27)\]with 2\.7 billion parameters takes 14 days on 96 A100 GPUs\.

The second approach leverages pre\-trained LLMs, avoiding training from scratch and significantly reducing training time\. Alongside techniques like structural pruning\[[1](https://arxiv.org/html/2606.04063#bib.bib7)\]and quantization\[[10](https://arxiv.org/html/2606.04063#bib.bib28)\], neural architecture search \(NAS\) has emerged as a key method for LLM compression\. However, current NAS applications for LLM compression face challenges, such as requiring resource\-intensive supernet training\[[6](https://arxiv.org/html/2606.04063#bib.bib8),[5](https://arxiv.org/html/2606.04063#bib.bib5)\]or updating only a small subset of candidate sub\-networks during supernet training\[[25](https://arxiv.org/html/2606.04063#bib.bib3),[19](https://arxiv.org/html/2606.04063#bib.bib11)\]\.

In contrast to prior works that explore only a limited portion of the search space\[[25](https://arxiv.org/html/2606.04063#bib.bib3),[19](https://arxiv.org/html/2606.04063#bib.bib11),[5](https://arxiv.org/html/2606.04063#bib.bib5)\], Our proposed NAS framework optimizes over the full defined discrete search space via a relaxation\. By directly optimizing architecture parameters under given constraints— without pre\-selection as in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\], our method explores a broader search space\.

Furthermore, unlike conventional approaches that apply uniform quantization post\-pruning, our method jointly optimizes both architectural configurations and layer\-specific non\-uniform quantization policies\. This enables the discovery of optimal pruning\-quantization combinations, achieving substantially lower memory usage and superior accuracy compared to state\-of\-the\-art methods\. Our main contributions are as follows\.

- •We introduce a novel differential weight\-entanglement supernet design together with a constrained differential optimization method for efficient compression of pre\-trained LLMs\. Our approach achieves superior accuracy and lower latency compared to the state\-of\-the\-art methods\.
- •We present the first unified NAS framework that simultaneously optimizes model architecture and layer\-wise quantization precision for LLM, addressing the longstanding limitation of treating pruning and quantization as separate steps\. Models found by our joint approach deliver up to1\.4×1\.4\\timesfaster inference than those produced by sequential NAS\-then\-quantization pipelines\.
- •We develop a novel vectorized implementation that significantly accelerates training of weight\-entanglement supernets for LLMs, reducing training time by up to 4×\\timescompared to the original approach\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]\.

## 2Related Work and Our Advancements

### 2\.1Weight\-Entanglement NAS

Sukthanker et al\.\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]introduce a weight superposition technique in TangleNAS, which consolidates all possible weight matrix configurations into a single weighted representation, with each configuration assigned a learned importance scalar\. This enables simultaneous training and architecture search in a single stage, offering the potential for reduced search costs\. However, applying TangleNAS directly to large language models \(LLMs\) presents significant challenges\.

- •The mixed\-operation approach from TangleNAS\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]is not suitable for compressing the depth dimension of LLMs\. It restricts depth reduction to dropping only the final consecutive blocks, whereas the importance of individual blocks in a pre\-trained foundation model varies substantially\. Removing a more critical block can lead to substantial accuracy degradation\.
- •The weight\-entanglement mechanism in TangleNAS was not optimized for GPU efficiency, making it impractical for training supernet at the scale required for LLMs \(see Section[3\.3](https://arxiv.org/html/2606.04063#S3.SS3)for further discussion\)\.

Our contribution: We propose an efficient weight\-entanglement\-based supernet design specifically tailored for compressing large language models\. Our framework significantly expands the search space by incorporating diverse quantization precision options across layers\. To address the first limitation of TangleNAS, we propose importance\-aware depth pruning that enables more flexible and effective depth compression \(detailed in Section[3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2)\)\. To overcome the second limitation, we develop a software\-level optimization that substantially accelerates supernet training, making it feasible to compress LLMs using only a single NVIDIA A100 80GB GPU within practical time budgets\.

### 2\.2Neural Architecture Search techniques for LLM compression

A notable feature of transformers is permutation equivariance, which allows reordering embedding features, MLP intermediate features, and attention heads without significantly impacting model accuracy\[[27](https://arxiv.org/html/2606.04063#bib.bib9)\]\. Leveraging this, a common preprocessing step involves ranking the components of a pre\-trained LLM by importance\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\. When selecting a sub\-network, we can then simply choose the first neurons or heads from the supernet\. This preprocessing method is widely used in the related work\[[5](https://arxiv.org/html/2606.04063#bib.bib5),[6](https://arxiv.org/html/2606.04063#bib.bib8),[25](https://arxiv.org/html/2606.04063#bib.bib3)\]and is also adopted in our study\.

LoNAS\[[19](https://arxiv.org/html/2606.04063#bib.bib11)\]andsubnet\-selection\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]adopted a two\-stage neural architecture search \(NAS\) approach to identify optimal sub\-architectures in a pre\-trained large language model \(LLM\)\. In the first stage, they train a supernet using LoRA\[[14](https://arxiv.org/html/2606.04063#bib.bib53)\]\. In the second stage, they employ a multi\-objective search to find sub\-architectures that optimize accuracy and performance metrics, such as latency and energy efficiency, with the supernet serving as an accuracy estimator for sub\-networks\.

Unlike LoNAS\[[19](https://arxiv.org/html/2606.04063#bib.bib11)\],subnet\-selectionmethod\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]incorporates a pre\-selection step before the NAS process\. They rank the model’s features and blocks by importance and select the top neurons, heads, or blocks when sampling sub\-networks\. They observed that random sampling, as used in LoNAS, introduces bias in the search space, where smaller sub\-architectures are updated more frequently than larger ones, complicating supernet training\. To address this and optimize the search space, they introduced a grid\-based sampling approach, dividing the search space intoKKpartitions and selecting the best candidate from each\. During supernet fine\-tuning, they randomly selectk≪Kk\\ll Ksub\-networks and use knowledge distillation to train these sub\-networks alongside the largest \(original\) network\. However, this heuristic approach can introduce a strong bias based on the initial selection criteria, potentially resulting in suboptimal architectures from the beginning\.

Our contribution: Previous techniques focus primarily on the often\-updated subnets while disregarding the bulk of those that are rarely or never updated, which may lead to missing the genuine optimal solution\. In this work, we present a novel differential neural architecture search \(NAS\) method designed for compressing pre\-trained large language models \(LLMs\)\. By leveraging weight entanglement style supernet, our differential supernet does not require randomly sampling architectures from search space, and thus avoids the skewness of architecture distribution as in LoNAS\. Meanwhile, our method does not rely on any heuristic to pre\-select architectures as insubnet\-selection\. In contrast, our method explores all conceivable candidate subnets during fine\-tuning, progressively converging toward the optimal structure\. Additionally, unlike earlier methods that lack support for quantization or require to apply compression and quantization techniques sequentially, our approach simultaneously optimizes architectural parameters and quantization precision across layers\.

## 3Method

### 3\.1Constrained Differential NAS

#### 3\.1\.1Problem Formulation

We approach the compression of large language models \(LLMs\) as a constrained optimization problem\. Our search spaceSSis parameterized byζ∈S\\zeta\\in S, governing the architectural structure in a fully differentiable manner similar as in\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]\. A sampled candidate network is drawn asζ^∼Pζ​\(S\)\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\), wherePζ​\(S\)P\_\{\\zeta\}\(S\)is a probability distribution over search spaceSS, parameterized byζ\\zeta\.ℒtrain\\mathcal\{L\}\_\{\\text\{train\}\}andℒval\\mathcal\{L\}\_\{\\text\{val\}\}denote the training and validation losses, respectively\. Consequently, the LLM compression task can be formulated as a bi\-level constrained optimization problem as described below\.

ζ∗=\\displaystyle\\zeta^\{\*\}=\\quadminζ⁡ℒval​\(w∗,ζ^\)\\displaystyle\\min\_\{\\zeta\}\\mathcal\{L\}\_\{\\text\{val\}\}\(w^\{\*\},\\hat\{\\zeta\}\)\(1a\)s\.t\.w∗=arg⁡minw⁡ℒtrain​\(w,ζ^\),\\displaystyle w^\{\*\}=\\arg\\min\_\{w\}\\mathcal\{L\}\_\{\\text\{train\}\}\(w,\\hat\{\\zeta\}\),Bmin<Fparams​\(ζdiscrete\)<Bmax\\displaystyle B\_\{\\text\{min\}\}<F\_\{\\text\{params\}\}\(\\zeta\_\{\\text\{discrete\}\}\)<B\_\{\\text\{max\}\}\(1b\)
Denotewwas weight of the pre\-trained model\. LetFparams​\(ζ^\)F\_\{\\text\{params\}\}\(\\hat\{\\zeta\}\)denotes the parameter count of the optimal neural architecture derived from the discretization step within the NAS procedure\. This discretization step entails applying thearg⁡max\\arg\\maxfunction to the architectural parametersζ\\zetain order to select the definitive architecture, yieldingζdiscrete=arg⁡max⁡\(ζ\)\\zeta\_\{\\text\{discrete\}\}=\\arg\\max\(\\zeta\)\. The involvement of thearg⁡max\\arg\\maxoperation rendersFparamsF\_\{\\text\{params\}\}non\-differentiable\. A simple method for approximatingFparams​\(ζdiscrete\)F\_\{\\text\{params\}\}\(\\zeta\_\{\\text\{discrete\}\}\)involves relaxing the hard constraint by calculating the expected value𝐄ζ^∼Pζ​\(S\)​\[Fparams​\(ζ^\)\]\\mathbf\{E\}\_\{\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\)\}\[F\_\{\\text\{params\}\}\(\\hat\{\\zeta\}\)\]\.

#### 3\.1\.2Constrained Optimization

We can transform the constrained optimization into an unconstrained one by adding a regularization term for the constraint[1b](https://arxiv.org/html/2606.04063#S3.E1.2)\. Particularly, the constraint[1b](https://arxiv.org/html/2606.04063#S3.E1.2)can be formalized with a pair of ReLU functions as follows:

Fconstraint=\\displaystyle F\_\{\\text\{constraint\}\}=\\;R​e​L​U​\(𝐄ζ^∼Pζ​\(S\)​\[Fparams​\(ζ^\)\]−Bmax\)\\displaystyle ReLU\(\\mathbf\{E\}\_\{\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\)\}\[F\_\{\\text\{params\}\}\(\\hat\{\\zeta\}\)\]\-B\_\{\\text\{max\}\}\)\+R​e​L​U​\(Bmin−𝐄ζ^∼Pζ​\(S\)​\[Fparams​\(ζ^\)\]\)\\displaystyle\+ReLU\(B\_\{\\text\{min\}\}\-\\mathbf\{E\}\_\{\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\)\}\[F\_\{\\text\{params\}\}\(\\hat\{\\zeta\}\)\]\)\(2\)The loss terms in Equation[1a](https://arxiv.org/html/2606.04063#S3.E1.1)can be reformulated as:

ℒtrain​\(w,ζ^\)=\\displaystyle\\mathcal\{L\}\_\{\\text\{train\}\}\(w,\\hat\{\\zeta\}\)=\\;𝐄ζ^∼Pζ​\(S\)​\[ℒCE​\(w,ζ^\)\]\\displaystyle\\mathbf\{E\}\_\{\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\)\}\[\\mathcal\{L\}\_\{\\text\{CE\}\}\(w,\\hat\{\\zeta\}\)\]\(3\)ℒval​\(w∗,ζ\)=\\displaystyle\\mathcal\{L\}\_\{\\text\{val\}\}\(w^\{\*\},\\zeta\)=\\;𝐄ζ^∼Pζ​\(S\)​\[ℒCE​\(w∗,ζ^\)\+η​Flatency​\(ζ^\)\]\\displaystyle\\mathbf\{E\}\_\{\\hat\{\\zeta\}\\sim P\_\{\\zeta\}\(S\)\}\[\\mathcal\{L\}\_\{\\text\{CE\}\}\(w^\{\*\},\\hat\{\\zeta\}\)\+\\eta F\_\{\\text\{latency\}\}\(\\hat\{\\zeta\}\)\]\+λ​Fconstraint\\displaystyle\+\\lambda F\_\{\\text\{constraint\}\}
whereLCEL\_\{\\text\{CE\}\}is the cross\-entropy loss andFlatency​\(ζ^\)F\_\{\\text\{latency\}\}\(\\hat\{\\zeta\}\)is the expected inference latency of the sampled architecture\. This latency term is computed as a probability\-weighted average of per\-choice latencies, which are retrieved from a pre\-calculated lookup table\. We denoteη\\etais a hyperparameter defining the trade\-off between validation lossLCEL\_\{\\text\{CE\}\}and inference latencyFlatencyF\_\{\\text\{latency\}\}\.λ\\lambdais a hyperparameter to control strength of the regularization termFconstraintF\_\{\\text\{constraint\}\}\. It is note\-worthy thatFlatencyF\_\{\\text\{latency\}\}can be substituted with other user\-defined metrics, such as energy consumption or memory usage\.

#### 3\.1\.3Pruning during supernet fine\-tuning

Assume we haveDDconfigurable architectural dimensions, such as the number of blocks, number of neurons, or number of heads per block\. For each architectural dimension, we haveCCdifferent choices\. Denotepcdp^\{d\}\_\{c\}as the probability of thect​hc^\{th\}choice in thedt​hd^\{th\}architectural dimension\. The architectural entropy is defined as:

H=−1D​∑d=1D∑c=1Cpcd​log⁡pcd\.\\displaystyle H=\-\\frac\{1\}\{D\}\\sum\_\{d=1\}^\{D\}\\sum\_\{c=1\}^\{C\}p^\{d\}\_\{c\}\\log p^\{d\}\_\{c\}\.\(4\)WhenH<ϵH<\\epsilon, indicating convergence to a single sub\-architecture, we prune all redundant branches, retaining only the optimal sub\-architecture for continued fine\-tuning\.

#### 3\.1\.4Knowledge Distillation

After the pruning step and the supernet has converged to an optimal sub\-network \(whenH<ϵH<\\epsilon\), we continue fine\-tuning it via a knowledge distillation approach, in which the largest subnet \(the original model\) acts as the teacher and the optimal sub\-network serves as the student, thereby boosting the sub\-network’s accuracy\.

### 3\.2Supernet design

#### 3\.2\.1Width dimensions

To differentiate between the width and depth aspects of the architecture, we denoteα\\alphaandβ\\betaas the respective parameters for these dimensions\. Consequently, the overall architectural parameters areζ=\{α,β\}\\zeta=\\\{\\alpha,\\beta\\\}\. The width aspects encompass elements like the hidden size, count of attention heads, head size, and intermediate size, while the depth aspect pertains to the count of transformer blocks\. Figure[1](https://arxiv.org/html/2606.04063#S3.F1)depicts our supernet architecture, emphasizing the width dimension\. Specifically, we introduce a mixed\-operation weight for every embedding layer, normalization layer, and linear layer\. This mixed weight is formed by a straightforward linear combination of possible sizes for the layer’s original weight, with contributions weighted by the sampled mixing coefficientsα^\\hat\{\\alpha\}\.

Drawing inspiration from DrNAS\[[7](https://arxiv.org/html/2606.04063#bib.bib12)\], we model the mixing weightsα^\\hat\{\\alpha\}, which encode the relative importance or selection probability of each width option, as random variables drawn from a Dirichlet distribution parameterized byα\\alpha\. We denoteαI​N\\alpha^\{IN\}andαO​U​T\\alpha^\{OUT\}as the architectural parameters governing the input and output dimension selections for a given layer, respectively\. For a linear layer with base weight matrixW0W\_\{0\}of sizeM×NM\\times N, letα^​iI​N\\hat\{\\alpha\}i^\{IN\}andα^​jO​U​T\\hat\{\\alpha\}j^\{OUT\}represent the sampled mixing weights for theii\-th input dimension choiceNiN^\{i\}and thejj\-th output dimension choiceMjM^\{j\}, respectively\. A sampling functionFs​a​m​p​l​e​\(W0,Ni,Mj\)F\_\{sample\}\(W\_\{0\},N^\{i\},M^\{j\}\)extracts the top\-left submatrix ofW0W\_\{0\}consisting of the firstMjM^\{j\}rows andNiN^\{i\}columns, then zero\-pads it to restore the originalM×NM\\times Nshape\. As in TangleNAS\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\], the mixed\-operation weightWm​i​x​e​dW\_\{mixed\}for the layer is calculated accordingly\.

![Refer to caption](https://arxiv.org/html/2606.04063v1/x1.png)Figure 1:An overview of mixed\-operation supernet design for width dimensions\.Wmixed=∑i∑jα^iIN​α^jOUT​Fsample​\(W0,Ni,Mj\)\\displaystyle W\_\{\\text\{mixed\}\}=\\sum\_\{i\}\\sum\_\{j\}\\hat\{\\alpha\}\_\{i\}^\{\\text\{IN\}\}\\hat\{\\alpha\}\_\{j\}^\{\\text\{OUT\}\}F\_\{\\text\{sample\}\}\(W\_\{0\},N^\{i\},M^\{j\}\)\(5\)with​α^iIN∼D​i​r​\(αiIN\),α^jOUT∼D​i​r​\(αjOUT\)\\displaystyle\\text\{with \}\\hat\{\\alpha\}\_\{i\}^\{\\text\{IN\}\}\\sim Dir\(\\alpha^\{\\text\{IN\}\}\_\{i\}\),\\hat\{\\alpha\}\_\{j\}^\{\\text\{OUT\}\}\\sim Dir\(\\alpha^\{\\text\{OUT\}\}\_\{j\}\)
The iterations in the loops of Equation[5](https://arxiv.org/html/2606.04063#S3.E5)are computationally independent\. Moreover, each iteration in the above equation requires a different padding size\(N−Ni,M−Mj\)\(N\-N^\{i\},M\-M^\{j\}\)\. Therefore, the computation of the above equation requires non\-uniform padding operations\. However, Pytorch does not support parallelism for such operations and requires a loop over eachiiandjjto run the slicing and padding operations\. Consequently, naively computing the above equation on GPU with Pytorch, as in TangleNAS\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]is expensive\. Further details on this issue are discussed in Section[3\.3](https://arxiv.org/html/2606.04063#S3.SS3)\.

#### 3\.2\.2Depth dimension

Unlike the width dimensions, applying the mixed\-operation approach to the depth dimension presents several challenges\. Figure[2\(a\)](https://arxiv.org/html/2606.04063#S3.F2.sf1)illustrates the mixed\-operation design for the depth dimension\. This mixed\-operation design requires to forward allLLblocks sequentially\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\], and thus, it imposes strict constraints on depth reduction\. Particularly, the final mixed\-operation output for the block sequence is computed as a weighted sum of the outputs, with weights governed by the architectural parametersβ\\beta\. By optimizingβ\\beta, whenβi→1\\beta\_\{i\}\\to 1, the mixed output is primarily influenced by the output of theit​hi^\{th\}block, effectively equivalent to sequentially forward all blocks up to theit​hi^\{th\}block and excluding all subsequent blocks\.

In a pre\-trained large language model \(LLM\), transformer blocks have different levels of importance, and removing more critical blocks can significantly impair model performance compared to removing less essential ones\. The block importance metric, first introduced by\[[20](https://arxiv.org/html/2606.04063#bib.bib10)\], quantifies a block’s sensitivity by measuring the cosine similarity between its input and output\. We adopt this metric to assess block importance in our study\. As illustrated on the right side of Figure[2](https://arxiv.org/html/2606.04063#S3.F2), pruning blocks based on their importance results in a lower validation loss for the compressed sub\-network compared to simply removing the final consecutive blocks\. This finding aligns with\[[23](https://arxiv.org/html/2606.04063#bib.bib56)\], who examined the validation loss of the Llama\-3 8B model when dropping 16 consecutive final blocks versus 16 non\-consecutive blocks selected by importance\. They found that dropping consecutive layers caused a significantly larger increase in validation loss compared to dropping non\-consecutive layers\.

We propose an importance\-aware depth pruning method to compress the depth dimension of large language models as in Algorithm[1](https://arxiv.org/html/2606.04063#alg1)\. Specifically, we modelβ\\betaas parameters of a categorical distributionC​a​t​\(β\)Cat\(\\beta\)\. In each training iteration, we sample the number of retained blocksL^∼C​a​t​\(β\)\\hat\{L\}\\sim Cat\(\\beta\)\. To account for block importance, we maintain an array of block indices sorted by the importance of the corresponding blocks\. For each sampledL^\\hat\{L\}, we perform the forward pass using only the topL^\\hat\{L\}most important blocks, bypassing the rest\. Whenβi→1\\beta\_\{i\}\\to 1, it indicates that we retain theiimost important blocks and discard the remainingL−iL\-ileast important ones\. To enable gradient updates forβ\\beta, we use a differentiable sampling technique, specifically the ReinMax method\[[18](https://arxiv.org/html/2606.04063#bib.bib50)\], a state\-of\-the\-art gradient estimation technique for categorical sampling that provides enhanced robustness, as opposed to the Gumbel\-Softmax trick described in\[[15](https://arxiv.org/html/2606.04063#bib.bib51)\]\. Figure[2\(b\)](https://arxiv.org/html/2606.04063#S3.F2.sf2)depicts our proposed design for the depth dimension\.

![Refer to caption](https://arxiv.org/html/2606.04063v1/x2.png)\(a\)The mixed\-op design
![Refer to caption](https://arxiv.org/html/2606.04063v1/x3.png)\(b\)Our design
![Refer to caption](https://arxiv.org/html/2606.04063v1/x4.png)

Figure 2:Left: The difference between the mixed\-op design \(a\) and our proposed design \(b\) for pruning depth dimension\. Right: Validation loss of compressed sub\-network from Llama\-3\.1\-8B model by dropping blocks with two schemes, namely dropping last consecutive blocks and dropping blocks by their importance\.Algorithm 1Importance\-Aware Probabilistic Depth Pruning0:Pre\-trained LLM with

NNtransformer blocks, learnable concentration parameters

β∈ℝN\\beta\\in\\mathbb\{R\}^\{N\}, block importance scores

I∈ℝNI\\in\\mathbb\{R\}^\{N\}
1:

s​o​r​t​e​d​\_​i​n​d​i​c​e​s←argsort​\(I\)sorted\\\_indices\\leftarrow\\text\{argsort\}\(I\)\{Indices of blocks sorted from most to least important\}

2:foreach training iterationdo

3:

L^∼C​a​t​\(β\)\\hat\{L\}\\sim Cat\(\\beta\)
4:

kept\_indices←sorted\_indices\[1:L^\]kept\\\_indices\\leftarrow sorted\\\_indices\[1:\\hat\{L\}\]
5:Sort

k​e​p​t​\_​i​n​d​i​c​e​skept\\\_indicesin ascending order

6:

i​n​p​u​t←input\\leftarrowinitial input to the first block

7:for

i=1i=1to

L^\\hat\{L\}do

8:if

i∈k​e​p​t​\_​i​n​d​i​c​e​si\\in kept\\\_indicesthen

9:

o​u​t​p​u​t←blocki​\(i​n​p​u​t\)output\\leftarrow\\text\{block\}\_\{i\}\(input\)
10:else

11:

o​u​t​p​u​t←i​n​p​u​toutput\\leftarrow input
12:endif

13:

i​n​p​u​t←o​u​t​p​u​tinput\\leftarrow output
14:endfor

15:Compute loss and backpropagate \(gradients flow to

β\\betavia ReinMax reparameterization\)

16:endfor

#### 3\.2\.3Weight Quantization

Fully fine\-tuning the supernet described above is memory\-intensive and impractical for deployment on a single NVIDIA A100 GPU\. To mitigate this, we utilize LoRA techniques to reduce memory usage\. Specifically, we incorporate LoRA adapters into all linear layers of the attention and MLP blocks\. LetQp​\(⋅\)Q\_\{p\}\(\\cdot\)represent the quantization function that quantizes an input matrix to precisionpp, withAAandBBdenoting the matrices of the LoRA adapters\. The standard approach to integrating LoRA with quantization involves quantizing the sum of the original weight and the LoRA adapter, as described in\[[4](https://arxiv.org/html/2606.04063#bib.bib25),[17](https://arxiv.org/html/2606.04063#bib.bib21)\]\. The quantized weight is formalized in Equation[6](https://arxiv.org/html/2606.04063#S3.E6)\.

W~=Qp​\(W0\+B​A\)\\displaystyle\\tilde\{W\}=Q\_\{p\}\(W\_\{0\}\+BA\)\(6\)
This design can eliminate mixed\-precision computation overhead during the inference phase\[[28](https://arxiv.org/html/2606.04063#bib.bib34)\]\. Specifically, we employ a unique pair of LoRA matrices A and B for each distinct weight quantization precision, rather than using a single pair for all precisions\. This method addresses discrepancies in precision ranges among varying precisions, which could otherwise lead to instability in the supernet fine\-tuning process\. LetPPdenote the set of weight precision\. The mixed\-op LoRA matrix for a precisionppis denoted asB​Ap,m​i​x​e​dBA\_\{p,mixed\}\. The quantized mixed\-op weightW~m​i​x​e​d\\tilde\{W\}\_\{mixed\}can be calculated as follows\.

W~mixed=∑p∈Pαpweight∗Qp​\(Wmixed\+B​Ap, mixed\)\\displaystyle\\tilde\{W\}\_\{\\text\{mixed\}\}=\\sum\_\{p\\in P\}\\alpha^\{\\text\{weight\}\}\_\{p\}\*Q\_\{p\}\(W\_\{\\text\{mixed\}\}\+BA\_\{\\text\{p, mixed\}\}\)\(7\)Denoteαpw​e​i​g​h​t\\alpha^\{weight\}\_\{p\}is the probability that the weight matrixWWbeing quantized to precisionpp\. One advantage of our approach is its compatibility with various quantization\-aware training \(QAT\) techniques\[[17](https://arxiv.org/html/2606.04063#bib.bib21),[4](https://arxiv.org/html/2606.04063#bib.bib25)\], enabling the integration of fine\-tuning and QAT to obtain the most effective architectures, we defer this direction for future research\. In this paper, we use Straight\-Through Estimator\[[2](https://arxiv.org/html/2606.04063#bib.bib32)\]to enable gradient flow through quantization operations\.

Note that the QA\-LoRA methods\[[28](https://arxiv.org/html/2606.04063#bib.bib34)\]are incompatible with our above design\. Specifically, QA\-LoRA alters the output shape of theQp\(\.\)Q\_\{p\}\(\.\)function to\[M,Ngroup\_size\]\[M,\\frac\{N\}\{\\text\{group\\\_size\}\}\]under low\-precision settings, whereas it retains the\[M,N\]\[M,N\]shape for 16\-bit precision, causing shape mismatches in the sum aggregation of Equation[7](https://arxiv.org/html/2606.04063#S3.E7)\.

To enhance processing speed through better GPU parallelism, we present a software enhancement in Section[3\.3](https://arxiv.org/html/2606.04063#S3.SS3)\. Regarding initialization, techniques for post\-training quantization can be employed, including GPTQ\[[11](https://arxiv.org/html/2606.04063#bib.bib19)\]or OmniQuant\[[22](https://arxiv.org/html/2606.04063#bib.bib20)\], among others\.

#### 3\.2\.4Activation Quantization

We further explore activation quantization with a similar design\. Particularly, the mixed activationXm​i​x​e​dX\_\{mixed\}of different precisions can be defined as follows:

Xmixed=∑p∈Pαpactivation​Qp​\(X\)\\displaystyle X\_\{\\text\{mixed\}\}=\\sum\_\{p\\in P\}\\alpha^\{\\text\{activation\}\}\_\{p\}Q\_\{p\}\(X\)\(8\)

### 3\.3Software Improvement: Vectorizing calculation of mixed\-op weights

The weight\-entanglement method requires multiple nested loops to calculate combined weights and low\-rank adaptations \(LoRA\) for each linear layer, which greatly slows down fine\-tuning\. Specifically, Equation[5](https://arxiv.org/html/2606.04063#S3.E5)for the aggregated weight matrix is computationally expensive due to loops over indicesiiandjj, plus zero\-padding that demands dynamic memory handling in each iteration\. The left side of Figure[3](https://arxiv.org/html/2606.04063#S3.F3)illustrates the original implementation and the implementation of our proposed software improvement for computing mixed weight\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]\.

To overcome this problem, we develop a method to vectorize the mixed weight calculation\. LetMiM\_\{i\}andNjN\_\{j\}be the input and output feature sizes linked to architectural parametersαiIN\\alpha\_\{i\}^\{\\text\{IN\}\}andαjOUT\\alpha\_\{j\}^\{\\text\{OUT\}\}\. We use a binary mask matrixGi,jG\_\{i,j\}\(sizeN×MN\\times M\) to indicate pairings between theii\-th input andjj\-th output dimensions, where each elementgn,mg\_\{n,m\}is:

gn,m=\{1if​n<Ni​and​m<Mj,0otherwise\.\\displaystyle g\_\{n,m\}=\\begin\{cases\}1&\\text\{if \}n<N\_\{i\}\\text\{ and \}m<M\_\{j\},\\\\ 0&\\text\{otherwise\}\.\\end\{cases\}To fix the loop inefficiencies, we rewrite Equation[5](https://arxiv.org/html/2606.04063#S3.E5)as:

Wmixed=W0⊙∑i∑jGi,j​αi,j\\displaystyle W\_\{\\text\{mixed\}\}=W\_\{0\}\\odot\\sum\_\{i\}\\sum\_\{j\}G\_\{i,j\}\\alpha\_\{i,j\}\(9\)Here,⊙\\odotis element\-wise multiplication andαi,j=αiI​N​αjO​U​T\\alpha\_\{i,j\}=\\alpha\_\{i\}^\{IN\}\\alpha\_\{j\}^\{OUT\}\. The masksGi,jG\_\{i,j\}can be precomputed at model startup and reused across blocks\. The sum involves multiplying a 3D tensor of allGGmasks by the scalar vector𝜶\\boldsymbol\{\\alpha\}, then summing element\-wise to create a 2D probabilistic mask that adjusts the base weightsW0W\_\{0\}\. This vectorized version avoids loops, boosting fine\-tuning speed\.

Although Equations[5](https://arxiv.org/html/2606.04063#S3.E5)and[9](https://arxiv.org/html/2606.04063#S3.E9)possess comparable complexity, our revised formulation supports parallel computation through broadcasting and element\-wise operations\. By precomputing the binary maskGGonly once at initialization, it avoids the repeated slicing and padding required for calculatingWm​i​x​e​dW\_\{mixed\}across layers in every training iteration\. We additionally apply this updated approach to the consolidated LoRA matrices presented in Equation[7](https://arxiv.org/html/2606.04063#S3.E7)\. While our technique incurs a minor overhead for retaining the binaryGGmatrix, approximately 3\.2 GB out of the 80GB available on an NVIDIA A100 GPU for the Llama\-3\.1\-8B model using theLlama3Spacesearch space \(outlined in Section[4\.4](https://arxiv.org/html/2606.04063#S4.SS4)\), this remains acceptable\. Furthermore, the overhead is determined exclusively by the search space dimensions and is unaffected by size of training input prompts\. In empirical tests with theLlama3Spacesearch space, this software enhancement provides up to 4\.3×\\timesimproved training throughput over the baseline weight\-entanglement method from\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]\. This optimization can be adopted for any transformer\-based architectures involving weight\-entangled search spaces\.

We evaluate the performance advantages of our software optimization strategy by benchmarking it against the standard weight\-entanglement technique described in\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]\. This assessment includes determining the average runtime per training sample per iteration, using a batch size of 8, for both strategies across different search space scales, which are determined by the quantity of candidate networks\. To vary the search space size, we adapt theLlama3Spacesearch space through modifications to architectural parameters, including the number of heads, embedding dimension, intermediate size, and head size\. As illustrated in the right side of Figure[3](https://arxiv.org/html/2606.04063#S3.F3), the speedup results demonstrate the clear superiority of our software\-driven enhancement compared to the baseline weight\-entanglement method\. Our approach delivers a 4\.3×\\timesreduction in training time for theLlama3Spacesearch space, at the expense of an additional 3\.2GB in memory usage\. This enhancement stems from the fact that larger search spaces amplify the sequential loop iterations in Equation[5](https://arxiv.org/html/2606.04063#S3.E5), leading to increased runtime burdens in the original setup\. In contrast, our vectorized method performs these operations concurrently, making it more resilient to expansions in search space size\. Moreover, the speedup intensifies further as the number of candidate options rises, though this is accompanied by a corresponding increase in overhead memory\. Therefore, when applying our method, it is crucial to consider the trade\-off between memory costs and training acceleration\.

![Refer to caption](https://arxiv.org/html/2606.04063v1/x5.png)

![Refer to caption](https://arxiv.org/html/2606.04063v1/x6.png)

Figure 3:\(Left\) Software Implementation for mixed\-op weight computation\. \(Right\) Speedup and overhead memory cost of our proposed methods compared to the original weight\-entanglement implementation\[[24](https://arxiv.org/html/2606.04063#bib.bib1)\]with different numbers of candidate networks \(sizes of search space\) running on a A100 GPU with 80GB memory\.

## 4Experiments

### 4\.1Baselines

We evaluate and compare our proposed approach withsubnet\-selection\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]and LoNAS\[[19](https://arxiv.org/html/2606.04063#bib.bib11)\], the state\-of\-the\-art NAS approach for LLM compression\. In addition, for a fair comparision, we also apply feature reordering based on their importance and sandwich training procedure when adapting LoNAS\.

Besides that, directly applying random search in search phase as in the original LoNAS paper\[[19](https://arxiv.org/html/2606.04063#bib.bib11)\]can lead to the dominance of small architectures in the result by the skewness in architectural distribution mentioned in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\. Therefore, in the search phase of LoNAS, we apply the same grid\-based partition strategy as in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]by splitting the search space intokkeven parts, then randomly searching the best architecture from each partition\.

### 4\.2Dataset

Calibration and Fine\-tuning Dataset:To ensure a fair comparison with other state\-of\-the\-art methods, we adopt the same experimental setup as described in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\. Specifically, we use the same calibration dataset from\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]for computing important score\. We utilize the Alpaca dataset\[[26](https://arxiv.org/html/2606.04063#bib.bib33)\]with approximately 52,000 instructions\. The fine\-tuning dataset is divided into two separate train and validation datasets during search phase of our method\. After the supernet converges to a single optimal sub\-architecture, we further fine\-tune this optimal sub\-architecture using the whole dataset\.

Evaluation Dataset:We evaluate and compare our method with other baselines with 7 diverse common\-reasoning tasks, namely BoolQ\[[8](https://arxiv.org/html/2606.04063#bib.bib37)\], PIQA\[[3](https://arxiv.org/html/2606.04063#bib.bib38)\], HellaSwag \(HS\)\[[29](https://arxiv.org/html/2606.04063#bib.bib39)\], WinoGrande \(WG\)\[[21](https://arxiv.org/html/2606.04063#bib.bib48)\]and ARC\[[9](https://arxiv.org/html/2606.04063#bib.bib41)\]\(ARC\-easy \(ARC\_E\) and ARC\-challenge \(ARC\_C\)\), and MMLU\[[13](https://arxiv.org/html/2606.04063#bib.bib42)\]\. Moreover, similar to experimental settings in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\], we perform 5\-shot evaluation for WinoGrande, 10\-shot for HellaSwag, 25\-shot for ARC\-challenge, 5\-shot evaluation for MMLU, and 0\-shot evaluation for remaining tasks\.

### 4\.3Latency Profiling

To evaluate and measure the inference latency of the Llama3 model across different weight\-activation quantization setups, we utilize several state\-of\-the\-art kernels tailored to each configuration\. Specifically, for W4A16 and W8A16 quantization, we adopt the Marlin kernel\[[12](https://arxiv.org/html/2606.04063#bib.bib30)\], which achieves top\-tier performance in these scenarios\. For alternative quantization configurations, we apply the method introduced in ABQ\-LLM\[[30](https://arxiv.org/html/2606.04063#bib.bib55)\], enabling support for arbitrary quantization in linear layers\. This profiling and data gathering process is conducted on an A100 GPU equipped with 80GB of memory\.

### 4\.4Search Space

We use Llama\-3\.1\-8B as the foundation LLM, however, it is remarkable that our method is adaptable to other foundation models\. For fair baseline comparisons, we assess our method against baselines within a search space calledLlama3Space\. For our quantization extension, we employ a distinct search space namedQLlama3Space\. Notably, we adopt a group size of 128 for quantization options inQLlama3Space, which requires us to exclude certain choices for hidden size, number of heads, and head dimensions that are not divisible by 128\. Details of both search spaces are provided in Alpaca[1](https://arxiv.org/html/2606.04063#S4.T1)\. For the sake of simplification, we denoteour\-no\-quantandour\-quantas our proposed method run withLlama3SpaceandQLlama3Spacesearch space, respectively\.

Table 1:Possible configurations of search space for Llama\-3\.1 8BDimensionsLlama3SpaceQLlama3SpaceHidden Dim\.\{2i\|i∈\[5,12\]\}\\\{2^\{i\}\|i\\in\[5,12\]\\\}\{2i\|i∈\[7,12\]\}\\\{2^\{i\}\|i\\in\[7,12\]\\\}Number of Heads\{8,16,32\}\\\{8,16,32\\\}\{16,32\}\\\{16,32\\\}Head Dim\.\{8,16,32,64,128\}\\\{8,16,32,64,128\\\}\{32,64,128\}\\\{32,64,128\\\}Intermediate Dim\.\{4096∗i\|i∈\[1\.0,2\.0,3\.0,3\.5\]\}\\\{4096\*i\|i\\in\[1\.0,2\.0,3\.0,3\.5\]\\\}\{4096∗i\|i∈\[1\.0,2\.0,3\.0,3\.5\]\}\\\{4096\*i\|i\\in\[1\.0,2\.0,3\.0,3\.5\]\\\}Number of Blocks\{1,…,32\}\\\{1,\.\.\.,32\\\}\{1,…,32\}\\\{1,\.\.\.,32\\\}Weight Bitwidth\{2,4,8\}\\\{2,4,8\\\}Activation Bitwidth\{2,4,8,16\}\\\{2,4,8,16\\\}\# of Candidates8×3×5×4×32=153608\\times 3\\times 5\\times 4\\times 32=153606×2×3×4×332×432×326\\times 2\\times 3\\times 4\\times 3^\{32\}\\times 4^\{32\}\\times 32
### 4\.5Experimental Details

For our method,subnet\-selection, and LoNAS, LoRA is set with a rank of 32, alpha of 16, and a dropout rate of 0\.05 during fine\-tuning\. We also pre\-process the pre\-trained LLM by reordering neurons and attention heads by their important scores as in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\. To compute the block importance, we follows the method as mentioned in\[[20](https://arxiv.org/html/2606.04063#bib.bib10)\]\. Unlike\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\], which applies LoRA solely to embedding and attention layers, we follow LoNAS\[[19](https://arxiv.org/html/2606.04063#bib.bib11)\]by applying LoRA to both attention and MLP layers\. Training settings for LoRA adapters are similar as settings in\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\.

In terms of the quantization approach, we utilize group\-wise quantization with a group size of 128\. Furthermore, we initialize the quantizers in our proposed supernet using OmniQuant\[[22](https://arxiv.org/html/2606.04063#bib.bib20)\]\. For quantizing models suggested by the baselines \(subnet\-selectionand LoNAS\), we similarly employ the OmniQuant method\. Notably, other post\-training quantization techniques can also be applied to initialize our supernet\.

### 4\.6Comparison to baselines

#### 4\.6\.1Without Quantization

Table[2](https://arxiv.org/html/2606.04063#S4.T2)presents a comparison of our proposed method against baseline approaches across common reasoning tasks, evaluated over four distinct parameter ranges\. Notably, our methodour\-no\-quantsurpasses state\-of\-the\-art baselines in both average accuracy and inference latency\. To illustrate the advantages of our approach, Figure[5](https://arxiv.org/html/2606.04063#S4.F5)shows Pareto fronts comparing our method to baselines, alongside the distribution of sub\-architectures across various model size ranges\. The distribution reveals a skew, with smaller models being far more prevalent than larger ones\. This observation align with observations in previous work\[[25](https://arxiv.org/html/2606.04063#bib.bib3)\]\. Specifically, in the 6–8 billion parameter range, only a few architectures are present\. When compared to two baselines, our method demonstrates superior performance across the 2–6 billion parameter range\. For larger models, the architectures suggested by our method align closely with those from thesubnet\-selectionmethod, primarily due to the limited number of available architectures in this range, allowing both methods to identify optimal configurations\. However, in the 2–6 billion parameter range, our method consistently outperforms all baselines, indicating its ability to identify superior compression configurations\. Besides that,subnet\-selectionalso outperforms LoNAS in this range of model size\. Therefore, for the experiment with quantization \(Section[4\.6\.2](https://arxiv.org/html/2606.04063#S4.SS6.SSS2)\), we apply quantization tosubnet\-selectionand use it as the baseline for comparison\. For small models \(<<2 billion parameters\), the performance gap between our method and the baselines is minimal, as highly compressed architectures lose significant information, resulting in reduced accuracy across all methods\.

Table 2:A comparison between baselines with our proposed methods for 7 common reasoning tasks, with 4 different ranges of model size, namely 2\-3 billion \(2\-3B\), 3\-4 billion \(3\-4B\), 4\-5 billions \(4\-5B\), and 5\-6 billions \(5\-6B\) parameters\. Remarkably, LoNAS\* are improved as mentioned in Section[4\.1](https://arxiv.org/html/2606.04063#S4.SS1)\.MethodParamsLatency\(ms\)ARC\_EARC\_CBoolQWGHSMMLUPIQAAcc\.\(%\)2B\-3Bsubnet\-selection2\.54B70\.9525\.6626\.6249\.8838\.5026\.3825\.7253\.7635\.22LoNAS\*2\.67B69\.8424\.4525\.7748\.8661\.8726\.0522\.9552\.2337\.45ours\-no\-quant2\.27B43\.1434\.1823\.9860\.8350\.6727\.4325\.3855\.2839\.683B\-4Bsubnet\-selection3\.28B81\.3128\.2025\.7748\.7837\.9527\.1524\.4954\.0335\.19LoNAS3\.28B81\.3124\.4927\.4749\.7238\.6926\.5824\.6551\.4134\.72ours\-no\-quant3\.37B40\.9038\.1323\.2160\.5552\.7229\.2125\.8656\.5340\.894B\-5Bsubnet\-selection4\.32B57\.4531\.9926\.8837\.8649\.4931\.5925\.0757\.7337\.22LoNAS\*4\.00B73\.7225\.2526\.1146\.9656\.6726\.8223\.7452\.8436\.91ours\-no\-quant4\.14B53\.9341\.3329\.7962\.6956\.2741\.9826\.9062\.6845\.935B\-6Bsubnet\-selection6\.06B87\.3861\.4137\.1270\.4365\.5959\.9630\.9870\.4656\.56LoNAS\*6\.84B99\.0042\.5931\.8357\.7763\.1251\.0024\.2664\.4247\.86ours\-no\-quant6\.06B87\.3859\.1842\.3272\.5465\.5160\.7749\.0770\.8260\.03![Refer to caption](https://arxiv.org/html/2606.04063v1/x7.png)Figure 4:Pareto fronts of our proposed method compared tosubnet\-selectionand LoNAS\. The second y\-axis is the architectural distribution across different ranges of model size forLlama3Spacesearch space\.
![Refer to caption](https://arxiv.org/html/2606.04063v1/x8.png)Figure 5:Pareto fronts ofour\-quantmethod compared to models compressed bysubnet\-selectionandour\-no\-quant, with 4 different quantization configurations, namely W4A4, W4A16, W8A8, and W8A16

#### 4\.6\.2With Quantization

In this experiment, we investigate theQLlama3Spacesearch space, which supports a wide range of quantization configurations\. For a fair comparison, we apply post\-training quantization to the architectures identified by thesubnet\-selectionandour\-no\-quantmethods using four quantization settings: W4A4, W8A8, W8A16, and W4A16\. These quantized models are then benchmarked against those produced by our proposed joint optimization approach, denotedour\-quant\. We exclude the LoNAS baseline from this analysis, assubnet\-selectionconsistently outperformed it in earlier experiments\. Figure[5](https://arxiv.org/html/2606.04063#S4.F5)presents the Pareto fronts of compressed architectures discovered by our\-quant, compared against those fromsubnet\-selectionandour\-no\-quantunder the various quantization settings\. The results clearly show that jointly optimizing both architectural configurations and quantization precisions enables our method to achieve superior efficiency–accuracy trade\-offs compared to sequential pipelines \(architecture search followed by quantization\)\. For example, at the same average accuracy of 40% across reasoning tasks, models compressed with our method achieve up to1\.4×1\.4\\timesfaster inference than those from competing baselines\. Alternatively, at a fixed inference latency of 30 ms, our models reach approximately 41% average accuracy on reasoning tasks, outperforming other baselines by roughly 6%\.

## 5Conclusions and Future Work

We present an effective compression approach for large language models \(LLMs\) using differential neural architecture search\. Notably, our non\-quantized method,ours\-no\-quant, outperforms state\-of\-the\-art approaches for non\-quantized configurations, while our quantized version,ours\-quant, which jointly optimizes both compression and quantization, surpasses sequential compression followed by quantization\. This suggests that jointly optimizing structural architecture and quantization yields more efficient compression configurations compared to treating them as separate problems\.

## References

- \[1\]S\. Ashkbooset al\.\(2024\)SliceGPT: compress large language models by deleting rows and columns\.InICLR,Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1)\.
- \[2\]Y\. Bengio, N\. Léonard, and A\. Courville\(2013\)Estimating or propagating gradients through stochastic neurons for conditional computation\.arXiv preprint arXiv:1308\.3432\.Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p2.7)\.
- \[3\]Y\. Bisk, R\. Zellers,et al\.\(2020\)PIQA: reasoning about physical commonsense in natural language\.InAAAI,Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[4\]Y\. Bondarenko, R\. Del Chiaro, and M\. Nagel\(2024\)Low rank quantization\-aware training for llms\.InICML,Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p1.4),[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p2.7)\.
- \[5\]R\. Cai, S\. Muralidharan,et al\.\(2024\)FLEXTRON: many\-in\-one flexible large language model\.InICML,pp\. 5298–5311\.Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1),[§1](https://arxiv.org/html/2606.04063#S1.p4.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p1.1)\.
- \[6\]R\. Caiet al\.\(2025\)LLaMaFlex: many\-in\-one llms via generalized pruning and weight sharing\.InICLR,Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p1.1)\.
- \[7\]X\. Chenet al\.\(2021\)DrNAS: dirichlet neural architecture search\.InICLR,Cited by:[§3\.2\.1](https://arxiv.org/html/2606.04063#S3.SS2.SSS1.p2.18)\.
- \[8\]C\. Clarket al\.\(2019\)BoolQ: exploring the surprising difficulty of natural yes/no questions\.InNAACL,Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[9\]P\. Clarket al\.\(2018\)Think you have solved question answering? try arc, the ai2 reasoning challenge\.arXiv:1803\.05457v1\.Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[10\]E\. Frantaret al\.\(2022\)Gptq: accurate post\-training quantization for generative pre\-trained transformers\.arXiv preprint arXiv:2210\.17323\.Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1)\.
- \[11\]E\. Frantaret al\.\(2023\)OPTQ: accurate post\-training quantization for generative pre\-trained transformers\.InICLR,Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p4.1)\.
- \[12\]E\. Frantaret al\.\(2025\)Marlin: mixed\-precision auto\-regressive parallel inference on large language models\.InPPoPP,pp\. 239–251\.Cited by:[§4\.3](https://arxiv.org/html/2606.04063#S4.SS3.p1.1)\.
- \[13\]D\. Hendryckset al\.\(2021\)Measuring massive multitask language understanding\.ICLR\.Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[14\]E\. J\. Huet al\.\(2022\)Lora: low\-rank adaptation of large language models\.\.ICLR\.Cited by:[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p2.1)\.
- \[15\]E\. Jang, S\. Gu, and B\. Poole\(2016\)Categorical reparameterization with gumbel\-softmax\.arXiv preprint arXiv:1611\.01144\.Cited by:[§3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2.p3.9)\.
- \[16\]M\. Javaheripiet al\.\(2023\)Phi\-2: the surprising power of small language models\.Microsoft Research Blog1\(3\),pp\. 3\.Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p2.1)\.
- \[17\]H\. Jeon, Y\. Kim, and J\. Kim\(2024\)L4Q: parameter efficient quantization\-aware fine\-tuning on large language models\.arXiv preprint arXiv:2402\.04902\.Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p1.4),[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p2.7)\.
- \[18\]L\. Liuet al\.\(2023\)Bridging discrete and backpropagation: straight\-through and beyond\.NeurIPS\.Cited by:[§3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2.p3.9)\.
- \[19\]J\. P\. Munozet al\.\(2024\)Lonas: elastic low\-rank adapters for efficient large language models\.InLREC\-COLING 2024,Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1),[§1](https://arxiv.org/html/2606.04063#S1.p4.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p2.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p3.2),[§4\.1](https://arxiv.org/html/2606.04063#S4.SS1.p1.1),[§4\.1](https://arxiv.org/html/2606.04063#S4.SS1.p2.1),[§4\.5](https://arxiv.org/html/2606.04063#S4.SS5.p1.1)\.
- \[20\]S\. Muralidharanet al\.\(2024\)Compact language models via pruning and knowledge distillation\.arXiv preprint arXiv:2407\.14679\.Cited by:[§3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2.p2.1),[§4\.5](https://arxiv.org/html/2606.04063#S4.SS5.p1.1)\.
- \[21\]K\. Sakaguchiet al\.\(2020\)Winogrande: an adversarial winograd schema challenge at scale\.InAAAI,Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[22\]W\. Shaoet al\.\(2024\)OmniQuant: omnidirectionally calibrated quantization for large language models\.InICLR,Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p4.1),[§4\.5](https://arxiv.org/html/2606.04063#S4.SS5.p2.1)\.
- \[23\]S\. T\. Sreenivaset al\.\(2024\)Llm pruning and distillation in practice: the minitron approach\.arXiv preprint arXiv:2408\.11796\.Cited by:[§3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2.p2.1)\.
- \[24\]R\. S\. Sukthankeret al\.\(2024\)Weight\-entanglement meets gradient\-based neural architecture search\.InAutoML,Cited by:[3rd item](https://arxiv.org/html/2606.04063#S1.I1.i3.p1.1),[1st item](https://arxiv.org/html/2606.04063#S2.I1.i1.p1.1),[§2\.1](https://arxiv.org/html/2606.04063#S2.SS1.p1.1),[Figure 3](https://arxiv.org/html/2606.04063#S3.F3),[§3\.1\.1](https://arxiv.org/html/2606.04063#S3.SS1.SSS1.p1.8),[§3\.2\.1](https://arxiv.org/html/2606.04063#S3.SS2.SSS1.p2.18),[§3\.2\.1](https://arxiv.org/html/2606.04063#S3.SS2.SSS1.p4.3),[§3\.2\.2](https://arxiv.org/html/2606.04063#S3.SS2.SSS2.p1.6),[§3\.3](https://arxiv.org/html/2606.04063#S3.SS3.p1.2),[§3\.3](https://arxiv.org/html/2606.04063#S3.SS3.p3.4),[§3\.3](https://arxiv.org/html/2606.04063#S3.SS3.p4.1)\.
- \[25\]R\. S\. Sukthanker, B\. Staffler, F\. Hutter, and A\. Klein\(2024\)Large language model compression with neural architecture search\.InNeurIPS,Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p3.1),[§1](https://arxiv.org/html/2606.04063#S1.p4.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p1.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p2.1),[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p3.2),[§4\.1](https://arxiv.org/html/2606.04063#S4.SS1.p1.1),[§4\.1](https://arxiv.org/html/2606.04063#S4.SS1.p2.1),[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p1.1),[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1),[§4\.5](https://arxiv.org/html/2606.04063#S4.SS5.p1.1),[§4\.6\.1](https://arxiv.org/html/2606.04063#S4.SS6.SSS1.p1.1)\.
- \[26\]R\. Taoriet al\.\(2023\)Stanford alpaca: an instruction\-following llama model\.GitHub\.Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p1.1)\.
- \[27\]H\. Xu, L\. Xiang, H\. Ye, D\. Yao, P\. Chu, and B\. Li\(2024\)Permutation equivariance of transformers and its applications\.InCVPR,Cited by:[§2\.2](https://arxiv.org/html/2606.04063#S2.SS2.p1.1)\.
- \[28\]Y\. Xuet al\.\(2024\)QA\-lora: quantization\-aware low\-rank adaptation of large language models\.InICLR,Cited by:[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p2.4),[§3\.2\.3](https://arxiv.org/html/2606.04063#S3.SS2.SSS3.p3.3)\.
- \[29\]R\. Zellerset al\.\(2019\)HellaSwag: can a machine really finish your sentence?\.InACL,Cited by:[§4\.2](https://arxiv.org/html/2606.04063#S4.SS2.p2.1)\.
- \[30\]C\. Zenget al\.\(2025\)Abq\-llm: arbitrary\-bit quantized inference acceleration for large language models\.InAAAI,Cited by:[§4\.3](https://arxiv.org/html/2606.04063#S4.SS3.p1.1)\.
- \[31\]P\. Zhanget al\.\(2024\)TinyLlama: an open\-source small language model\.arXiv preprint arXiv:2401\.02385\.Cited by:[§1](https://arxiv.org/html/2606.04063#S1.p2.1)\.