Cross-Attention Speculative Decoding

TL;DR

This paper introduces Budget EAGLE, a cross-attention-based Transformer decoder for speculative decoding that simplifies architecture, enhances training efficiency, and maintains competitive inference speedups compared to existing models.

Contribution

The paper presents Budget EAGLE, the first cross-attention-based SD model that eliminates auxiliary components, along with Two-Stage Block-Attention Training for stable, efficient training.

Findings

Beagle achieves comparable inference speedups to state-of-the-art models.

Beagle demonstrates higher training efficiency across multiple LLMs.

The architecture simplifies speculative decoding without sacrificing performance.

Abstract

Speculative decoding (SD) is a widely adopted approach for accelerating inference in large language models (LLMs), particularly when the draft and target models are well aligned. However, state-of-the-art SD methods typically rely on tightly coupled, self-attention-based Transformer decoders, often augmented with auxiliary pooling or fusion layers. This coupling makes them increasingly complex and harder to generalize across different models. We present Budget EAGLE (Beagle), the first, to our knowledge, cross-attention-based Transformer decoder SD model that achieves performance on par with leading self-attention SD models (EAGLE-v2) while eliminating the need for pooling or auxiliary components, simplifying the architecture, improving training efficiency, and maintaining stable memory usage during training-time simulation. To enable effective training of this novel architecture, we propose Two-Stage Block-Attention Training, a new method that achieves training stability and convergence efficiency in block-level attention scenarios. Extensive experiments across multiple LLMs and datasets show that Beagle achieves competitive inference speedups and higher training efficiency than EAGLE-v2, offering a strong alternative for architectures in speculative decoding.