Cross-Attention Speculative Decoding
This work addresses the problem of accelerating inference in large language models for AI practitioners, offering a more efficient and generalizable alternative to existing speculative decoding methods.
The paper tackles the complexity and generalization issues in speculative decoding for large language models by introducing Budget EAGLE (Beagle), a cross-attention-based Transformer decoder that achieves performance on par with leading self-attention models like EAGLE-v2, while simplifying the architecture and improving training efficiency.
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.