讓大模型的訓練和推理,比更快還更快!谷歌2022年終總結第四彈
雖然谷歌的Bard翻車了,但谷歌的AI實力仍然不容小覷。
自開年以來,由Jeff Dean領銜的Google Research年終總結系列「Google Research, 2022 & beyond」一直在持續更新,最近也是更新到了第四期。
本期以「提升模型效率」為主題,一起看看谷歌工程師們都想出了哪些ideas!
運行效率成關鍵
在過去十年里,深度學習呈現爆炸式發展,很大程度上是由于新算法和體系結構的融合、數據量的顯著增加以及計算能力的提高。
相比十年前,人工智能和機器學習模型變得更加巨大、更復雜,具有更深且更復雜的網絡結構、更多的參數、訓練時用到了更多的數據,共同促進了機器學習歷史上一些最具變革性的成果。
隨著這些模型越來越多地部署在生產和業務應用程序中,模型的推理效率和運行成本已經從一個次要因素變成了一個主要的限制因素。
Google在這方面的應對措施就是繼續在機器學習效率方面投入巨資,主要解決以下四個難題:
1、高效的模型架構(Efficient Architecture)
2、數據效率(Data Efficiency)
3、訓練效率(Training Efficiency)
4、推理效率(Inference Efficiency)
除了效率之外,模型還面臨著圍繞真實性、安全性、隱私性和時效性(freshness)等諸多難題。
這篇文章將重點介紹一系列Google Research研究公司為應對上述挑戰而開發的新算法。
高效的模型架構
一個最基本的問題是:是否有更好的方法來參數化模型以提高效率?
2022年,Google Research專注于新技術,通過檢索上下文增強模型,混合專家,使變壓器(位于大多數大型機器學習模型的核心)更有效,從而注入外部知識。
上下文增強模型(Context-augmented models)
為了追求更高的質量和效率,神經模型可以從大型數據庫或可訓練的記憶庫中的外部context中得到增強。
通過利用檢索到的上下文,神經網絡不需要在其內部參數中記憶大量的世界知識,就能夠獲得更好的參數效率、可解釋性和真實性。
在「上下文增強語言建模的解耦上下文處理」一文中,研究人員探索了一種基于解耦編解碼器架構將外部上下文合并到語言模型中的簡單架構。

論文鏈接:?https://arxiv.org/abs/2210.05758?
該模型能夠顯著地節省計算量,同時在自回歸語言建模和開放領域問題回答任務方面給出了有競爭力的結果。
預訓練的大語言模型(LLM)通過對大訓練集的自監督消耗了大量的信息,但目前還不清楚這些模型的「世界知識」是如何與輸入的context相互作用的。

通過知識感知微調(knowledge aware fine-tuning, KAFT) ,研究人員將反事實上下文和無關上下文結合到標準的監督數據集中,增強了 LLM 的可控性和魯棒性。
探索模塊化深度網絡的問題之一是如何設計具有相應計算模塊的概念數據庫,研究人員提出了一個理論架構,將「記憶事件」(remember events)以草圖(sketches)的形式存儲在一個外部 LSH 表中,其中包括一個pointers模塊來處理草圖。
上下文增強模型的另一個難題是從大型數據庫中快速檢索信息的加速器。
研究人員開發了一個基于 TPU 的最近鄰搜索算法,該算法與 TPU 的性能模型保持一致,并為預期召回提供分析保證,從而達到最佳性能。
搜索算法通常涉及大量的超參數和設計選擇,這使得它們很難在新任務上進行優化,所以研究人員進一步提出了一種新的約束優化算法來自動調整超參數,將期望的成本或召回作為輸入,該算法產生的調整在經驗上非常接近速度召回的Pareto frontier,并在標準基準上給出領先的性能。
混合專家模型
混合專家(MoE, Mixture-of-experts)模型已被證明是增加神經網絡模型容量而不會過度增加其計算成本的有效手段。MoEs 的基本思想是從多個專家子網絡中構建一個網絡,其中每個輸入由一個合適的專家子網絡處理。
因此,與標準的神經網絡相比,MoEs 只調用整個模型的一小部分,從而提高了語言模型應用程序(如 GLaM)的效率。

決定哪些專家應該被激活(be active)以參與到某一特定輸入取決于路由功能(routing function),路由的設計非常具有挑戰性,因為開發者的預期是每個專家模型都是恰如其分的,不會出現利用不足和過度的情況。

在最近的一項工作中,研究人員提出了專家選擇路由(Expert Choice Routing) ,一種新的路由機制,它不是將每個輸入token分配給 top-k 專家,而是反過來將每個專家分配給 top-k token,可以自動調節專家的負載平衡,同時自然地允許多個專家處理輸入token

論文鏈接:?https://openreview.net/pdf?id=jdJo1HIVinI?
高效的Transformers
Transformer是當下最流行的序列到序列的模型,從視覺到自然語言的理解,在一系列具有挑戰性的任務中已經展現了非常強大的性能。
這類模型的一個核心組成部分是注意力層(attention layer),通過計算「查詢」(query)和「鍵」(key)之間的相似度來構建一個適當的「值」(value)加權組合。雖然性能強,但注意力機制的計算效率卻不高,復雜度通常為輸入序列長度的二次方。
隨著Transformer規模的不斷擴大,其中一個重要問題的研究十分有價值:是否有任何自然發生的結構或模式的學習模型,可以解了注意力有效的原理。
對此,Google Research研究了中間 MLP 層的學習嵌入,發現它們非常稀疏,比如,T5-大模型只有 < 1% 的非零項,其稀疏性進一步表明,可以在不影響模型性能的情況下降低FLOP

論文鏈接:?https://arxiv.org/pdf/2210.06313.pdf?
研究人員最近提出了 Treeformer模型,一種依賴于決策樹的標準注意力計算的替代方案,直觀地說,它可以快速識別與查詢相關的一小部分鍵,并且只在這個集合上執行注意力操作。根據經驗,Treeformer 可以使注意力層的 FLOP 減少30倍。
同時研究人員還引入了序貫注意( sequential attention),一種將注意和貪婪算法相結合的可微特征選擇方法,該技術經過證明,可以直接無縫遷移到大規模模型。
另一種提升Transformer效率的方法是在注意層加速softmax的計算。
在low-rank approximation of the softmax kernel研究基礎上,研究人員提出了一類新的隨機特征,提供了softmax kernel的首個「positive and bounded」隨機特征近似,并且在序列長度上是計算線性的。

論文鏈接:?https://arxiv.org/abs/2205.15317?
并且還提出了首個涵蓋多個注意力masking的機制,如因果編碼和相對位置編碼。
訓練效率
有效的優化方法是現代機器學習應用的基石,在大規模環境中尤為重要。
在這種情況下,即使是像Adam這樣的一階自適應方法也往往需要大量的計算量,而且訓練的穩定性也會變得非常困難。
此外,這些方法往往與神經網絡的體系結構無關,沒有考慮模型架構內的結構信息,導致訓練效率低下,這也促使新技術能夠更有效地優化現代神經網絡模型。
Google Research開發了一些新的基于模型架構的訓練技術,例如,用于訓練Transofmre網絡,包括新的規模不變的Transofmre網絡和新的clipping方法,當與原版隨機梯度下降(SGD)結合時,可以加快訓練速度。

論文鏈接:?https://arxiv.org/pdf/2202.00980.pdf?
使用這種方法,研究人員首次實現了僅使用簡單的 SGD 就能有效地訓練 BERT,而無需自適應(adaptivity)。
此外,研究人員提出了一種新的方法LocoProp,在使用與一階優化器相同的計算和內存資源的情況下,獲得類似于二階優化器的性能。
LocoProp 通過將神經網絡分解為一個層的組合,從而獲得了神經網絡的模塊化視圖。然后允許每一層都有自己的損失函數以及輸出目標和重量調整器。通過這種設置,經過適當的向前向后梯度傳遞后,LocoProp 繼續對每個層的「局部損失」執行并行更新。

論文鏈接:?https://proceedings.mlr.press/v151/amid22a.html?
事實上,這些更新在理論上和經驗上都類似于高階優化器,在深度自編碼器基準測試中,LocoProp 獲得了與高階優化器相當的性能,同時速度明顯更快。
像 SGD 這樣的優化器的一個關鍵假設是,每個數據點都是從一個分布中獨立和完全相同地采樣的,不過在強化學習等實際應用環境中,很難滿足該假設,因為模型(或智能體)必須從根據自己的預測生成的數據中學習。
研究人員提出了一種新的算法方法,稱為帶有反向經驗重放的 SGD,可以在線性動力系統(linear dynamical system)、非線性動力系統和強化學習的 Q 學習等幾種情況下找到最優解。

論文鏈接:?https://arxiv.org/abs/2103.05896?
此外,這種方法的一個改進版本IER ,經過實驗被證明是最先進的,并且是各種流行的 RL 基準測試上最穩定的經驗重播技術。
數據效率
深度神經網絡嚴重依賴于大型數據集,隨之而來的還有存儲成本和潛在的安全/隱私問題,在這些數據集上訓練現代深度神經網絡也會帶來很高的計算成本。
解決這個問題的一個有前景的方法是數據子集選擇,其中學習器的目標是從大量的訓練樣本中找到信息量最大的子集,以接近(甚至改進)整個訓練集的訓練。
研究人員分析了一個子集選擇框架設計工作與任意模型族在一個實際的batch設置,在這種設置下,學習器可以一次采樣一個樣本,訪問上下文和真實標簽,但為了限制開銷,一旦選擇了足夠大的一批樣本,只能更新其狀態,即,進一步訓練模型權重。
并基于此開發了一種稱為 IWeS 的算法,該算法通過重要抽樣來選擇樣本,其中分配給每個樣本的抽樣概率是基于對先前選定的批次訓練的模型的熵。文中提供了一個理論分析,證明了泛化和抽樣率的界限。

論文鏈接:?https://arxiv.org/pdf/2301.12052.pdf?
訓練大型網絡的另一個問題是,它們可能對部署時看到的訓練數據和數據之間的分布變化高度敏感,特別是在處理數量有限的訓練數據時,這些數據可能不包括所有的部署時間情景。
最近的一項研究假設「極端簡單性偏差」是神經網絡這種脆弱性背后的關鍵問題,其最新工作使這一假設成為可行,導致了兩種新的互補方法DAFT 和 FRR ,結合起來提供了顯著更強大的神經網絡。特別是,這兩種方法使用對抗微調以及逆特征預測,提升了學習網絡的健壯性。

論文鏈接:?https://arxiv.org/pdf/2006.07710.pdf?
推理效率
增加神經網絡的規模已被證明在提高其預測準確性方面有驚人的效果,然而,在現實世界中利用這些優勢是具有挑戰性的,因為大型模型的推理成本可能過高了,該問題也促使策略提高服務效率,而不犧牲準確性。
研究人員提出了實現該目標的不同策略,特別是基于知識蒸餾和自適應計算的策略。
蒸餾
蒸餾是一種簡單而有效的模型壓縮方法,極大地擴展了大型神經模型的潛在適用性,并且在廣告推薦等一系列實際應用中被證明是非常有效的。
大多數蒸餾的用例都涉及到將基本模型直接應用到給定的領域,對于何時以及為什么應該這樣做只有有限的理解。Google的研究著眼于根據特定的環境對蒸餾進行裁剪,并系統性地研究了決定蒸餾成功的因素。
在算法方面,通過仔細建模教師模型提供的標簽中的噪聲,研究人員開發了一個原則性的方法來重新加權的訓練例子,和一個穩健的方法來抽樣的數據子集有教師標簽。

論文鏈接:https://arxiv.org/abs/2210.06711
在「教師指導訓練」的過程中,研究人員提出了一個新的蒸餾框架: 不是被動地使用教師來標注一個固定的數據集,而是主動地使用教師來指導信息樣本的選擇來進行標注,這使得蒸餾過程在有限的數據或長尾設置中更有效。

論文鏈接:?https://arxiv.org/abs/2208.06825?
同時還研究了從交叉編碼器(dual-encoder, 例如 BERT)到因子雙編碼器(dual-encoder)的新方法,這也是對(query, document)對的相關性進行評分的重要設置。

論文鏈接:https://proceedings.mlr.press/v162/menon22a/menon22a.pdf
文中研究了交叉編碼器和雙重編碼器之間性能差距的原因,指出這可能是泛化的結果,而非雙編碼器的容量限制。
精心構造的蒸餾損耗函數可以緩解這種情況,縮小交叉編碼器和雙編碼器之間的性能差距。
隨后,在 EmbedDistil進一步研究了通過匹配來自教師模型的嵌入來進一步改進雙編碼器精餾。該策略還可用于從大到小的雙編碼器模型中提取信息,繼承和凍結教師的文檔嵌入被證明是非常有效的。

論文鏈接:?https://arxiv.org/abs/2301.12005?
在理論方面提供了一個新的視角,通過監督復雜性的蒸餾,衡量學生如何能夠預測教師的標簽。

論文鏈接:?https://arxiv.org/abs/2301.12245?
利用神經正切核(NTK)理論得出了一些概念性的結論,比如,能力差距可能會影響蒸餾,因為這樣的教師的標簽可能出現類似于純粹的隨機標簽的學生。

論文鏈接:?https://arxiv.org/abs/2301.12923?
同時進一步證明,蒸餾導致學生欠擬合的點也是對于教師模型比較難的問題,從直覺上講,這可能有助于學生將其有限的能力集中在那些可以合理建模的樣本上。
自適應計算
雖然蒸餾是降低推理成本的有效方法,但它在所有樣本中的作用是一致的(uniformly),直觀地說,一些「簡單」的樣本可能本質上需要比相對更「難」樣本更少的計算。
自適應計算的目標是設計能夠進行這種依賴于樣本的計算的機制。
Confident Adaptive Language Modeling(CALM)為 T5等基于Transformer的文本生成器引入了受控的早期退出功能。

論文鏈接:?https://arxiv.org/abs/2207.07061?
在這種形式的自適應計算中,模型動態修改每個解碼步驟使用的Transformer層數,其中early-exit gate使用帶有決策閾值的置信度量,該置信度量經過校準以滿足統計性能保證。
這樣,模型只需要為最具挑戰性的預測計算完整的解碼器層堆棧,更簡單的預測只需要計算幾個解碼器層。在實踐中,該模型平均使用約三分之一的層進行預測,在保持同等水平的生成質量的同時,產生2-3倍的加速。
一個常用的自適應計算機制包括兩個或多個基本模型的級聯,其關鍵問題是決定是簡單地使用當前模型的預測,還是將預測推遲到下游模型,學習何時推遲需要設計一個合適的損失函數,它可以利用適當的信號作為推遲決策的監督。
Google Research系統性地研究了現有的損失函數,證明由于標簽平滑的隱式應用,它們可能不適合訓練樣本,文中還展示了可以通過延遲規則的事后訓練來緩解這種情況,這種訓練不需要以任何方式修改模型內部。

論文鏈接:?https://openreview.net/pdf?id=_jg6Sf6tuF7?
對于檢索應用程序,標準的語義搜索技術對大型模型生成的每個嵌入使用固定的表示,也就是說,不管下游任務及其相關的計算環境或約束,表示的大小和能力大多是固定的。
Matryoshka 表示學習引入了根據部署環境調整表示的靈活性,強制表示在它的坐標中有一個自然的排序,這樣對于資源受限的環境,只使用表示的最高的幾個坐標;而對于更豐富和精度關鍵的設置,可以使用表示的更多坐標。

論文鏈接:?https://openreview.net/pdf?id=9njZa1fm35?
當結合標準近似最近鄰搜索技術,如掃描神經網絡,MRL 能夠提供多達16倍的低計算相同的召回率和準確度量。
總結
大型機器學習模型在多個領域都展現出具有變革性的結果,但訓練和推理的效率正在成為使這些模型在現實世界中切實可行的關鍵需求。
通過開發新的基礎技術,Google Research已經在使大型機器學習模型高效方面進行了大量投資,這也需要持續性的努力,未來將繼續探索核心挑戰,使機器學習模型更加健壯和高效。



























