PaperlayerView exact revision on arXiv

Source-linked AI summary

Thanos: A Block-wise Pruning Algorithm for Efficient Large Language Model Compression

Ivan Ilin, Peter Richtarik

arXiv:2504.05346v1cs.LGcs.AIcs.CLcs.PF

TL;DR

Large language models require substantial memory and computation, making deployment costly. Thanos is a post-training, block-wise pruning method that needs no retraining; it surpasses competing methods in structured pruning and outperforms existing methods on several models in unstructured pruning.

  • Problem

    Large language models have large memory footprints and high inference complexity, requiring many powerful, expensive GPUs to train and run.

    Source: §1 · p. 1

  • Method

    Thanos is a post-training pruning method that removes weights block-wise without requiring retraining or training from scratch.

    Source: §1 · p. 1, §1.3 · p. 2

  • Results

    Thanos surpasses other widely used methods in structured pruning and achieves lower WikiText perplexity than SparseGPT, Wanda, and Magnitude pruning on TinyLlama 1.3B and LLaMA-3 models from 3B to 70B parameters.

    Source: §5.2 · p. 7

  • Takeaways & Limitations

    Thanos offers pruning as an alternative to quantization for reducing model size without retraining.

    Source: §1 · p. 1

    Unless stated otherwise, α=0.1; this lowers semi-structured sparsity from 0.5 to 0.45, while structured sparsity remains unchanged.

    Source: §5.1 · p. 7

Abstract

from arXiv · show

This paper presents Thanos, a novel weight-pruning algorithm designed to reduce the memory footprint and enhance the computational efficiency of large language models (LLMs) by removing redundant weights while maintaining accuracy. Thanos introduces a block-wise pruning strategy with adaptive masks that dynamically adjust to weight importance, enabling flexible sparsity patterns and structured formats, such as $n:m$ sparsity, optimized for hardware acceleration. Experimental evaluations demonstrate that Thanos achieves state-of-the-art performance in structured pruning and outperforms existing methods in unstructured pruning. By providing an efficient and adaptable approach to model compression, Thanos offers a practical solution for deploying large models in resource-constrained environments.

1. Introduction

LLMs’ memory and inference costs motivate compression, and Thanos addresses this with post-training, block-wise weight pruning. The paper also extends the method to outlier-row structured pruning and n:m sparsity.

1. Introduction

  • LLMs’ large memory footprint and inference complexity make deployment costly, motivating compression through quantization or pruning.The paper notes that training and running these models requires many powerful, expensive GPUs.Source: §1 · p. 1

1.1. The structure of LLMs

  • LLMs contain sequential transformation blocks with linear layers, and the paper focuses on pruning their weights because they represent most model parameters.Post-training pruning uses a pretrained, optimized model and a small calibration set; the paper describes processing blocks sequentially, with a forward pass before pruning and another after pruning to prepare the next block’s input.Source: §1.1 · p. 1, §1.2 · pp. 1–2

1.3. Contribution of this work

  • Thanos is a block-wise pruning algorithm that updates blocks of weights without significantly compromising performance.The paper presents it as a way to optimize and quickly update a weight block.Source: §1.3 · p. 2
  • Thanos retains a fraction of outlier rows during structured pruning, a strategy reported to improve perplexity and zero-shot results over competing baselines.Prior work cited in the paper indicates that a small number of outlier weights are important for inference quality.Source: §1.3 · p. 2
  • Thanos extends to semistructured n:m sparsity, with the paper highlighting performance gains on NVIDIA Ampere GPUs, particularly for 2:4 sparsity.The cited format is described as widely adopted.Source: §1.3 · p. 2
  • The paper evaluates Thanos on language-modeling tasks and reports that it significantly outperforms existing methods for structured pruning.The supplied contribution passage ends before specifying the remaining scope of this result.Source: §1.3 · p. 2

2. Formulation of the problem

The formulation casts pruning a linear layer as minimizing its output change subject to a discrete sparsity constraint. Because exact optimization is difficult, the paper motivates mask selection and weight reconstruction, and develops the pruning solution from a single-weight case.

2.1. Objective function

  • The pruning objective minimizes the difference between a linear layer’s outputs before and after pruning, using its weights and calibration inputs.The paper frames this as a layer-wise compression subproblem.Source: §2.1 · p. 2
  • The mask M marks weights to eliminate, and the sparsity ratio p measures the fraction of weights removed.A mask entry of 1 indicates that the corresponding weight is eliminated.Source: §2.1 · p. 2
  • The formulation constrains pruned weights to zero and fixes the number removed using a discrete sparsity constraint.The floor function is used because the number of removed parameters must be an integer.Source: §2.1 · pp. 2–3
  • Exact optimization is NP-hard because of the discrete constraints, making exact solutions for large layers computationally infeasible.The paper says existing methods therefore rely mostly on approximation techniques.Source: §2.1 · p. 3
  • A common strategy first selects a pruning mask, then optimizes the remaining weights with the mask fixed, reducing the latter step to linear least squares.The mask may be chosen using a metric such as weight magnitude.Source: §2.1 · p. 3

2.2. Pruning of a single parameter of the layer

  • The paper simplifies the constrained problem by first solving the case of removing a single weight from the layer.The chosen weight is set to zero while other mask entries remain zero.Source: §2.2 · p. 3
  • The change in the pruned matrix’s kth row is denoted Δk:, and the single-weight solution can be applied iteratively until the target sparsity is reached.The procedure removes the least important weights one at a time; the accompanying complexity table assumes c=b and ignores sequence length a, with B as Thanos’s block-size hyperparameter.Source: §2.2 · p. 3, Table 1 · p. 3

3. Background and related work

The background contrasts magnitude-based pruning with data-aware methods that use calibration data, then describes Hessian-based approaches and their approximations. SparseGPT and Wanda are presented as related sequential and diagonal-Hessian strategies, respectively.

3.1. Magnitude pruning

  • Magnitude pruning removes individual weights whose magnitudes fall below a threshold.The paper identifies it as a simple, widely used sparsification method.Source: §3.1 · p. 3
  • The paper focuses on data-aware pruning because methods using small calibration datasets are reported to balance sparsity and performance effectively.The passage says such methods consistently yield superior results for LLMs.Source: §3.1 · p. 3

3.2. OBD and OBS

  • OBD and OBS are cited as classical approaches to the pruning problem; the paper introduces SOBS as a loss function derived from its solution.The passage identifies H = 2XX^T as the Hessian of the objective, while OBD simplifies the pruning metric by assuming a diagonal Hessian.Source: §3.2 · p. 4

3.3. SparseGPT

  • SparseGPT prunes parameters sequentially from left to right and considers parameters on the right, keeping the Hessian constant for all rows.Its block-wise procedure updates the mask, adjusts remaining weights, and refines the Hessian while selecting weights using SOBS.Source: §3.3 · p. 4

3.4. Wanda

  • Wanda approximates the intensive Hessian-inverse computation by assuming a diagonal Hessian, paralleling the relationship between OBD and OBS.Under this assumption, its pruning metric simplifies.Source: §3.4 · p. 4

4. Thanos pruning

Thanos prunes multiple weights with block-wise updates and dynamically adjusted masks. It extends this approach to unstructured, structured, and semi-structured n:m sparsity.

4.1. Updating several weights at a time

  • Removing multiple weights at once accounts for cumulative adjustment effects, but searching all possible weight combinations creates a computational bottleneck.With s weights removed from a row of b weights, there are b!/(s!(b−s)!) candidate combinations for each row.Source: §4 · p. 4, §4.1 · pp. 4–5

4.2. Pruning metric

  • Thanos uses Wanda’s pruning metric, which is optimal for removing one weight at a time without adjusting the remaining weights.Thanos first selects a mask, then optimizes adjustments to the weights within that mask.Source: §4.2 · p. 5, §4.1 · p. 5

4.3. Pruning in a block-wise manner

  • Dividing W into smaller blocks makes weight updates computationally feasible and lets the pruning mask adapt as pruning progresses.For each local block, the algorithm updates the row and remaining weights to its right.Source: §4.3 · p. 5

4.4. Unstructured sparsity

  • For unstructured sparsity, Thanos dynamically constructs a global residual mask from the number of weights already pruned and the target sparsity level.This allows pruning beyond local block-wise or row-wise sparsity patterns.Source: §4.4 · p. 5

4.5. How to get indices of weights for removal

  • The mapping ϕ converts each binary mask row into the indices of its nonzero elements, identifying weights designated for removal.For example, ϕ(1, 0, 0, 1, 1) = (1, 4, 5).Source: §4.5 · p. 6

4.7. Structured sparsity; 4.7.1. OUTLIER ROWS (DEFINITION OF α)

  • Structured pruning removes whole columns, with an optional selective strategy that retains a fraction of rows as outliers.The outlier fraction α determines how many rows are not pruned; α = 0 prunes all rows, while α = 1 leaves all rows unpruned.Source: §4.7 · p. 6, §4.7.1 · p. 6

4.7. Structured sparsity; 4.8. Semi-structured n:m sparsity

  • Thanos is significantly faster than SparseGPT in structured pruning and is faster for structured sparsity on models up to 3 billion parameters.The semi-structured setting also streamlines pruning because each row has a uniform number of weights designated for removal.Source: §4.7 · p. 6, §4.8 · p. 7

4.8. Semi-structured n:m sparsity

  • Thanos extends to n:m sparsity, which requires exactly n zeros in every block of m consecutive weights and supports hardware-optimized patterns such as 2:4.In this setting, uniform per-row pruning counts streamline operations and reduce computational overhead.Source: §4.8 · p. 7

5. Experiments

The experiments evaluate Thanos across LLM families, language modeling, and zero-shot tasks. They report strong structured-pruning results, with block-size choices varying by sparsity pattern.

5.1. Models and evaluation.

  • Evaluation covers OPT and LLaMA-family models, seven zero-shot tasks, and WikiText-2 validation perplexity, with shared C4 calibration data for methods that require it.The calibration set contains 128 sequences with context length sampled from C4.Source: §5.1 · p. 7

5.2. Perplexity evaluation

  • The authors attribute structured-pruning gains to cumulative weight changes from removing multiple columns that cannot be approximated by summing individual-column changes.The passage describes this as a reason for Thanos's stronger structured-sparsity performance.Source: §5.2 · pp. 7–8
  • Thanos consistently surpasses other widely used methods in structured pruning and is more efficient than SparseGPT.The paper describes the advantage as especially relevant at high sparsity levels.Source: §5.2 · p. 7, Table 2 · p. 8
  • For WikiText perplexity, Thanos beats Magnitude pruning, SparseGPT, and Wanda on TinyLlama 1.3B and LLaMA-3 models from 3B to 70B.Lower perplexity indicates better language-modeling performance.Source: §5.2 · p. 7, Table 2 · p. 8

5.3. Zero-shot evaluation on different LLM’s tasks

  • In zero-shot evaluation, Thanos outperforms competing methods in structured and semi-structured sparsity and remains competitive in unstructured sparsity.It achieves the best results for LLaMA-2 70B and LLaMA-3 models from 1B to 3B parameters.Source: Table 3 · p. 8, §5.3 · p. 8

5.4. Experiments with different blocksizes

  • Unstructured pruning uses block size B = 128, while structured patterns such as 4:8 and 2:4 use B = 512.The authors selected B = 128 after observing little perplexity change across unstructured block sizes; larger blocks were more effective for structured pruning.Source: §5.1 · p. 7, §5.4 · p. 8

6. Conclusion and Limitations

The section points readers to Appendix A for the paper’s conclusion and limitations.

6. Conclusion and Limitations

  • The paper directs readers to Appendix A for its concluding comments and limitations.Source: §6 · p. 8

A.1. Conclusion

Thanos combines block-wise pruning with adaptive masks and outlier-row detection. The authors report advantages in structured and semi-structured sparsity, alongside competitive performance and efficiency.

A.1. Conclusion

  • Thanos performs block-wise weight pruning with adaptive masks that adjust dynamically during pruning.The authors say this flexibility extends beyond local block or row sparsity constraints and helps retain generalization as sparsity increases.Source: §A.1 · p. 12
  • Across language-modeling tasks and model sizes, Thanos significantly outperforms other approaches in structured and semi-structured sparsity while maintaining competitive performance and improving efficiency over SparseGPT.The authors describe the method as performance-preserving across a wide range of model sizes and sparsity levels.Source: §A.1 · p. 12
  • In structured sparsity, simultaneous multi-weight updates address cumulative changes that cannot be approximated by summing individual weight changes.The authors contrast this with SparseGPT’s column-by-column pruning of one weight per row at a time.Source: §A.1 · p. 12
  • Thanos’s structured-pruning approach incorporates outlier-row detection and is reported to improve perplexity and zero-shot performance over baselines.Source: §A.1 · p. 12

A.2. Limitations

Thanos has computational costs that grow with pruning demands, and its relative performance varies by sparsity setting; the appendix also includes a notation reference.

A.2. Limitations

  • Solving multiple linear systems per weight block makes Thanos computationally costly, especially for larger matrices or higher sparsity, and this burden grows with model size.The authors identify the resulting time and resource requirements as a limiting factor.Source: §A.2 · p. 12
  • SparseGPT often surpasses Thanos on zero-shot evaluation in unstructured sparsity, and the authors suggest possible challenges at very large scale or fine-grained sparsity.They attribute the possible challenges to adaptable weight updates that do not always yield optimal performance in those configurations.Source: §A.2 · p. 12
  • Wanda is generally less competitive than Thanos and SparseGPT at lower sparsity levels, yet remains robust in unstructured sparsity for LLaMA-2 and LLaMA-3.Source: §A.2 · p. 12

B. Table of Frequently Used Notation

  • Appendix B includes Table 4 as a summary of the paper’s key notation.Source: Table 4 · p. 13

C. Experiments with different blocksize B

The experiments vary Thanos’s block size when pruning TinyLlama-1.1B. Perplexity is relatively stable across block sizes for unstructured sparsity, while structured sparsity benefits from larger blocks.

C. Experiments with different blocksize B

  • For unstructured sparsity, perplexity remained relatively stable across tested block sizes, while structured sparsity pruning was more effective with larger blocks.The experiments pruned TinyLlama-1.1B using B values from 8 to 4096; B = 128 was selected for most unstructured experiments to balance efficiency and performance.Source: C. Experiments with different… · p. 14
  • The block-size experiments covered B values from 8 to 4096 on TinyLlama-1.1B.Source: C. Experiments with different… · p. 14, Table 5 · p. 14

D. Full Zero-Shot evaluation on different LLM’s tasks

This section presents zero-shot evaluation tables for OPT, LLaMA-2, and LLaMA-3 across four sparsity settings, alongside the pruning procedure used to select rows and weights. It derives the OBS metric and notes that OBD simplifies it by assuming a diagonal Hessian.

D.1. Unstructured 50% sparsity; D.2. Structured 30% sparsity

  • Tables 6–11 report zero-shot evaluations of different methods for OPT, LLaMA-2, and LLaMA-3, covering unstructured 50% sparsity and structured 30% sparsity.The captions provide no scores.Source: Table 6 · p. 14, Table 7 · p. 15, Table 8 · p. 15

D.3. Semi-structured 4:8 sparsity; D.4. Semi-structured 2:4 sparsity

  • Tables 12–17 report zero-shot evaluations of different methods for OPT, LLaMA-2, and LLaMA-3 at semi-structured 4:8 and 2:4 sparsity.The captions identify the models and sparsity formats, but provide no scores or comparisons.Source: Table 12 · p. 17, Table 13 · p. 18, Table 14 · p. 19

E.1. General pruning algorithm

  • The generic pruning procedure processes model blocks sequentially, prunes each linear layer to a fixed sparsity ratio, then forwards the pruned block to prepare inputs for the next block.Layers within a block are pruned independently using the same procedure.Source: §E.1 · p. 21

E.2. Objective function; E.3. Constrained optimization problem

  • The objective minimizes the change in a layer’s output after pruning, but its fixed-sparsity constraint makes mask selection discrete and the constrained problem NP-hard.The paper separates finding a pruning mask from adjusting the remaining weights.Source: §E.2 · p. 22, §E.3 · pp. 22–23

E.4. Pruning of a single parameter of the layer; E.5. Objective function for d calibration samples

  • The single-weight solution is applied greedily and iteratively, removing weights until the target sparsity is reached; with multiple calibration samples, the loss averages results across samples.The single-weight simplification underlies the iterative pruning approach.Source: §E.4 · p. 23, §E.5 · p. 25

F.1. Magnitude pruning

  • Magnitude pruning removes weights below a magnitude threshold without calibration data, but its performance degrades at high sparsity, motivating data-aware pruning with small calibration datasets.Magnitude pruning is described as computationally efficient, while data-aware methods are reported to balance sparsity and performance better for large models.Source: §F.1 · pp. 25–26

F.2. OBD and OBS

  • The OBD/OBS derivation approximates the loss near zero using a Taylor expansion and notes that the Hessian is shared across weight-matrix rows and is symmetric positive definite.The derivation assumes pretrained weights are at or near a local optimum.Source: §F.2 · pp. 26–27
  • The optimal δ and λ are found by setting the Lagrangian’s partial derivatives with respect to both variables to zero.The derivation assumes an invertible Hessian and uses its symmetry and positive definiteness to obtain the optimal weight change.Source: §F.2 · p. 27
  • OBS chooses the row and weight index with the smallest pruning metric, then updates that row using the optimal weight change.The text defines the saliency metric for the whole weight matrix and applies the update rule to the selected row.Source: §F.2 · p. 28
  • OBD assumes a diagonal Hessian, simplifying its inverse and the resulting pruning metric.The text identifies OBS as the neural-network pruning method described in the Optimal Brain Surgeon paper.Source: §F.2 · p. 28

F.3. SparseGPT

This section describes SparseGPT’s sequential pruning with adaptive block masks and Hessian-based updates, alongside the approaches used by Thanos and Wanda. It also explains how block-wise updates propagate pruning effects and how block size trades in-block communication against computational cost.

F.3. SparseGPT

  • SparseGPT’s row-specific Hessians change as weights are pruned; computing c masked Hessians at O(b^3) each yields overall O(cb^3) complexity.SparseGPT avoids this issue by pruning sequentially from left to right and considering only unseen weights to the right, keeping the Hessian the same across rows.Source: §F.3 · pp. 28–29
  • SparseGPT selects the p% smallest-SOBD weights in each local block and dynamically adapts its mask as pruning changes the importance of remaining weights.For pruned weights, it applies an OBS inverse-Hessian update to compensate by adjusting remaining weights, including weights to the right of the current block.Source: §F.3 · p. 29, Figure 5 · p. 30

F.4. Wanda

  • Wanda prunes the p% least important weights in each row using the product of weight magnitude and the corresponding input-column norm, without iterative mask updates or weight adjustments.The resulting mask is constructed in one step, and the selected weights are removed while the remaining weights stay unchanged.Source: §F.4 · p. 32

G. Thanos pruning

  • Thanos considers updating more than one weight at a time, extending beyond SparseGPT’s column-by-column removal of one weight per row.SparseGPT is identified as the comparison method in the passage introducing this motivation.Source: G. Thanos pruning · p. 32

G.1. Updating several weights at a time

  • Solving for the optimal set of s weights to prune entails a combinatorial search and is described as impractical for LLMs.With a predefined pruning mask, the task instead separates into selecting the mask and optimizing adjustments to its selected parameters.Source: §G.1 · p. 34

G.2. Pruning in a block-wise manner

  • Thanos partitions the weight matrix into local blocks to make weight-adjustment systems computationally manageable, then updates remaining weights to the right of each block.The method proceeds from the leftmost block and solves a system of linear equations for every row in each block.Source: §G.2 · pp. 34–35, Figure 7 · p. 34
  • The update rule enables in-block communication and propagates pruning effects to weights outside the current block, including remaining parameters to its right.Together, these interactions efficiently spread pruning effects through the matrix.Source: §G.2 · p. 35
  • Block size B must balance computational cost against the benefits of in-block communication; solving each weight-matrix row’s linear system requires O(B^3) operations.The passage recommends keeping B small enough to avoid high costs but large enough to gain from in-block communication.Source: §G.2 · p. 35

G.3. Pruning metric

This section describes Thanos’s pruning metric and adaptive mask construction across unstructured, structured, and semi-structured settings, along with permutation, index-extraction, complexity, and implementation procedures. It also covers padding for batched solvers on non-uniform rows and reports that block-wise pruning is faster than SparseGPT on models smaller than 1 billion parameters.

G.3. Pruning metric

  • Wanda’s metric is optimal for removing one weight at a time when the remaining weights are not tuned.The loss for each candidate weight is evaluated to select the minimum-loss removal; the computation can exploit the localized effect of removing a single weight.Source: §G.3 · pp. 35–36

G.4.1. UNSTRUCTURED SPARSITY

  • Thanos constructs a global residual mask iteratively, avoiding fixed local-block or row-wise sparsity constraints while adapting to weight updates.At each step it selects a local block from the residual mask, then updates weights and constructs the next residual mask according to the remaining pruning target.Source: §G.4.1 · pp. 37–38, Figure 8 · p. 37

G.4.2. STRUCTURED SPARSITY

  • For structured pruning, Thanos removes whole columns and can exclude outlier rows, addressing columns containing weights that are difficult to remove.It selects columns by their induced loss and uses row and column permutations to arrange pruning targets, then applies inverse permutations to restore the original ordering.Source: §G.4.2 · p. 38

G.4.4. COLUMN AND ROW PERMUTATIONS

  • Permutation matrices order columns by column loss and rows by row loss, enabling pruning in the desired order and restoration of the original matrix ordering.Columns with the smallest losses are placed first; rows are similarly reordered according to their losses.Source: §G.4.4 · pp. 38–39, §G.4.2 · p. 38

G.4.5. STRUCTURED n:m SPARSITY

  • Thanos extends to n:m sparsity, which places exactly n zeros in each block of m consecutive weights and can use NVIDIA Ampere’s optimized 2:4 implementation.This setting uses local block-structured masks rather than a global residual mask; uniform removals per row avoid matrix padding and keep linear-system dimensions consistent.Source: §G.4.5 · p. 39

G.4.6. HOW TO GET INDICES OF WEIGHTS FOR REMOVAL

  • A mapping from each mask row to the indices of its nonzero entries extracts the weight indices designated for removal.For example, (1, 0, 0, 1, 1) maps to (1, 4, 5).Source: §G.4.6 · p. 40

G.6. Complexity

  • The complexity estimate accounts for Hessian computation and inversion, pruning-metric calculation and selection, and weight updates for each block.The section also points to a comparison of computational complexities across the methods considered.Source: §G.6 · p. 40

H.1. Efficient batched operations

  • To reduce implementation overhead, Thanos replaces nested loops with NumPy and PyTorch broadcasting operations that leverage underlying C++ optimizations.The text identifies varying per-row removal counts, and thus linear systems of varying sizes, as a core challenge.Source: §H.1 · p. 41
  • Padding each row’s removal-index vector to the maximum row size enables batched linear solvers while accommodating non-uniform rows.The maximum size is rmax := max_i ||M_i:||_F^2; shorter index vectors are padded with ones.Source: §H.1 · p. 42
  • For models smaller than 1 billion parameters, Thanos is faster than SparseGPT because it prunes parameters in blocks rather than individually.The passage attributes this efficiency to reduced computation time for smaller-scale problems in this parameter range (Figure 9a).Source: §H.1 · pp. 42–43
  • Thanos’s block-wise pruning reduces computation time by processing parameters in blocks instead of pruning them individually.The passage identifies this block-wise approach as the reason for Thanos’s speed advantage over SparseGPT on models smaller than 1 billion parameters.Source: §H.1 · p. 43

H.2. GPU memory limitations

Thanos segments linear-system solves into vertical blocks to fit GPU memory, and its structured n:m pruning is more efficient and faster than SparseGPT for models up to 3 billion parameters.

H.2. GPU memory limitations

  • Batched solves across all rows can exceed GPU memory, even on NVIDIA A100 GPUs with 40 or 80 GB, for models as small as OPT-1.3B.Thanos subdivides columns into vertical blocks sized to fit GPU memory.Source: §H.2 · p. 43
  • Processing vertical blocks sequentially makes larger matrices feasible on GPUs while limiting memory consumption.Thanos solves the equations for rows in one block before proceeding to the next vertical block.Source: §H.2 · p. 43

H.3. Accelerated structured n:m sparsity

  • Structured patterns such as 2:4 and 4:8 prune a uniform number of weights per row, eliminating the need for padding adjustments.Source: §H.3 · p. 43
  • For structured sparsity, Thanos is more efficient and faster than SparseGPT for models up to 3 billion parameters.The paper attributes the gain to block-wise pruning aligned with structured patterns.Source: §H.3 · p. 43

Related papers

Loading 2504.05346v1…