Model Optimizations

Quantization reduces the numerical precision of model weights and activations from floating-point (typically FP32) to lower-bit representations (INT8, INT4, or even binary). This directly reduces model size, memory bandwidth requirements, and computational cost, while leveraging integer arithmetic units that are faster and more energy-efficient on most hardware.

Post-Training Quantization (PTQ)

Post-training quantization converts a pre-trained FP32 model to lower precision without retraining. The simplest approach uniformly maps the FP32 value range to the target precision. More sophisticated methods like per-channel quantization and dynamic range adjustment minimize the accuracy impact. PTQ is attractive because it requires no additional training, but it can struggle with models that have wide activation ranges.

Quantization-Aware Training (QAT)

Quantization-aware training simulates quantization during training, allowing the model to adapt its weights to the constraints of lower precision. During training, fake quantization nodes insert quantization and dequantization operations in the forward pass while using straight-through estimators for gradients. QAT typically achieves better accuracy than PTQ, especially at very low bit widths.

LLM Quantization: GPTQ, AWQ, and Beyond

Recent advances in quantization for large language models have pushed the boundaries to 4-bit and even 2-bit precision. Techniques like GPTQ, AWQ, and SqueezeLLM use sophisticated calibration methods to maintain quality at extreme compression ratios. These methods have made it possible to run large language models on consumer hardware, democratizing access to powerful AI capabilities.

Pruning removes unnecessary parameters from neural networks, creating sparse models that require less computation and memory. The key insight is that neural networks are typically over-parameterized, and many weights contribute little to the output. Identifying and removing these redundant parameters can yield significant efficiency gains.

Unstructured Pruning

Unstructured pruning removes individual weights based on magnitude or other criteria, creating irregular sparsity patterns. While this can achieve high compression ratios, the irregular memory access patterns make it difficult to accelerate on standard hardware. Specialized sparse computation libraries and hardware are needed to realize the theoretical speedups.

Structured Pruning

Structured pruning removes entire neurons, channels, or attention heads, producing models that maintain dense computation patterns. This is more hardware-friendly because standard dense matrix operations can be used on the smaller model. Channel pruning for CNNs and head pruning for Transformers are the most common structured pruning approaches.

The Lottery Ticket Hypothesis

The Lottery Ticket Hypothesis, proposed by Frankle and Carlin, suggests that dense networks contain sparse subnetworks (winning tickets) that can achieve comparable accuracy when trained in isolation. This finding has profound implications for understanding why neural networks work and has inspired new approaches to training efficient models from the start.

A randomly-initialized, dense neural network contains a subnetwork that is initialized such that — when trained in isolation — it can match the test accuracy of the original network after training for at most the same number of iterations.

Knowledge distillation transfers the learned representations from a large, complex teacher model to a smaller, more efficient student model. Rather than training the student directly on hard labels, distillation trains it to match the soft probability distributions produced by the teacher, which contain richer information about class relationships.

The Distillation Loss

The standard distillation loss combines the student's cross-entropy loss on ground truth labels with a KL divergence loss between the teacher's and student's softened output distributions. The temperature parameter controls the softness of the distributions: higher temperatures produce softer distributions that transfer more information but may introduce noise.

Feature-Based and Attention Distillation

Feature-based distillation extends the concept by matching intermediate representations rather than just outputs. The student is trained to mimic the teacher's internal feature maps at selected layers. Attention transfer, which matches attention maps between teacher and student, is another popular variant for Transformer models.

Distillation for Large Language Models

Distillation has become a critical technique in the large language model era. Models like DistilBERT achieve 97% of BERT's performance with 40% fewer parameters by distilling during pre-training. For deployment on resource-constrained devices, distillation provides a principled way to compress knowledge from models that are too large to serve directly.

Operator fusion combines multiple sequential operations into a single fused kernel, reducing memory traffic and kernel launch overhead. For example, fusing a convolution, batch normalization, and ReLU activation into a single operation avoids writing and reading intermediate results from memory, significantly improving performance.

Graph-Level Optimizations

Graph-level optimizations analyze the computational graph to identify fusion opportunities, eliminate redundant operations, and optimize memory allocation.

• Constant folding — Pre-compute operations whose inputs are all known at compile time (e.g., shape calculations).
• Dead code elimination — Remove computations whose outputs are never used downstream.
• Layout transformation — Convert tensor memory layouts (NCHW vs. NHWC) to match hardware preferences.
• Common subexpression elimination — Detect and share identical computations to avoid redundant work.
• Memory planning — Schedule tensor lifetimes to reuse memory, reducing peak memory consumption.

ML Compilers

ML compilers like TensorRT, Apache TVM, and XLA automate these optimizations for specific hardware targets. TensorRT, for example, can automatically fuse operations, select optimal kernel implementations, and apply quantization for NVIDIA GPUs. These compilers can achieve 2-5x speedup over naive framework execution.

Custom Kernels

Custom kernel development using CUDA, Triton, or other low-level languages enables optimizations that are impossible at the graph level. Flash Attention is a prominent example: by rewriting attention as a single fused kernel that is aware of the GPU memory hierarchy, it achieves both lower memory usage and faster execution than the standard multi-operation implementation.

class="tok-comment"># Triton: custom fused add + ReLU kernel
class="tok-decorator">@triton.jit
def fused_add_relu_kernel(x_ptr, y_ptr, out_ptr, n,
                          BLOCK: tl.constexpr):
    pid = tl.program_id(class="tok-number">0)
    offsets = pid * BLOCK + tl.arange(class="tok-number">0, BLOCK)
    mask = offsets < n
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    result = tl.maximum(x + y, class="tok-number">0.0)  class="tok-comment"># Fused add + ReLU
    tl.store(out_ptr + offsets, result, mask=mask)

In practice, the most effective optimization pipelines combine multiple techniques rather than relying on a single approach. A typical pipeline might start with architecture-level efficiency (choosing an efficient base architecture), apply knowledge distillation during training, follow with quantization-aware training, and finish with operator fusion during deployment.

Order of Operations Matters

The order of optimization matters because techniques interact. Pruning before quantization is generally better than the reverse because pruning can remove parameters that would have caused quantization errors. Similarly, distillation-then-quantization typically outperforms quantization-then-distillation because the distilled model starts from a stronger quality baseline.

1. Architecture design — Choose an efficient base architecture (MobileNet, EfficientNet) or use NAS.
2. Knowledge distillation — Train a smaller student from a larger teacher model.
3. Pruning — Remove redundant parameters (structured pruning preferred for hardware compatibility).
4. Quantization-aware training — Fine-tune with simulated quantization for the target bit width.
5. Post-training quantization — Apply final quantization calibration if QAT was not used.
6. Operator fusion and graph optimization — Use an ML compiler (TensorRT, TVM) for the target hardware.

Benchmarking on Target Hardware

Measuring the compound effect of optimizations requires careful benchmarking on the target hardware. Theoretical speedup predictions (e.g., "4x from INT8 quantization") rarely match real-world measurements due to memory bottlenecks, kernel launch overhead, and data loading constraints. End-to-end profiling on the target device is essential.

Deployment-Specific Strategies