STAT: Shrinking Transformers After Training

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Flynn, Megan, Wang, Alexander, Alvarez, Dean Edward, De Sa, Christopher, Damle, Anil
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917681686380544
author Flynn, Megan
Wang, Alexander
Alvarez, Dean Edward
De Sa, Christopher
Damle, Anil
author_facet Flynn, Megan
Wang, Alexander
Alvarez, Dean Edward
De Sa, Christopher
Damle, Anil
contents We present STAT: a simple algorithm to prune transformer models without any fine-tuning. STAT eliminates both attention heads and neurons from the network, while preserving accuracy by calculating a correction to the weights of the next layer. Each layer block in the network is compressed using a series of principled matrix factorizations that preserve the network structure. Our entire algorithm takes minutes to compress BERT, and less than three hours to compress models with 7B parameters using a single GPU. Using only several hundred data examples, STAT preserves the output of the network and improves upon existing gradient-free pruning methods. It is even competitive with methods that include significant fine-tuning. We demonstrate our method on both encoder and decoder architectures, including BERT, DistilBERT, and Llama-2 using benchmarks such as GLUE, Squad, WikiText2.
format Preprint
id arxiv_https___arxiv_org_abs_2406_00061
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle STAT: Shrinking Transformers After Training
Flynn, Megan
Wang, Alexander
Alvarez, Dean Edward
De Sa, Christopher
Damle, Anil
Machine Learning
Artificial Intelligence
Computation and Language
We present STAT: a simple algorithm to prune transformer models without any fine-tuning. STAT eliminates both attention heads and neurons from the network, while preserving accuracy by calculating a correction to the weights of the next layer. Each layer block in the network is compressed using a series of principled matrix factorizations that preserve the network structure. Our entire algorithm takes minutes to compress BERT, and less than three hours to compress models with 7B parameters using a single GPU. Using only several hundred data examples, STAT preserves the output of the network and improves upon existing gradient-free pruning methods. It is even competitive with methods that include significant fine-tuning. We demonstrate our method on both encoder and decoder architectures, including BERT, DistilBERT, and Llama-2 using benchmarks such as GLUE, Squad, WikiText2.
title STAT: Shrinking Transformers After Training
topic Machine Learning
Artificial Intelligence
Computation and Language
url https://arxiv.org/abs/2406.00061