大型語言模型中的快速推理解碼的循環起草者

Recurrent Drafter for Fast Speculative Decoding in Large Language Models

March 14, 2024

作者: Aonan Zhang, Chong Wang, Yi Wang, Xuanyu Zhang, Yunfei Cheng

cs.AI

摘要

本文介紹了一種改進的推測解碼方法，旨在增強服務大型語言模型的效率。我們的方法充分利用了兩種已建立的技術的優勢：經典的雙模型推測解碼方法和較新的單模型方法Medusa。受Medusa的啟發，我們的方法採用了單模型策略進行推測解碼。然而，我們的方法通過採用單一、輕量級的草稿頭部，具有循環依賴設計，本質上類似於經典推測解碼中使用的小型草稿模型，但不涉及完整Transformer架構的複雜性。由於循環依賴，我們可以使用束搜索快速過濾掉草稿頭中的不需要的候選項。該方法結合了單模型設計的簡單性，避免了在Medusa中僅用於推斷的創建數據依賴樹注意結構的需求。我們在幾個流行的開源語言模型上實證了所提出方法的有效性，並對採用此方法涉及的權衡進行了全面分析。

English

In this paper, we introduce an improved approach of speculative decoding aimed at enhancing the efficiency of serving large language models. Our method capitalizes on the strengths of two established techniques: the classic two-model speculative decoding approach, and the more recent single-model approach, Medusa. Drawing inspiration from Medusa, our approach adopts a single-model strategy for speculative decoding. However, our method distinguishes itself by employing a single, lightweight draft head with a recurrent dependency design, akin in essence to the small, draft model uses in classic speculative decoding, but without the complexities of the full transformer architecture. And because of the recurrent dependency, we can use beam search to swiftly filter out undesired candidates with the draft head. The outcome is a method that combines the simplicity of single-model design and avoids the need to create a data-dependent tree attention structure only for inference in Medusa. We empirically demonstrate the effectiveness of the proposed method on several popular open source language models, along with a comprehensive analysis of the trade-offs involved in adopting this approach.