近日，一篇匿名提交給自然語言處理頂會(huì) ACL 的論文《 N-grammer: Augmenting Transformers with latent n-grams 》中，研究者受到統(tǒng)計(jì)語言建模的啟發(fā)，通過從文本序列的離散潛在表示構(gòu)建 n-gram 來增強(qiáng)模型，進(jìn)而對(duì) Transformer 架構(gòu)進(jìn)行了一個(gè)簡(jiǎn)單而有效的修改，稱為 N-grammer。

具體地，N-grammer 層通過在訓(xùn)練期間將潛在 n-gram 表示合并到模型中來提高語言模型的效率。由于 N-grammer 層僅在訓(xùn)練和推理期間涉及稀疏操作，研究者發(fā)現(xiàn)具有潛在 N-grammer 層的 Transformer 模型可以匹配更大的 Transformer，同時(shí)推理速度明顯更快。在 C4 數(shù)據(jù)集上對(duì)語言建模的 N-grammer 進(jìn)行評(píng)估表明，本文提出的方法優(yōu)于 Transformer 和 Primer 等基準(zhǔn)。

論文地址：https://openreview.net/pdf?id=GxjCYmQAody

N-grammer 層

在網(wǎng)絡(luò)高層次上，該研究引入了一個(gè)簡(jiǎn)單的層，該層基于潛在 n-gram 用更多的內(nèi)存來增強(qiáng) Transformer 架構(gòu)。一般來說，N-grammer 層對(duì)于任意 N-gram 來說已經(jīng)足夠了，該研究?jī)H限于使用 bi-gram，以后將會(huì)研究高階 n-gram。這個(gè)簡(jiǎn)單的層由以下幾個(gè)核心操作組成：

• 給定文本的 uni-gram 嵌入序列，通過 PQ （Product Quantization）推導(dǎo)出離散潛在表示序列；
• 推導(dǎo)潛在序列 bi-gram 表示；
• 通過哈希到 bi-gram 詞匯表中查找可訓(xùn)練的 bi-gram 嵌入；
• 將 bi-gram 嵌入與輸入 uni-gram 嵌入相結(jié)合。

此外，當(dāng)提到一組離散項(xiàng)時(shí)，該研究使用符號(hào) [m] 表示集合{0,1，···，m−1}。

序列的離散潛在表示

第一步，N-grammer 層從給定的輸入嵌入序列學(xué)習(xí) Codebook，獲得具有乘積量化（Product Quantization，PQ）（Jegou 等人，2011 年）的離散潛在表示的并行序列。輸入嵌入是一個(gè) uni-gram 嵌入序列 x ϵ R^( l×h×d )，其中 l 是序列長(zhǎng)度，h 是頭數(shù)量，d 是每個(gè)頭嵌入維度。該研究在 R^ k×h×d 中學(xué)習(xí)了一個(gè) Codebook c，通過相同的步驟，該研究選取距離輸入嵌入最小的 code book ID，形成序列 x 的離散潛在表示 z ϵ[k]^l×h 的并行序列：

離散潛在表示 Bi-gram ID

第二步是將離散潛在表示 z 轉(zhuǎn)換為 bi-gram ID b ϵ [k^2 ]^( l×h )。它們通過組合來自前一個(gè)位置的 uni-gram 潛在 ID z，然后在當(dāng)前位置形成潛在 bi-gram ID：

其中 k 是 codebook 大小，這直接將離散潛在序列從詞匯空間[k] 映射到潛在 bi-gram 詞匯空間 [k^2 ] 。

構(gòu)建 bi-gram 表示

第三步是構(gòu)建序列 bi-gram 潛在表示 b?？紤]所有的 k^2 bi-gram，并通過對(duì)每個(gè)這樣的 bi-gram 嵌入來增強(qiáng)模型。在實(shí)踐中，對(duì)于 uni-gram 詞匯為 32,000 的機(jī)器翻譯模型壓縮，在不犧牲質(zhì)量的情況下，需要將 187 個(gè) token 聚類為 k = 212 個(gè) cluster。在這種情況下，需要考慮所有的 bi-gram，涉及構(gòu)建一個(gè)包含 1600 萬行的嵌入表。由于所構(gòu)建的表仍然很大，該研究通過對(duì)每個(gè)頭使用單獨(dú)的哈希函數(shù)，將潛在 bi-gram ID 映射到大小為 v 的較小的 bi-gram 詞匯表。

更準(zhǔn)確地講，該研究有一個(gè)潛在 bi-gram 嵌入表 B ϵ R^v×h×d_b，其中 v 為 bi- gram 詞匯，d_b 為 bi-gram 嵌入維度。然后將文本序列 bi-gram 嵌入構(gòu)建為：

與嵌入進(jìn)行結(jié)合

最后一步是將 uni-gram 嵌入 x ϵ R^(l×h×d)與潛在 bi-gram 嵌入 y∈R^(l×h×db)相結(jié)合，形成文本序列新表示。bi-gram 嵌入和 uni-gram 嵌入都是獨(dú)立的層歸一化(LN)，然后沿著嵌入維度連接兩者以產(chǎn)生 w = [LN(x), LN(y)] ϵ R^l×h×(d+db) ，并將其作為輸入傳遞給 Transformer 網(wǎng)絡(luò)的其余部分。

實(shí)驗(yàn)結(jié)果

該研究在 C4 數(shù)據(jù)集上將 N-grammer 模型與 Transformer 架構(gòu)（Vaswani 等人，2017 年）以及最近提出的 Primer 架構(gòu)（So 等人，2021 年）進(jìn)行了比較。其中，該研究使用 Adam 優(yōu)化器，所有模型的學(xué)習(xí)率為 10^-3，而對(duì)于 n-gram 嵌入表，學(xué)習(xí)率為 10^-2。

下表 1 比較了 N-grammer、Primer 和 Transformer 模型，其中基線 Transformer 模型有 16 層和 8 個(gè)頭，模型維度為 1024。研究者在 TPU v3 上以 256 的批大小和 1024 的序列長(zhǎng)度訓(xùn)練所有模型。研究者對(duì) N-grammer 模型進(jìn)行了消融研究，bi-gram 嵌入維度大小從 128 到 512 不等。由于添加 n-gram 嵌入增加了可訓(xùn)練參數(shù)的數(shù)量，該研究還在表 1 中訓(xùn)練了兩個(gè)大基線（Transformer-L 和 Primer-L），它們的參數(shù)順序與 N-grammer 模型相同。然而，與較大的 Transformer 模型不同，N-grammer 的訓(xùn)練和推理成本與嵌入層中的參數(shù)數(shù)量不成比例，因?yàn)樗鼈円蕾囉谙∈璨僮鳌?/p>

該研究還測(cè)試了一個(gè)簡(jiǎn)單版本的 N-grammer，研究者直接從 uni-gram 詞匯表（3.3 節(jié)中的）而不是從潛在表示中計(jì)算 n-gram（3.1 節(jié)的）。由表 1 可知，它對(duì)應(yīng)于在 clusters 列中沒有條目的 N- grammer。