In this blog, we present the key techniques to gain AI efficiency, meaning models that are:
- Faster: Accelerate inference times through advanced optimization techniques
- Smaller: Reduce model size while maintaining quality
- Cheaper: Lower computational costs and resource requirements
- Greener: Decrease energy consumption and environmental impact
For this, Pruna provides an open-source toolkit that simplifies scalable inference, requiring just a few lines of code to optimize your models in each of the mentioned aspects.
So first, let’s take a quick look at an overview of these techniques, and then we’ll dive deeper into each one.
Optimization Techniques
To get started, we created a high-level overview of the different techniques implemented in Pruna. This list can be further enriched; however, it provides a solid basis for your understanding.
| Technique | Description | Impacts |
|---|---|---|
| Batching | Groups multiple inputs together to be processed simultaneously, improving computational efficiency and reducing overall processing time. | Speed (✅), Memory (❌), Accuracy (~) |
| Caching | Stores intermediate results of computations to speed up subsequent operations, reducing inference time by reusing previously computed results. | Speed (✅), Memory (⚠️), Accuracy (~) |
| Speculative Decoding | Speculative decoding speeds up AI text generation by having a small, fast model predict several tokens at once, which a larger model then verifies, creating an efficient parallel workflow. | Speed (✅), Memory (❌), Accuracy (⚠️) |
| Compilation | Compilation optimizes the model with instructions for specific hardware. | Speed (✅), Memory (➖), Accuracy (~) |
| Distillation | Trains a smaller, simpler model to mimic a larger, more complex model. | Speed (✅), Memory (✅), Accuracy (❌) |
| Quantization | Reduces the precision of weights and activations, lowering memory requirements. | Speed (✅), Memory (✅), Accuracy (❌) |
| Pruning | Removes less important or redundant connections and neurons, resulting in a sparser, more efficient network. | Speed (✅), Memory (✅), Accuracy (❌) |
| Recovering | Restores the performance of a model after compression. | Speed (⚠️), Memory (⚠️), Accuracy (🟢) |
✅(improves), ➖(stays the same), ~(could worsen), ❌(worsens)
Technique requirements and constraints
Before we continue, note that each one of these techniques and their underlying implementation algorithms has specific requirements and constraints. Some techniques can only be applied on specific hardware, like GPUs, or models like LLMs or image generation models. Others might require a tokenizer, processor, or dataset to function. Lastly, not all techniques can be used interchangeably, and therefore have compatibility limitations.
The Optimization Techniques
We will now dive a bit deeper into different optimization techniques. Although we will dive a bit deeper into the various techniques and their underlying algorithms, we will not be going into the nitty-gritty details, and keep it high level, and for each technique, highlight one of the fundamental underlying algorithms that has been implemented in the Pruna library.
Batching AI model inference
Batching groups multiple inputs together to be processed simultaneously, improving computational efficiency and reducing overall processing time. Instead of processing one prompt at a time, the GPU processes multiple prompts in parallel, maximizing hardware utilization. This significantly increases throughput since modern GPUs are designed for parallel computation. Batching reduces the per-example computational overhead and allows for better distribution of fixed costs across multiple inputs, thus often increasing the throughput.
For batching, we implemented WhisperS2T, which works on top of Whisper models. It intelligently batches smaller speech segments and is designed to be faster than other implementations, boasting a 2.3X speed improvement over WhisperX and a 3X speed boost compared to HuggingFace Pipeline with FlashAttention 2 (Insanely Fast Whisper).
Caching intermediate results
Caching stores intermediate results of computations to speed up subsequent operations, reducing inference time by reusing previously computed results. For transformer-based LLMs, this typically involves storing key-value pairs from previous tokens to avoid redundant computation. When generating text token by token, each new token can reuse cached computations from previous tokens rather than recomputing the entire sequence. This dramatically improves inference efficiency, especially for long-context applications. However, caching goes beyond only saving KV computations and can be used in multiple places for LLMs and image generation models.
For caching, we implemented DeepCache, which works on top of diffuser models. DeepCache accelerates inference by leveraging the U-Net blocks of diffusion pipelines to reuse cached high-level features. The nice thing is that it is training-free and almost lossless, while accelerating models 2X to 5X.
Speculative decoding with parallelizing generation
Speculative decoding improves the efficiency of language model inference by parallelizing parts of the generation process. Instead of generating one token at a time, a smaller, faster draft model generates multiple candidate tokens in a single forward pass. The larger, more accurate model then verifies or corrects these tokens in parallel, allowing for faster token generation without significantly sacrificing output quality. This approach reduces the number of sequential steps required, thereby lowering overall latency and accelerating inference. It’s essential to note that the effectiveness of speculative decoding depends on the alignment between the draft and target models, as well as the chosen parameters, such as batch size and verification strategy.
For speculative decoding, we have not implemented any algorithms. Yet! Stay tuned to discover our future speculative decoding algorithms.
Compilation for specific hardware
Compilation optimizes the model for specific hardware by translating the high-level model operations into low-level hardware instructions. Compilers like NVIDIA TensorRT, Apache TVM, or Google XLA analyze the computational graph, fuse operations where possible, and generate optimized code for the target hardware. This process eliminates redundant operations, reduces memory transfers, and leverages hardware-specific acceleration features, resulting in faster inference times and lower latency. It is essential to note that each combination of model/hardware will have a different optimal compilation setup.
For compilation, we implemented Stable-fast, which works on top of diffuser models. Stable-fast is an optimization framework for Image-Gen models. It accelerates inference by fusing key operations into optimized kernels and converting diffusion pipelines into efficient TorchScript graphs.
Distillation for smaller models
Distillation trains a smaller, simpler model to mimic a larger, more complex model. The larger “teacher” model produces outputs that the smaller “student” model learns to replicate, effectively transferring knowledge while reducing computational requirements. This technique preserves much of the performance and capabilities of larger models while significantly reducing parameter count, memory usage, and inference time. Distillation can target specific capabilities of interest rather than general performance.
For distillation, we implemented Hyper-SD, which works on top of diffusion models. Hyper-SD is a distillation framework that segments the diffusion process into time-step groups to preserve and reformulate the ODE trajectory. By integrating human feedback and score distillation, it enables near-lossless performance with drastically fewer inference steps.
Quantization for lower precision
Quantization reduces the precision of weights and activations, lowering memory requirements by converting high-precision floating-point numbers (FP32/FP16) to lower-precision formats (INT8/INT4). It reduces model size, memory bandwidth requirements, and computational complexity. Modern quantization techniques, such as Post-Training Quantization (PTQ) and Quantization-Aware Training (QAT), minimize accuracy loss while achieving substantial efficiency gains. Hardware accelerators often have specialized support for low-precision arithmetic, further enhancing performance.
For quantization, we implemented Half-Quadratic Quantization (HQQ), which works on top of any model. HQQ utilizes fast and robust optimization techniques for on-the-fly quantization, eliminating the need for calibration data and making it applicable to any model. This algorithm is adapted explicitly for diffuser models.
Pruning away redundant neurons
Pruning removes less important or redundant connections and neurons, resulting in a sparser, more efficient network. Various pruning strategies exist, including magnitude-based pruning (removing the smallest weights) and lottery ticket hypothesis approaches (finding sparse subnetworks). Key design choices typically involve deciding which structure to prune (e.g., weight, neuron, blocks) and determining how to score structures (e.g., using weight magnitude, first-order, or second-order information). Pruning can significantly reduce model size (often by 80-90%) with minimal performance degradation when done carefully. Sparse models require specialized hardware or software support to realize computational gains.
For pruning, we implemented structured pruning, which works on top of any model. Structured pruning removes entire units like neurons, channels, or filters from a network, leading to a more compact and computationally efficient model while preserving a regular structure that standard hardware can easily optimize.
Recovering performance with training
Recovering is special since it allows for improving the performance of compressed models. After compression, it restores the performance of a model through techniques like finetuning or retraining. After aggressive pruning, models typically experience some performance degradation, which can be mitigated by additional training steps. This recovery phase allows the remaining parameters to adapt and compensate for the compression. Approaches for efficient recovery include learning rate rewinding, weight rewinding, and gradual pruning with recovery steps between pruning iterations. The recovery process helps achieve optimal trade-offs between model size and performance.
For recovering, we implemented text-to-text PERP, which works on top of text generation models. This recoverer is a general-purpose PERP recoverer for text-to-text models using norm, head, and bias finetuning and optionally HuggingFace’s LoRA. Similarly, we support text-to-image PERP for other image generation models.
What’s next?
This blog provided a brief introduction to each of these categories, but there are many more nuances, techniques, and implementations that we will highlight in upcoming blogs. The cool thing is that each of these techniques has been implemented in the open-source Pruna library and is ready for you to experiment with!
Enjoy the Quality and Efficiency!
Want to take it further?
- Compress your own models with Pruna and give us a ⭐ to show your support!
- Explore our materials collection, or dive into our courses.
- Join the conversation and stay updated in our Discord community.
Top comments (0)