Compress and Sparsify LLM

Sparsity Fundamentals and Pruning

Sparsity, in the context of neural networks, refers to the property where many model parameters (weights) are zero or near-zero, meaning they contribute little or nothing to the output. Exploiting sparsity through pruning allows for model compression and potentially faster inference.

What is Sparsity?

A dense model utilizes all its parameters for computation. A sparse model, however, contains a significant proportion of zero-valued weights. This can be achieved intentionally through pruning or emerge from specific architectural designs like MoE. Sparsity reduces the effective number of parameters, leading to lower memory requirements and potentially faster computation if hardware/software can leverage it.

Pruning Techniques

Pruning involves systematically removing less important weights or structures from a trained model, often followed by fine-tuning to recover any lost accuracy.

• Magnitude-based Pruning: The simplest form, where weights with the smallest absolute values are set to zero. It's effective but often results in unstructured sparsity, which can be hard to accelerate on standard hardware.
• Structured Pruning: Removes entire structures like neurons, channels, or even attention heads. This creates regular patterns of sparsity that are easier to leverage for speedups on hardware like GPUs and CPUs.
• Lottery Ticket Hypothesis: Suggests that dense networks contain smaller subnetworks ("winning tickets") that, when trained in isolation from initialization, can match the performance of the full network. Finding these tickets involves iterative pruning and rewinding weights.

Architectural Sparsity: Mixture of Experts (MoE)

Instead of pruning a dense model, MoE introduces sparsity architecturally by activating only parts of the network for each input.

Traditional Transformer vs. MoE

Traditional transformer models are dense; all parameters in layers like the feed-forward network are used for every input token. MoE replaces dense feed-forward layers with multiple parallel "expert" networks. Only a few selected experts process each token, making computation sparse while allowing for a massive total number of parameters.

Key Components of MoE Architecture

• Experts: Smaller neural networks (typically feed-forward layers) that specialize in different aspects of the data.
• Gating Mechanism (Router): A small network that determines which expert(s) (usually 1 or 2) should process the current input token based on its representation. The router's efficiency and load balancing are critical.

The Role of Gating Mechanisms

The gating network dynamically routes each token to the most relevant expert(s). Its goal is to ensure experts are utilized effectively (load balancing) and that the routing decisions lead to high model quality. Common gating mechanisms use simple linear layers followed by softmax or variants like Top-K routing.

GShard: Scaling MoE across Devices

GShard was a significant step in scaling MoE models for large-scale systems like machine translation. It introduced techniques to efficiently partition the experts and the gating mechanism across multiple accelerator devices (TPUs), enabling models with hundreds of billions of parameters while managing communication overhead.

• GShard Architecture: Experts are sharded across devices, while the gating mechanism can be replicated or sharded. GShard optimized routing and communication patterns for distributed training and inference. It demonstrated MoE's effectiveness beyond single-machine setups.
• Evaluation and Results: GShard achieved state-of-the-art results in machine translation with significantly fewer computational resources (FLOPs) per token compared to dense models of similar quality, showcasing the efficiency benefits of sparse activation.

CoLT5: Conditional MoE for Text Generation

CoLT5 (Conditional Computation Learning for T5) explored applying MoE principles within the context of the T5 encoder-decoder architecture, focusing on efficient pre-training and fine-tuning for text generation tasks. It investigated different strategies for routing and expert placement.

• CoLT5 Architecture: Integrated MoE layers into the T5 framework, experimenting with routing mechanisms and how expert outputs are combined. Addressed challenges like training stability and representation quality in generative MoEs.
• Evaluation and Results: CoLT5 demonstrated strong performance on various NLP benchmarks, often outperforming dense T5 models of similar computational cost (FLOPs) during inference.
• Ablations & Limitations: Studies analyzed the impact of the number of experts, routing strategies, and capacity factors. Limitations included potential training instabilities and the complexity of tuning MoE hyperparameters compared to dense models.

Quantization of Large Language Models

Quantization reduces the memory footprint and computational cost of LLMs by representing weights and/or activations with lower-precision numbers (e.g., 8-bit integers) instead of the standard 32-bit or 16-bit floating-point formats.

Quantifying LLMs

This involves mapping the high-precision floating-point values of model parameters (and sometimes activations during inference) to a smaller set of low-precision values. This significantly reduces model size (e.g., 8-bit models are ~4x smaller than 32-bit) and can accelerate computation on hardware with specialized low-precision units.

Key Quantization Techniques

• LLM.int8(): A technique that enables 8-bit integer matrix multiplication for the linear layers in transformers while identifying and preserving outlier activation values in higher precision (FP16). This maintains model performance remarkably well for models >6B parameters with minimal code changes.
• QLoRA (Quantized Low-Rank Adaptation): An efficient fine-tuning technique. It quantizes a pre-trained model to very low precision (e.g., 4-bit NormalFloat) to save memory, then fine-tunes it by adding and training small, low-rank adapter layers. This drastically reduces the memory required for fine-tuning large models. (Note: 8-bit optimizers, like those in `bitsandbytes`, complement techniques like QLoRA by reducing optimizer state memory during training/fine-tuning).
• BitNet (e.g., BitNet b1.58): A more recent and aggressive approach proposing 1-bit LLMs where weights are constrained to {-1, +1} (or {-1, 0, +1} in variants like b1.58). This promises extreme compression and potential energy efficiency but requires specific training recipes and architectural adjustments (like replacing LayerNorm). Its practicality and performance across diverse tasks are still under active research as of early 2025.