如何高效訓(xùn)練 Transformer？微軟亞研院開源TorchScale工具包

近期，微軟亞洲研究院從深度學(xué)習(xí)基礎(chǔ)理論出發(fā)，研發(fā)并推出了 TorchScale 開源工具包。TorchScale 工具包通過采用 DeepNet、Magneto 和 X-MoE 等最先進(jìn)的建模技術(shù)，可以幫助研究和開發(fā)人員提高建模的通用性和整體性能，確保訓(xùn)練模型的穩(wěn)定性及效率，并允許以不同的模型大小擴(kuò)展 Transformer 網(wǎng)絡(luò)。

如今，在包括語音、自然語言處理（NLP）、計(jì)算機(jī)視覺（CV）、多模態(tài)模型和 AI for Science 等領(lǐng)域的研究中，Transformer 已經(jīng)成為一種通用網(wǎng)絡(luò)結(jié)構(gòu)，加速了 AI 模型的大一統(tǒng)。與此同時(shí)，越來越多的實(shí)踐證明大模型不僅在廣泛的任務(wù)中能產(chǎn)生更好的結(jié)果、擁有更強(qiáng)的泛化性，還可以提升模型的訓(xùn)練效率，甚至衍生出新的能力。因此，學(xué)術(shù)界和產(chǎn)業(yè)界都開始追求更大規(guī)模的模型。

然而隨著模型的不斷擴(kuò)大，其訓(xùn)練過程也變得更加困難，比如會出現(xiàn)訓(xùn)練不收斂等問題。這就需要大量的手動調(diào)參工作來解決，而這不僅會造成資源浪費(fèi)，還會產(chǎn)生不可預(yù)估的計(jì)算成本。

與其揚(yáng)湯止沸，不如釜底抽薪。微軟亞洲研究院從深度學(xué)習(xí)基礎(chǔ)理論出發(fā)，創(chuàng)新推出了 TorchScale 工具包，并已將其開源。TorchScale 是一個(gè) PyTorch 庫，允許科研和開發(fā)人員更高效地訓(xùn)練 Transformer 大模型。同時(shí)，它有效地提升了建模的性能和通用性，提高了 Transformer 的穩(wěn)定性和訓(xùn)練效率。

TorchScale GitHub 頁面：

https://github.com/microsoft/torchscale

“我們希望通過 TorchScale 的系列工作從更底層出發(fā)做一些基礎(chǔ)性的研究創(chuàng)新，通過數(shù)學(xué)或者理論上的指導(dǎo)和啟發(fā)，在 Transformer 模型擴(kuò)展的工作中取得更好的效果，而不是單純的調(diào)參或僅從工程層面部分緩解某些問題。TorchScale 能夠支持任意的網(wǎng)絡(luò)深度和寬度，實(shí)驗(yàn)驗(yàn)證它可以輕松擴(kuò)大模型規(guī)模，而且只需要幾行代碼就能夠?qū)崿F(xiàn)多模態(tài)模型的訓(xùn)練?！蔽④泚喼扪芯吭鹤匀徽Z言計(jì)算組首席研究員韋福如表示。

據(jù)了解，TorchScale 主要從三個(gè)方面幫助科研人員克服了擴(kuò)展 Transformer 大模型時(shí)的困難：

• DeepNet：提升模型的穩(wěn)定性。
• Magneto：提升模型的通用性。
• X-MoE：提升模型訓(xùn)練的高效性。

圖1：TorchScale 解決大模型在

穩(wěn)定性、通用性、高效性上面臨的問題

DeepNet：讓Transformer訓(xùn)練深度超過1000層

盡管近年來模型參數(shù)的數(shù)量越來越大，已經(jīng)從百萬級擴(kuò)展到萬億級，但參數(shù)的深度卻一直受限于 Transformer 訓(xùn)練的不穩(wěn)定性。為了解決這一問題，一些科研人員嘗試通過更好的初始化或架構(gòu)來提升 Transformer 的穩(wěn)定性，但這也只能讓 Transformer 在百層級別的深度下保持穩(wěn)定。

微軟亞洲研究院的研究員們發(fā)現(xiàn)，模型輸出的劇烈變化是導(dǎo)致模型不穩(wěn)定的重要原因。為此，研究員們在殘差連接處使用了一種新的歸一化函數(shù)——DeepNorm。新的函數(shù)由理論推導(dǎo)而來，可以把模型輸出的變化限制在常數(shù)范圍內(nèi)。這種方法只需要改變幾行代碼，就可以大幅提升 Transformer 的穩(wěn)定性。通過引入新的 DeepNorm 函數(shù)，研究員們訓(xùn)練了超深的 Transformer 網(wǎng)絡(luò) DeepNet，在保證模型穩(wěn)定的同時(shí)，可以將模型深度擴(kuò)展到1000層以上。

DeepNorm 同時(shí)具備 Post-LN 的性能和 Pre-LN 的訓(xùn)練穩(wěn)定性。這個(gè)新方法或?qū)⒊蔀?Transformer 的首選替代方案，它不僅適用于深模型，更適用于大模型。值得一提的是，與具有120億參數(shù)的48層模型相比，微軟亞洲研究院32億參數(shù)的200層模型在100多個(gè)語言、超10000個(gè)語言對和130億個(gè)文本對的多語言機(jī)器翻譯實(shí)驗(yàn)中實(shí)現(xiàn)了5 BLEU 的提升。在大規(guī)模多語言翻譯任務(wù)上，隨著 DeepNet 模型深度從10層擴(kuò)展至100層和1000層，模型也獲得了更高的 BLEU 值。

圖2：隨著模型深度從10層擴(kuò)展至100層和1000層，DeepNet 有效提升了多語言翻譯結(jié)果

韋福如說，“此前，科研人員在訓(xùn)練更大規(guī)模的模型時(shí)，往往需要投入大量的精力在模型調(diào)參上，無形中增加了實(shí)驗(yàn)成本，有的模型在訓(xùn)練中途就無法繼續(xù)下去了，即使給模型打上補(bǔ)丁也還是會影響模型性能。DeepNet 可以幫助科研人員大幅降低調(diào)參的負(fù)擔(dān)，在提升模型性能的同時(shí)降低實(shí)驗(yàn)成本?！?/span>

Magneto：真正實(shí)現(xiàn)多模態(tài)模型架構(gòu)統(tǒng)一

跨語言、視覺、語音和多模態(tài)的模型在模型結(jié)構(gòu)上走向大一統(tǒng)的趨勢如今已經(jīng)愈發(fā)明顯。具體而言，從 NLP 領(lǐng)域開始，Transformer 已成為 AI 各領(lǐng)域的主流結(jié)構(gòu)。然而，盡管都使用了 Transformer，但不同模態(tài)任務(wù)的模型結(jié)構(gòu)在具體實(shí)現(xiàn)時(shí)仍存在顯著差異。例如，GPT 和 ViT 模型采用了 Pre-LN Transformer，而 BERT 和機(jī)器翻譯模型使用的是 Post-LN 來獲得更好的性能。更重要的是，對于多模態(tài)模型，不同輸入模態(tài)的最優(yōu) Transformer 變體通常是不同的。以微軟亞洲研究院推出的多模態(tài)預(yù)訓(xùn)練模型 BEiT-3為例，其使用 Post-LN 對于視覺部分是次優(yōu)的，而 Pre-LN 對于語言部分是次優(yōu)的。

要想讓多模態(tài)預(yù)訓(xùn)練真正實(shí)現(xiàn)大一統(tǒng)就需要一個(gè)統(tǒng)一的架構(gòu)，該架構(gòu)需要在不同任務(wù)和模態(tài)上都能有良好的性能表現(xiàn)。另外，如之前所述，Transformer 架構(gòu)訓(xùn)練的穩(wěn)定性也是一個(gè)痛點(diǎn)。微軟亞洲研究院的研究員們意識到，通用模型的開發(fā)需要更基礎(chǔ)的 Transformer，即 Foundation Transformer。首先它的建模能夠作為各種任務(wù)和模式的統(tǒng)一架構(gòu)，這樣就可以使用相同的主干而無需反復(fù)魔改。其通用的設(shè)計(jì)原則也應(yīng)該支持多模態(tài)基礎(chǔ)模型的開發(fā)，在不犧牲性能的前提下將統(tǒng)一的 Transformer 用于各種模態(tài)。其次，它的網(wǎng)絡(luò)結(jié)構(gòu)應(yīng)能夠保障訓(xùn)練的穩(wěn)定性，從而降低基礎(chǔ)模型大規(guī)模預(yù)訓(xùn)練的難度。

為了實(shí)現(xiàn)這些目標(biāo)，微軟亞洲研究院的研究員們提出了一個(gè) Foundation Transformer——Magneto。在 Magneto 中，研究員們引入了 Sub-LN，為每個(gè)子層（即多頭自注意力和前饋網(wǎng)絡(luò)）添加了額外的 LayerNorm，并且提出了一種新的初始化方法，為從根本上提高訓(xùn)練的穩(wěn)定性提供了理論保證。

通過對 Magneto 在不同任務(wù)和模態(tài)上的評測，包括掩碼語言建模（即BERT）、因果語言建模（例如GPT）、機(jī)器翻譯、掩碼圖像建模（即BEiT）、語音識別和視覺語言預(yù)訓(xùn)練（即BEiT-3），結(jié)果顯示在下游任務(wù)上，Magneto 顯著優(yōu)于各種 Transformer 變體。此外，得益于訓(xùn)練穩(wěn)定性的提高，Magneto 還允許使用更高的學(xué)習(xí)率來進(jìn)一步提高結(jié)果。

圖3：Magneto 在語言、圖像、

語音和多模態(tài)任務(wù)上的實(shí)驗(yàn)結(jié)果

X-MoE：優(yōu)于基線SMoE模型，助力模型高效訓(xùn)練

在有關(guān)大模型訓(xùn)練的研究中，除了將網(wǎng)絡(luò)深度做得更深和將寬度即隱藏維度擴(kuò)大以外，還可以利用混合專家系統(tǒng)（Mixture of Experts, MoE）。盡管 MoE 可以在諸如語言模型和視覺表示學(xué)習(xí)等廣泛問題上獲得更好的性能，但也會導(dǎo)致更高的計(jì)算成本，這促使越來越多的科研人員開始探索稀疏混合專家模型（Sparse Mixture-of-Experts, SMoE）。SMoE 主要通過構(gòu)建稀疏激活的神經(jīng)網(wǎng)絡(luò)來增加模型容量。在不顯著增加計(jì)算開銷的情況下 SMoE 模型在各種任務(wù)（包括機(jī)器翻譯、圖像分類和語音識別）上的性能都優(yōu)于稠密模型。

在 SMoE 模型中，路由機(jī)制發(fā)揮著重要的作用。給定輸入 token，路由機(jī)制會測量每個(gè) token 與專家之間的相似度分?jǐn)?shù)，然后再根據(jù)路由得分將 token 分配給最匹配的專家。因此，近年來許多研究都集中在如何設(shè)計(jì) token 專家分配算法上。然而，微軟亞洲研究院的研究員們發(fā)現(xiàn)，當(dāng)前的路由機(jī)制傾向于以專家為中心來推動隱藏表示聚類，這容易引起表征坍塌（Representation Collapse），損害模型性能。

為了緩解現(xiàn)有的路由機(jī)制引起的表征坍塌問題，微軟亞洲研究院的研究員們提出了新的方法 X-MoE，為 SMoE 模型引入了一種簡單而有效的路由算法。具體來說，區(qū)別于現(xiàn)有 SMoE 模型直接使用隱藏向量進(jìn)行路由，X-MoE 先將隱藏向量投射到低維空間中，再對 token 表示和專家表示進(jìn)行 L2 歸一化，來測量低維超球面上的路由分?jǐn)?shù)。此外，研究員們還提出了軟專家門（soft expert gate），以學(xué)習(xí)控制專家的激活。

圖4：X-MoE 流程圖

微軟亞洲研究院的研究員們對這一新方法在跨語言模型預(yù)訓(xùn)練任務(wù)上進(jìn)行了評測。實(shí)驗(yàn)結(jié)果表明，在語言建模和微調(diào)性能方面，基于 X-MoE 的模型始終優(yōu)于基線 SMoE 模型。實(shí)驗(yàn)分析還表明，與 SMoE 基線相比，X-MoE 方法有效緩解了表征坍塌問題。該方法在預(yù)訓(xùn)練和微調(diào)期間也實(shí)現(xiàn)了更一致的路由行為，證實(shí)了 X-MoE 路由算法的有效性。

“隨著技術(shù)的持續(xù)演進(jìn)，大模型的訓(xùn)練不僅僅是工程層面的工作。我們應(yīng)該從基礎(chǔ)研究的角度出發(fā)，探索下一代 Transformer 網(wǎng)絡(luò)架構(gòu)。與此同時(shí)，在 AI 模型大一統(tǒng)趨勢的推動下，我們更應(yīng)該追求同一結(jié)構(gòu)來支持不同模態(tài)的輸入，并在不同語言和模態(tài)的任務(wù)上獲得良好的性能。通過理論指導(dǎo)讓模型變得更大、更穩(wěn)定、更通用?！表f福如說。