LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Elhoushi, Mostafa, Shrivastava, Akshat, Liskovich, Diana, Hosmer, Basil, Wasti, Bram, Lai, Liangzhen, Mahmoud, Anas, Acun, Bilge, Agarwal, Saurabh, Roman, Ahmed, Aly, Ahmed A, Chen, Beidi, Wu, Carole-Jean
Format: Preprint
Published: 2024
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866913552177037312
author Elhoushi, Mostafa
Shrivastava, Akshat
Liskovich, Diana
Hosmer, Basil
Wasti, Bram
Lai, Liangzhen
Mahmoud, Anas
Acun, Bilge
Agarwal, Saurabh
Roman, Ahmed
Aly, Ahmed A
Chen, Beidi
Wu, Carole-Jean
author_facet Elhoushi, Mostafa
Shrivastava, Akshat
Liskovich, Diana
Hosmer, Basil
Wasti, Bram
Lai, Liangzhen
Mahmoud, Anas
Acun, Bilge
Agarwal, Saurabh
Roman, Ahmed
Aly, Ahmed A
Chen, Beidi
Wu, Carole-Jean
contents We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.
format Preprint
id arxiv_https___arxiv_org_abs_2404_16710
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
Elhoushi, Mostafa
Shrivastava, Akshat
Liskovich, Diana
Hosmer, Basil
Wasti, Bram
Lai, Liangzhen
Mahmoud, Anas
Acun, Bilge
Agarwal, Saurabh
Roman, Ahmed
Aly, Ahmed A
Chen, Beidi
Wu, Carole-Jean
Computation and Language
Artificial Intelligence
Machine Learning
We present LayerSkip, an end-to-end solution to speed-up inference of large language models (LLMs). First, during training we apply layer dropout, with low dropout rates for earlier layers and higher dropout rates for later layers, and an early exit loss where all transformer layers share the same exit. Second, during inference, we show that this training recipe increases the accuracy of early exit at earlier layers, without adding any auxiliary layers or modules to the model. Third, we present a novel self-speculative decoding solution where we exit at early layers and verify and correct with remaining layers of the model. Our proposed self-speculative decoding approach has less memory footprint than other speculative decoding approaches and benefits from shared compute and activations of the draft and verification stages. We run experiments on different Llama model sizes on different types of training: pretraining from scratch, continual pretraining, finetuning on specific data domain, and finetuning on specific task. We implement our inference solution and show speedups of up to 2.16x on summarization for CNN/DM documents, 1.82x on coding, and 2.0x on TOPv2 semantic parsing task. We open source our code and checkpoints at https://github.com/facebookresearch/LayerSkip.
title LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2404.16710