Don't Be Greedy, Just Relax! Pruning LLMs via Frank-Wolfe
This work offers a more efficient pruning method for LLMs, reducing computational and storage requirements, though it is incremental as it builds on existing layer-wise pruning approaches.
The paper tackles the problem of pruning large language models (LLMs) by addressing the combinatorial optimization challenge in layer-wise pruning, proposing a method based on convex relaxation and the Frank-Wolfe algorithm that reduces per-layer pruning error and outperforms baselines on GPT architectures.
Pruning is a common technique to reduce the compute and storage requirements of Neural Networks. While conventional approaches typically retrain the model to recover pruning-induced performance degradation, state-of-the-art Large Language Model (LLM) pruning methods operate layer-wise, minimizing the per-layer pruning error on a small calibration dataset to avoid full retraining, which is considered computationally prohibitive for LLMs. However, finding the optimal pruning mask is a hard combinatorial problem and solving it to optimality is intractable. Existing methods hence rely on greedy heuristics that ignore the weight interactions in the pruning objective. In this work, we instead consider the convex relaxation of these combinatorial constraints and solve the resulting problem using the Frank-Wolfe (FW) algorithm. Our method drastically reduces the per-layer pruning error, outperforms strong baselines on state-of-the-art GPT architectures, and remains memory-efficient. We provide theoretical justification by showing that, combined with the convergence guarantees of the FW algorithm, we obtain an approximate solution to the original combinatorial problem upon rounding the relaxed solution to integrality.