選自arxiv
作者：楊植麟、Thang Luong等
機器之心編譯
參與：魔王、杜偉
2017 年，楊植麟等人提出一種解決 Softmax 瓶頸的簡單有效的方法——Mixture of Softmaxes（MoS）。但該方法成本高昂，於是最近楊植麟等人再次瞄準 softmax 瓶頸問題，提出兼顧表達能力和高效性的新方法 Mixtape。
論文連結：https://papers.nips.cc/paper/9723-mixtape-breaking-the-softmax-bottleneck-efficiently.pdf
softmax 瓶頸限制了神經語言模型的表達能力（expressiveness）。Mixture of Softmaxes (MoS) 是解決該理論局限的有效方法，但與 softmax 相比，MoS 無論在內存還是時間上都成本較高。
來自 CMU 和谷歌大腦的楊植麟、Thang Luong、Ruslan Salakhutdinov 和 Quoc Le 提出了一種新方法 Mixtape，該輸出層利用三項新技術——logit 空間向量門控（logit space vector gating）、sigmoid 樹分解（sigmoid tree decomposition）和門控共享（gate sharing），更高效地打破了 softmax 瓶頸。
研究者在語言建模和機器翻譯等四個基準數據集上進行了實驗，結果表明 Mixtape 層與 MoS 層性能相當，且 Mixtape 的效率是後者的 3.5-10.5 倍。在 10-30K 的詞彙量下，使用 Mixtape 的網絡僅比基於 softmax 的網絡慢 20%-34%，其困惑度和翻譯質量均優於 softmax。
softmax 帶給我們的苦與樂
大量神經網絡使用 softmax 作為標準輸出層，包括大部分神經語言模型。但是，正如楊植麟等人在
之前研究 [19]
中所指出的，softmax 限制了神經語言模型的表達能力，因為它將輸出表示限制在低秩，這不足以建模自然語言的複雜性。該局限叫做「softmax 瓶頸」。
為了打破這一瓶頸，[19] 提出新方法 Mixture of Softmaxes (MoS)，它將離散潛在變量引入到輸出層中，通過 log-sum-exp 非線性變換使對數概率矩陣變為高秩。但是，MoS 的內存和時間成本均高於 softmax，這使得它在計算資源有限的情況下實際應用性減弱。
MoS →Mixtape
為了降低 MoS 的計算成本，最近楊植麟等人提出了一種高效解決 softmax 瓶頸的新型輸出層 Mixtape。Mixtape 可作為額外層嵌入到任意現有網絡的交叉熵損失函數之前。它不像 MoS 在概率空間中部署標量混合（scalar mixture），而是在 logit 空間中應用向量門控機制，以避免使用多個成本高昂的 softmax。
此外，Mixtape 還使用了另外兩項新技術來進一步降低計算成本。
首先，向量門控機製成本高昂，因為該機制需要我們對詞彙表中的每一個詞計算 softmax 門控。為此，研究者提出了 sigmoid 樹分解技術，將 softmax 概率門控分布分解為一個深度為 2 的二叉樹結構，每一個分支所包含的概率值部分由 sigmoid 函數決定。sigmoid 樹分解更加高效，因為它避免了 softmax 中的減除運算。
另一項技術是門控共享，即對所有低頻詞共享門控值，得到部分高秩的表示。該技術在不影響性能的情況下，節約了一定量的內存和計算資源，因為即使沒有門控共享，低頻詞的門控值通常也很難準確估計。
Mixtape 有多強大
Mixtape 結合了以上三項技術，效率較 MoS 有顯著提升，而且 Mixtape 在四個基準數據集上的性能堪比甚至超過後者。在正常詞彙量情況下（如 10K-30K），基於相同的批大小，Mixtape 層的速度是 MoS 層的 1.6-11.5 倍；基於相同的內存，Mixtape 層的速度是 MoS 的 3.5-10.5 倍。
在正常詞彙量情況下，基於相同的批大小，使用 Mixtape 的網絡速度僅比使用 softmax 的網絡慢 5%-18%；基於相同的內存，前者僅比後者慢 20%-34%。在包含 100K token 的大詞彙量情況下，基於 Mixtape 的網絡的速度僅比基於 softmax 的網絡慢 60%。
Mixtape 和 MoS 在困惑度和翻譯質量方面均優於 softmax。有趣的是，這些基準數據集各自的詞彙量從 10K 到 100K 不等，其輸入表示也不盡相同（包括單詞和 BPE 子詞），這表明 Mixtape 對不同輸入都很高效，具備穩健性。
Mixtape 高效解決 softmax 瓶頸的奧秘
softmax 瓶頸問題的定義。
2017 年，CMU 楊植麟等人提出了解決 Softmax 瓶頸的簡單有效方法 MoS。近期，楊植麟等人為了解決效率低下問題，再次提出了新方法 Mixtape，該方法既能和 MoS 一樣學習高秩表示，其效率又高於 MoS。
如前所述，Mixtape 運用了三項新技術，接下來我們就來查看其細節。
logit 空間向量門控
MoS 成本最高的部分是計算 K 個 softmax，如果我們能夠只用一個 softmax 計算最終概率分布，就可以節約大量計算資源。我們很容易想到將混合（mixture）從概率空間移到 logit 空間，即在 softmax 運算之前混合表示，從而得到條件分布
。但是，正如論文 [19] 所述，該公式會導致低秩表示，因為它仍然使用公式 (1) 中的矩陣分解。
於是，該研究做了一個小小的修改：在 logit 空間中使用混合運算，從而得到高秩表示。其關鍵思想是：使用向量門控機制，而不是標量混合。也就是說，該方法未對每一個 token 使用共享混合權重集，而是對不同的 token 應用不同的權重集。使用向量門控後，條件分布的公式可寫為：
但是，在 Mixtape 取得高效率的路上還有一個障礙。對於每一個 context-token 對 (c, x)，先驗 π_c,x,k 需要執行歸一化操作，而這需要對每一對的先驗概率執行 softmax 運算。
Sigmoid 樹分解
為了高效計算先驗 π_c,x,k，研究者不使用 softmax，而是提出一種新技術——將 softmax 分布分解為 sigmoid 函數樹結構。具體而言，計算 (K 1) 個 sigmoid 輸出，並利用它們定義樹分支的概率。例如，當 K = 4 時，先驗被定義為：
其中 γ_ 表示 sigmoid 概率，σ 表示 sigmoid 函數。
這就是 sigmoid 樹分解。這種分解能夠通過 (K 1) 個 sigmoid 函數完全恢復 K-way 概率分布。使用 sigmoid 函數可以移除 softmax 中的減除運算，更加高效。
令 g_c 作為語境 c 下的 d_1 維度最後一層隱藏狀態，則預激活先驗（pre-activation prior）l_ 的計算公式為：
其中 v_x ∈ R^(d_2)、U_k ∈ R^(d_2×d_1)、u_k ∈ R^(d_1)、b_x,k ∈ R 是模型參數。d_2 是表示門控嵌入大小的超參數，通常比正常的詞嵌入大小 d 要小。語境嵌入可以通過以下公式得到：
其中 H_k ∈ R^(d×d_1) 是模型參數。
門控共享
通過上述兩個方法已經可以得到高效的高秩模型，但仍然存在改進空間。研究者觀察到，我們仍需要對詞彙表中的每個 token 計算門控先驗，而這成為影響效率的瓶頸。但是，由於缺乏訓練樣本，我們很難估計低頻 token 的門控先驗，因此學習低頻 token 的門控先驗可能只是對算力的浪費。基於此，研究者提出門控共享，即對所有低頻詞共享相同的門控先驗。具體而言，對於低頻 token x，預激活門控先驗被定義為：
使用門控共享後，研究者可使用共享門控先驗來混合語境嵌入 h_c,k，然後再與 token 嵌入 w_x 相乘，由於低頻 token 無需存儲門控 logits 從而節省了內存空間。門控共享還加速計算，因為所有低頻 token 僅計算一組門控先驗。
Mixtape 奧秘匯總
Mixtape 層可總結為：
給出最後一層的隱藏狀態 g_c，使用公式 (5) 計算語境嵌入 h_c,k；對每個高頻 token x，使用公式 (4) 計算預激活門控先驗 l_c,x,k；對於所有低頻 token，使用公式 (6) 計算預激活門控先驗 l_c,x,k；使用 sigmoid 樹分解，計算公式 (3) 中的門控先驗 π_c,x,k；使用向量門控，利用公式 (2) 獲得下一個 token 的概率。
Mixtape 層的架構如下圖所示：
圖 1：Mixtape 層架構圖。
實驗
實驗包括三部分：
Mixtape 層打破 softmax 瓶頸，從而改進了當前最優的機器翻譯系統；研究者對比了 Mixtape、MoS 和 softmax 的困惑度、翻譯質量、速度和內存約束，證明 Mixtape 能夠在效果和效率之間做好權衡；控制變量實驗證明了門控共享的優勢。
表 1：在 WMT 英德和英法語言對數據上的性能對比。Mixtape 在這兩項任務上分別使用了 2 億和 8 億參數。
表 2：數據集統計數據概況。「PTB」和「1B」分別表示 Penn Treebank 數據集和 One Billion Word 數據集。
表 3：模型在 Penn Treebank 上的困惑度和訓練時間對比情況。
表 4：模型在 One Billion Word 數據集上的困惑度和訓練時間對比情況。
表 5：模型在 WMT』14 英德語言對數據上的 BLEU 值和訓練時間對比。
表 6：模型在 WMT』14 英法語言對數據上的 BLEU 值和訓練時間對比。