Sparsity Limit to Prune Large Language Models for on-Device AI Assistants: Llama-2 as an Example

1. Introduction

Large language models (LLMs) have shown impressive performance on various tasks and signficant impact on human society [1,2,3]. However, they require substantial computational resources and energy consumption. To make LLMs more accessible to individual consumers, a recent trend is to train moderate-sized models such as Llama [4,5], ChatGLM [6], and Falcon [7]. Despite being relatively smaller, these models still contain billions of parameters, posing challenges for the development of personalized and privacy-preserving AI assistants that can operate locally on user devices like laptops or mobile phones. For instance, the smallest model in the Llama-2 series has 7 billion parameters and requires over 20 GB memory, making it impractical to run on a mobile phone with limited RAM.

Recent advances in neuroscience and deep learning offer promising solutions by leveraging sparse connectivity. Biological neural networks and neuronal activities in the brain are generally sparse [8,9,10,11,12,13], in contrast to the dense connections and activations widely adopted in LLMs, including the QKV matrices and the attention matrices of the Transformer architecture [14]. Notably, a recent computational model found that in mammalian olfactory system, sparse inter-hemispheric projections can align two olfactory cortical neural representations of environmental odors when the number of olfactory cortical neurons is sufficiently high [15]. Similarly, theoretical studies revealed that under suitable conditions, sparse Transformers only need $\mathcal{O}\left(n\right)$ connections per attention layer to approximate any sequence-to-sequence function, rather than $\mathcal{O}\left(n^2\right)$, where n is the number of tokens [16]. In practice, the deep learning community have developed various pruning and sparsification skills to obtain sparse models [17].

Consequently, pruning and sparsifying LLMs has gained popularity, leading to the development of diverse methods [18,19,20]. For instance, [18] utilized gradient information to perform structural pruning of LLMs. Motivated by the empirical observations of emergent hidden-state features with large magnitudes in LLMs [19,21,22] developed a simple but effective pruning approach on a per-output basis: the weights were firstly multiplied by the corresponding inputs, and those with smallest output magnitudes were pruned. Additionally, [20] proposed a two-step structured pruning procedure involving targeted structured pruning based on precedent models and dynamic batch loading that adjusts training data proportions based on respective losses.

Our research focuses on developing local LLMs as AI assistants on user devices, prioritizing lightweight models over perfect accuracy. Therefore, we adopted the simple yet effective method from [19] to prune the Meta-developed Llama-2 7B model [4], explored the sparsity limit of pruning, and evaluated the performance at the sparsity limit, to facilitate the development of personalized AI assistants on user devices. Our contributions are two-fold:

• We demonstrated that $60\%$ of the weights in the Llama-2 7B model could be pruned without notable decrease in language modeling capability, as evidenced by the wikitext perplexity metrics.
• We examined the user experience of daily dialogue and query handling at this sparsity limit to assess the feasibility of on-device AI assistants. The pruned model generated fluent and helpful answers but with factual inaccuracies, raising intriguing theoretical questions about the nature of pruning or sparsification. We hypothesized that these inaccuracies might result from forgetting or hallucination, proposed a simple protocol to distinguish between the two mechanisms, and discussed future directions to improve the pruned models.

In summary, our work paves the way for the development of lightweight and local LLMs suitable for user devices, balancing performance with resource efficiency.

2. Methods

We adopted the pruning method on a per-output basis detailed in [19]. As shown in Figure 1, unlike the magnitude-based and layer-wise pruning, this method examines the contribution of each weight ${\mathbf{W}}_{ij}$ to a specific output i by multiplying it with corresponding inputs, i.e., the contribution of weight ${\mathbf{W}}_{ij}$ to output i is estimated as

$$\begin{array}{c}\hfill {\mathbf{O}}_{ij}=|{\mathbf{W}}_{ij}|·||{\mathbf{X}}_j{\left|\right|}_2,\forall j,\end{array}$$

(1)

where $\left|\right|{\mathbf{X}}_j{\left|\right|}_2$ is the ${\ell }_2$ norm of the j-th input dimension vector concatenated across different inputs. Denote $s\%$ as the desired sparsity ratio, then for output i, we will prune $s\%$ of the weights ${\mathbf{W}}_{ij}$ among j based on the sorted ${\mathbf{O}}_{ij}$ values. In this paper, we varied $s\%$ from 0 to $90\%$, to figure out the sparsity limit of pruned LLMs.

Secondly, among the different pruning strategies explored in [19], we adopted the unstructured pruning approach rather than the structured N:M sparsity pattern, where N out of M consecutive weights connected to an output were pruned [23]. Although the latter can utilize the NVIDIA GPU sparse tensor cores to accelerate matrix calculations, [19] demonstrated that the unstructured pruning generally performed better and was more robust.

Finally, following [19], we used the same calibration data (the C4 training set [24]) to estimate input distribution, calculate the contribution to output ${\mathbf{O}}_{ij}$, and facilitate pruning.

3. Results

3.1. The sparsity limit to prune LLMs, evaluated by wikitext perplexity

We varied the sparsity ratio from 0 to $90\%$, to probe the sparsity limit of pruning the Llama-2 7B model. Following [18,19,22], wikitext perplexity was used to measure the language modeling capability of pruned models. As shown in Figure 2, the perplexity increases slowly for sparsity ratios from 0 to $60\%$, remaining the same order ($\lesssim 10$). However, from 60% to 90%, the perplexity grows exponentially. This dramatic change at a sparsity ratio of $60\%$ suggests that, with current pruning method, $60\%$ is roughly the sparsity limit. Furthermore, we confirmed that despite an increase in wikitext perplexity, the $60\%$-sparsity model with fewer than half of the connections has a perplexity less than twice that of the full model (Figure A1). Altogether, these results showed that Llama-2 7B model could be pruned up to $60\%$; in other words, only $40\%$ connections are required for relatively intact function.

3.2. Pruned sparse LLM as AI assistants: user experience examination

Next, we investigated the user experience of daily dialogue with the pruned Llama-2 model at the 60% sparsity limit, in order to evaluate the potential of pruned sparse LLMs as future on-device AI assistants. Specifically, to simulate daily dialogue of mobile phone users, we posed several queries to the $60\%$-sparsity pruned models. Two example queries are:

• “Tell me about Boston”;
• “Describe the Python programming language, in terms of its syntax, history, user experience, and popularity”.

As shown in Figure A2, the $60\%$-sparsity pruned model can generate fluent, meaningful, and helpful answers that are indistinguishable from those of the full Llama-2 7B model. For fairness, ChatGPT-4o was used to rate the answers of both models, and the $60\%$-sparsity pruned model obtained good scores that are slightly lower than the full model (6 vs 7, and 7 vs 8 out of 10, Figure A3–Figure A6).

However, we noticed that the $60\%$-sparsity model sometimes produces factually incorrect statements (the red texts in Figure A2). For example, when talking about Python language, it wrongly alleged that “(Python) was created in 1990 by Guido van Rosenberg”, while the full model correctly pointed out that “(Python) was created by Guido van Rossum and first released in 1991”. Similarly, for the Boston-related query, the answer from the $60\%$-sparsity model also contained many factual inaccuracies, as labeled by the red texts in Figure A2 and identified by ChatGPT-4o (Figure A3).

This observation raised interesting theoretical questions about the nature of pruning (sparsification) and the mechanism of generating factual inaccuracies. With some weights zeroed during pruning, are these factual mistakes due to forgetting (e.g., the pruned model forgot thus could not find the name of Guido van Rossum as the Python creator)? Or, the pruned model still remembered this person’s name and his role as the Python inventor, but did not find and generated some hallucinations due to the missing weights? Similar phenomenon named “catastrophic forgetting” has been observed and studied in the context of transfer learning [25,26] but less explored in pruning. In Discussion, we proposed a simple protocol to distinguish between the two mechanisms (forgetting vs hallucination) in future studies, as well as possible solutions to remedy this inaccuracy issue.

Last but not least, as mentioned in [18], the pruned model sometimes generated meaningless sentences, repetitive tokens, and even mixed multilingual characters. Also, we confirmed the importance of prompt engineering, as in many other unpruned LLMs (see the review [27]). For instance, if we rephrase the Python-related question as “Can you explain briefly to me what is the Python programming language?”, the pruned model occasionally repeat this question and end the dialogue. In Discussion, we provided some potential methods to mitigate these problems.

4. Discussion

Summary. We adopted a simple yet effective pruning approach to sparsify the Llama-2 7B model, and found that up to $60\%$ weights could be pruned without significant decrease in the model performance. At this sparsity limit, the pruned model was able to generate fluent and helpful answers to daily-life queries, demonstrating its potential for local and lightweight AI assistants on user devices. Meanwhile, some factual inaccuracies were identified in the answers of the pruned model, whose mechanism and remedies will be studied in the future.

Future directions. To distinguish forgetting vs hallucination and account for the factual inaccuracies generated by the pruned model, notice that the key difference is whether it still remembers the true answers. Therefore, we propose to test the pruned model with appropriate prompts after identifying the factual mistakes. For example, we may ask the $60\%$-sparsity model who is the Python creator or who is Guido van Rossum. If it correctly assigned Guido van Rossum as the Python creator instead of “Guido van Rosenberg”, we may conclude that this pruned model still remembers Guido van Rossum as the Python creator but simply did not find this name and generated the above hallucination. In reality, it is possible that both mechanisms coexist after pruning. Nevertheless, this attempt may serve as a small but useful step towards understanding the nature of pruning.

To eliminate these factual errors and improve the pruned models, we will adopt the fine-tuning methods such as that described in [28]. Another possible route is to perform layer-wise pruning with different sparsity ratios: currently, all layers are pruned evenly with the same sparsity ratio; however, it has been realized that different layers in deep neural networks are not equal, and generally, deeper layers can be pruned more [29,30,31,32,33]. Therefore, in the future, we will firstly calculate the significance of various layers as stated in [32,33], then assign the corresponding sparsity ratio to each layer for more efficient pruning.

In order to develop on-device AI assistants, we will also try more updated moderate-sized LLMs such as the newly-released Llama-3 8B model [5]. Moreover, we will employ various techniques to further reduce the model size while boosting the model performance [34]. For instance, LLM quantization has been an active research area, which stores the parameters with fewer digits to reduce the model size and save resources [22,35,36]. Another example is low-rank factorization which decomposes the weight matrices into smaller matrices [37]. Altogether, this work represents a small yet promising step towards on-device AI assistants, particularly when combined with various fine-tuning, layer-specific pruning, and LLM compression techniques in the long run.