DDTree透過單次區塊擴散前向傳遞建構草稿樹，提升推測解碼加速比

DDTree透過單次區塊擴散前向傳遞建構草稿樹，提升推測解碼加速比。

DDTree（Diffusion Draft Tree）是一種新型推測解碼方法，從單次區塊擴散前向傳遞產生每個位置的token分佈，直接建構草稿樹，並以樹狀注意力在單次目標模型前向傳遞中驗證多條可能延續路徑。相較於vanilla DFlash僅驗證單一路徑，DDTree在固定節點預算下選擇最有潛力的分支，保留目標模型輸出分佈不變，並在所有60種資料集/模型/溫度組合中實現更高加速比，峰值達HumanEval上的8.22x，以及MATH-500上的10.73接受長度。

核心創新

DDTree解決vanilla DFlash的關鍵限制：雖然DFlash以單次區塊擴散前向傳遞產生整個草稿區塊，實現state-of-the-art性能並超越EAGLE-3等自迴歸草稿器，但僅驗證單一草稿路徑，浪費了每個位置的豐富分佈資訊。

• DDTree直接從DFlash的每個位置邊際分佈（marginal distribution）建構草稿樹，每個深度i的節點代表該位置候選token。
• 使用固定節點預算B，透過簡單的最佳優先堆疊演算法（best-first heap algorithm）選擇最可能匹配目標模型的延續，證明此樹最大化草稿模型下的期望接受長度（作為目標模型的代理目標）。
• 驗證階段採用祖先專用注意力遮罩（ancestor-only attention mask），讓目標模型在單次前向傳遞中評分整個樹，維持無損（lossless）特性，完全保留目標模型輸出分佈。

運作流程

每個解碼回合從前一bonus token b開始，b為目標模型已產生但尚未前向傳遞的保證token，草稿器可依賴其身分但不含其特徵。

• 步驟1：執行單次區塊擴散草稿器，產生接下來L個位置的每個位置分佈qi（邊際分佈，非路徑條件分佈）。
• 步驟2：在節點預算B下建構草稿樹，最大化因式分解分佈下的期望接受長度，證明最佳解為前綴閉包（prefix-closure）下的最高機率前綴集合。
• 步驟3：編譯樹為目標模型輸入張量，以樹狀注意力執行單次前向傳遞。
• 步驟4：樹狀走訪（tree walk）：依目標模型解碼規則選擇下個token，若匹配樹中子節點則接受前綴，直至不匹配，攜帶首個不匹配token作為下一回合bonus token。

與既有方法的比較

DDTree建基於DFlash的低草稿延遲優勢，超越傳統樹狀驗證與平行草稿方法。

• 相較SpecInfer與Medusa的樹狀驗證，DDTree專為區塊擴散設計，避免自迴歸草稿的多前向傳遞開銷。
• 優於EAGLE家族與OPT-Tree的自迴歸樹建構，每深度需單獨草稿前向傳遞，計算開銷更高。
• 區別於DART的連續性修剪（需外部N-gram分數與運行時N-gram trie），DDTree直接從DFlash per-position機率建樹，無需輔助評分，具明確代理目標且證明最優。

實驗結果與權衡

DDTree在多模型規模與領域一致超越vanilla DFlash，峰值加速比8.22x（HumanEval），接受長度10.73（MATH-500）。

• 在所有60種資料集/模型/溫度組合中，DDTree加速比均高於DFlash。
• MATH-500案例研究顯示明確權衡：節點預算B增大時接受長度穩步上升，但加速比在中間B峰值後下降，因驗證器成本主導。
• 接受長度直方圖轉移：短接受變稀少，全區塊接受大幅增加，強化長前綴機率。

理論基礎

代理目標為草稿模型因式分解近似下的期望接受長度，命題1證明其為前綴機率的可加和，命題2證明最佳前B高機率前綴自動形成有效樹。

• 區塊擴散平行預測每個位置qi，條件僅前區塊上下文，非區塊內先前token，DDTree精準利用此邊際分佈建樹。
• 樹建構高效，無需目標模型條件機率，直接最大化代理期望接受token數。

實作與再現

程式碼基於Hugging Face Transformers庫實現，支援CUDA PyTorch環境。

• 安裝：pip install -r requirements.txt。
• 基準測試：bash run_benchmark.sh，輸出於runs/，日誌於logs/。
• 再現論文圖表：python3 plot_results.py生成圖；python3 make_latex_table.py生成LaTeX表格。
• 專案頁：https://liranringel.github.io/ddtree；論文：https://liranringel.github.io/ddtree/DDTree.pdf；程式庫：https://github.com/liranringel/ddtree。

推測解碼背景與意義

推測解碼透過輕量草稿器提議多token，目標模型平行驗證，解決自迴歸模型逐token前向傳遞的延遲瓶頸，特別在大模型時代關鍵。區塊擴散如DFlash以單次傳遞產生整區塊，成本低且準確高，但傳統僅單路徑驗證限制接受長度。DDTree擴展此基礎，證明區塊擴散分佈可高效轉化為樹狀探索，提升端到端加速，置於領先推測解碼方法之列。作者Liran Ringel與Yaniv Romano邀請從事推測解碼或擴散草稿者提供回饋。