TLDR: A new research paper introduces a novel structured pruning method for Large Language Models (LLMs) that is inspired by functional brain networks in the human brain. By treating an LLM as a ‘digital brain’ and applying Independent Component Analysis (ICA) layer-by-layer, the method identifies and preserves crucial ‘functional networks’ of neurons. This approach aims to maintain the LLM’s performance while significantly reducing its size and computational requirements. Experimental results show that this brain-inspired method outperforms several existing pruning techniques on various benchmark datasets and models like LLaMA2-7B-chat-hf and Vicuna-7B-v1.5.
Large Language Models (LLMs) are incredibly powerful, but their massive size, often containing billions of parameters, leads to significant computational costs for both training and deployment. This makes them challenging to use efficiently in real-world applications. To address this, researchers often employ compression techniques, with ‘pruning’ being a prominent method. Pruning involves removing redundant parts of the model while trying to maintain its performance.
There are two main types of pruning: unstructured and structured. Unstructured pruning targets individual weights, essentially setting them to zero, but this doesn’t always translate to faster inference or reduced memory use without specialized hardware. Structured pruning, on the other hand, removes entire architectural elements like attention heads or neuron groups, which directly reduces the model’s size and improves efficiency on standard hardware.
Current structured pruning methods typically assess the importance of individual structural units and remove those deemed less important. However, many of these methods overlook the crucial interactions and collaborations among artificial neurons. This oversight can disrupt the LLM’s overall functional architecture, leading to a decline in performance after pruning.
A Brain-Inspired Approach to LLM Pruning
Inspired by the striking similarities between artificial neural networks and functional neural networks in the human brain, a new study proposes a novel structured pruning method. This approach aims to prune LLMs by identifying and preserving ‘functional networks’ within them. The core idea is to treat an LLM as a ‘digital brain’ and decompose it into these functional networks, much like how functional brain networks are identified in neuroimaging data.
The researchers draw parallels between the signals of neurons in LLMs and BOLD (Blood-Oxygen Level Dependent) signals measured in fMRI (functional Magnetic Resonance Imaging) data, which are used to infer neuronal activity in the human brain. Just as fMRI signals reflect the activity of biological neurons, LLM neuron outputs represent internal state changes as the model processes information. Both exhibit temporal dynamics.
To identify these functional networks, the study adapts Independent Component Analysis (ICA), a data-driven method widely used in neuroscience to extract functional brain networks from high-dimensional fMRI data. Specifically, they use canonical ICA (CanICA), a spatial ICA method. This technique assumes that observed signals are linear combinations of statistically independent source signals, where the spatial patterns of these components depict functional networks.
Methodology and Implementation
Given the immense number of artificial neurons in LLMs, applying traditional ICA globally across the entire model would be computationally prohibitive. To overcome this, the researchers propose a ‘layer-by-layer ICA-based local pruning’ method. This means CanICA is applied independently to the neuron signals of each layer of the LLM, particularly focusing on the Multi-Layer Perceptron (MLP) layers, which are crucial for introducing non-linearity and expressive power.
After applying CanICA to each layer, functional networks are identified. These networks are represented by a source signal matrix, where each row corresponds to a network and each column to an associated neuron. By applying a threshold to this matrix, ‘mask matrices’ are obtained, which pinpoint the critical neurons that must be preserved during pruning.
To ensure that all important neurons are retained, especially since neuron associations can vary with different calibration data, the researchers divided calibration data samples into multiple groups. CanICA was run multiple times, and the mask matrices from different groups were aggregated using an OR operation to create a final ‘global mask matrix’ for pruning.
Experimental Validation and Results
The proposed method was tested on popular LLMs, including LLaMA2-7B-chat-hf, Vicuna-7B-v1.5, and ChatGLM3-6B-base. Performance was evaluated using the lm-evaluation-harness framework, assessing perplexity on the Wikitext-2 dataset and zero-shot classification on benchmarks like PIQA, HellaSwag, Winogrande, OpenBookQA, ARC-Easy, and ARC-Challenge.
The results demonstrate that this brain-inspired pruning approach successfully identifies and locates functional networks and key neurons, leading to efficient model pruning. At a 20% pruning rate, the method (CanICA) often outperformed other state-of-the-art structured pruning methods like LLM-Pruner, SliceGPT, FLAP, and Shortened LLaMA across various tasks. For instance, on Vicuna-7B-v1.5, it achieved the best performance in 5 out of 7 tasks compared to several competitors. Similar strong performance was observed on LLaMA2-7B-chat-hf and ChatGLM3-6B-base, even with its slightly different architecture.
The study also explored the impact of hyperparameters, such as the number of independent components in ICA, finding optimal performance around 128 or 256 components. Furthermore, it was observed that the method could achieve better pruning results with fewer calibration samples compared to some existing methods.
Also Read:
- Streamlining AI: How PCA Makes Large Language Models Leaner and Smarter
- Streamlining LLM Reasoning: Compressing Thought Chains with Step Entropy
Future Directions
While promising, the current method has limitations. It performs pruning layer-by-layer, a ‘local’ strategy. Future work could explore a ‘global’ strategy by identifying functional networks across all layers, which might better maintain functional coherence. Additionally, the ICA algorithm used is linear; exploring non-linear deep neural networks like auto-encoders for identifying functional networks could capture more complex dependencies among neurons. This research highlights a fascinating intersection where insights from neuroscience can significantly advance artificial intelligence, fostering mutual inspiration and co-development between the two fields.
For more details, you can read the full research paper here.
