Efficient Model Pruning for Large Language Models: Wanda's Simplified Approach

Large Language Models (LLMs) have transformed natural language processing, but their immense size poses computational challenges. Wanda, a novel pruning method, offers an efficient solution.

Introduction:

Large Language Models (LLMs) have revolutionized the field of natural language processing, enabling remarkable achievements in various language-related tasks. However, their staggering size, with billions of parameters, demands enormous computational resources for training and deployment. Reducing this computational cost without sacrificing performance is a pressing challenge.

Motivation:

Wanda, or "Pruning by Weights and Activations," is inspired by a recent revelation in the world of LLMs: the emergence of large-magnitude features. These features, previously unnoticed, play a critical role in model behavior. Traditional pruning methods, like magnitude pruning, fail to account for this phenomenon. Wanda, however, leverages these features to guide its pruning strategy.

The Math Behind Wanda:

At the heart of Wanda lies a straightforward yet powerful pruning metric that combines weight magnitudes and input activation norms. This metric, designed explicitly for LLMs, quantifies the importance of each weight in the network.

The pruning score for a weight Wij that connects input j to output i is defined as follows:

Sij = |Wij| · ∥Xj∥?

Here,

• |Wij| represents the absolute value of the weight magnitude.
• ∥Xj∥? calculates the ?? norm of input feature activations, aggregated across all tokens.

Code Implementation:

Wanda pruning can be implemented efficiently in PyTorch with the following code snippet:

import torch

# W: weight matrix (C_out, C_in);
# X: input matrix (N * L, C_in);
# s: desired sparsity level, between 0 and 1;
def prune(W, X, s):
    # Calculate the importance metric for each weight
    metric = W.abs() * X.norm(p=2, dim=0)
    
    # Sort the weights based on their importance metric
    _, sorted_idx = torch.sort(metric, dim=1)
    
    # Determine how many weights to prune based on the desired sparsity level
    pruned_idx = sorted_idx[:, :int(W.size(1) * s)]
    
    # Scatter the pruned weights to zero
    W.scatter_(dim=1, index=pruned_idx, src=0)
    
    # Return the pruned weight matrix
    return W        

Why Wanda?

Wanda pruning, as described in the paper "A Simple and Effective Pruning Approach for Large Language Models," offers several key features that make it a valuable technique for reducing the computational cost of Large Language Models (LLMs):

• Simplicity: Wanda's pruning approach is conceptually simple. It calculates weight importance based on a straightforward metric that combines weight magnitudes and input activation norms.
• Efficiency: Wanda can be integrated seamlessly into a single forward pass of an LLM during inference. This one-pass approach is efficient and does not require iterative weight updates or complex weight reconstruction procedures.
• No Retraining Required: Unlike many other pruning methods, Wanda does not require retraining of the pruned LLM. The sparse LLM obtained after pruning is immediately usable without further fine-tuning.
• Per-Output Pruning: Wanda prunes weights on a per-output basis, considering the importance of each weight for individual model outputs. This per-output pruning granularity is crucial for effectively pruning LLMs and prevents imbalanced pruning across different output features.
• Balanced Magnitude and Activation Consideration: Wanda's pruning metric balances weight magnitudes and input activation norms. This ensures that weights connected to large-magnitude features are preserved, even if their own magnitudes are small. It takes into account the specific characteristics of LLMs, where large-magnitude features can emerge.
• Robustness: The pruning metric used in Wanda is empirically shown to be robust and can be estimated with a modest number of calibration samples, making it practical even without access to the original training data.
• Potential for Structured Pruning: While initially designed for unstructured sparsity, Wanda can be extended to structured N: M sparsity, where at most N out of every M contiguous weight is allowed to be non-zero. This flexibility allows for different pruning patterns.

Conclusion:

Wanda pruning is a promising approach to reduce the computational cost of Large Language Models. It leverages a unique pruning metric that takes into account both weight magnitudes and input activation norms. What makes Wanda particularly appealing is its simplicity and efficiency. It can significantly reduce the size of LLMs without requiring retraining or complex weight updates.

In a world where LLMs continue to play a central role in natural language processing, Wanda's contribution to sparsity in these models is invaluable. As researchers and practitioners explore ways to make LLMs more accessible and practical, Wanda stands as a promising solution to tackle the computational challenges posed by these behemoths.

By Kirouane Ayoub