How can representation dimension dominate structurally pruned LLMs?

TL;DR

This work examines how structured pruning of LLMs, via SliceGPT, affects model functionality with a focus on representation dimension $d$. It shows that $d$ dominates linear transformations and the forward flow in transformer blocks, enabling explicit analytical relations between sparsity $s$ and performance metrics: $rac{\ln \texttt{PPL}_0(D)}{\ln \texttt{PPL}(D)} = 1-s$ for perplexity and $\ln \frac{\texttt{acc}(D)}{\texttt{acc}_0(D)} \propto 1-s$ for accuracy. These relations are validated empirically on LLaMa-3-8B-Instruct and Phi-3-mini-4k-Instruct across WikiText2 and several multiple-choice tasks, highlighting a near-inverse relationship between perplexity and accuracy under dimension reduction. The results connect mechanistic interpretability with pruning theory, offering precise, dimension-driven guidance for pruning hyperparameters and suggesting broader applicability beyond SliceGPT. The study thus advances understanding of how subnetworks governed by representation dimension influence both predictions and overall performance in pruned LLMs.

Abstract

Pruning assumes a subnetwork exists in the original deep neural network, which can achieve comparative model performance with less computation than the original. However, it is unclear how the model performance varies with the different subnetwork extractions. In this paper, we choose the representation dimension (or embedding dimension, model dimension, the dimension of the residual stream in the relevant literature) as the entry point to this issue. We investigate the linear transformations in the LLM transformer blocks and consider a specific structured pruning approach, SliceGPT, to extract the subnetworks of different representation dimensions. We mechanistically analyse the activation flow during the model forward passes, and find the representation dimension dominates the linear transformations, model predictions, and, finally, the model performance. Explicit analytical relations are given to calculate the pruned model performance (perplexity and accuracy) without actual evaluation, and are empirically validated with Llama-3-8B-Instruct and Phi-3-mini-4k-Instruct.

Figures (3)

• Figure 1: The linear transformations before (solid arrows) and after (dashed arrows) pruning in the model. The original model maps the input embeddings $E_{\texttt{in}}$ from the input space (illustrated with the shading grey ellipse) to the output embedding $E_{\texttt{out}}$ in the output space, through a series of linear transformations (i.e. those defined by the weight matrices ${\bf W}_1, \ldots,{\bf W}_N$ ). After pruning, ${\bf W}_i$ is converted to ${\bf W}^{'}_i$ ($i \in \{1,2,\ldots, N\}$), and the original output $E_{\texttt{out}}$ is shifted to $E^{'}_{\texttt{out}}$. The final predictions ($y$ and $\hat{y}$) are generated (normally non-linearly) from $E_{\texttt{out}}$ and $E^{'}_{\texttt{out}}$. Denote the prediction domain of the pruned model as $\hat{Y}$ and that of the unpruned model $Y$, the mapping $g:Y\rightarrow\hat{Y}$ shifts the model performance.
• Figure 2: Activation flow in the LLM transformer. Linear transformations are defined by weight matrices like ${\bf W}_i$, and non-linear transformations are represented with teletype font. The detailed shapes of the weight matrices are clarified in \ref{['tab:shape']}.
• Figure 3: Fitted evaluation results. The fitting coefficients and errors are in \ref{['sec:error']}.

Theorems & Definitions (3)

• Proposition 1
• Proposition 2
• Proposition 3