MOONSHOT : A Framework for Multi-Objective Pruning of Vision and Large Language Models

TL;DR

MOONSHOT is a flexible framework that enhances neural network pruning by jointly optimizing multiple objectives, leading to better compression and performance on large models like Llama and Vision Transformers.

Contribution

It introduces a multi-objective pruning approach that extends existing methods, improving efficiency and effectiveness without retraining, applicable to billion-parameter models.

Findings

Reduces perplexity by up to 32.6% on Llama-3.2 with 2:4 sparsity.

Improves zero-shot accuracy by up to 4.9 points across seven benchmarks.

Enhances ImageNet accuracy by over 5 points at 70% sparsity.

Abstract

Weight pruning is a common technique for compressing large neural networks. We focus on the challenging post-training one-shot setting, where a pre-trained model is compressed without any retraining. Existing one-shot pruning methods typically optimize a single objective, such as a layer-wise reconstruction loss or a second-order Taylor approximation of the training loss. We highlight that neither objective alone is consistently the most effective across architectures and sparsity levels. Motivated by this insight, we propose MOONSHOT, a general and flexible framework that extends any single-objective pruning method into a multi-objective formulation by jointly optimizing both the layer-wise reconstruction error and second-order Taylor approximation of the training loss. MOONSHOT acts as a wrapper around existing pruning algorithms. To enable this integration while maintaining scalability to billion-parameter models, we propose modeling decisions and introduce an efficient procedure for computing the inverse Hessian, preserving the efficiency of state-of-the-art one-shot pruners. When combined with state-of-the-art pruning methods on Llama-3.2 and Llama-2 models, MOONSHOT reduces C4 perplexity by up to 32.6% at 2:4 sparsity and improves zero-shot mean accuracy across seven classification benchmarks by up to 4.9 points. On Vision Transformers, it improves accuracy on ImageNet-1k by over 5 points at 70% sparsity, and on ResNet-50, it yields a 4-point gain at 90% sparsity.