A Proximal Operator for Inducing 2:4-Sparsity
This addresses the accuracy loss in sparse models for AI hardware acceleration, but it is incremental as it builds on existing sparsity methods.
The paper tackles the problem of decreased accuracy in models with 2:4 sparsity by deriving a regularizer and proximal operator to find better sparsity masks, improving over previous state-of-the-art on models up to 13B parameters and matching performance on 70B models.
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.