Amphista: Bi-directional Multi-head Decoding for Accelerating LLM Inference (2025)

Abstract: Large Language Models (LLMs) inherently use autoregressive decoding, which lacks parallelism in inference and results in significantly slow inference speed. While methods such as Medusa constructs parallelized heads, they lack adequate information interaction across different prediction positions. To overcome this limitation, we introduce Amphista, an enhanced speculative decoding framework that builds upon Medusa. Specifically, Amphista models an *Auto-embedding Block* capable of parallel inference, incorporating bi-directional attention to enable interaction between different drafting heads. Additionally, Amphista integrates *Staged Adaptation Layers*, which ensure a seamless transition of semantic information from the target model’s autoregressive inference to the drafting heads’ non-autoregressive inference, effectively achieving paradigm shift and feature fusion. Experimental results on Vicuna models using MT-Bench and Spec-Bench demonstrate that Amphista achieves substantial acceleration while maintaining generation quality. On MT-Bench, Amphista delivers up to **2.75×** speedup over vanilla autoregressive decoding and **1.40×** over Medusa on Vicuna 33B in wall-clock time.

## introduction generative large language models (llms) have made remarkable advances in language processing by scaling the transformer decoder block, offering a potential pathway toward artificial general intelligence (agi) @xcite . however, the autoregressive nature of nexttoken prediction and the large parameter size of foundational models result in low inference efficiency, marked by high latency per token and low throughput per second during decoding. in this context, acceleration during inference has become a burgeoning research area. speculative de-figure 1 : top-1/5 accuracy for different heads of medusa and amphista. we perform testing with randomly sampled 5% sharegpt conversation data. amphista far outperforms medusa in terms of head accuracy, especially for the latter two heads. nonetheless, as illustrated in figure 1 , we find that the prediction accuracy of separately independent medusa heads is relatively low, which progressively worsens and adversely impacts acceleration performance in downstream tasks. to mitigate these inaccuracies stemming from the absence of feature dependencies while maintaining parallel inference, we first introduce the autoembedding block, which integrates a bi-directional self-attention mechanism @xcite . this structure not only allows earlier heads to attend to subsequent ones, but more importantly, enables backward heads to leverage information from preceding heads. this enhancement allows drafting heads to better capture contextual information, thereby improving the acceptance rate of their predictions. moreover, in the multi-step drafting framework, this non-autoregressive structure achieves lower drafting latency compared to an autoregressive approach. additionally, we identify a significant gap between the autoregressive target model and the non-autoregressive draft model in their prediction paradigms. to bridge this discrepancy and further enhance feature representations across different drafting heads, we propose the staged adaptation layers. these layers serve as an intermediary module to facilitate feature integration and transformation between the target model and draft heads. once adopted, semantically enriched features are passed through mlp activations and fed into the auto-embedding block. this enhances the bi-directional attention mechanism's ability to fuse features across heads, ultimately boosting acceptance rates and reducing wall-clock time. lastly, to further align the draft model with the target model with minimal computational cost, we incorporate the sampled token from the target model's latest prediction into the staged adaptation layers. this critically integrated information harmonizes amphista with the target model, yielding a significant improvement in performance. on mt-bench, amphista achieves up to 2.75× speedup over vanilla autoregressive decoding and 1.40× over medusa on vicuna 33b, as consistently high accuracy (see figure 1 ). to summarize, our contributions are as follows: ## preliminaries in this section, we introduce some preliminary background related to our work as follows: speculative decoding. speculative decoding has been successfully applied to llm decoding algorithm recently @xcite . the core idea is to leverage a small, lowerquality model (draft model) together with a large, higher-quality model (target model) to accelerate token generation. concretely, in each decoding step, the algorithm first uses the draft model to autoregressively generate a sequence of future tokens. these drafted tokens are then verified by the target model in a single forward pass. during the verification process, a certain strategy is applied to determine which tokens are accepted by the target model and which are rejected and discarded. previous work @xcite has theoretically and empirically demonstrated that the token output distribution of speculative decoding is consistent with the autoregressive generation of original target model, but with fewer decoding steps, thus enhancing generation efficiency. medusa decoding. medusa decoding @xcite represents an efficient speculative decoding algorithm based on the draft-and-verify principle, inheriting principles from blockwise parallel decoding @xcite . specifically, medusa integrates independent mlp layers, called drafting heads, with the target model to form a unified architecture. in each decoding step, the target model's lm_head samples the next token, while the next-i mlp heads predict tokens at subsequent positions. these drafted tokens are then verified by the target model's forward pass to decide their acceptance. by leveraging lightweight mlp layers, medusa strikes an effective balance between computational efficiency and prediction accuracy, leading to substantial acceleration. hydra @xcite , which is a subsequent state-of-the-art optimization based on medusa, transforms the independent mlp heads into sequentially dependent mlp heads, further enhancing the predictive accuracy. tree attention. tree attention @xcite enables parallel computation of attention scores for multiple draft candidates. medusa uses a tree causal mask, allowing each target model hidden states 𝑳𝑴 𝑯𝒆𝒂𝒅𝟏 𝑳𝑴 𝑯𝒆𝒂𝒅𝟐 𝑳𝑴 𝑯𝒆𝒂𝒅𝒌 … (auto-embedding block) (staged adaptation layers) input: 𝑻 𝟏 , 𝑻 𝟐 , … , 𝑻 𝒏 𝑻𝒏+𝟐 token ( ) 𝑻𝒏+𝟑 token ( ) 𝑻𝒏+𝒌+𝟏 token (×) : accepted × : rejected target model amphista: non-autoregressive decoding 𝒎𝒍𝒑𝟏 𝒎𝒍𝒑⌊k/𝟐⌋ 𝒎𝒍𝒑 ⌊ 𝒌 𝟐 +𝟏⌋ 𝒎𝒍𝒑𝒌 bidirectional self attention transformer decoder positional embedding transformer decoder mlp block 𝑳𝑴 𝑯𝒆𝒂𝒅 attention block 𝑻 𝒏+𝟏 token ( ) adapted hidden states adapted hidden states … … decoder×n figure 2: the framework of amphista decoding. our method improves medusa in two folds: (1) we introduce staged adaptation layers, consisting of a group of causal transformer decoder layers built upon the target model, to adapt the target model's hidden states and the sampled token in two stages. this module ensures that the adapted features contain richer contextual information, supporting multiple-token predictions rather than focusing solely on the immediate next-token prediction. (2) we introduce an auto-embedding block, which is a bi-directional transformer encoder module with positional encoding. this block allows each head to attend to others, fostering cooperative predictions and thereby enhancing the speculative accuracy during the drafting stage. ## amphista the overview of amphista is shown in figure 2 . building its pipeline upon target model, amphista contains two main modules: (1) staged adaptation layers. they are causal transformer decoder layers that adapt the target model's hidden states and sampled token embedding in two stages, each focusing on different drafting positions. this adaptation process results in hidden states that are enhanced with position-aware contextual information, improving overall prediction accuracy, especially for the latter steps. (2) auto-embedding block. it is a transformer encoder module that conducts bidirectional self-attention computations among the representations of different draft heads, allowing each head can be attended by the others. this facilitates collaborative prediction among these heads, thereby improving overall prediction accuracy. hidden states contain semantically rich information, there are still differences in the representation requirements between the target model and the draft heads. specifically, the hidden states of the target model are trained only for predicting the next token, while draft heads need more contextual and positon-aware hidden states to perform multi-step speculation. to address this problem, medusa-2 applies lora (hu et al., 2021) for joint training of the target model and draft heads, which may compromise the generality on downstream tasks. hydra employs a single prefix layer for all positions, lacking targeted adaptation for different positions. ## related work increasing techniques have been proposed to enhance the inference speed of large language models (llms), covering aspects of system hardware, model architecture, and decoding algorithms. a significant branch of these techniques is model compression, which includes methods such as model quantization @xcite @xcite , pruning @xcite @xcite , and distillation @xcite @xcite . additionally, techniques like kv-cache @xcite , flash-attention @xcite , and early exiting @xcite @xcite have also significantly reduced inference overhead. another important line is speculative decoding, which our work is based on. it can be broadly categorized into two types. the first treats the target model and draft model separately and independently, involving the use of a small language model @xcite @xcite @xcite , external database, or n-grams pool @xcite @xcite to generate candidate token sequences or token trees @xcite , which the llm then verifies. the second type views the draft model as a dependent approximation of the target model, deriving the draft model directly from the target model or building additional modules on top of the target model for drafting @xcite @xcite @xcite @xcite . unlike these approaches, we propose a novel method using an auto-embedding block combined with staged adaptation layers to further enhance acceleration. ## conclusion we propose amphista, an efficient nonautoregressive speculative decoding framework that accelerates inference through parallel decoding and improves alignment between target and draft models via feature adaptation. amphista integrates two core components: the auto-embedding block, leveraging bi-directional self-attention for collaborative speculation among drafting heads, and staged adaptation layers, transforming target model semantics for multi-step predictions. additionally, amphista exploits sampled tokens to further optimize alignment. extensive experiments confirm the superiority of amphista, showcasing the promise of non-autoregressive methods in speculative decoding.