多詞元預測頭如何加速推理而不影響輸出品質

Mellum2.1 Thinking MTP 模型的核心變化在於加入了一個與主模型共同訓練的額外 Transformer 區塊,稱為多詞元預測(MTP)頭。這個組件的功能是預測剛剛採樣出的詞元之後的下一個詞元。

在推理過程中,引擎將 MTP 頭的預測結果作為投機性草稿(speculative draft)。主模型會逐一驗證這些草稿詞元,由於最終輸出仍由主模型決定,因此生成內容保持不變。

根據 JetBrains 提供的數據,這種機制能使解碼速度提升 1.4 至 1.8 倍,且約有 90% 的草稿詞元會被接受並直接採用,從而減少實際需要從頭生成的步驟。

模型檔案結構與主要推理框架的啟用方式

MTP 頭的權重儲存在名為 model-mtp.safetensors 的檔案中,並使用 mtp. 作為前綴標識。在 config.json 設定檔裡,透過 num_nextn_predict_layers: 1 來宣告此功能的存在,其餘模型結構與標準版 Mellum2.1 Thinking 完全一致。

對於使用 vLLM 0.33 或更新版本的使用者,需在啟動服務時加入參數 --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' 以啟用投機性解碼。

若使用 transformers 5.20 或更新版本,則在呼叫 generate 函數時需設定 use_mtp=True。至於 llama.cpp 使用者,應從 GGUF 儲存庫下載對應檔案,並使用 --spec-type draft-mtp --spec-draft-n-max 1 參數進行執行。