Source-linked AI summary
Structured Pruning of Large Language Models
Ziheng Wang, Jeremy Wohlwend, Tao Lei
TL;DR
Large language models are costly because their size and latency continue to grow, motivating the question of whether they need to remain large. The paper introduces FLOP, which adaptively prunes low-rank components and achieves strong compression, competitive performance, and speedups across language-modeling settings. It also applies FLOP to adaptive embeddings and softmax layers.
Problem
Growing model size and latency make large language models slow and expensive, motivating study of model capacity through compression.
Method
FLOP parameterizes weight matrices with adaptive low-rank factorization and learns which rank-1 components to remove during training.
Results
FLOP outperforms block-structured and unstructured pruning baselines while achieving significant compression and speedups; one language model loses 0.8 perplexity at 50% compression and exceeds 2x speed-up.
Takeaways & Limitations
Structured low-rank pruning reduces the overhead of large language models while preserving performance across evaluated language-modeling and language-understanding settings.
Takeaways & Limitations
The l0 regularization approach alone does not reliably control final model size, which can vary with learning rates or pruning schedules.
Abstract
from arXiv · showhide
Large language models have recently achieved state of the art performance across a wide variety of natural language tasks. Meanwhile, the size of these models and their latency have significantly increased, which makes their usage costly, and raises an interesting question: do language models need to be large? We study this question through the lens of model compression. We present a generic, structured pruning approach by parameterizing each weight matrix using its low-rank factorization, and adaptively removing rank-1 components during training. On language modeling tasks, our structured approach outperforms other unstructured and block-structured pruning baselines at various compression levels, while achieving significant speedups during both training and inference. We also demonstrate that our method can be applied to pruning adaptive word embeddings in large language models, and to pruning the BERT model on several downstream fine-tuning classification benchmarks.
1 Introduction
Large language models are increasingly costly because of their size and latency. FLOP uses adaptive low-rank factorization to prune them structurally, preserving dense computation while achieving strong compression, performance, and speedups.
- Motivation: Unstructured pruning reduces model size but produces sparse matrices that are difficult to accelerate on common hardware.Structured pruning improves speed but often performs worse than unstructured pruning at the same parameter budget.
- Method: FLOP parameterizes weight matrices with adaptive low-rank factorization and removes rank-1 components during training.The method applies to any matrix multiplication and uses magnitude pruning or improved l0 regularization.
- Adaptive Embeddings: FLOP dynamically learns embedding dimensions for different word clusters, improving the parameter-reduction versus performance trade-off for large vocabulary layers.Input embeddings and softmax output layers can contain most language-model parameters when vocabulary size is large.
- Results: FLOP outperforms block-structured and unstructured pruning baselines while improved l0 regularization further improves performance in most cases.The method is evaluated on Wiki-103, Enwiki8, and GLUE using recurrent networks and Transformers.
- Results: 50% compression costs only 0.8 perplexity, while training and inference achieve over 2x speed-up without additional hardware or software requirements.The result is reported for a large word-level language model with adaptive embeddings.
2 Related Work
Prior compression work includes pruning, distillation, and quantization, with pruning spanning unstructured and structured approaches. FLOP differs through low-rank parameterization and augmented-Lagrangian control of compression.
- Compression Methods: Model compression research is commonly organized around weight pruning, knowledge distillation, and quantization.These techniques have been applied across architectures and NLP applications.
- Combining Techniques: Compression techniques can be combined, and FLOP is compatible with quantization and distillation because it applies to matrix multiplications.The cited related work notes that combining techniques can achieve greater compression.
- Pruning Methods: Prior pruning methods include magnitude-based unstructured pruning, dropout, and structured pruning with thresholding or l1 regularization.FLOP instead uses low-rank parameterization for compression.
3 Background
The pruning problem introduces variables that determine which weights or blocks remain, with the effective model size measured by the l0 norm. A differentiable re-parameterization enables joint optimization of pruning variables and model parameters.
- Pruning Formulation: Pruning is modeled by associating binary variables with individual weights or blocks, such as weight-matrix columns.The pruned parameters are represented by multiplying model parameters with the pruning variables.
- Pruning Objective: The l0 norm counts retained parameters or blocks and therefore measures the effective size of the pruned model.The regularization coefficient λ encourages sparsity during training.
- Pruning Objective: The training objective minimizes expected loss plus l0 regularization over the pruned model parameters.Direct optimization is difficult because the pruning variables are discrete and have exponentially many configurations.
- Differentiable Optimization: The re-parameterization trick relaxes binary pruning variables into differentiable continuous variables generated from uniform samples through a learnable inverse CDF.Gradient optimization can then update both model parameters and re-parameterization parameters.
- Differentiable Optimization: The Hard Concrete stretch-and-rectify process places probability mass near 0 and 1, providing a relaxation of Bernoulli pruning variables.The expected l0 regularization can be computed in closed form and remains differentiable.
4 Method
FLOP uses low-rank factorization to prune rank-1 components while retaining dense matrix structure, and extends this strategy to adaptive embeddings with explicit compression control.
- Structured Pruning using Factorization: Structured pruning accelerates computation by removing parameter groups, but its restrictive patterns can reduce performance relative to unstructured pruning.Column pruning is described as producing speedups during training and inference while often achieving lower performance than unstructured pruning.
- Structured Pruning using Factorization: FLOP reparameterizes each weight matrix as a product of two smaller matrices and prunes rank-1 components through learned diagonal masks.The factorization is W = PQ, with each component formed by a column of P and a row of Q.
- Structured Pruning using Factorization: Unlike unstructured pruning, low-rank pruning preserves dense matrices, avoiding sparse-weight indices and specialized computation support.After pruning, the remaining factors are smaller but dense, enabling simple matrix multiplications on common hardware.
- Pruning Adaptive Embedding and Softmax Layer: FLOP automatically learns separate embedding dimensions for word clusters instead of requiring their reduced dimensions to be manually specified.The method uses a diagonal mask for each cluster and adjusts each cluster’s parameter budget during training.
- Augmented Lagrangian Method: The augmented Lagrangian enforces an equality between expected model size and a target size during pruning training.Its multipliers are jointly updated, and the updates increase training loss unless the equality constraint is met.
- Training and Inference: Shared stochastic pruning masks reduce the active computation in training, while expected mask values are converted into a deterministic inference mask.At inference, the largest expected mask values are retained to match the desired l0 norm.
5 Experimental Setup
The evaluation covers recurrent and Transformer-XL language modeling, adaptive embedding settings, and BERT classification, using matched pruning baselines and training configurations.
- Tasks: The Wiki-103 setup uses a 12-layer SRU with tied adaptive embedding and softmax layers, approximately half of whose parameters belong to adaptive layers.The base model contains 100M parameters in total.
- Tasks: The Transformer-XL setup introduces pruning in self-attention and feedforward matrices while preserving the unfactorized model’s parameter count before pruning.The model is a 12-layer base model evaluated on Enwik8.
- Baselines: The baselines include fixed-ratio factorization, neuron pruning, magnitude-based unstructured pruning, and factorized magnitude-based pruning.These alternatives isolate the contributions of low-rank factorization and l0 pruning.
- Experimental Controls: All methods use the same learning-rate and dropout configurations, while pruning schedules and Lagrangian-variable learning rates are tuned separately.The study compares unstructured, structured, and factorization-based pruning baselines.
6 Results
FLOP consistently improves compression results across recurrent, character-level, Transformer-XL, and downstream fine-tuning settings. It preserves performance while substantially reducing parameters and, in language modeling, can also deliver strong compact-model results.
- Word-level Language Model: A 50% parameter reduction with FLOP-l0 raises Wiki-103 test perplexity only from 24.5 to 25.3.The unpruned SRU base model has test perplexity 24.5; FLOP-l0 loses 0.8 perplexity at 50% compression.
- Word-level Language Model: FLOP adaptively reduces dimensions most aggressively for less-frequent words in adaptive embeddings.Figure 1 breaks parameter usage down across word-frequency clusters and model components.
- Char-level Language Model: A 70% compression level yields perplexity 1.25 versus 1.24 for the uncompressed character-level model.FLOP-l0 obtains the best performance across the reported pruning levels and nearly matches the uncompressed model.
- Transformer-XL models: On Enwiki8, FLOP-l0 outperforms the compared pruning methods and reaches 1.17 BPC with 4M parameters.The comparison includes FAC, unstructured AGP, FLOP-AGP, and FLOP-l0 at 80% compression, plus FLOP variants at 90%.
- BERT on Classification Tasks: On downstream fine-tuning, the method conserves nearly 99% of performance while reducing parameters by 35%.The attainable compression level is limited because embedding layers constitute a significant portion of the remaining parameters.
7 Analysis
The analysis tests whether factorization-based pruning preserves performance more effectively than input-feature pruning and measures its computational speedups and pruning dynamics.
- Analysis: The analysis examines factorization, performance retention, computational speed, and the training dynamics of pruning decisions.The section introduces an analysis of several aspects of the method.
- Factorization: A larger unfactorized NP-l0 model with 90% more parameters reaches 1.25 BPC, on par with the factorized base model.The comparison addresses whether factorization benefits arise only from using a larger initial hidden size.
- Factorization: Factorization-based pruning retains relative model performance more effectively than input feature pruning across four compression levels.Table 5 reports BPC and performance loss relative to the uncompressed model.
- Speed analysis: Inference speedups range from 1.5x to 2.2x on CPUs, while GPU training speedups reach 2.4x at the tested compression levels.These measurements are reported for FLOP and require no special hardware or software requirements.
- Learning dynamics: HardConcrete parameters become bimodal during training, separating parameters that are pruned from those retained with high certainty.Negative α indicates likely pruning, positive α indicates retention, and magnitude reflects decision certainty.
8 Conclusion
The paper presents adaptive low-rank factorization as a generic structured pruning method and evaluates it on large language models, finding speedups and compression with minimal performance loss.
- 8 Conclusion: The method provides significant speedups and compression rates on large models while losing minimal performance compared with methods including unstructured magnitude pruning.The conclusion frames these results as reducing the overhead of large language models and informing understanding of model capacity.
A.1 Optimization details
FLOP trains factorized models with warmup, gradually increases the pruning target, and jointly updates model parameters and Lagrangian multipliers while sampling shared batch masks.
- Optimization schedule: FLOP first trains the factorized model for warmup epochs and then begins pruning, while other pruning baselines use the same warmup process.The FAC baseline instead trains a smaller factorized model directly from scratch.
- Optimization schedule: The target size t is gradually increased at a linear rate toward the desired maximum size tmax.The sparsity is specified at each pruning iteration using an annealing schedule with m controlling the number of annealing steps.
- Optimization updates: Model parameters and Lagrangian multipliers receive joint gradient updates at every iteration, with a separately tuned multiplier learning rate.The multipliers are initialized to zero at the start of training.
- Mask sampling: A pruning mask z is sampled for each training batch and shared across its examples, allowing computation to focus on parameters active in that batch.The shared mask supports selecting currently active parameters for smaller computation.
A.2 Experimental Details
The experiments use Wiki-103, Enwik8, and Transformer-XL or SRU configurations with dataset-specific model sizes, factorization dimensions, training schedules, and hardware settings.
- Datasets: Experiments use standard train/dev/test splits from Wiki-103, Enwik8, and GLUE benchmarks.Training configurations are described in the experimental details.
- SRU: The Enwik8 SRU uses six layers, hidden size 3056, initial factorization dimension 512, 30 warmup epochs, and up to 100 pruning epochs.The model is trained with batch size 64 and unroll length 256.
- SRU: The Wiki-103 SRU uses 12 layers, hidden dimension 2048, factorization dimension 512, adaptive embedding dimensions 1024, 256, and 64, plus 50 warmup and 100 pruning epochs.Batch sizes are 64 or 96, with an unroll length of 256.
- Optimization settings: SRU runs use inverse-square-root learning-rate scheduling, while AGP and l0 regularization hyperparameters are tuned separately.The initial factor l0 is set to 2 or 3 for model parameters, and multiplier learning rates are tuned from 3 through 6.
- Transformer-XL: Transformer-XL pruning runs use up to 300k iterations, learning rate 0.00025, batch size 32, and four GPUs per run.The 12-layer base model uses up to 200k iterations, learning rate 0.0003, batch size 48, and eight GPUs.