PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Stephen Zhang, Vardan Papyan

arXiv:2409.13652v3cs.LGcs.AI

TL;DR

Large transformers are costly to train and deploy, while pruning methods often lose performance as compression increases. OATS uses input-activation second moments to decompose weights into sparse and low-rank terms, and reports strong performance across models and modalities.

  • Problem

    Large foundation models are costly to train and deploy, and existing pruning methods often lose performance as compression increases.

    Source: §1 · p. 1

  • Method

    OATS scales weights using input-activation second moments and approximates them as a sum of sparse and low-rank matrices.

    Source: §2 · p. 2, §2.3 · p. 3

  • Results

    Across models, modalities, compression rates, and performance metrics, OATS consistently outperforms prior state of the art; at 50% compression, its five-shot MMLU accuracy margin on Phi-3 Mini is 5.42%.

    Source: §3.2 · p. 5, §7 · p. 10

  • Takeaways & Limitations

    Vision-transformer visualizations indicate that sparse and low-rank components capture different image segments.

    Source: §7 · p. 10

    The visualized component behaviors are not fully understood; further investigation is needed to identify the mechanisms driving them.

    Source: §5 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Large transformers are costly to train and deploy, and pruning methods lose performance as compression increases. OATS addresses this with outlier-scaled sparse-plus-low-rank decomposition, evaluated across language and vision models.

1 INTRODUCTION

  • Billions-parameter transformers succeed across domains but are computationally expensive to train and deploy.Source: §1 · p. 1
  • As compression rises, existing pruning methods consistently lose performance (Yin et al., 2024a), while structured pruning can trade away more accuracy for acceleration (Chen et al., 2022).Source: §1 · p. 1

1.1 CONTRIBUTIONS

  • OATS approximates weights as sparse plus low-rank matrices and scales them by the second moment of their input embeddings to emphasize outliers.Source: §1.1 · p. 2
  • Across Phi-3, Llama-3, ViT, and DINOv2, OATS achieves state-of-the-art performance and improves CPU inference acceleration over unstructured pruning.Source: §1.1 · p. 2
  • In vision transformers, sparse and low-rank components focus on different image regions, effectively segmenting images into distinct areas.Source: §1.1 · p. 2

2 THE OATS ALGORITHM

OATS approximates transformer weights with sparse and low-rank terms using alternating thresholding. It scales weights according to input-activation second moments, then selects the decomposition rank and sparsity from two hyperparameters.

2 THE OATS ALGORITHM

  • OATS frames each transformer weight matrix as a robust-PCA-style sum of sparse and low-rank matrices.Source: §2 · p. 2

2.1 ALTERNATING THRESHOLDING

  • Alternating thresholding iteratively estimates the low-rank term by truncated SVD and the sparse term by hard-thresholding, storing the low-rank term as two factors.Source: §2.1 · p. 2

2.2 ALTERNATIVE SPARSITY PATTERNS

  • Row-wise thresholding keeps the largest-magnitude entries per row and has been reported to improve performance (Sun et al., 2024b); N:M sparsity can exploit NVIDIA sparse tensor cores (Mishra et al., 2021).Source: §2.2 · p. 3

2.3 INCORPORATING OUTLIER INFORMATION

  • OATS scales W by a diagonal second-moment matrix derived from calibration activations, emphasizing outlier features in reconstruction; it decomposes WD≈S+L and applies D^-1 afterward.Source: §2.3 · p. 3

2.4 OATS PARAMETERS

  • Compression rate ρ and rank ratio κ determine the retained nonzeros and low-rank rank; κ specifies the share of compressed-layer parameters assigned to the low-rank term.Source: §2.4 · p. 3

3 EXPERIMENTS ON LARGE LANGUAGE MODELS

OATS is evaluated against post-training pruning methods on language-model accuracy and perplexity, with additional ablations and hardware comparisons. It outperforms prior methods across the reported compression settings, while the studies examine design choices, hyperparameters, and speed.

3.1 EXPERIMENT SETUP

  • Evaluations cover Phi-3 and Llama-3 on five-shot MMLU, eight zero-shot tasks, and WikiText-2, with SparseGPT, Wanda, and DSNoT as pruning baselines.Calibration uses 128 C4 sequences of length 2048; compression rates are 0.3, 0.4, 0.5, and 0.6, with OWL ratios at 0.6.Source: §3.1 · p. 4

3.2 RESULTS

  • At 50% compression, OATS surpasses prior pruning algorithms on five-shot MMLU by 5.42% on Phi-3 Mini, 2.52% on Phi-3 Medium, 2.86% on Llama-3 8B, and 2.03% on Llama-3 70B.OATS outperforms prior methods across all reported compression rates, with a widening gap as compression increases.Source: Table 2 · p. 5, §3.2 · p. 5
  • At 50% compression, OATS has a 2.05% advantage over prior pruning algorithms on Phi-3 Mini averaged across eight zero-shot tasks.The improvement over prior methods increases with compression.Source: Table 3 · p. 5, §3.2 · p. 5
  • At 50% compression, OATS reduces WikiText-2 perplexity by 8.49% on Phi-3 Medium, 8.99% on Phi-3 Mini, 9.04% on Llama-3 8B, and 9.30% on Llama-3 70B.Table 4 compares perplexity, for which lower values are better.Source: Table 4 · p. 6, §3.2 · p. 6
  • At 60% compression using OWL ratios, OATS exceeds prior methods on five-shot MMLU by 6.39% on Phi-3 Mini, 5.61% on Phi-3 Medium, and 4.98% on Llama-3 8B.This evaluation uses the higher compression rate of 60%.Source: Table 5 · p. 6, §3.2 · p. 6

3.3 STUDIES AND HYPERPARAMETER EXPLORATION

  • At 40% compression and a 20% rank ratio on Phi-3 Mini, second-moment scaling and row-wise pruning each significantly improve OATS performance.Table 6 reports the ablation results for these two design choices.Source: Table 6 · p. 6, §3.3 · p. 6
  • Rank ratios of 25–30% give the best performance; gains from increasing iterations are sharp through 20 iterations, then level off and saturate around 80.Figure 1 examines zero-shot and five-shot accuracy as rank ratio and iteration count vary.Source: Figure 1 · p. 7, §3.3 · p. 7

3.4 HARDWARE SPEEDUP

  • At 40% compression, OATS is 1.37× faster on CPU than unstructured pruning; with 2:8 sparsity, it exceeds prior methods by 6.34% on Phi-3 Mini at 50% compression and by 9.2% on Llama-3 8B at 50%.OATS uses 2:8 structured sparsity on the sparse term, compared with 2:4 in the prior methods; varying the rank ratio allows flexible compression-performance trade-offs.Source: §3.4 · pp. 7–8, Figure 2 · p. 8

4 EXPERIMENTS ON VISION TRANSFORMERS

The vision experiments compare OATS with SparseGPT, Wanda, and DSNoT on ImageNet top-1 accuracy, finding a 0.41% drop for DinoV2 at 50% compression.

4 EXPERIMENTS ON VISION TRANSFORMERS

  • Experiments cover supervised ViT-Base and self-supervised DinoV2-Giant, with their model sizes and training settings differing.ViT-Base has 86.6M parameters and DinoV2-Giant has 1.14B parameters.Source: §4 · p. 8
  • The comparison measures ImageNet validation top-1 accuracy, using the same 2,048 training images for calibration across pruning experiments.OATS uses rank ratio κ=20% and N=80 iterations; embedding and classifier layers are excluded from compression.Source: §4 · p. 8
  • 0.41% top-1 accuracy drop for DinoV2 when compressed by 50% using OATS.Table 8 reports ImageNet validation accuracy.Source: Table 8 · p. 8, §4 · p. 8

5 VISUALIZING AND INTERPRETING THE DECOMPOSITION

The paper uses attention rollout to examine sparse-only and low-rank-only information flow in compressed vision transformers, finding that the components attend to distinct image regions.

5 VISUALIZING AND INTERPRETING THE DECOMPOSITION

  • Attention rollout isolates sparse-only and low-rank-only information flow; Figure 3 visualizes the two settings in a standard transformer block.Each setting zeros one component and propagates inputs through the other.Source: §5 · p. 9, Figure 3 · p. 9
  • Figure 4 shows rollout for COCO images passed to ViT-B compressed by 50%, with rank ratio 20%.The figure presents attention rollout visualizations for various Microsoft COCO images.Source: §5 · p. 9, Figure 4 · p. 9
  • Rollouts reveal three patterns: sparse terms capture subjects while low-rank terms capture backgrounds; components split one subject’s features; or isolate different subjects.The visualizations show the components attending to distinct image areas, effectively segmenting them.Source: §5 · pp. 9–10
  • The observed patterns offer initial insight, but their underlying mechanisms require further investigation.Source: §5 · p. 10

6 RELATED WORKS

Related work situates OATS alongside pruning, sparse and low-rank transformer methods, Robust PCA algorithms, and studies of pruning’s interpretability and performance effects.

Connection with Wanda

  • At κ=0, OATS reduces to Wanda’s scaled hard-thresholding pruning step (Sun et al., 2024b).The paper describes Wanda as a special case of OATS under this rank-ratio setting.Source: §6 · p. 10

Sparse and Low-Rank Approximation in Transformers

  • Sparse and low-rank transformer work spans theoretical results on logits (Zhao et al., 2024c), attention approximation (Chen et al., 2021a), and LoRAP pruning (Li et al., 2024).Source: §6 · p. 10

Structured Pruning and Low-Rank Adaptation

  • LoSparse, LoRAPrune, and APT combine structured pruning with trained low-rank adapters, unlike OATS’s retraining-free sparse-plus-low-rank decomposition (Li et al., 2023; Zhang et al., 2024a; Zhao et al., 2024a).Source: §6 · p. 10

Robust PCA Algorithms

  • Robust PCA alternatives include convex relaxation and gradient descent on low-rank factors; OATS uses alternating thresholding for simplicity, while other algorithms remain a future direction.The cited alternatives include Zhou et al. (2010), Yi et al. (2016), and Tong et al. (2021).Source: §6 · p. 10

Pruning and Interpretability

  • The authors postulate that OATS’s low-rank term may mitigate pruning harms; larger gaps over prior methods at higher compression suggest it plays a critical role.The cited discussion connects those gaps to Tables 2 and 5.Source: §6 · p. 10

7 CONCLUSION

The conclusion summarizes OATS’s training-free sparse–low-rank transformer compression, its reported performance across models and settings, and how its inference focus differs from prior approaches.

7 CONCLUSION

  • OATS consistently outperforms prior state of the art across performance metrics, compression rates, models, and modalities while improving CPU speed-up.The evaluation spans large-scale transformers across language and vision tasks.Source: §7 · p. 10
  • Without retraining, OATS decomposes model weight matrices into sparse and low-rank components after scaling weights by input-embedding second moments and applying alternating thresholding.The approach is inspired by prior work on the emergence of outlier features.Source: §7 · p. 10
  • Vision-transformer visualizations show sparse and low-rank structures capturing different image segments, revealing their potential for transformer efficiency, performance, and interpretability.The paper presents this as the first demonstration of sparse and low-rank decompositions for large-scale transformers.Source: §7 · p. 10

A.1 ADDITIONAL RELATED WORKS

  • OATS avoids reconstruction-error minimization because it can degrade large-transformer performance around outlier features; it instead uses a simpler formulation than Yu et al. (2017).The passage says pruning weights within outlier channels should be avoided; this issue was not relevant to Yu et al.’s pre-transformer work.Source: §A.1 · pp. 18–19
  • Vision-transformer pruning methods are not comparable to OATS here because prominent approaches often require parameter training, whereas OATS requires none.The cited literature includes Zhu et al. (2021), Chen et al. (2021b), Chavan et al. (2022), Yu et al. (2022a;b), and Yu & Xiang (2023); its models are mostly much smaller than the LLMs studied.Source: §A.1 · p. 19
  • Unlike SLOPE, which trains low-rank components during pre-training, OATS approximates pre-trained weights with robust PCA for training-free inference acceleration.SLOPE adds low-rank components alongside N:M sparsity in the final pre-training phase and trains them on the model loss.Source: §A.1 · p. 19
  • SLIM combines pruning and quantization, using average absolute values and one pruning/quantization step followed by singular-value thresholding, unlike OATS’s L2 norm and alternating thresholding.Both methods use a scaling technique and a low-rank term to help restore performance after compression.Source: §A.1 · p. 19

A.2 TIME COMPLEXITY AND WALL-CLOCK TIME FOR OATS

OATS's compression time depends on transformer depth, iteration count, and singular-value-thresholding cost, and is higher than prior pruning algorithms. The authors report that reducing iterations preserves performance gains, while parallel compression can reduce per-block time.

A.2 TIME COMPLEXITY AND WALL-CLOCK TIME FOR OATS

  • OATS has time complexity O(LNα), where L is transformer-block count, N is iteration count, and α is the cost of singular-value thresholding.Table 9 reports wall-clock time for one alternating-thresholding iteration on one transformer block across the compressed models.Source: §A.2 · p. 19, Table 9 · p. 20
  • OATS requires more wall-clock time than prior pruning algorithms, but compression is performed once before deployment and can be parallelized across layers in a block.Compressing a Llama-3 70B block in parallel across four NVIDIA A40 GPUs takes 71.10 seconds.Source: §A.2 · p. 20
  • With 20 iterations—one quarter of the iterations—OATS still outperforms all prior pruning algorithms across performance metrics.The experiment compresses Llama-3 70B by 50% with a rank ratio of 0.3; fewer iterations also reduce total wall-clock time.Source: §A.2 · p. 20, Table 10 · p. 20

A.3 USING A ROBUST SCALING MATRIX

These ablations examine the scaling matrix and the order and form of thresholding in OATS. The reported findings favor outlier-sensitive scaling and the original thresholding order; the magnitude-based sparse-component alternative is described as an experiment without a result in these passages.

A.3 USING A ROBUST SCALING MATRIX

  • Robust scaling produces significantly worse performance than OATS's outlier-sensitive scaling, which captures feature scale.The comparison uses Phi-3-Mini at 50% compression and a 25% rank ratio.Source: §A.3 · p. 20, Table 11 · p. 20

A.4 SWITCHING THE ORDER OF THRESHOLDING

  • Switching to hard-thresholding before singular-value thresholding remains competitive but falls short of OATS's original order.The ablation uses Phi-3-Mini at 40% compression and a 20% rank ratio.Source: §A.4 · p. 21, Table 12 · p. 21

A.5 MAGNITUDE-BASED PRUNING FOR THE SPARSE COMPONENT

  • The magnitude-based pruning ablation tests whether the low-rank term can capture all outlier information while hard-thresholding selects the sparse term without scaling.The described setup is evaluated on Phi-3-Mini at 40% compression with a 20% rank ratio.Source: §A.5 · p. 21, Table 13 · p. 21

A.6 ADDITIONAL CPU SPEEDUP EXPERIMENTS

The appendix adds CPU-speedup, hyperparameter, and low-rank-term comparisons. It notes that OATS and unstructured pruning yield comparable CPU speedups in the additional experiment, while structure-specific optimizations may offer further gains.

A.6 ADDITIONAL CPU SPEEDUP EXPERIMENTS

  • OATS and unstructured pruning deliver comparable CPU speedups in the additional 256-token Phi-3 Medium 15B experiment.The comparison measures throughput and speedup relative to dense models; the authors expect further gains from optimizations that exploit OATS's sparse-plus-low-rank structure.Source: §A.6 · pp. 21–22, Table 14 · p. 21

A.7 ADDITIONAL HYPERPARAMETER TESTS FOR OATS

  • Additional hyperparameter configurations are tested for OATS on Phi-3 Mini and Llama-3 8B.The passage identifies Table 15 as reporting these experiments but does not state their outcomes.Source: §A.7 · p. 22, Table 15 · p. 22

A.8 PERFORMANCE GAP BETWEEN OATS AND WANDA

  • The appendix compares OATS with Wanda to assess the impact of including a low-rank term.Table 16 is identified as reporting the performance gaps, without results stated in these passages.Source: §A.8 · p. 22, Table 16 · p. 22

A.9 QWEN 2.5 EXPERIMENTS

The appendix reports OATS benchmarks on Qwen 2.5 3B Instruct, specifies the evaluated MMLU subject subset, and describes how attention-rollout visualizations are generated.

A.9 QWEN 2.5 EXPERIMENTS

  • OATS is benchmarked against prior pruning algorithms on Qwen 2.5 3B Instruct with rank ratio 0.2 and 80 iterations.Table 17 presents the benchmarks.Source: §A.9 · p. 22, Table 17 · p. 23

A.10 MMLU SUBJECTS

  • The evaluated MMLU subjects include College Computer Science, aligning with the subset used in Ashkboos et al. (2024).The referenced subset is available in their TransformerCompression codebase.Source: §A.10 · p. 23

A.11 ATTENTION ROLLOUT: DETAILS

  • Attention-rollout visualizations average attention matrices across heads and discard the lowest 40% of attention pixels.Discarding the lowest-valued pixels was inspired by Gil (2021).Source: §A.11 · p. 23

A.12 ZERO-SHOT TASK-SPECIFIC PERFORMANCE

The appendix points to task-specific zero-shot results and an alternative Wikitext-2 perplexity evaluation, then documents implementation and hyperparameter settings, including DSNoT variants for language and vision models.

A.12 ZERO-SHOT TASK-SPECIFIC PERFORMANCE

  • Table 18 presents task-specific performance for the zero-shot evaluations reported in Section 3.2 and Appendix A.14.The supplied passage does not state the individual task results.Source: §A.12 · p. 23, Table 18 · p. 24

A.13 WIKITEXT-2 PERPLEXITY

  • Table 19 reports Wikitext-2 perplexities using the alternative normalization and dataset segmentation adopted by prior works.Perplexity is lower-is-better; the implementation is sourced from Sun et al. (2024b).Source: §A.13 · p. 24, Table 19 · p. 25

A.14 PRUNING IMPLEMENTATION AND HYPERPARAMETERS

  • The experiments use implementations based on the listed codebases and Huggingface Transformers for language models and vision transformers.The supplied passages do not enumerate the codebase names in this point.Source: §A.14 · p. 25

A.14.1 SPARSEGPT HYPERPARAMETERS

  • SparseGPT uses blocksize 128 and Hessian dampening of 0.01, with a ladder of 0.1 used only for non-positive-definiteness issues in Cholesky decomposition.These settings apply across experiments.Source: §A.14.1 · p. 25

A.14.2 DSNOT HYPERPARAMETERS

  • DSNoT uses SparseGPT- or Wanda-generated initial masks, 50 iterations, and an update threshold of 0.1; Tables 20 and 21 report language-model and vision results by initial method.Table 20 gives LLM performance metrics, while Table 21 gives ImageNet validation accuracy.Source: §A.14.2 · pp. 25–26, Table 20 · p. 26, Table 21 · p. 26

Related papers

Loading 2409.13652v3…