핵심 개념
本文提出了一種基於上下文感知的助手模型選擇方法,通過在推理過程中動態選擇最合適的助手模型,以提升大型語言模型在資源受限環境下的推理速度,並在多個任務和模型架構上驗證了該方法的有效性。
摘要
大型語言模型 (LLM) 雖然應用廣泛,但在資源受限的環境下使用仍然存在困難,其不斷增長的規模也增加了使用的門檻。其中一個顯著的問題是自回歸生成過程中的高延遲,使得大型 LLM 的使用依賴於先進的計算基礎設施。助手解碼方法利用較小的草稿模型指導較大的目標模型生成,有助於緩解這一問題,但其性能取決於兩個模型之間的一致性。因此,如果草稿模型在某些領域的能力相對於目標模型不足,則性能可能會下降。或者,可以利用多個草稿模型來更好地覆蓋目標模型的專業知識,但是當有多個黑盒草稿模型可用時,在不了解其構造細節的情況下選擇助手模型可能會很困難。為了更好地理解這個決策問題,我們將其視為一個上下文老虎機問題,其中策略必須根據上下文選擇草稿模型。我們證明,即使沒有草稿模型的先驗知識,僅根據獨立的草稿/目標模型的輸出創建離線數據集,並根據這些輸出的對齊程度訓練策略,就可以在候選模型有效的情況下加速多個領域的性能。進一步的結果表明,這在具有多個助手解碼候選模型的各種設置中都成立,突出了其靈活性以及此類決策可以發揮的積極作用。
引言
隨著 Transformer 的出現,大型語言模型 (LLM) 時代來臨,並發展出能夠以驚人的人類方式進行推理和行動的 LLM。然而,使用資源密集型模型和技術仍然是一個先決條件,因此,已經開發並應用了一些方法來緩解與這些模型的實際可用性相關的問題。隨著時間的推移,一個取得持續進步的主要領域是文本生成的自動回歸解碼方面,其中每個新標記的生成都需要通過模型進行完整的推理過程,這沒有充分利用注意力機制和現代加速器(例如 GPU、TPU)並行計算的能力。
解決這一問題的一個越來越流行的方法是推測解碼。在推測解碼中,通過最小化高延遲順序計算的數量並用更便宜的計算代替它們來減少延遲。並非直接從較大的模型中採樣,而是通過接受-拒絕採樣使用來自較小且更便宜的模型的樣本來近似採樣。具體來說,一個小的草稿模型自動回歸地生成文本,然後由一個更大的目標模型並行驗證。因此,大型模型不需要重複生成文本,而是在真正無能為力的時候通過糾正輸出來指導小型模型。這可以減少對大型 LLM 的調用次數,從而節省時間和內存。然而,這需要兩個模型,以及它們在生成能力上有一定的相似性,才能使這種方法獲得顯著的加速。雖然存在一些方法可以規避其中一些需求,但這些方法通常受到額外微調需求的限制(在資源受限的環境中很難實現)或生成質量下降的限制。由於規模和成本之間存在明顯的權衡,如果草稿模型與目標模型很好地對齊,那麼這種方法將非常有效。然而,雖然可以確保最終輸出遵循目標分佈,但選擇不充分的草稿模型可能會由於大量的拒絕而導致缺乏加速。雖然其他方法允許改變輸出分佈偏移以進一步加速推理,但此類偏移在許多高風險場景中可能會產生問題。從這個角度來看,擁有多個草稿模型(每個模型都適用於不同的設置)有助於在不降低生成質量的情況下加速推理。通過動態選擇草稿模型,可以在多個領域實現加速,而額外成本卻很小。然而,這需要學習如何在給定上下文的情況下選擇最佳草稿選項,這就引入了一個需要解決的決策問題。
那麼如何學習這個決策過程呢?我們首先將其視為一個上下文老虎機問題,目標是讓策略根據給定的查詢選擇一個草稿模型。這也需要獎勵,策略可以從中學習估計和比較可以採取的不同行動的理想程度。為此,我們使用一個離線過程來估計不同草稿模型與目標模型在一組訓練樣本上的上下文對齊程度,從而構建一個數據集,定義目標模型對特定草稿候選模型的偏好。這使我們能夠訓練一個策略,在不知道草稿模型更多細節的情況下考慮這些偏好。通過在推理時部署此策略,就可以權衡這些偏好,從而加快目標模型的生成速度。我們進一步表明,當使用自推測解碼時,該策略非常有用,其中草稿模型是目標模型參數的一個子集。
方法
假設一個大型目標模型 Me 產生了很大的端到端延遲,我們希望避免這種延遲。推測解碼旨在通過使用草稿模型來近似目標模型來解決延遲問題。然而,如前所述,草稿模型必須與目標模型相似,否則採樣分佈將過於不同,無法提高速度。因此,雖然草稿模型可以提供幫助,但只有當它們的知識分佈與目標模型相似時,它們才是可靠的。因此,如果目標模型具有多種專業知識,那麼僅使用一個草稿模型可能無法很好地服務於一般情況。但是,通過在任何給定場景中動態選擇不同的草稿模型,只要決策者有能力且高效,就可以觀察到每個草稿模型的優勢。
當面對查詢 q 時,從多個唯一的候選模型中選擇一個草稿模型可能會導致性能的不同,具體取決於所選的選項。從上下文老虎機的角度來看,q 是一個上下文,其中有 k 個臂,每個臂返回一個獨立的獎勵 r。每個臂對應一個不同的草稿模型,其獎勵是通過推測解碼生成輸出序列所需的時間。因此,每個臂可以為每個 q 產生不同的獎勵。然後,目標是學習一個策略 π(·|q),對於任何給定的上下文 q,該策略可以從臂中選擇可以產生最大獎勵的臂。從推測解碼的角度來看,目標是選擇其能力與給定查詢的目標最匹配的草稿模型,因為這將最小化目標模型必須被調用的次數。
隨機選擇草稿模型可能會顯著增加延遲,因此以樣本高效的方式學習做出正確的決策非常重要。雖然理想的獎勵是真實/觀察到的加速,但如果與草稿模型的對齊程度未知,這可能會很昂貴。因此,可能需要一個更便宜的代理。然而,有兩個因素直接影響真實獎勵:1)目標模型和草稿模型之間的對齊程度,以及 2)草稿模型的大小。這提供了一種收集策略訓練數據的替代方法:使用草稿模型自動回歸地生成文本,並計算與目標模型輸出的對齊分數,然後根據草稿模型的大小調整這些分數。接下來,我們將描述如何收集數據以離線訓練策略。
給定一組查詢 Q = {qi}n i=1,我們為目標模型 oe i 以及每個候選草稿模型 {oj i}k j=1 生成基於每個 qi 的輸出。然後,我們使用相似性度量來計算每個候選輸出的分數
sj i = f(oe i, oj i)
作為衡量目標模型和候選模型在 qi 上對齊程度的一種方法。還可以結合推理速度得分。例如,如果我們將特定草稿模型的推理速度的某個相對度量視為 cj i,則可以將分數調整為加權和
sj i = α · f(oe i, oj i) + (1 −α) · cj i
其中 α ∈ [0, 1] 權衡這兩個組成部分。
實驗結果
學習選擇草稿模型
在我們的第一個實驗中,我們使用 T5 編碼器-解碼器模型。作為目標模型,我們使用經過指令微調的 Flan-T5-XXL,而我們的草稿候選模型是公開可用的 T5-Small 模型,一個是基本版本,另一個是在文本摘要上進行了微調。我們在翻譯(IWSLT2017 英德)和文本摘要(XSUM)上進行評估。
表 1 比較了相同大小的草稿模型在其專業領域的差異。雖然每個模型在其知識領域(T5-Small 的英德翻譯和 T5-Small-XSum 的文本摘要)內都能顯著加速生成,但在其專業領域之外使用時,它們基本上沒有幫助,甚至有害,如 T5-Small 在 XSUM 上的速度降低了 1%,以及 T5-Small-XSum 在英德翻譯上使用時的速度降低了 17%。相比之下,該策略確保了兩個領域內的加速,而決策的延遲可以忽略不計。
這突出了使用策略的一些直接好處,即它可以在不知道任何關於草稿候選模型本身的明確信息的情況下,識別出適合上下文的正確草稿模型。相反,僅從每個草稿模型和目標模型中生成採樣輸出就足以通過使用計算出的獎勵來發展區分不同領域的通用能力。
平衡質量和速度
同樣重要的是,草稿模型的使用成本必須相對於目標模型足夠低。這促使我們進行了第二個實驗,該實驗僅在 XSUM 上進行評估,但比較了在大小和目標模型對齊方面不同的草稿候選模型。比較了多個草稿模型:Flan-T5-Small(8000 萬參數)、上述相同的 T5-Small(6000 萬參數)模型和 Flan-T5-Base(2.2 億參數)。
表 2 顯示了使用這些不同的草稿候選模型和 Flan-T5-XXL 目標模型通過推測解碼獲得的加速比。儘管較大的草稿模型可能比較小的草稿模型與目標模型更好地對齊,但使用它們會導致延遲,最終效率可能更低。因此,在候選模型之間做出決策時,平衡對齊和效率是一個需要考慮的問題。圖 2 顯示了離線訓練如何幫助實現這一點。通過設置 α 在 §2.3 中定義的目標之間變化,其中我們根據草稿模型的大小使用固定的推理成本,我們觀察到動態策略如何最終適應 α 選擇設置的偏好。
例如,隨著 α 接近 1,無論質量如何,策略都越來越傾向於選擇最小的草稿模型。同時,α →0 表明越來越傾向於選擇與目標生成最匹配的模型。這證明了使用這種不同獎勵的加權方案所帶來的通用靈活性,同時也證明了即使是更簡單的推理懲罰代理也足以正確地平衡這兩者。
需要學習多少個例子才能區分?
還需要考慮決策者需要多少個例子才能正確學習區分不同的例子。為此,我們研究了隨著用於訓練策略 πθ 的例子數量的增加,策略學習使用離線數據集中標註的分數以顯示可見的速度提升的速度。
我們重複使用第一個實驗中的模型,但隨著用於訓練策略 πθ 的例子數量的增加,我們跟踪解碼速度。
正如我們在圖 3 中觀察到的那樣,學習選擇正確模型的速度相當快,因為訓練不到 10000 個例子就足以達到與在整個離線數據集(包含近 40 萬個例子)上訓練相當的性能水平。這一結果證明了該方法的普遍效率,因為從少量例子中收集輸出並訓練策略顯示出相當強的泛化能力。
自動回歸生成作為一種選擇
在某些情況下,草稿模型將毫無用處,在這種情況下,自動回歸地使用目標模型仍然是最合理的选择。
為此,我們嘗試觀察向決策者提供此選項如何影響我們之前的實驗。我們重複表 1 中的相同實驗,但允許我們的策略學習選擇自動回歸地生成。為了避免輸出完全匹配,我們從目標模型中採樣輸出,並根據貪婪輸出來評分。由於目標模型相對於草稿模型的規模很大,我們使用 α = 0.5 來平衡大小和質量得分。
表 3 顯示添加自動回歸選項不會降低解碼速度或生成質量,實際上可能會加速整體解碼。這可能是因為某些上下文對於所有候選草稿模型來說都是分佈外的;在這些情況下,自動回歸生成可以顯著提高效率。
我們進一步驗證了該策略是否可以在草稿模型無用的情況下忽略它們。我們通過在任務中添加 GSM8K 來進行實驗,只有目標模型與之對齊。由於兩個草稿模型都無法加速此任務的推理,因此理想情況下,該策略應避免為來自此設置的例子起草。由於 GSM8K 比英德翻譯和 XSUM 小得多,因此我們減少了例子數量以匹配所有數據集的大小。
表 4 顯示自動回歸生成(不出所料)優於輔助生成。然而,使用策略顯示出與自動回歸生成相當的速度,這表明它學會了忽略草稿模型,因為每個模型的貪婪輸出之間存在明顯的差異。
推廣到多任務草稿模型
為了證明該方法對更一般設置的適用性,特別是在草稿模型可能勝任多項任務的情況下,我們進一步將基於策略的選擇方法應用於 SpecBench,使用 Vicuna-33B 作為目標模型,並使用較小的草稿模型(表 5)。鑑於 SpecBench 的規模(480 個例子,平均分為 6 個任務),我們僅將其用作測試集。為了訓練我們的策略,我們使用了提取 SpecBench 例子的原始任務數據集,並從每個數據集中採樣相同數量的例子(2000 個)。對於 MT-BENCH,總共只有 80 個例子,這些例子都包含在測試集中,但我們通過替換採樣來用於訓練集。因此,此任務的結果可能過於自信。由於 Vicuna 模型是僅解碼器 Transformer,因此我們將句子表示調整為輸入序列的最終隱藏表示。我們的結果表明,我們從 T5 架構中得到的初步發現仍然成立,這表明這種基於策略的訓練方法既穩健又可以推廣到不同的設置。
自我草稿的消融研究
儘管輔助解碼有很多好處,但草稿依賴於以下小型草稿模型的可用性:
(1) 與目標模型共享詞彙表。
(2) 在感興趣的任務上與目標模型對齊。
這樣的模型可能很難獲得,導致自我草稿,其中草稿模型存在於目標模型中。為了探索此設置的差異,我們進行了額外的消融研究。
我們使用 LLaMA-2-13B-Chat 模型,並使用提前退出,遵循 Kavehzadeh 等人的方法,對所有退出使用單個語言建模頭。雖然存在其他方法,但這些方法可能具有組合數量的潛在草稿選項,並且需要在推理過程中預先確定路徑流。同時,對並行解碼使用額外語言建模頭的方法需要額外的參數,這些參數可能與資源約束相衝突,或者降低生成質量。
在 ALPACA 和 TRIVIAQA 上分別進行的實驗結果(表 6)表明,儘管中間層草稿導致解碼速度降低,但使用策略可以最大程度地減少性能損失,特別是在存在自動回歸選項的情況下,這突出了所提出的離線策略學習方法也具有自我草稿的潛力。這表明使用策略仍然是回退到最有效解碼選項的有用方法。此外,當考慮沒有自動回歸選項的情況時,我們注意到策略方法能夠在 ALPACA 上恢復到與最佳情況中間草稿模型相似的性能。雖然 TRIVIAQA 並非如此,但这可能是由於該數據集中答案較短造成的。
討論
LLM 路由
LLM 在一系列任務中都表現出了非凡的能力,但它們的成本和能力存在很大差異。廣義上講,能力越強的模型往往比能力較弱的模型更昂貴。這在現實世界中部署 LLM 時會導致一個困境——將所有查詢都路由到最大、能力最強的模型會產生最高質量的響應,但成本可能很高,而將查詢路由到較小的模型可以節省成本,但可能會導致響應質量較低。同樣,並非所有模型都同樣適合同一組任務,這意味著路由到最合適的模型也非常重要。
我們的工作與這種模型路由的概念(即根據查詢選擇最佳模型)有很多相似之處。特別是,可以將草稿模型集視為一組子網絡,類似於專家混合 (MoE) 風格的範例。同時,該策略充當正確子網絡的路由器。更先進的路由技術已被探索作為一種利用野外存在的眾多 LLM 的方法,但尚未在推測解碼等下游設置中廣泛使用。
自適應推測解碼
推測解碼方法需要使用許多預定義的超參數,這些超參數會顯著影響加速,即使是很小的變化也會產生顯著的影響。最近的工作已經開始探索如何解耦這個過程,例如通過動態選擇在每個解碼步驟中生成的草稿標記的數量。Kavehzadeh 等人進一步討論了為每個實例動態選擇模型,但是由於需要在提前退出時在生成後計算置信度得分,因此他們的方法僅限於他們的特定設置。
雖然我們沒有引入新的解碼算法,但我們首次嘗試通過根據輸入在多個草稿模型之間切換來使推測解碼具有自適應性。然而,可能需要更複雜的自適應級別,因為每個解碼步驟可能不同,因此可能需要在整個過程中仔細調整不同的超參數,以便最大程度地提高加速。
輔助解碼的決策
輔助解碼可能需要做出多個決定。其中之一是確定每個步驟要解碼的理想草稿標記數量。另一個與如何拒絕標記有關,通常使用貪婪或基於採樣的標記匹配啟發式算法。然而,在強制執行特定選擇時存在權衡,這需要進一步研究以更好地理解如何調整此類技術。
這項工作建議在解碼過程開始時(即在過程中開始時)添加一個額外的決策,前提是存在多個草稿選項。雖然我們將自己限制在一個更獨立的環境中進行更完整的分析,但有多種方法可以讓這些方法在一個更大的管道中共存。然而,由於理解不同選擇如何以及如何共同影響整體報告結果並非易事,因此此類工作留待將來探索。
衡量輸出之間的對齊程度
我們觀察到,標記級別的相似性得分對於訓練決策者是有效的,這可以歸因於輔助解碼本身依賴於匹配輸出的標記級別分佈。因此,如果草稿模型的貪婪解碼輸出與目標輸出高度相似,則邏輯空間中的概率分佈之間的相似度也越高,這在採樣時會導致更少的拒絕。
然而,此類指標有其局限性,因為它們主要捕獲文本的表面元素,其中分佈的邊際差異對輸出文本有很大影響。此外,不同的指標可能會過擬合特定的任務,因此需要更好地衡量草稿/目標對齊,這有望通過設計更好的指標本身或通過學習比較不同粒度級別的特徵(例如,目標和草稿邏輯與文本輸出)來更好地估計獎勵,從而訓練出更好的策略。此外,語義含義也可能發揮重要作用,因為具有顯著結構的輸出可能仍然具有相同的含義,而這是標記級別的相似性指標無法充分捕捉到的。
推測解碼作為近似推理
推測解碼可以類比為一種近似推理形式,其中由於使用感興趣的模型執行推理的難處理性,因此使用近似方法來學習模型的估計值。雖然訓練草稿模型等同於執行變分推理(即用代理近似難處理的分佈),但这可能很昂貴。因此,僅訓練策略可以看作是對一組固定分佈進行加權,使其成為目標模型的更好代理。
一些工作進一步嘗試從這個角度研究推測解碼。特別是,Zhou 等人通過使用 KL 散度損失來構建草稿模型,有效地基於草稿輸出的似然信息構建了目標模型的後驗分佈。同時,Liu 等人探索了當例子分佈發生變化時相同的技術,構建了一個可以適應不斷變化的用戶輸入的草稿模型。這種設置也可能受益於多個草稿模型,其中以查詢為條件可以使草稿模型更有效地適應以更好地推廣到未見過的設置。
托管多個草稿模型
此方法的一個重要方面與需要與專家模型一起託管多個草稿模型有關。這可能會產生額外的成本,特別是當專家模型和選定的草稿模型不在同一設備上時。雖然自我草稿等方法避免了這個問題,並且創建最小規模草稿模型的可能性通常減輕了對過度內存使用的擔憂,但需要考慮的一個特殊方面是硬件級別的優化,這些優化可以最好地使選定的草稿模型以最快的速度加載,從而避免由於設備之間的數據傳輸帶來的帶寬限制而導致的額外延遲。
結論
這項工作首次嘗試將輔助生成集成到存在多個黑盒草稿候選模型的環境中。當沒有關於哪個草稿候選模型最適合輔助解碼給定例子的先驗知識時,該問題可以建模為上下文老虎機問題,其目標是估計每個草稿選項的未知獎勵。
我們的工作表明,離線強化學習提供了一種有效的方法,可以在此設置中學習區分可用選項並在各種例子中提供加速解碼,並提供了一種從模型中收集離線數據以供學習的邏輯方法。我們的結果和消融研究表明,使用這種方法學習策略可以適應一般的偏好,同時考慮決策過程中更複雜的方面,突出了其穩健性。此外,這種方法是可擴展的,並且對引入更多草稿模型或刪除草稿模型具有魯棒性,為無法獲得唯一優越的草稿模型的設置提供了一種可行的替代方案。
然而,仍然存在需要進一步發展的領域。例如,以在線方式學習可能會使該方法更廣泛地適用。或者,探索如何在每個解碼步驟(而不是每個例子)中動態選擇草稿模型,以及將此工作方向與試圖在每個步驟中自適應地選擇推測長度的方向相結合,是將我們的發現與當前工作相結合的可行方法,希望能從所有方法中獲益。
局限性
這項工作有一些局限性,這些局限性定義了未來工作的範圍。
草稿模型和數據域的選擇 結果可能源於域/任務之間存在的明顯界限。在這些界限不明確的環境中,結果可能会有所不同。然而,技術限制和缺乏足夠的預訓練模型進行比較使得立即探索這一點變得困難。
額外的存儲和內存 使用多個獨立草稿的模型需要額外的內存,當在這方面存在明確的約束時(自我草稿由於使用單個模型而避免了這一點),這可能更難管理。此外,在某些特定情況下,例如當輸入/輸出序列非常長或需要許多離線例子時,收集離線數據集可能會很困難。
自我草稿 我們在一個不進行任何與語言模型本身明確相關的參數的額外訓練的環境中工作,無論這些參數是現有參數還是由於該方法而添加的新參數。雖然我們探索的方法也可以應用於這些情況,但計算限制使得目前很難嚴格地進行此類研究,因此我們將其留給未來的研究。
道德聲明
本文討論了為推測解碼動態選擇多個黑盒草稿模型的概念,提出了一種離線強化學習方法來自適應地選擇良好的草稿模型以提供幫助。我們的結果與模型的解碼速度有關,這不太可能導致道德問題或對此類結果的錯誤解讀。
통계
使用較小的草稿模型 T5-Small 和 Flan-T5-Small,在 XSUM 數據集上進行推測解碼,分別獲得了 1.21 倍和 1.09 倍的加速比。
在 IWSLT2017 英德翻譯和 XSUM 數據集上,使用基於策略的草稿模型選擇方法,與單獨使用任何一個草稿模型相比,都能獲得更好的加速效果。
訓練策略所需的例子數量相對較少,僅使用不到 10000 個例子就能達到與使用整個離線數據集(包含近 40 萬個例子)訓練相當的性能水平。
在 SpecBench 基准測試中,使用 Vicuna-33B 作為目標模型,並使用 Vicuna-68m、Vicuna-160m 和 LLaMA-68m 作為草稿候選模型,基於策略的選擇方法在多個任務上都取得了與最佳草稿模型相當或更好的加速效果。