網友盛贊：Transformer引領年度論文的簡化版本來了

從大模型的根源開始最佳化。

Transformer 架構可以說是近期深度學習領域許多成功案例背後的主力。建構深度Transformer 架構的簡單方法是將多個相同的Transformer 「區塊」（block）依序堆疊起來，但每個「區塊」都比較複雜，由許多不同的元件組成，需要以特定的排列組合才能實現良好的性能。

自從 2017 年 Transformer 架構誕生以來，研究者們基於其推出了大量衍生研究，但幾乎沒有改動過 Transformer “塊”。

那麼問題來了，標準 Transformer 區塊是否可以簡化？

在最近的一篇論文中，來自 ETH Zurich 的研究者討論瞭如何在不影響收斂特性和下游任務表現的情況下簡化 LLM 所必需的標準 Transformer 區塊。基於訊號傳播理論和經驗證據，他們發現可以移除一些部分，例如殘差連接、歸一化層（LayerNorm）、投影和值參數以及MLP 序列化子區塊（有利於並行佈局），以簡化類似GPT 的解碼器架構以及編碼器式BERT 模型。

研究者探討了在不影響訓練速度的情況下，是否可以移除所涉及的元件，並對Transformer區塊進行哪些架構修改。

論文連結：https://arxiv.org/pdf/2311.01906.pdf

Lightning AI 創辦人、機器學習研究者Sebastian Raschka 將這項研究稱為自己的「年度最愛論文之一」：

但也有研究者質疑：「這很難評，除非我看過完整的訓練過程。如果沒有歸一化層，也沒有殘差連接，如何能在大於1 億參數的網路中進行擴展？」

#Sebastian Raschka 表示同意：「是的，他們試驗的架構相對較小，這是否能推廣到數十億參數的Transformer 上還有待觀察。」但他仍然表示這項工作令人印象深刻，並認為成功移除殘差連接是完全合理的（考慮到其初始化方案）。

對此，圖靈獎得主Yann LeCun 的評價是：「我們只是觸及了深度學習架構領域的皮毛。這是一個高維度空間，因此體積幾乎完全包含在表面中，但我們只觸及了表面的一小部分。」

#為什麼需要簡化Transformer 區塊？

研究者表示，在不影響訓練速度的前提下簡化 Transformer 區塊是一個有趣的研究問題。

首先，現代神經網路架構設計複雜，包含許多元件，而這些不同元件在神經網路訓練動態中所扮演的角色，以及它們之間如何相互作用，人們對此尚不清楚。這個問題關乎深度學習理論與實務之間的差距，因此非常重要。

訊號傳播理論（Signal propagation）已被證明具有影響力，因為它能夠激勵深度神經網路架構中的實際設計選擇。訊號傳播研究了初始化時神經網路中幾何資訊的演化，透過跨輸入的分層表徵的內積來捕捉，在訓練深度神經網路方面取得了許多令人印象深刻的成果。

然而，目前理論只考慮初始化時的模型，而且往往只考慮初始前向傳遞，因此無法揭示深度神經網路訓練動態的許多複雜問題，例如殘差連接對訓練速度的助益。雖然訊號傳播對修改動機至關重要，但研究者表示，他們不能僅從理論上就得出簡化的 Transformer 模組，還要依靠經驗見解。

在實際應用方面，考慮到目前訓練和部署大型 Transformer 模型的高昂成本，Transformer 架構的訓練和推理管線的任何效率提升都代表著巨大的潛在節約意義。如果能夠透過移除非必要元件來簡化 Transformer 模組，既能減少參數數量，又能提高模型的吞吐量。

這篇論文也提到，移除殘差連接、值參數、投影參數和序列化子區塊之後，可以同時做到在訓練速度和下游任務表現方面與標準 Transformer 相符。最終，研究者將參數量減少了 16%，並觀察到訓練和推理時間的吞吐量增加了 16%。

如何簡化 Transformer 區塊？

研究者結合訊號傳播理論和經驗觀察，介紹如何從 Pre-LN 模組出發，產生最簡單的 Transformer 區塊（如下圖）。

在論文第四章的每一個小節，作者分別介紹如何在不影響訓練速度的情況下每次刪除一個區塊組件。

這一部分的所有實驗都在CodeParrot 資料集上使用了一個18-block 768-width 的因果僅解碼器類別GPT 模型，這個資料集足夠大，因此當作者處於單一訓練epoch 模式時，泛化差距非常小（見圖2），這使得他們可以專注於訓練速度。

刪除殘差連結

#研究者首先考慮刪除注意力子區塊中的殘差連結。在公式（1）的符號中，這相當於將 α_SA 固定為 0。簡單地移除注意力殘差連接會導致訊號退化，即秩崩潰（rank collapse），從而導致可訓練性差。在論文 4.1 部分，研究者詳細解釋了他們的方法。

刪除投影/ 值參數

#從圖3 可以得出結論，完全移除值和投影參數W^ V、W^P 是可能的，而且每次更新的訓練速度損失最小。也就是說，當β_V = β_P = 0 和identity 初始化的

#時，在相同的訓練步數後，本研究基本上能達到Pre-LN 區塊的性能。在這種情況下，在整個訓練過程中都有 W^V = W^P = I，即值和投影參數是一致的。作者在 4.2 節介紹了詳細方法。

刪除MLP 子區塊殘差連接

與上述幾個模組相比，刪除MLP 子區塊殘差連接要更具挑戰性。與先前的研究一樣，作者發現，在使用 Adam 時，如果沒有 MLP 殘差連接，透過訊號傳播使活化更加線性仍會導致每次更新訓練速度的顯著下降，如圖 22 所示。

他們也嘗試了 Looks Linear 初始化的各種變體，包括高斯權重、正交權重或恆等權重，但都無濟於事。因此，他們在整個工作中使用標準活化（例如 ReLU）和 MLP 子塊中的初始化。

作者轉向並行 MHA 和 MLP 子區塊的概念，這在幾個近期的大型 transformer 模型中已被證明很受歡迎，例如 PALM 和 ViT-22B。並行 transformer 區塊如下圖所示。

作者在論文 4.3 節詳細介紹了移除 MLP 子區塊殘差連接的具體操作。

刪除歸一化層

最後一個被刪除的是歸一化層，這樣就得到了圖 1 右上角的最簡塊。從訊號傳播初始化的角度來看，作者可以在本節簡化的任何階段移除歸一化層。他們的想法是，Pre-LN 區塊中的歸一化會隱式地降低殘差分支的權重，而這種有利的效果可以透過另一種機制在沒有歸一化層的情況下複製：要麼在使用殘差連接時明確降低殘差分支的權重，要麼將注意力矩陣偏向identity / 將MLP 非線性轉換為「更」線性。

由於作者在修改過程中考慮了這些機制（例如降低 MLP β_FF 和 Shaped Attention 的權重），因此無需進行歸一化處理。作者在第 4.4 節介紹了更多資訊。

實驗結果

深度擴展

#鑑於訊號傳播理論通常關注很大的深度，而這種情況下通常會出現訊號退化。因此一個很自然的問題就是，本文的簡化 transformer 區塊所提升的訓練速度是否也能擴展到更大的深度？

從圖6 可以觀察到，將深度從18 個區塊擴展到72 個區塊後，本研究的模型和Pre-LN transformer 的表現都得到了提高，這表明本研究中的簡化模型不僅訓練速度更快，還能利用更大的深度所提供的額外能力。事實上，在使用歸一化時，本研究中的簡化區塊和 Pre-LN 的每次更新軌跡在不同深度下幾乎沒有差異。

BERT

接下来，作者展示了他们的简化块性能除了适用于自回归解码器之外，还适用于不同的数据集和架构，以及下游任务。他们选择了双向仅编码器 BERT 模型的流行设置，用于掩蔽语言建模，并采用下游 GLUE 基准。

如图 7 所示，在 24 小时运行时内，与（Crammed）Pre-LN 基线相比，本研究的简化块可以媲美掩蔽语言建模任务的预训练速度。另一方面，在不修改值和投影的情况下删除残差连接再次导致训练速度的显著下降。在图 24 中，作者提供了 microbatch 步骤的等效图。

此外，在表 1 中，研究者发现他们的方法在 GLUE 基准上经过微调后，性能与 Crammed BERT 基准相当。

他们在表 2 中对下游任务进行了细分。为了进行公平比较，他们使用了与 Geiping & Goldstein (2023) 相同的微调协议（5 个 epoch、各任务超参数恒定、dropout regularisation）。

效率提升

在表 1 中，研究者还详细列出了使用不同 Transformer 块的模型在掩蔽语言建模任务中的参数数量和训练速度。他们以预训练 24 小时内所采取的 microbatch 步骤数与基线 Pre-LN Crammed BERT 的比率计算了速度。结论是，模型使用的参数减少了 16%，SAS-P 和 SAS 的每次迭代速度分别比 Pre-LN 块快 16% 和 9%。

可以注意到，在这里的实现中，并行块只比 Pre-LN 块快 5%，而 Chowdhery et al.（2022 ）观察到的训练速度则快 15%，这表明通过更优化的实现，整个训练速度有可能进一步提高。与 Geiping & Goldstein（2023 年）一样，此处实现也使用了 PyTorch 中的自动算子融合技术 (Sarofeen et al., 2022)。

更长的训练

最后，考虑到当前在更多数据上长时间训练较小模型的趋势，研究者讨论了简化块在长时间训练后是否仍能达到 Pre-LN 块的训练速度。为此，他们在 CodeParrot 上使用图 5 中的模型，并使用 3 倍 token 进行训练。准确地说，是在批大小为 128、序列长度为 128 的情况下进行了约 120K 步（而不是 40K 步）的训练，这将导致约 2B 个 token。

从图 8 可以看出，当使用更多的 token 进行训练时，简化的 SAS 和 SAS-P 代码块的训练速度仍然与 PreLN 代码块相当，甚至优于 PreLN 代码块。

更多研究细节，可参考原论文。

以上是網友盛贊：Transformer引領年度論文的簡化版本來了的詳細內容。更多資訊請關注PHP中文網其他相關文章！