Model Compression & Quantization

Deploying large models on resource-constrained devices requires compression techniques that reduce model size, computation, and latency while preserving accuracy. This module covers pruning, distillation, quantization, and parameter-efficient fine-tuning with mathematical foundations.

1. Taxonomy of Compression Methods

Model Compression
├── Pruning
│   ├── Unstructured (weight-level)
│   ├── Structured (filter/channel/layer)
│   ├── Dynamic (at inference time)
│   └── Lottery Ticket Hypothesis
├── Quantization
│   ├── Post-Training Quantization (PTQ)
│   ├── Quantization-Aware Training (QAT)
│   ├── INT8, INT4, NF4
│   └── GPTQ, AWQ, GGUF
├── Knowledge Distillation
│   ├── Classical (Hinton)
│   ├── Feature-based
│   ├── Relation-based
│   └── Self-distillation
└── Parameter-Efficient Fine-Tuning
    ├── LoRA / QLoRA
    ├── Adapters
    ├── Prefix Tuning
    └── Prompt Tuning

2. Pruning

2.1 Unstructured Pruning

Remove individual weights below a threshold:

where is the pruning threshold. The pruned model:

Magnitude pruning (Han et al., 2015): Remove weights with smallest absolute value.

Optimal Brain Damage (LeCun et al., 1989): Use second-order information:

where is the diagonal of the inverse Hessian. Prune weights with largest (smallest cost increase).

2.2 Structured Pruning

Remove entire filters, channels, or layers:

Filter pruning: Remove filter if:

Channel pruning: Remove channel if its corresponding filter's L1-norm is small.

Taylor expansion pruning:

where is the -th eigenvalue of the Hessian.

2.3 Lottery Ticket Hypothesis

Frankle & Carbin (2019): A dense network contains a sparse subnetwork ("winning ticket") that can match the full network's accuracy.

Algorithm:

1. Train full network to convergence
2. Prune smallest-magnitude weights
3. Reset remaining weights to initialization
4. Retrain from initialization
5. Repeat

Convergence: The winning ticket converges with the same learning rate as the original network, suggesting the initialization matters more than the architecture.

2.4 Gradual Magnitude Pruning (GMP)

where is the sparsity schedule from 0% to target sparsity .

3. Knowledge Distillation

3.1 Classical Distillation (Hinton et al., 2015)

A small student network learns from a large teacher network :

Soft labels: Teacher logits with temperature :

Distillation loss:

where:

• : Cross-entropy with ground truth labels
• : KL divergence between teacher and student soft labels
• scaling: Compensates for gradient magnitude change with temperature
• : Balance hyperparameter

Temperature analysis: Higher produces softer probability distributions, revealing more information about the teacher's knowledge about inter-class relationships.

3.2 Feature-Based Distillation

Match intermediate representations:

where are projection functions that match dimensions.

FitNets (Romero et al., 2015): Use hint layers to match intermediate features.

3.3 Attention Transfer

Match attention maps:

3.4 Self-Distillation

The network distills knowledge from itself:

where and are predictions from different layers or branches.

4. Quantization

4.1 Uniform Quantization

Map continuous values to discrete levels:

where is the scale factor and is the zero-point.

Scale and zero-point:

4.2 Symmetric Quantization

4.3 Quantization Error

Bound: (absolute error per element)

Relative error: (smaller for larger magnitudes)

Signal-to-quantization-noise ratio (SQNR):

4.4 Mixed-Precision Quantization

Assign different bit-widths to different layers:

where is the number of parameters in layer and is the total bit budget.

4.5 Post-Training Quantization (PTQ)

Calibration: Run a small calibration set through the model to determine activation ranges:

4.6 Quantization-Aware Training (QAT)

Simulate quantization during training:

Use straight-through estimator (STE) for gradient:

5. Advanced Quantization Methods

5.1 GPTQ: Generalized Post-Training Quantization

GPTQ (Frantar et al., 2022) uses optimal brain quantization:

where is the quantization error.

Block quantization: Process columns at a time to reduce memory.

5.2 AWQ: Activation-Aware Weight Quantization

Protect salient channels based on activation magnitudes:

where with .

5.3 GGUF/GGML Quantization

Used by llama.cpp for efficient CPU inference:

Block-wise quantization: Quantize blocks of 32 weights with shared scale:

• Q4_0: 4-bit, block size 32
• Q4_1: 4-bit with minimum value
• Q5_0: 5-bit quantization
• Q8_0: 8-bit quantization

6. Low-Rank Adaptation (LoRA)

6.1 LoRA: Low-Rank Adaptation

LoRA (Hu et al., 2022) freezes the pretrained weights and injects trainable low-rank matrices:

where , , and .

Forward pass:

where is the scaling factor and is the rank.

Initialization: , .

Parameter count: (original) vs. (LoRA). For , this is of original.

6.2 QLoRA

QLoRA (Dettmers et al., 2023) combines 4-bit quantization with LoRA:

1. Quantize pretrained model to 4-bit NF4
2. Add LoRA adapters in 16-bit
3. Train only LoRA parameters

NF4 (NormalFloat4): Optimal 4-bit quantization for normally distributed weights:

where and are chosen to minimize quantization error for .

Double quantization: Quantize the quantization constants:

6.3 LoRA Variants

LoRA+ (Hayou et al., 2024): Use different learning rates for and :

with .

rsLoRA (Kalajdzievski, 2023): Rank-stabilized LoRA:

DoRA (Liu et al., 2024): Weight-Decomposed Low-Rank Adaptation:

where is a learnable magnitude vector.

AdaLoRA (Zhang et al., 2023): Adaptive rank allocation per layer:

where is the importance score of layer .

7. Implementation

import torch
import torch.nn as nn
import torch.nn.functional as F

class LoRALinear(nn.Module):
    def __init__(self, in_features, out_features, rank=16, alpha=1.0):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features, bias=True)
        self.linear.weight.requires_grad = False
        self.linear.bias.requires_grad = False
        
        self.lora_A = nn.Parameter(torch.randn(rank, in_features) * 0.01)
        self.lora_B = nn.Parameter(torch.zeros(out_features, rank))
        self.scaling = alpha / rank
    
    def forward(self, x):
        h = self.linear(x)
        h += (x @ self.lora_A.T @ self.lora_B.T) * self.scaling
        return h

def freeze_model(model):
    for param in model.parameters():
        param.requires_grad = False

def add_lora(model, target_modules=['q_proj', 'v_proj'], rank=16):
    for name, module in model.named_modules():
        if any(target in name for target in target_modules):
            if isinstance(module, nn.Linear):
                lora = LoRALinear(
                    module.in_features, 
                    module.out_features, 
                    rank=rank
                )
                lora.linear.weight = module.weight
                lora.linear.bias = module.bias
                parent = get_parent(model, name)
                setattr(parent, name.split('.')[-1], lora)
    return model

class KnowledgeDistillationLoss(nn.Module):
    def __init__(self, temperature=4.0, alpha=0.5):
        super().__init__()
        self.T = temperature
        self.alpha = alpha
    
    def forward(self, student_logits, teacher_logits, labels):
        # Soft target loss
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.T, dim=1),
            F.softmax(teacher_logits / self.T, dim=1),
            reduction='batchmean'
        ) * (self.T ** 2)
        
        # Hard target loss
        hard_loss = F.cross_entropy(student_logits, labels)
        
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

class StructuredPruner:
    def __init__(self, model, pruning_ratio=0.3):
        self.model = model
        self.ratio = pruning_ratio
    
    def compute_filter_importance(self, layer):
        if isinstance(layer, nn.Conv2d):
            return layer.weight.data.abs().sum(dim=[1, 2, 3])
        elif isinstance(layer, nn.Linear):
            return layer.weight.data.abs().sum(dim=1)
        return None
    
    def prune_conv(self, layer, importance, mask):
        num_keep = int(len(importance) * (1 - self.ratio))
        _, indices = torch.sort(importance, descending=True)
        keep = indices[:num_keep]
        
        new_weight = layer.weight.data[keep]
        new_bias = layer.bias.data[keep] if layer.bias is not None else None
        
        pruned = nn.Conv2d(
            layer.in_channels, num_keep, 
            layer.kernel_size, layer.stride, layer.padding
        )
        pruned.weight.data = new_weight
        if new_bias is not None:
            pruned.bias.data = new_bias
        
        return pruned, keep

class QuantizedLinear(nn.Module):
    def __init__(self, in_features, out_features, bits=8):
        super().__init__()
        self.bits = bits
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.scale = None
        self.zero_point = None
    
    def quantize_weight(self):
        if self.bits == 8:
            w_max = self.weight.data.abs().max()
            self.scale = w_max / 127
            self.zero_point = 0
            w_q = torch.clamp(
                torch.round(self.weight.data / self.scale), -128, 127
            )
            return w_q * self.scale
        return self.weight.data
    
    def forward(self, x):
        w = self.quantize_weight()
        return F.linear(x, w, self.bias)

8. SVG: Pruning Before/After

9. SVG: Quantization Levels

10. SVG: Knowledge Distillation Diagram

11. LoRA Architecture

12. Comparison of Compression Methods

13. Open Problems

• Extreme quantization: 2-bit and ternary quantization for edge deployment
• Structured sparsity: Hardware-friendly sparse patterns (N:M sparsity)
• Combining methods: Joint pruning + quantization + distillation
• LLM-specific: Quantizing mixture-of-experts, attention layers, KV-cache
• Adaptive compression: Dynamic compression based on input complexity
• Theoretical guarantees: Understanding when compression preserves performance

References

1. Han, S., Pool, J., Tran, J., & Dally, W. (2015). Learning both Weights and Connections for Efficient Neural Networks. NeurIPS.
2. Hinton, G., Vinyals, O., & Dean, J. (2015). Distilling the Knowledge in a Neural Network. arXiv:1503.02531.
3. Frankle, J., & Carbin, M. (2019). The Lottery Ticket Hypothesis. ICLR.
4. Frantar, E., Ashkboos, S., Hoefler, T., & Alistarh, D. (2022). GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers. ICLR.
5. Hu, E. J., Shen, Y., Wallis, P., et al. (2022). LoRA: Low-Rank Adaptation of Large Language Models. ICLR.
6. Dettmers, T., Pagnoni, A., Holtzman, A., & Zettlemoyer, L. (2023). QLoRA: Efficient Finetuning of Quantized LLMs. NeurIPS.
7. Liu, S.-Y., Wang, C.-Y., Yin, H., et al. (2024). DoRA: Weight-Decomposed Low-Rank Adaptation. arXiv.