Large Language Model Efficiency: Pruning, Quantization, and Adaptation Strategies
Efficiency, Scaling, and Capacity Models
- Foundational Notions of Model Capacity: The concept of capacity models serves as the fundamental building block for scaling large language models (LLMs).
- Scaling and Performance: Significant research indicates that increasing compute, alongside the use of transformer models and vast amounts of data, leads to improved performance on the models themselves and their downstream capabilities.
- Emergence in LLMs: There is a phenomenon known as "emergence," where certain capabilities only become available once model sizes surpass a specific capacity threshold. Below this threshold, performance on these tasks is non-existent or minimal. While there are caveats to this analysis, the distinction between small and large model capabilities remains valid.
- Optimization Strategies for Deployment:
* Distillation: A method to push information into smaller models to mimic the performance of larger, pre-trained, and post-trained models.
* Pruning: A technique to make large models more efficient by identifying parameters that contribute minimally to the overall output and removing them.
* Weight Grooming Signals: Pruning relies on signals to determine which parameters are the "logical units of computation" that can be removed.
* Hardware Coordination: Effective pruning requires tight coordination with the underlying hardware. If pruned weights (zeros) are distributed uniformly across a weight matrix in a simple linear layer, little efficiency is gained unless the hardware is specifically arranged to take advantage of sparsity.
* Unit Pruning: Instead of individual weights, entire parts of the network or matrices can be removed to optimize dense matrix multiplications.
* Quantization: Reducing the precision of the underlying hardware's numerical representation.
* Standard Precision: Operations are typically done in FP32 (Floating Point 32).
* The Approximation Hypothesis: Since the entire LLM is essentially an approximation of a function, exact values are often unnecessary. Reducing precision from FP32 to lower formats can save significant compute without drastically affecting output.
* Quantization-Aware Training: Suddenly reducing precision in a full-precision model can break the alignment between layers. To recover losses, practitioners apply quantization during the forward pass of training while treating it as a differentiable operation in the backward pass.
* Mathematical Constraint: Quantization is naturally a non-differentiable operation, making it impossible to apply standard gradient backpropagation (backprop) directly through it without these compensatory training simple techniques.
- Reducing Real-Time Compute:
* Self-Attention Optimization: Reducing computation for long-context tasks by optimizing the self-attention mechanism.
* Early Termination: Avoiding certain layers or terminating computation early through the network hierarchy to save resources.
Training Side Efficiency and Adaptation
- Transitioning to Training: While deployment focuses on optimizing an existing big model, there is a separate challenge in adapting models to target domains (e.g., specialized code generation for a specific new language like Rust).
- The Challenge of 70 Billion Parameters: Training even a "small" model by modern standards, such as a 70,000,000,000 parameter model, is problematic with regular training methods.
- Motivations for Efficiency:
* Cost reduction for large corporations.
* Accessibility for small companies.
* Sustainability and environmental impact ("saving the planet").
- Adapting Pre-trained Models:
* The Setup: A base model consists of a series of layers with a language modeling head at the output for generation.
* Full Fine-Tuning: Involves touching all parameters in the model during training.
* Top-k Layer Fine-Tuning: Only the top k layers are updated while the rest of the model remains frozen.
* If k=0, only the output layer is trained.
* Limitations: This creates a disconnect because frozen lower layers do not adapt to provide information in the way the changing upper layers require.
* Memory Efficiency: Fine-tuning fewer parameters reduces the number of gradients that must be kept in memory. Memory requirements vary by optimization algorithm, with some requiring multiple rounds of previous gradient aggregations.
Adapting with Adapters and Prefix Tuning
- Adapters: Small, task-specific layers inserted into the fixed transformer architecture.
* Mechanism: Transformers consist of self-attention and feed-forward layers. Adapters adjust the representations coming out of these components for a specific task without permanently changing the base parameters.
* Bottleneck Architecture: An adapter projects a representation of dimension d down to a smaller dimension r, applies a nonlinearity, and then projects it back to dimension d.
* Dimensional Integrity: The dimensions must remain consistent at the entry and exit points to ensure the subsequent frozen layers can process the signal.
- Prefix Tuning and Steering:
* Concept: Learning task-specific continuous tokens (prefixes) that steer the model's internal representations during the generation process.
* Implementation: Adding approximately 100 task-specific tokens to the start of an input. The transformer attends to these tokens, influencing the representation created for the actual input text.
* Efficiency Aspect: This saves parameter updates because only the prefix tokens are trained; the entire transformer is frozen. However, it does not save memory during training because all model activations must still be stored for backpropagation.
Low-Rank Adaptation (LoRA) and QLoRA
- Low-Rank Adaptation (LoRA): The current standard for adapting large models.
* The Weight Update Hypothesis: Gradient descent updates can be viewed as adding a matrix to the original weight matrix.
* Mathematical Decomposition: Instead of updating the full weight matrix W of size dimesd, LoRA uses two low-rank matrices, A and B, such that the update ΔW=BimesA.
* Parameter Count: For a rank r, the number of trainable parameters is 2imesrimesd, which is significantly smaller than d2.
* Zero Latency Inference: After training, the matrices B and A can be multiplied and added back into the original weights Wext(frozen)+BimesA=Wext(updated), resulting in a model with the same architecture and no additional inference cost.
- QLoRA: A variation that combines quantization with LoRA. The original model is quantized to a format like 4ext−bitInt, and then the low-rank adapters (B and A) are trained on top in a higher precision.
- Multi-Task LoRA: For models handling diverse instructions, multiple adapters (Bk and Ak) can be used. Gating mechanisms or "mixture of experts" logic can determine which adapter should contribute to a specific input instance.
Managing Long Context and Agentic Trajectories
- The Agentic Problem: Agents operating in environments accumulate trajectories of actions, thoughts, and observations. As the trajectory grows, the number of tokens increases, making self-attention (which is quadratic in complexity) extremely expensive.
- Computational Hardness: Attending to 20 past events is significantly harder and more expensive than attending to 2.
- Context Compression:
* Summarization: Periodically using an LLM to compress the trajectory into a smaller summary state. Performance often improves with compression (e.g., from an accuracy of 56 to higher).
* Selective Compression: Strategies include summarizing the whole history or only environmental observations while keeping the agent's internal "reasoning" thoughts intact.
- External Memory Systems:
* Instead of keeping everything in the LLM's active context window, information is stored in an external memory (like a file or database).
* Retrieval-Generation Loop: The agent queries the external memory, retrieves relevant past data (e.g., a previous login attempt), and uses it to generate the next action. This mimics a human checking a worksheet.
- KV Caching and Policy: To avoid re-computing keys and values for every token, systems use caching. Efficient systems employ predictive caching policies to determine which tokens to keep in high-bandwidth memory and which to evict.
Systems and Hardware Interaction
- Flash Attention: A system-level optimization that reorganizes the computation of the attention matrix to exploit memory hierarchies, avoiding the need to fully enumerate the massive attention matrix in memory.
- Sparsity in Attention:
* Research shows that attention is often "sparse," meaning many tokens have near-zero attention values.
* Sparsity Thresholds: Experiments show that accuracy can be maintained even if the attention matrix is 90% sparse (meaning 90% of the values are effectively zero).
- Hardware and Tensor Cores:
* Modern hardware like FPGAs and GPUs use Tensor Cores for matrix multiplication.
* Block-Based Sparsity: To exploit sparsity, indices must be sorted and packaged such that non-zero blocks are sent to processing units efficiently.
* If zeros are distributed too uniformly, hardware cannot skip the operations, and no speedup is realized.
Questions & Discussion
- Nonlinearity Placement in Adapters: A student asked why the nonlinearity is placed inside the bottleneck rather than at the end. The instructor explained that without the nonlinearity during the transformation, the adapter would reduce to a single linear layer and lose the capacity for task-specific adaptation.
- Frozen Layer Disconnect: There was a discussion on whether frozen layers with large-magnitude vectors might conflict with new adapters. The response was that the adapter must learn to transform its output back into the space expected by the subsequent frozen layers. This is often managed via specific initializations like zero-initialization or identity-initialization.
- Training Speed vs. Epochs: In response to a query about training iterations, the instructor noted that fewer parameters generally mean fewer training epochs and less data are needed, but there is a trade-off regarding how deeply the model learns concepts versus heuristics.
- Adapter Placement: A participant asked about the benefit of prefix tuning over simply adding layers at the end. The instructor clarified that prefix tuning leverages the transformer's built-in attention mechanism to extract information more naturally from the task-specific guidance, acting like a "continuous" version of prompting (zero-shot instructions).