Cross-Attention Speculative Decoding

2025-05-30

AI Summary

Existing speculative decoding methods predominantly rely on tightly coupled self-attention Transformer decoders, resulting in architectural complexity and poor generalizability. This paper proposes Beagle—the first lightweight speculative decoding framework based exclusively on cross-attention mechanisms, fully decoupling the backbone language model from auxiliary modules. Our key contributions are: (1) introducing a novel cross-attention architecture that replaces conventional self-attention, eliminating the need for auxiliary decoders; and (2) proposing a two-stage block-level attention training strategy to jointly optimize training stability and convergence efficiency. Experiments demonstrate that Beagle achieves inference throughput comparable to EAGLE-v2, with higher training efficiency, constant memory footprint, and seamless transferability across diverse LLMs and datasets. These advances significantly enhance the generality, practicality, and deployment flexibility of speculative decoding.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Computer Vision: Large Vision Models

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

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.

Problem

• Reduces complexity in speculative decoding architectures
• Eliminates need for auxiliary components in SD models
• Improves training efficiency and memory stability

Innovation

• Cross-attention-based Transformer decoder SD model
• Two-Stage Block-Attention Training method
• Simplified architecture without auxiliary components