A Proximal Operator for Inducing 2:4-Sparsity

TL;DR

This paper introduces a novel regularizer and proximal operator to induce 2:4 sparsity in neural networks, improving pruning efficiency and accuracy, especially for large language models, by exploiting local feature correlations.

Contribution

We derive a proximal operator for 2:4 sparsity that enhances model pruning, achieving state-of-the-art results on large language models with minimal accuracy loss.

Findings

Improved pruning of models up to 13B parameters.

Matched state-of-the-art performance on 70B models.

Efficient solution for 2:4 sparsity regularization.

Abstract

Recent hardware advancements in AI Accelerators and GPUs allow to efficiently compute sparse matrix multiplications, especially when 2 out of 4 consecutive weights are set to zero. However, this so-called 2:4 sparsity usually comes at a decreased accuracy of the model. We derive a regularizer that exploits the local correlation of features to find better sparsity masks in trained models. We minimize the regularizer jointly with a local squared loss by deriving the proximal operator for which we show that it has an efficient solution in the 2:4-sparse case. After optimizing the mask, we use maskedgradient updates to further minimize the local squared loss. We illustrate our method on toy problems and apply it to pruning entire large language models up to 70B parameters. On models up to 13B we improve over previous state of the art algorithms, whilst on 70B models we match their performance.