Qwen-2.5-1.5b-mtp-speculative

Qwen2.5-1.5B + Multi-Token-Prediction Speculative Decoding Heads

This model augments Qwen2.5-1.5B with two lightweight Multi-Token-Prediction (MTP) heads for self-speculative decoding — the model drafts multiple future tokens per forward pass using its own auxiliary heads, then verifies them against its own backbone, requiring no separate draft model.

• Head 0 predicts the token at t+2 (one step beyond the backbone's own t+1 prediction)
• Head 1 predicts the token at t+3, conditioned on Head 0's own prediction

At inference time, both heads are fully autoregressive and self-conditioned: Head 0 consumes the backbone's own greedy/sampled token embedding, and Head 1 consumes Head 0's own predicted token embedding — never a ground-truth "oracle" token, since that information isn't available yet during real drafting. See Architecture below for details.

Architecture

Each MTP head (MTPModule) is a small transformer block that fuses two signals:

1. Previous hidden state — either the backbone's final hidden state (for Head 0) or the previous head's own output hidden state (for Head 1)
2. Shifted token embedding — the embedding of the token at the head's target offset

fused = concat(LayerNorm(prev_hidden), LayerNorm(shifted_embedding))
hidden = Linear(fused)
hidden = CausalTransformerEncoderLayer(hidden)
logits = LMHead(LayerNorm(hidden))

Each head's lm_head weight is initialized as a copy of the backbone's own lm_head weight, giving the heads a strong prior toward the backbone's output distribution before any training happens. Both the hidden-state branch and the embedding branch are independently normalized (via separate LayerNorm layers) before being concatenated and projected — normalizing both branches, rather than just one, keeps their scales comparable going into the fusion projection.

Training Process

Training proceeded in two stages, both using a streamed mixture of FineWeb-Edu (45%), open-web-math (20%), and CodeSearchNet (35%).

Stage 1 — Causal MTP head training (hard-label cross-entropy)

The MTP heads were trained from scratch on top of a partially-unfrozen backbone (last 6 transformer layers + final norm + lm_head unfrozen; all earlier layers frozen). Each head was trained with a standard next-token cross-entropy loss against the ground-truth token at its target offset (teacher-forced embeddings), with the backbone's own next-token loss included in the same joint objective so the unfrozen backbone layers keep adapting alongside the new heads. This stage establishes the heads' basic capacity to predict multiple tokens ahead before any distillation signal is introduced.

Stage 2 — KL-distillation against the backbone's own distribution

The key insight motivating this stage: the acceptance test used at real inference time (rejection sampling / greedy-match verification) compares a head's draft distribution against the backbone's own output distribution at that position — it does not care whether the draft matches the ground-truth corpus token, only whether it matches what the backbone itself would have predicted. Hard-label cross-entropy against the corpus is only a proxy for this; it's optimizing a different target than what acceptance rate actually measures.

Stage 2 resumes from the Stage 1 checkpoint and adds a forward-KL distillation loss per head: the backbone's own softmax distribution (temperature-softened, stop-gradient teacher) at the matching future position becomes the training target for that head's output distribution, directly optimizing what the acceptance test checks. A small hard-label cross-entropy term is retained alongside the KL term as a regularizer, to keep the heads from drifting toward a degenerate distribution early in training.

def distillation_kl_loss(student_logits, teacher_logits, temperature):
    student_log_probs = F.log_softmax(student_logits.float() / temperature, dim=-1)
    teacher_probs      = F.softmax(teacher_logits.float().detach() / temperature, dim=-1)
    kl = F.kl_div(student_log_probs, teacher_probs, reduction="batchmean", log_target=False)
    return kl * (temperature ** 2)

loss = (
    main_loss
    + kl_weight_h0 * kl_h0 + ce_weight_h0 * loss_h0_ce
    + kl_weight_h1 * kl_h1 + ce_weight_h1 * loss_h1_ce
)

Teacher/student index alignment: head k (shift = k+1) at local position i predicts the token at global position i+shift+1. Since main_logits[:, j, :] predicts token j+1, setting j+1 = i+shift+1 gives j = i+shift — so the teacher slice for head k is main_logits[:, shift:-1, :], length-aligned to match that head's own logits.

Benchmarks

Evaluated out-of-sample — prompts drawn from the same streaming data mixture used in training, but skipped well past any record the training loop could have consumed, guaranteeing zero train/test overlap. All numbers below use float16 (found empirically faster than bfloat16 on this hardware — see note below) with KV-cached, single-parallel-verify-per-round self-speculative decoding.

Greedy decoding (do_sample=False)

Sampling (do_sample=True, temperature=0.5)

Sampling mode uses the standard speculative-decoding rejection test — a draft is accepted with probability min(1, p_target(x)/p_draft(x)), and on rejection the next token is drawn from the normalized residual distribution max(0, p_target - p_draft). This is provably distributionally equivalent to sampling directly from the backbone at the same temperature, so sampling-mode speedup comes "for free" without altering the output distribution.

Hardware note: bfloat16, despite being architecturally supported, was measurably slower than float16 for this workload on the T4 GPU used for benchmarking (0.68x vs 2.15x+ speedup under otherwise identical settings) — a known T4-specific throughput characteristic, not a correctness issue. If deploying on newer datacenter GPUs (A100/H100), re-benchmark both dtypes before choosing.

Usage

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "Harish241412/Qwen-2.5-1.5b-mtp-speculative",
    torch_dtype=torch.float16,   # benchmark this vs bfloat16 on your own hardware
    device_map="cuda:0",
    trust_remote_code=True,      # required: this repo ships a custom architecture
)
tokenizer = AutoTokenizer.from_pretrained("Harish241412/Qwen-2.5-1.5b-mtp-speculative")

inputs = tokenizer("def calculate_factorial(n):", return_tensors="pt").input_ids.to("cuda:0")

# Greedy self-speculative decoding
outputs, metrics = model.generate_speculative(inputs, max_new_tokens=80, do_sample=False)

# Sampling self-speculative decoding
outputs, metrics = model.generate_speculative(
    inputs, max_new_tokens=80, do_sample=True, temperature=0.5
)

print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print(metrics)  # {'total_rounds': ..., 'both_heads_accepts': ..., ...}

generate_speculative currently supports batch size 1 only — per-sequence KV cache management for batched speculative decoding (as done via paged attention in serving stacks like vLLM/TGI) is not implemented in this reference version.

Limitations

• Batch size 1 only for generate_speculative.
• Accept rate (and therefore realized speedup) varies with prompt domain — code-like, boilerplate-heavy prompts tend to draft more successfully than open-ended prose, since the training mixture skews toward code/technical text.
• The eager PyTorch reference implementation here is for benchmarking and validation; production deployment (e.g. via vLLM/SGLang with a compiled speculative-decoding engine) would be expected to realize a larger fraction of the theoretical speedup ceiling implied by the acceptance rate.
• Only 2 MTP heads / 2 lookahead positions are supported in this release.