Don't Be Greedy, Just Relax! Pruning LLMs via Frank-Wolfe

TL;DR

This work addresses the challenge of pruning large language models without full retraining by turning the combinatorial mask selection into a convex optimization problem over the mask's hull and solving it with the projection-free Frank-Wolfe algorithm. The proposed sparsefw method accounts for interactions between weights, supports unstructured and semi-structured sparsity, and achieves substantial reductions in per-layer pruning error with strong empirical gains on modern GPT architectures. The authors provide a data-dependent theoretical bound linking the relaxed FW solution to an approximate integral mask, and demonstrate practical improvements in perplexity and zero-shot accuracy, albeit with limitations from local-global objective mismatch that can be mitigated by fixing a fraction of salient weights. Overall, sparsefw offers a scalable, theoretically grounded alternative to greedy pruning for large-scale transformers, with meaningful implications for efficiency in deployment.

Abstract

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.

Figures (4)

• Figure 1: Visualization of $\mathcal{C}_k$ for $d_{\text{out}}=3$, $d_{\text{in}}=1$. Left: $k=1$, Right: $k=2$.
• Figure 2: LLaMA-3.1-8B pruned to 60% unstructured sparsity with using Wanda warmstart with 256 samples. This figure shows the relative reduction in pruning error (y-axis) for each matrix type (see legend for colors) for all layers of the model (x-axis) compared to the warmstart mask.
• Figure 3: LLaMA-3.1-8B pruned to $2{:}4$ sparsity using . Left: Perplexity over the number of iterations per layer with 256 samples. Right: Perplexity over the number of calibration samples with 2000 iterations per layer. The solid curve represents the mean over multiple random seeds, the shaded regions represent the min-max range.
• Figure 4: LLaMA-3.1-8B optimized towards 60% unstructured sparsity with using 256 calibration samples. Lightly colored curves show the results individual matrices; the solid curve is their median. Left: Relative pruning error reduction versus iterations for continuous and thresholded masks. Right: Average threshold residual (mean $\ell_1$ distance between continuous and thresholded masks) versus iterations.

Theorems & Definitions (3)

• Lemma 1: Informal
• Lemma 2: Formal statement of \ref{['lem:sparsefw-optimality-informal']}
• proof : Proof of \ref{['lem:sparsefw-optimality-formal']}