這項由高通AI研究院的Ramchalam Kinattinkara Ramakrishnan、Zhaocong Yuan等七位研究人員共同完成的研究，發(fā)表于2025年7月3日，論文編號為arXiv:2507.02659v1。感興趣的讀者可以通過arXiv平臺獲取完整論文內(nèi)容。這項研究解決了一個在人工智能領(lǐng)域頗為頭疼的問題：如何讓一個小巧的AI模型為各種不同的大型AI模型提供高效的"草稿服務(wù)"。
當我們使用ChatGPT、Claude這樣的大型語言模型時，它們需要逐字逐句地生成回答，就像一個作家在稿紙上一個字一個字地寫作。這個過程很慢，特別是在手機或其他移動設(shè)備上使用時更是如此。為了解決這個問題，研究人員想出了一個巧妙的辦法：讓一個小而快的"草稿員"模型先快速寫出初稿，然后讓大模型來檢查和修正這個初稿。這就像是讓一個速記員先快速記錄，然后讓專業(yè)編輯來潤色一樣。
然而，現(xiàn)實中存在一個棘手的問題。不同的AI模型就像來自不同國家的人，它們使用著不同的"詞匯表"。一個專門為Llama模型訓(xùn)練的草稿員，無法直接為Qwen或其他模型提供草稿服務(wù)，因為它們對同一個詞匯的理解方式不同。這就好比一個習慣了美式英語的速記員，突然要為一個只懂英式英語的編輯工作，兩者之間的詞匯差異會造成很多誤解。
高通AI研究院的團隊提出了一個名為OmniDraft的解決方案，它的核心思想是創(chuàng)建一個"萬能翻譯官"，讓同一個小型草稿模型能夠為任何大型目標模型提供服務(wù)。這個方案包含了三個巧妙的創(chuàng)新。
一、跨詞匯表的智能翻譯系統(tǒng)
研究團隊首先解決的是不同模型之間的"語言障礙"問題。他們設(shè)計了一個叫做"n-gram緩存"的翻譯系統(tǒng)，這個系統(tǒng)就像是一個智能詞典，能夠記住不同模型之間的詞匯對應(yīng)關(guān)系。
傳統(tǒng)的做法是只處理兩個模型詞匯表中完全相同的詞匯，這就像兩個人只能用共同認識的詞匯交流，大大限制了交流的豐富性。而OmniDraft的n-gram緩存更加聰明，它能夠處理更復(fù)雜的對應(yīng)關(guān)系。比如，草稿模型可能將"snowflake"（雪花）分解為"snow"、"f"、"la"、"ke"四個部分，而目標模型可能將其識別為"snow"和"flake"兩個部分。n-gram緩存能夠?qū)W會這種對應(yīng)關(guān)系，將草稿模型的四個片段正確地組合成目標模型能理解的兩個詞匯。
這個過程就像一個經(jīng)驗豐富的翻譯官，不僅能翻譯單個詞匯，還能理解不同語言中詞匯組合的方式。當草稿模型提出一系列詞匯片段時，翻譯系統(tǒng)會查看緩存，看看這些片段是否能組合成目標模型更喜歡的形式。如果找到了匹配的組合，就會將多個小片段合并成一個完整的詞匯，大大提高了被目標模型接受的可能性。
更重要的是，這個緩存系統(tǒng)是動態(tài)學(xué)習的。每當系統(tǒng)遇到新的詞匯對應(yīng)關(guān)系時，都會將其記錄下來，供將來使用。這就像一個翻譯官在工作中不斷積累經(jīng)驗，遇到的對應(yīng)關(guān)系越多，翻譯能力就越強。
二、在線混合蒸餾訓(xùn)練
解決了翻譯問題后，研究團隊面臨的第二個挑戰(zhàn)是如何讓草稿模型更好地理解目標模型的"思維方式"。他們開發(fā)了一種叫做"在線混合蒸餾"的訓(xùn)練方法。
這個過程可以比作師傅帶徒弟的學(xué)習方式。草稿模型（徒弟）在實際工作中觀察目標模型（師傅）的表現(xiàn)，然后調(diào)整自己的行為來更好地配合師傅。具體來說，當目標模型接受了草稿模型的建議時，草稿模型會記住這次成功的經(jīng)驗；當目標模型拒絕建議并給出修正時，草稿模型也會從這次"糾錯"中學(xué)習。
傳統(tǒng)的訓(xùn)練方法通常是離線進行的，就像學(xué)生在考試前突擊復(fù)習一樣。而OmniDraft采用的是在線學(xué)習方式，更像是邊工作邊學(xué)習的學(xué)徒制。這種方法的優(yōu)勢在于，草稿模型能夠根據(jù)具體的使用場景和用戶數(shù)據(jù)不斷調(diào)整自己，而不是一成不變地使用固定的知識。
混合蒸餾的"混合"體現(xiàn)在訓(xùn)練方法的靈活性上。對于可以直接對應(yīng)的詞匯，系統(tǒng)使用一種叫做"逆向KL散度"的方法來對齊兩個模型的概率分布，這就像讓徒弟學(xué)習師傅對同一個問題的判斷方式。對于需要通過n-gram緩存翻譯的詞匯，系統(tǒng)則使用"最大似然估計"的方法，重點提高這些詞匯被正確預(yù)測的概率。
研究團隊還引入了一個動態(tài)權(quán)重參數(shù)λ，用來平衡這兩種訓(xùn)練方式的重要性。這個參數(shù)可以根據(jù)實際情況調(diào)整，比如當遇到的翻譯詞匯較多時，可以增加翻譯相關(guān)訓(xùn)練的權(quán)重；當直接對應(yīng)的詞匯較多時，則增加概率對齊訓(xùn)練的權(quán)重。
三、自適應(yīng)草稿長度調(diào)整
OmniDraft的第三個創(chuàng)新是智能的草稿長度調(diào)整機制。這個機制就像一個經(jīng)驗豐富的秘書，能夠根據(jù)不同情況調(diào)整匯報的詳細程度。
在實際應(yīng)用中，草稿模型需要決定每次應(yīng)該提供多少個詞匯建議。提供太少的建議可能無法充分利用加速的潛力，而提供太多的建議則可能導(dǎo)致大部分被拒絕，反而浪費計算資源。傳統(tǒng)的做法是使用固定的草稿長度，但這顯然不夠靈活。
OmniDraft引入了一個"接受率預(yù)測頭"，這個小型神經(jīng)網(wǎng)絡(luò)能夠預(yù)測每個詞匯建議被目標模型接受的可能性。基于這些預(yù)測，系統(tǒng)會動態(tài)計算繼續(xù)提供更多建議的風險。如果預(yù)測顯示后續(xù)建議被拒絕的概率很高，系統(tǒng)就會提前停止，避免浪費計算資源。
這個預(yù)測機制使用了一種叫做"sigmoid函數(shù)"的數(shù)學(xué)工具來估計接受概率，然后計算所有建議中至少有一個被拒絕的總體概率。當這個概率超過預(yù)設(shè)的閾值時，系統(tǒng)就會停止生成更多建議。這就像一個精明的銷售員，能夠判斷客戶的興趣程度，在合適的時候結(jié)束推銷。
在在線學(xué)習環(huán)境中，這個預(yù)測頭面臨著一個特殊的挑戰(zhàn)：隨著草稿模型不斷改進，詞匯被接受的概率也在變化，這意味著預(yù)測頭需要同步調(diào)整。研究團隊提出了兩種解決方案。
第一種是"聯(lián)合訓(xùn)練"方法，讓草稿模型和預(yù)測頭同時更新。這種方法簡單直接，但可能因為兩個組件的學(xué)習速度不同而產(chǎn)生不穩(wěn)定性。第二種是"交替訓(xùn)練"方法，為預(yù)測頭維護一個更大的數(shù)據(jù)緩沖區(qū)，包含歷史數(shù)據(jù)，這樣可以提供更穩(wěn)定的訓(xùn)練環(huán)境。實驗表明，交替訓(xùn)練方法通常能獲得更好的性能。
四、實驗驗證與性能表現(xiàn)
為了驗證OmniDraft的有效性，研究團隊進行了大量的實驗測試。他們選擇了一個僅有68M參數(shù)的Llama小模型作為草稿員，并測試了它與多個不同大型模型的配合效果，包括Llama3-8B、Qwen2-7B和Vicuna-7B。
實驗涵蓋了四個不同的任務(wù)領(lǐng)域。在數(shù)學(xué)推理任務(wù)中，他們使用了GSM8K數(shù)據(jù)集，這個數(shù)據(jù)集包含了各種小學(xué)數(shù)學(xué)應(yīng)用題。在編程任務(wù)中，他們結(jié)合了MBPP和HumanEval兩個代碼生成數(shù)據(jù)集。在文本生成方面，他們使用了Alpaca指令跟隨數(shù)據(jù)集。在文本摘要任務(wù)中，他們采用了XSum新聞?wù)獢?shù)據(jù)集。
實驗結(jié)果令人印象深刻。在跨詞匯表的場景中，傳統(tǒng)的直接映射方法（SpDDM）幾乎無法工作，接受率通常在0.1左右，加速比甚至低于1，這意味著不僅沒有加速反而變慢了。而使用OmniDraft的LDM（直接映射訓(xùn)練）方法能將接受率提升到0.2-0.4的范圍，加速比達到1.2-1.6倍。
當加入n-gram損失項后，性能進一步提升。LDM + λLN-gram方法在所有任務(wù)上都表現(xiàn)出色，接受率通常能達到0.2-0.4，加速比在1.2-1.7倍之間。特別值得注意的是，在GSM8K數(shù)學(xué)推理任務(wù)上，無論是Llama3-8B還是Qwen2-7B作為目標模型，都能獲得最大的加速效果，這可能是因為數(shù)學(xué)推理任務(wù)具有更強的結(jié)構(gòu)性和可預(yù)測性。
研究團隊還測試了使用LoRA（Low-Rank Adaptation）技術(shù)的效果。LoRA是一種參數(shù)高效的微調(diào)方法，只需要更新模型的一小部分參數(shù)。實驗顯示，即使使用LoRA這種"輕量級"的訓(xùn)練方式，OmniDraft仍然能夠獲得顯著的性能提升，雖然效果略低于全參數(shù)微調(diào)，但對于資源受限的邊緣設(shè)備來說，這種方案提供了很好的性能和效率平衡。
在自適應(yīng)草稿長度調(diào)整的實驗中，研究團隊發(fā)現(xiàn)了一些有趣的現(xiàn)象。聯(lián)合訓(xùn)練方法雖然能夠獲得更高的接受率，但在某些任務(wù)上的加速比反而不如交替訓(xùn)練方法。這表明高接受率不一定直接轉(zhuǎn)化為更好的加速效果，可能是因為聯(lián)合訓(xùn)練方法容易低估接受概率，導(dǎo)致過早停止生成建議。