Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

OATS: Outlier-Aware Pruning Through Sparse and Low Rank Decomposition

Stephen Zhang, Vardan Papyan

TL;DR

OATS tackles the challenge of compressing large transformers without retraining by decomposing weight matrices into a sparse part and a low-rank part, guided by outlier information from input activations. The method applies an alternating thresholding procedure to solve a Robust PCA objective, augmented by a diagonal scaling matrix derived from the second moment of inputs. Across both large language models (Phi-3, Llama-3) and vision transformers (ViT, DinoV2), OATS achieves state-of-the-art performance at up to $60\%$ compression and provides substantial CPU speedups by blending structured and unstructured pruning. Visual analyses show the sparse and low-rank components attend to different image regions, suggesting complementary roles in information flow. Overall, OATS demonstrates that sparse+low-rank decompositions capture meaningful, outlier-driven structure in transformers, enabling efficient deployment with minimal performance loss.

Abstract

The recent paradigm shift to large-scale foundation models has brought about a new era for deep learning that, while has found great success in practice, has also been plagued by prohibitively expensive costs in terms of high memory consumption and compute. To mitigate these issues, there has been a concerted effort in post-hoc neural network pruning techniques that do not require costly retraining. Despite the considerable progress being made, existing methods often exhibit a steady drop in model performance as the compression increases. In this paper, we present a novel approach to compressing large transformers, coined OATS, that utilizes the second moment information in the input embeddings to decompose the model weights into a sum of sparse and low-rank matrices. Without any retraining, OATS achieves state-of-the-art performance when compressing models by up to $60\%$ on large language models such as Llama-3 and Phi-3 and vision transformers such as ViT and DINOv2 while delivering up to $1.37\times$ the CPU acceleration versus a model that was comparably pruned.

OATS: Outlier-Aware Pruning Through Sparse and Low Rank Decomposition

TL;DR

OATS tackles the challenge of compressing large transformers without retraining by decomposing weight matrices into a sparse part and a low-rank part, guided by outlier information from input activations. The method applies an alternating thresholding procedure to solve a Robust PCA objective, augmented by a diagonal scaling matrix derived from the second moment of inputs. Across both large language models (Phi-3, Llama-3) and vision transformers (ViT, DinoV2), OATS achieves state-of-the-art performance at up to compression and provides substantial CPU speedups by blending structured and unstructured pruning. Visual analyses show the sparse and low-rank components attend to different image regions, suggesting complementary roles in information flow. Overall, OATS demonstrates that sparse+low-rank decompositions capture meaningful, outlier-driven structure in transformers, enabling efficient deployment with minimal performance loss.

Abstract

The recent paradigm shift to large-scale foundation models has brought about a new era for deep learning that, while has found great success in practice, has also been plagued by prohibitively expensive costs in terms of high memory consumption and compute. To mitigate these issues, there has been a concerted effort in post-hoc neural network pruning techniques that do not require costly retraining. Despite the considerable progress being made, existing methods often exhibit a steady drop in model performance as the compression increases. In this paper, we present a novel approach to compressing large transformers, coined OATS, that utilizes the second moment information in the input embeddings to decompose the model weights into a sum of sparse and low-rank matrices. Without any retraining, OATS achieves state-of-the-art performance when compressing models by up to on large language models such as Llama-3 and Phi-3 and vision transformers such as ViT and DINOv2 while delivering up to the CPU acceleration versus a model that was comparably pruned.
Paper Structure (55 sections, 14 equations, 5 figures, 21 tables, 2 algorithms)

This paper contains 55 sections, 14 equations, 5 figures, 21 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Contributions
  3. The OATS Algorithm
  4. Alternating Thresholding
  5. Alternative Sparsity Patterns
  6. Row-Wise Thresholding
  7. N:M Sparsity
  8. Incorporating Outlier Information
  9. OATS Parameters
  10. Experiments on Large Language Models
  11. Experiment Setup
  12. Models and Tasks
  13. Pruning Benchmarks
  14. Calibration Data
  15. Layer-Wise Compression Rates
  16. ...and 40 more sections

Figures (5)

  • Figure 1: The effect of varying the rank ratio and number of iterations on zero-shot and five-shot accuracy.
  • Figure 2: Experiments evaluating OATS with 2:8 structured sparsity on the sparse terms against 2:4 sparsity of state-of-the pruning algorithms. The rank ratio for OATS is varied to capture the performance across different compression rates.
  • Figure 3: A visualization of how the attention rollout is computed to isolate the contribution of the sparse terms versus low-rank terms given by the OATS algorithm.
  • Figure 4: Attention rollout visualization applied to various images on the Microsoft COCO dataset.
  • Figure : AlternatingThresholding