ai ai/agenticai llm ai/books/designingllmapplication
Tldr
- K-V caching avoids repeated attention computation by reusing cached keys and values.
- Early exit saves compute by stopping before the final Transformer layer.
- Knowledge distillation creates a smaller model that approximates a larger teacher.
- Speculative decoding uses a small draft model to propose tokens for a larger model to verify.
- Parallel/multi-token decoding reduces the sequential bottleneck of autoregressive generation.
- Quantization reduces memory usage by representing numbers with fewer bits.
- Every optimization involves trade-offs involving accuracy, latency, memory, storage, or implementation complexity.
Overview
LLMs are expensive to deploy and run because of their large compute, memory, and energy requirements. Inference is especially challenging on edge devices.
Two major inference bottlenecks are:
-
Autoregressive generation: decoder-only LLMs generate one token at a time, creating a sequential limitation.
-
Long sequences: as input length increases, the compute required increases quadratically.
Inference optimization mainly targets three goals:
-
Reduce compute
-
Speed up decoding
-
Reduce storage requirements
1. Techniques for Reducing Compute
K-V Caching
Key idea: cache and reuse the keys (K) and values (V) from Transformer self-attention instead of repeatedly computing them, reducing repeated work during Transformer decoding.
This is useful when prompts contain repeated information, such as:
-
Repeated system instructions
-
Product manuals
-
API documentation
-
RAG context
-
Few-shot examples
Benefits
-
Reduces duplicated computation.
-
Can significantly reduce inference time.
-
Can reduce cost when providers charge cached tokens only once.
-
Few-shot caching can sometimes be a lightweight alternative to fine-tuning.
Trade-off
Caching exchanges compute for storage. K-V caches can become very large for long sequences, so providers may limit cache lifetime or charge based on caching duration.
Practical point: Put cacheable content at the beginning of the prompt (the prompt prefix).
Early Exit
Key idea: stop inference before the final Transformer layer and use an intermediate layer’s representation as the output.
Sequence-level early exit
The entire input sequence exits at the same layer.
The exit layer can be:
-
Fixed in advance, or
-
Dynamically selected based on the input.
Adapters/classifiers can be trained at different layers to decide whether the model should exit. Examples mentioned include FastBERT and HashEE.
Token-level early exit
Different tokens in the same sequence can exit at different layers.
-
More fine-grained than sequence-level early exit.
-
More complex to implement.
-
Can be more effective, but is slower than sequence-level early exit.
-
CALM (Confident Adaptive Language Modeling) is an example.
Trade-off
Early exit reduces compute but can reduce model performance. Dynamic selection of an appropriate exit layer can help minimize this degradation.
Dynamic inference
Inference computation can be adapted to the input.
Mixture-of-Experts (MoE) is an important example: a routing function selects only a subset of expert modules for each input, reducing computation.
Knowledge Distillation
Goal: create a smaller student model that approximates a larger teacher model, enabling faster inference.
The process has two main stages:
-
Prepare distillation data.
-
Train the student model.
Ways to prepare distillation data
-
Unsupervised generation: prompt the teacher to generate responses for the desired capability.
-
Data augmentation: give the teacher seed examples and have it generate similar examples.
-
Intermediate representations: use activations or logits from the teacher (white-box distillation).
-
Teacher feedback: the teacher evaluates student outputs and provides preferences or improvement instructions.
-
Self-teaching: the same model generates and refines its own outputs to create the distillation set.
A relatively small, high-quality dataset can be sufficient; the source notes that about 1,000 high-quality examples can be enough for a strong distillation set.
Distillation methods
-
Supervised fine-tuning: train the student to match teacher outputs.
-
KL divergence: minimize the difference between teacher and student output probability distributions.
-
Internal representation similarity: align teacher and student representations, potentially layer by layer.
-
Reinforcement learning: train a reward model from the distillation data and optimize the student against that reward.
Weak-to-Strong Generalization
A weaker teacher can generate weak labels that are then used to fine-tune a stronger student.
The stronger student’s pretrained representations allow it to learn useful capabilities even from weaker supervision.
Important limitations
Knowledge distillation can cause:
-
Capability degradation
-
Catastrophic forgetting
Careful evaluation is therefore necessary.
If teacher weights are unavailable, the source notes that supervised fine-tuning is the available distillation approach; white-box techniques require access to internal representations.
2. Techniques for Accelerating Decoding
Autoregressive LLMs generate one token after another, so decoding is inherently sequential.
Speculative Decoding
A small draft model generates several candidate tokens. The larger main model then evaluates those candidates together and accepts or rejects them.
The more draft tokens accepted, the more useful the draft model is.
Key metrics
-
Token acceptance rate: percentage of draft-model tokens accepted by the main model.
-
Decoding speedup: reduction in latency compared with ordinary autoregressive decoding.
Constructing draft models
Possible approaches include:
-
DistillSpec: distill the draft model from the main model.
-
Self-speculative decoding: use a subset of the main model’s layers as the draft model.
-
Retrieval-based speculative decoding (REST): retrieve commonly occurring phrases, prefixes, boilerplate, or quoted text from external repositories instead of generating them from scratch.
Parallel Decoding
Instead of generating only one token/output component at a time, independent portions can be generated simultaneously.
Two approaches include:
-
Multiple model instances: split independent parts of an output and generate them in parallel.
-
Multi-token decoding: use the same model to predict multiple future tokens.
Medusa
Medusa adds extra decoding heads that predict subsequent tokens:
-
Standard head → next token
-
Additional heads → later tokens
This enables multiple tokens to be proposed at once.
3. Techniques for Reducing Storage Needs
Quantization
Quantization converts numerical representations from higher precision to lower precision, reducing memory requirements.
Common formats discussed:
| Format | Main characteristic |
|---|---|
| FP32 | 32-bit floating point; high precision |
| FP16 | 16-bit floating point; lower range and precision |
| BF16 | 16-bit format with FP32-like range but lower precision |
| INT8 | 8-bit integer; much smaller storage requirement |
FP32 uses 32 bits (4 bytes) per number.
Converting FP32 to FP16 is lossy and reduces both range and precision. BF16 preserves approximately the same numerical range as FP32 while sacrificing precision.
Quantizing FP32 values to INT8 can reduce memory requirements by roughly a factor of 4, at the cost of additional precision loss.
Integer Quantization
Two main approaches:
Symmetric quantization
-
Original zero maps to integer zero.
-
A common method is absmax quantization.
-
The absolute maximum value is mapped to the largest positive/negative integer value.
Asymmetric quantization
-
Zero does not have to map to integer zero.
-
The original minimum and maximum are mapped to the minimum and maximum values of the target integer range.
Outliers
Outliers can negatively affect quantization. One approach is clipping, where extreme values are limited to the maximum/minimum representable values.
4. Weight vs. Activation Quantization
Quantization is typically applied after training.
Weight quantization
Generally easier because weights are known ahead of time, allowing the system to determine:
-
Ranges
-
Outliers
-
Scaling factors
Activation quantization
More difficult because activations are produced dynamically during inference.
Two approaches:
-
Dynamic scaling: calculate statistics such as ranges and outliers during inference.
-
Static scaling: estimate statistics using a calibration dataset.
Static scaling can speed up inference but may introduce more quantization errors.
Quick Comparison
| Goal | Techniques | Main trade-off |
| Reduce compute | K-V caching, early exit, knowledge distillation | Storage or performance/capability trade-offs |
| Speed up decoding | Speculative decoding, parallel decoding | Added system/model complexity |
| Reduce storage | Quantization | Reduced numerical precision/range |