Speculative Decoding with MoE Models

Autoregressive generation faces a persistent latency problem, even when memory-saving techniques like offloading and quantization are applied to address the static footprint of a Mixture of Experts model. Each token is generated sequentially, requiring a full forward pass through the large MoE model, a process limited by memory bandwidth and computational overhead. Speculative decoding offers a powerful method to reduce this latency by fundamentally changing how tokens are generated.

The strategy relies on a simple observation: a large, powerful model is slow, while a small, less powerful model is fast. Speculative decoding uses two models in tandem:

1. The Target Model: Our large, high-performance MoE model. It is accurate but has high latency per token.
2. The Draft Model: A much smaller, faster model. This could be a distilled version of the MoE or simply a standard Transformer with far fewer parameters. Its purpose is speed, not perfection.

Instead of the MoE model generating one token at a time, the fast draft model generates a "draft" of several future tokens. The large MoE model then checks this entire draft in a single, parallel forward pass, which is significantly faster than multiple sequential passes.

The Speculative Decoding Loop

The process operates in a loop that generates tokens until a stopping condition is met. Each iteration involves a draft-and-verify cycle.

1. Drafting: The small draft model takes the current sequence and autoregressively generates a short continuation of $k$ tokens. For example, if the current text is "The quick brown fox," the draft model might propose "jumps over the." This step is very fast.
2. Verification: The large MoE model receives the original sequence plus the $k$ drafted tokens ("The quick brown fox jumps over the"). It performs one forward pass to compute the probability distributions for the next token at each position.
3. Acceptance: We compare the draft model's predictions against the target MoE's outputs. Starting from the first drafted token ("jumps"), we check if it's what the MoE model would have likely chosen.
   • If the tokens match (or pass a sampling check), we accept the token and move to the next one in the draft ("over").
   • This continues until we find a mismatch or the draft is fully accepted.
4. Correction: The first time a token is rejected, we discard it and all subsequent tokens in the draft. We then sample a new, correct token using the more accurate probability distribution calculated by the MoE model at that position.
5. Continuation: The process repeats from the end of the newly extended sequence.

The following diagram illustrates this workflow. The point is that the expensive Target MoE Model is invoked only once per cycle to validate multiple tokens in parallel, amortizing its high latency.

The speculative decoding process. A fast draft model generates multiple candidate tokens, which are then validated in a single pass by the slower, more accurate MoE model.

Performance Gains and Implementation Details

The speedup from speculative decoding is directly related to the average number of tokens accepted per cycle. If the draft model is a good approximation of the target MoE, it can consistently generate several correct tokens, allowing the system to bypass multiple sequential and expensive forward passes. For example, if on average three tokens are accepted per cycle, the wall-clock time for generation can be reduced by a factor approaching three.

A high-level implementation would follow this logic:

# Pseudocode for Speculative Decoding
def speculative_decode(prompt, target_moe_model, draft_model, max_len, k):
    tokens = tokenize(prompt)

    while len(tokens) < max_len:
        # 1. Draft k tokens using the fast model
        draft_tokens = []
        context = tokens
        for _ in range(k):
            next_token = draft_model.generate(context, 1)
            draft_tokens.append(next_token)
            context.append(next_token)

        # 2. Verify all k+1 tokens with one call to the target model
        # The input includes the original sequence plus the draft
        target_logits = target_moe_model.forward(tokens + draft_tokens)

        # 3. Compare and accept/reject
        accepted_count = 0
        for i in range(k):
            # Check if draft token matches the most likely token from the target model
            draft_token_i = draft_tokens[i]
            target_distribution_i = softmax(target_logits[len(tokens) + i - 1])

            if is_accepted(draft_token_i, target_distribution_i):
                tokens.append(draft_token_i)
                accepted_count += 1
            else:
                # 4. Correct with a sample from the target model's distribution
                corrected_token = sample(target_distribution_i)
                tokens.append(corrected_token)
                break # Exit verification loop

        # If all draft tokens were accepted, sample one final token
        if accepted_count == k:
            final_distribution = softmax(target_logits[-1])
            final_token = sample(final_distribution)
            tokens.append(final_token)

    return detokenize(tokens)

Note: The acceptance function is_accepted can be a simple greedy check (i.e., argmax(target_distribution_i) == draft_token_i) or a more sophisticated sampling method like rejection sampling to preserve the original model's output distribution.

Trade-offs

While effective, speculative decoding requires careful tuning and introduces its own set of trade-offs.

• Draft Model Choice: The ideal draft model is a balancing act. It must be significantly faster than the MoE model, but accurate enough to achieve a high acceptance rate. A model that is too small will produce low-quality drafts, leading to frequent rejections and minimal speedup. A good candidate for a draft model is often a distilled version of the target MoE, as it is explicitly trained to mimic the larger model's behavior.

• Memory Overhead: This technique requires loading both the target MoE and the draft model into GPU memory. While the draft model is small, this adds to the already substantial memory pressure of serving a large MoE. This trade-off of memory for latency must be evaluated based on the available hardware.

• Draft Length ($k$): The number of tokens to draft, $k$, is a critical hyperparameter. A larger $k$ provides a higher potential speedup but also increases the probability of an early rejection, which wastes the computation spent on generating the later part of the draft. The optimal $k$ depends on the task and the quality of the draft model. For highly predictable tasks like code generation, a larger $k$ may be effective. For more creative or complex tasks, a smaller $k$ is often safer.

By intelligently combining a fast but imperfect draft with a slow but accurate verification, speculative decoding effectively parallelizes the autoregressive process. It is a powerful tool for reducing the end-to-end latency of MoE inference, making these large, sparse models more practical for interactive applications.