🤖 machine learning

Beyond KV Reconstruction: Functional Reconstruction for MLA Draft Models in Speculative Decoding

This paper proposes a "Functional Reconstruction" method for speculative decoding that optimizes converted Multi-head Latent Attention (MLA) draft models to reproduce the post-output-projection behavior of their original MHA/GQA counterparts, thereby significantly improving token acceptance rates without requiring retraining or verifier supervision.

Original authors: Weiye Shi, Fanxu Meng, Muhan Zhang

Published 2026-07-31

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Weiye Shi, Fanxu Meng, Muhan Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to send a secret message across a crowded room. In the world of artificial intelligence, large language models are like brilliant but heavy-footed giants. When they write a story or solve a math problem, they have to remember everything they've said so far. This memory is called a "Key-Value (KV) cache." Think of it as a giant, ever-growing backpack the giant carries. As the story gets longer, the backpack gets heavier and heavier, slowing the giant down because it takes so much energy just to carry it around.

To fix this, scientists invented a new trick called Multi-head Latent Attention (MLA). Instead of carrying the whole heavy backpack, the giant now carries a tiny, compressed "summary note" (a latent state) that holds the same information but takes up way less space. It's like swapping a suitcase full of clothes for a single, smart postcard that tells you exactly what to wear. This makes the giant much faster and lighter.

However, there's a catch. Most of the smartest giants we have today were trained to carry the heavy backpacks. To make them use the new postcards, we have to "convert" them. But sometimes, this conversion is like translating a book into a different language and accidentally changing the meaning of a few key words. The giant can still walk, but if you ask it to guess the next word in a sentence, it might guess the wrong one. This paper explores a specific problem: when we try to use these converted giants to speed up writing by guessing ahead (a technique called "speculative decoding"), the conversion errors cause the guesses to fail, slowing everything down again. The researchers found a way to "re-tune" the conversion so the giant guesses correctly again, without needing to retrain the whole thing from scratch.

The Problem: The "Good Enough" Translation

Let's say you have a master chef (the original AI model) who knows exactly how to cook a perfect dish. You want to hire a sous-chef (the "draft" model) to guess the next ingredient so the master chef can cook faster. If the sous-chef guesses right, the master chef just nods and keeps going. If the sous-chef guesses wrong, the master chef has to stop, correct them, and start over, which wastes time.

Now, imagine you take a master chef who only knows how to cook with a massive, heavy wok (the old "KV cache" method) and force them to use a tiny, lightweight pan (the new "MLA" method). You can do this conversion without buying a new chef, but the new pan changes how the heat hits the food. The chef might still be able to cook, but their "taste" for the next ingredient might be slightly off.

The researchers found that when they used these converted chefs as the "sous-chefs" to guess ahead, they were guessing wrong far too often. The conversion process introduced tiny errors in how the chef processed information. In a normal cooking scenario, these errors might be invisible. But in the high-speed game of "guess the next word," even a tiny taste difference causes the master chef to reject the guess, turning a speed-up into a slowdown.

The Solution: Functional Reconstruction

The paper proposes a clever fix called Functional Reconstruction. Instead of trying to rebuild the chef from the ground up (which would take forever and cost a fortune), the researchers treat the converted chef like a musical instrument that is slightly out of tune.

Here is how they do it:

The Setup: They take the converted chef (the MLA model) and the original master chef (the frozen GQA model).
The Test: They feed both chefs the exact same ingredients (calibration data) and ask them to cook the same dish.
The Tuning: They look at the final plate (the output) from both chefs. If the converted chef's plate tastes even a tiny bit different, they tweak the specific knobs on the converted chef's pan (the query and key projections) until the taste matches the master chef's plate perfectly.

Crucially, they do this without asking the master chef to grade the final dish against a recipe book. They just want the converted chef to mimic the master chef's process exactly. This is called "functional reconstruction" because they are fixing the function (the output behavior) rather than just the structure (the size of the pan).

What They Found

The researchers tested this on 192 different scenarios, mixing different types of chefs (Llama and Qwen models), different conversion tools, and different tasks like writing code, answering questions, or summarizing news.

The Good News: In 37 out of 64 matched situations, this "tuning" made a big difference. The converted chefs started guessing the next word correctly much more often, which means the whole system could write faster. In some cases, the speed-up was significant, with acceptance rates jumping by over 4 percentage points.
The Neutral News: In 26 cases, the tuning didn't change much. The chefs were already doing okay, or the errors were too small to matter.
The Bad News: In just one case, the tuning actually made things slightly worse. This tells us that while the method is powerful, it isn't magic; it can't fix every single problem, especially if the initial conversion was a disaster.

Why This Matters

The most important takeaway from this paper is that converting a model and making it a good "guessing" model are two different jobs. Just because a model can be converted to save memory doesn't mean it will be good at helping other models work faster.

The researchers showed that you don't need to retrain the whole model or use a super-complex verification system to fix this. You just need to "re-tune" the converted model to match the original one's behavior on a small set of test data. This is a fast, efficient way to get the best of both worlds: the memory savings of the new technology and the accuracy of the old, proven models.

However, the authors are careful to note that this isn't a magic wand that solves everything. If the initial conversion is too broken, or if the computer system running the model (the "backend") is incompatible, this tuning can't fix it. But for many common cases, it's a simple, effective tool to make AI faster and smarter.

Technical Summary: Beyond KV Reconstruction – Functional Reconstruction for MLA Draft Models in Speculative Decoding

Problem Statement

Multi-head latent attention (MLA) has emerged as a critical architecture for long-context Large Language Model (LLM) inference, replacing the growing key-value (KV) cache with compact latent states to reduce memory traffic. While existing methods (e.g., TransMLA, MHA2MLA) enable the conversion of pre-trained Multi-Head Attention (MHA) or Grouped-Query Attention (GQA) models into MLA without costly retraining, these structural conversions often fail when the converted model is used as a draft model in speculative decoding.

The core issue is that speculative decoding relies on high agreement between the draft model's token proposals and the target model's verification. Direct conversion from MHA/GQA to MLA introduces two primary sources of error that degrade this agreement:

Low-rank factorization errors: The compression of KV projections into a lower-dimensional latent space cannot exactly reproduce the original stacked K/V maps when the original rank exceeds the latent rank.
RoPE handling errors: The decomposition of Rotary Position Embeddings (RoPE) during conversion introduces residuals in the retained positional paths.

While these errors might be tolerable for standalone generation, they cause significant drops in token acceptance rates during speculative decoding, turning the draft model into a computational overhead rather than an accelerator. Existing conversion methods focus on structural validity (cache efficiency) rather than functional fidelity (proposal agreement).

Methodology

The paper proposes Functional Reconstruction, an end-to-end (E2E), converter-agnostic post-conversion procedure designed to optimize draft quality without altering the converted cache structure or requiring verifier supervision.

Core Mechanism

The method treats MLA draft construction as a functional reconstruction problem rather than a cache-compression problem. It operates in three stages:

Initialization: A standard converter (TransMLA or MHA2MLA) provides an initial cache-capable MLA checkpoint.
Training Objective: The method trains only the query and KV projections introduced by the converter. The output projection ( $W_O$ $W_{O}$ ) and the original MHA/GQA blocks remain frozen.
- For a given calibration hidden state $h_\ell$ , the system computes the output of the frozen original GQA block ( $A^{orig}_\ell(h_\ell)$ ) and the trainable converted MLA block ( $A^{MLA}_\ell(h_\ell)$ ).
- The loss function minimizes the masked mean-squared error (MSE) between the post-output-projection responses of the two blocks:
  $\mathcal{L}^{(\ell)}_{E2E} = \mathbb{E}_b \left[ \frac{\| M_b \odot (A^{MLA}_\ell(H^b_\ell) - \text{sg}[A^{orig}_\ell(H^b_\ell)]) \|_F^2}{\sum_t M_{b,t}} \right]$
  where $M_b$ is a mask for valid tokens and $\text{sg}$ denotes stop-gradient.
Scope: The objective covers the entire attention function from layer input to post-output-projection output, jointly addressing low-rank KV conversion errors and RoPE handling errors.

Key Characteristics

Converter-Agnostic: The same objective refines checkpoints from different converters (TransMLA, MHA2MLA) while preserving their specific cache topologies.
No Verifier Supervision: It requires no target model logits, activations, or benchmark labels, relying solely on calibration hidden states from the original model.
Inference-Free: It introduces no new parameters, branches, or cache states during inference.

Key Contributions

Distinction of Objectives: The paper establishes that structural conversion (ensuring a valid MLA parameterization) and drafting (ensuring proposal agreement for speculative decoding) are distinct objectives. A model can be structurally valid but functionally poor as a draft.
Training-Time Optimization: It demonstrates that draft quality can be optimized entirely at training time using a functional objective that matches the original model's attention behavior, without needing verifier supervision.
General Applicability: The method is shown to be converter-agnostic, improving or maintaining performance across different converters (TransMLA, MHA2MLA), model families (Llama, Qwen), and backends (Hugging Face, vLLM).

Experimental Results

The authors evaluated 192 configurations spanning four draft-target pairs (Llama-1B/3B $\to$ Llama-8B; Qwen-1.5B/3B $\to$ Qwen-7B), two converters, two backends, and four tasks (HumanEval, Alpaca, Natural Questions, CNN/DailyMail).

Acceptance Rates: Compared to a "Partial RoPE" baseline (a converter-local reconstruction method), Functional Reconstruction yielded material acceptance improvements in 37 of 64 matched task cells (using a 0.5 percentage point tolerance).
- 26 cells remained practically unchanged.
- 1 cell showed a material decrease.
Throughput: Among the 37 cells with improved acceptance, 12 also showed material throughput increases. The method introduces no inference overhead, so throughput gains are driven by higher acceptance rates.
Interaction Effects: The results revealed a "crossover" pattern:
- Llama + TransMLA: Consistently improved.
- Qwen + MHA2MLA: Consistently improved.
- Llama + MHA2MLA: Mostly unchanged (suggesting the initial conversion was already functionally close).
- Qwen + TransMLA: Mixed results, with some backend sensitivity.
Recovery Limits: While the method recovers a significant portion of the acceptance lost during conversion, it does not always fully restore the performance of the original unconverted GQA draft, particularly when the initial conversion or backend implementation introduces severe mismatches (e.g., tensor parallelism semantics).

Significance and Claims

The paper argues that MLA conversion and MLA draft construction should be treated as related but separate tasks.

Hierarchical Success: The authors define three levels of success: structural validity (executable parameters), functional fidelity (matching attention maps), and draft utility (speculative acceptance). Current converters often satisfy the first but fail the latter two.
Functional Interface: The significance of Functional Reconstruction lies in placing the optimization objective at a shared functional interface (post- $W_O$ output) rather than assuming every converter leaves the same error distribution. It acts as a "training-time control knob" to refine the draft model within the fixed constraints of the latent rank and cache budget.
Practical Implication: The work suggests a three-level workflow for future MLA pipelines: (1) establish cache-efficient parameters, (2) apply functional reconstruction to optimize the attention map, and (3) validate via matched acceptance. It emphasizes that while functional reconstruction is a general refinement, it cannot compensate for fundamental mismatches in latent rank or backend implementation (e.g., tensor parallelism semantics).

The paper concludes that while Functional Reconstruction materially improves acceptance in many scenarios, achieving full speculative speedup still requires high-quality initial converters and optimized MLA kernels.