小模型性能飽和、表現(xiàn)不佳，根源是因為Softmax?
語言建模領域的最新進展在于在極大規(guī)模的網(wǎng)絡文本語料庫上預訓練高參數(shù)化的神經(jīng)網(wǎng)絡。在實踐中，使用這樣的模型進行訓練和推斷可能會成本高昂，這促使人們使用較小的替代模型。然而，已經(jīng)觀察到較小的模型可能會出現(xiàn)飽和現(xiàn)象，表現(xiàn)為在訓練的某個高級階段性能下降并趨于穩(wěn)定。
最近的一篇論文發(fā)現(xiàn)，這種飽和現(xiàn)象可以通過較小模型的隱藏維度與目標上下文概率分布的高秩之間的不匹配來解釋。這種不匹配通過著名的 softmax 瓶頸現(xiàn)象影響了這些模型中使用的線性預測頭的性能。
論文鏈接：???https://arxiv.org/pdf/2404.07647.pdf??
本文在不同設置下衡量了 softmax 瓶頸的影響，并發(fā)現(xiàn)基于小于 1000 個隱藏維度的模型往往在預訓練的后期采用退化的潛在表征，從而導致評估性能降低。
簡介
表征退化問題是影響用于文本數(shù)據(jù)的自監(jiān)督學習方法等多種模態(tài)的常見現(xiàn)象。對語言模型的中間表征進行的許多觀察揭示了它們的低角度可變性（或各向異性），或者在訓練過程中出現(xiàn)的異常維度。然而，這些觀察大多是針對維度與 BERT 或 GPT-2 系列模型相當?shù)南鄬^小規(guī)模的模型進行的。
這些模型通常由一個神經(jīng)網(wǎng)絡 f_θ 組成，該神經(jīng)網(wǎng)絡接受 token 序列：
并在 R^d 中生成一個相對低維的上下文表征，其中 d 是模型的隱藏維度。然后它們依賴于一個語言建模頭，該頭部產(chǎn)生上下文 token 概率的對數(shù)。語言建模頭的常見選擇是一個線性層，其參數(shù)為 W ∈ R^(V×d)，其中 V 是可能 token 的數(shù)量。因此得到的下一個 token 概率分布是
其中 σ 是 softmax 函數(shù)。
在語言建模領域，當前的趨勢在于擴展引入了 GPT-2 的生成預訓練方法，這意味著在巨大的網(wǎng)絡文本語料庫上訓練由數(shù)十億參數(shù)組成的神經(jīng)模型。然而，訓練和應用這些高參數(shù)化模型會引發(fā)能源和硬件相關(guān)的問題，這需要尋求通過較小的模型實現(xiàn)類似性能水平的方法。
然而，對 Pythia 模型套件的評估表明，將小型模型訓練在非常大的語料庫上可能會導致飽和，表現(xiàn)為在預訓練后期性能下降。本文通過表征退化的視角探討了這種飽和現(xiàn)象，并發(fā)現(xiàn)這兩種現(xiàn)象之間存在著強烈的相關(guān)性，同時進一步證明了表征退化在小型模型的語言建模頭中發(fā)生，并在理論和實證上展示了線性語言建模頭如何成為基于小隱藏維度的架構(gòu)的性能瓶頸。
語言模型飽和現(xiàn)象
本文首先驗證了確實可以觀察和量化 Pythia 檢查點的性能飽和，因為它們是一系列模型尺寸的唯一發(fā)布的中間檢查點。本文測量了從它們的預訓練數(shù)據(jù)集（即 The Pile）中隨機抽取的 5 萬個 token 的 Pythia 檢查點的交叉熵。
在圖 1a 中可以清楚地看到，連 4.1 億參數(shù)的模型都遇到了飽和現(xiàn)象，表現(xiàn)為在高級訓練階段域內(nèi)損失的增加。
在圖 1b 中，本文根據(jù) Hoffmann et al. (2022) 的方法，對從 4.1 億參數(shù)開始的模型的數(shù)據(jù)點進行了擬合，只優(yōu)化模型相關(guān)的常數(shù)（A 和 α），同時重用所有其他值（B = 410.7，β = 0.28，E = 1.69）。這里回顧了 Hoffmann et al. (2022) 給出的參數(shù)計數(shù) N 和 token 計數(shù) T 之間的關(guān)系：
本文發(fā)現(xiàn)最佳參數(shù)為 A = 119.09 和 α = 0.246。作者展示了與最佳和最終檢查點相對應的 token 計數(shù)的擬合曲線?？梢杂^察到，最終檢查點的性能平均低于外推值約 8％。損失最?。ㄗ罴眩z查點由于學習率冷卻不完全，預計會低于外推法，但其表現(xiàn)僅低于外推法約 4%。
在用于語言模型評估工具（LM Evaluation Harness）評估的數(shù)據(jù)集中，也觀察到了類似的性能飽和現(xiàn)象，如表 1 所示。
性能飽和是秩飽和（Rank Saturation）
規(guī)模各向異性
各向異性是是在各種小型語言模型中觀察到的一種常見的表征退化形式，它包括特定層中表征分布的角度可變性降低。之前的研究（Ethayarajh, 2019; Godey et al., 2024）注意到，小型變形語言模型的幾乎所有層都是各向異性的。衡量向量表征集合 H 中各向異性的常用方法是平均余弦相似度：
然而，目前尚不清楚各向異性是否會影響具有超過 10 億參數(shù)的模型。為了解決這個問題，本文計算了一系列模型中間表征在層間的平均余弦相似度；即 GPT-2，OPT，Pythia 和 Gemma。本文使用了 The Pile 的子樣本，因為假設該數(shù)據(jù)集的領域包括或匹配這些套件中使用的預訓練數(shù)據(jù)集的領域。
在圖 2 中，可以觀察到，大多數(shù) Transformer 模型的大多數(shù)層在某種程度上都是各向異性的，而不論其規(guī)模如何。然而，在最后一層中似乎存在一個二分現(xiàn)象，其中模型要么幾乎是各向同性的，要么是高度各向異性的。本文注意到這種二分現(xiàn)象與 Pythia 套件的飽和現(xiàn)象之一相一致，其中只有包含 1.6 億個或更少參數(shù)的模型受到最后一層各向異性的影響。
本文研究了 Pythia 套件中各向異性的訓練動態(tài)，并將其與圖 3 中的飽和現(xiàn)象進行比較。
圖 3 清晰地展示了性能飽和現(xiàn)象的出現(xiàn)與模型最后一層表征中各向異性出現(xiàn)之間的明顯相關(guān)性。它還顯示了在訓練過程中，各向異性在飽和點附近會突然增加。在這里觀察到，在特定的領域內(nèi)語料庫中，模型在飽和時迅速失去性能，并且似乎永遠無法完全從這種爆炸中恢復過來。
奇異值飽和
平均余弦相似度是衡量分布均勻性的有價值的指標，但包含其他指標可以幫助更好地捕捉某些流形的復雜性。此外，它只關(guān)注語言模型的輸出嵌入，而不關(guān)注它們的權(quán)重。本節(jié)通過研究語言建模頭的奇異值分布來擴展本文的分析，以將實證觀察與本文的理論發(fā)現(xiàn)聯(lián)系起來。
圖 4 展示了沿訓練過程中最終預測層權(quán)重 W 的奇異值分布：
圖 4 揭示了一種特定的頻譜飽和模式，它與性能飽和現(xiàn)象大致同時發(fā)生。圖中顯示，奇異值分布在訓練過程中逐漸變平，幾乎達到均勻性，然后突然演變?yōu)榧夥宸植?，最大奇異值相對其他分布較高。
為了更準確地量化這種行為，本文使用奇異熵度量，計算為歸一化奇異值分布與均勻分布之間的 Kullback-Leibler 散度。
圖 5 顯示了使用少于 4.1 億個參數(shù)的模型與使用較大參數(shù)的模型的奇異分布演變方式不同。小型模型的頭部看到它們的奇異值分布逐漸變得更加均勻，直到它們突然退化為止，這再次與語言模型性能下降相關(guān)。較大模型的奇異值分布趨于更穩(wěn)定，并且在整個訓練過程中沒有顯示出明顯的單調(diào)模式。
Softmax 瓶頸與語言維度
自然語言的固有維度
直觀地說，上文中觀察到的奇異值分布飽和現(xiàn)象只適用于較小的模型，這就對 LM 頭的優(yōu)化所涉及的維度提出了質(zhì)疑。本節(jié)建議根據(jù)經(jīng)驗測量 LM 頭的秩的臨界值，并估計該頭的輸出應該匹配的上下文概率分布的維度。
為了經(jīng)驗性地測量線性頭部秩的影響，本文提出在預訓練的上下文表征上訓練秩受限的頭部，這些上下文表征來自高參數(shù)化語言模型。為了控制最大秩 r，考慮形式為 W = AB ∈ R^(V×d) 的頭部，其中 A ∈ R^(V×r) 和 B ∈ R^(r×d) 的系數(shù)從 N（0,1）中抽?。╠ 是模型的隱藏維度）。這種 W 矩陣的秩受參數(shù) r ∈ [1, d] 的限制對一系列值進行了掃描。
通過凍結(jié)語言模型，并在大約 1.5 億 token 上訓練秩受限的頭部，同時調(diào)整學習速率以適應可訓練參數(shù)的數(shù)量。
在圖 6 中可以觀察到，無論模型大小如何，當語言建模頭 W 的秩低于 1000 時，困惑度開始明顯下降。這暗示了對于具有更大隱藏維度的模型來說，頭部不是主要的性能瓶頸，但對于具有較小隱藏維度的模型來說，它可能會獨立于輸出表征的質(zhì)量而損害性能。
另一個有趣的因素是估計數(shù)據(jù)本身固有的維度。為了避免與特定歸納偏差相關(guān)的可能影響，本文在覆蓋范圍各異的幾個數(shù)據(jù)集上訓練了樸素的 5-gram 語言模型（IMDb，Wikitext，以及 The Pile），使用了兩種不同詞匯量的分詞器（Llama-2 為 30k tokens，Pythia 為 50k tokens）。給定 C 個觀察到的 5-gram，本文考慮矩陣 W ∈ R^(C×V)，其中每行是給定 4 個 token 上可能 token 的概率分布，并計算它們的奇異值分布，如 Terashima et al. (2003) 所述。
圖 7 報告了 W-error，根據(jù) Eckart-Young-Mirsky 定理預測的秩為 d 的矩陣 W 的最小近似誤差（見引理 5.2），并將其歸一化為 W 的 Frobenius 范數(shù)。
理論瓶頸
同時，W 的估計秩與隱藏維度的常規(guī)數(shù)量級相比也不可忽視。這里將從理論角度分析理想線性語言建模頭的維度與性能之間的聯(lián)系。
本節(jié)旨在確定上下文分布固有維度與可歸因于語言模型輸出表征的較低維度而產(chǎn)生的性能瓶頸之間的正式聯(lián)系。為此構(gòu)想了一個在理想上下文表征上優(yōu)化的語言建模頭，探討了其譜特性與在相同表征上訓練低秩頭時產(chǎn)生的性能差距之間的關(guān)系。
更多研究細節(jié)，可查看原論文。
本文轉(zhuǎn)自 機器之心 ，作者：機器之心
