ON-LSTM：用有序神經元表達層次結構

原文鏈接：https://spaces.ac.cn/archives/6621

本文轉載自 蘇劍林[1]在科學空間上關於ON-LSTM非常好的解析, 中間加入一些細節的說明, 以期讀者在理解原理之餘, 對其實現也有概念. 下面開始…

前言

今天介紹一個有意思的LSTM變種：ON-LSTM，其中“ON”的全稱是“Ordered Neurons”，即有序神經元，換句話說這種LSTM內部的神經元是經過特定排序的，從而能夠表達更豐富的信息。ON-LSTM來自文章《Ordered Neurons: Integrating Tree Structures into Recurrent Neural Networks》，顧名思義，將神經元經過特定排序是爲了將層級結構（樹結構）整合到LSTM中去，從而允許LSTM能自動學習到層級結構信息。這篇論文還有另一個身份：ICLR 2019的兩篇最佳論文之一(不知道是不是跟作者來自Mila有關係…)，這表明在神經網絡中融合層級結構（而不是純粹簡單地全向鏈接）是很多學者共同感興趣的課題。

筆者留意到ON-LSTM是因爲機器之心的介紹，裏邊提到它除了提高了語言模型的效果之外，甚至還可以無監督地學習到句子的句法結構！正是這一點特性深深吸引了我，而它最近獲得ICLR 2019最佳論文的認可，更是堅定了我要弄懂它的決心。認真研讀、推導了差不多一星期之後，終於有點眉目了，遂寫下此文。

在正式介紹ON-LSTM之後，我忍不住要先吐槽一下這篇文章實在是寫得太差了，將一個明明很生動形象的設計，講得異常晦澀難懂，其中的核心是$\tilde{f_{t}}$和$\tilde{i_{t}}$的定義，文中幾乎沒有任何鋪墊就貼了出來，也沒有多少詮釋，開始的讀了好幾次仍然像天書一樣…總之，文章寫法實在不敢恭維～

1. 背景: 回顧LSTM

首先來回顧一下普通的LSTM(不考慮層數)。用常見的記號，普通的LSTM寫爲：

如果熟悉了神經網絡本身，其實這樣的結構沒有什麼神祕的，$f_{t}$, $i_{t}$, $o_{t}$就是三個單層全連接模型，輸入是歷史信息$h_{t−1}$和當前信息$x_{t}$，用sigmoid激活，因爲sigmoid的結果在0～1之間，所以它們的含義可以詮釋爲“門（gate）”，分別稱爲遺忘門、輸入門、輸出門。不過我個人覺着gate這個名字是不夠貼切的，“valve（閥門）”也許更貼切些。

有了門之後，$x_{t}$被整合爲$\hat{c_{t}}$，然後通過∘運算（對應逐位相乘，有時候也記爲⊗）與前面的“門”結合起來，來對$c_{t−1}$和$\hat{c_{t}}$進行加權求和。

Notice: 這裏, 相比最一般的RNN(SimpleRNN), LSTM在層間傳遞中多了一個細胞狀態($c_{t}$)[2]:

下圖是原博主繪製的圖, 可以看到相比RNN, LSTM在層內傳遞(同一層不同時間步)的時候多了一個細胞狀態($c_{t}$)[2].

2. 背景: 語言和序信息

在常見的神經網絡中，神經元通常都是無序的，比如遺忘門$f_t$是一個向量，向量的各個元素的位置沒有什麼規律。如果把LSTM運算過程中涉及到的所有向量的位置按照同一方式重新打亂，權重的順序也相應地打亂，然後輸出結果可以只是原來向量的重新排序（考慮多層的情況下，甚至可以完全不變），信息量不變，不影響後續網絡對它的使用。

換言之，LSTM以及普通的神經網絡都沒有用到神經元的序信息，ON-LSTM則試圖把這些神經元排個序，並且用這個序來表示一些特定的結構，從而把神經元的序信息利用起來。

ON-LSTM的思考對象是自然語言。一個自然句子通常能表示爲一些層級結構，這些結構如果人爲地抽象出來，就是我們所說的語法信息，而ON-LSTM希望能夠模型在訓練的過程中自然地學習到這種層級結構，並且訓練完成後還能把它解析出來（可視化），這就利用到了前面說的神經元的序信息。（曾經做過的相關研究《最小熵原理（三）：“飛象過河”之句模版和語言結構》[3]）

爲了達到這個目標，我們需要有一個層級的概念，層級越低代表語言中顆粒度越小的結構，而層級越高則代表顆粒度越粗的結構，比如在中文句子中，“字”可以認爲是最低層級的結構，詞次之，再上面是詞組、短語等。層級越高，顆粒度越粗，那麼它在句子中的跨度就越大。

用原文的圖示就是：

3. ON-LSTM

上面的最後一句“層級越高，顆粒度越粗，那麼它在句子中的跨度就越大”很容易理解，但它對於ON-LSTM的設計有着指導作用。

首先，這要求我們在設計ON-LSTM的編碼時能區分高低層級的信息；其次，這也告訴我們，高層級的信息意味着它要在高層級對應的編碼區間保留更久（不那麼容易被遺忘門過濾掉），而低層級的信息則意味着它在對應的區間更容易被遺忘。

3.1 設計：分區間更新

有了這個指導之後，我們可以着手建立。假設ON-LSTM中的神經元都排好序後，向量$c_t$的index越小的元素，表示越低層級的信息，而index越大的元素，則表示越高層級的信息。然後，ON-LSTM的門結構和輸出結構依然和普通的LSTM一樣：

相比正常的LSTM, ON-LSTM的核心就在於: $\hat{c_{t}}$到$c_{t}$的更新方式變化了.

接下來，初始化一個全零的$c_{t}$，即沒有任何記憶，或者想象爲一個空的U盤。然後，我們將歷史信息和當前輸入按一定規律存入到$c_{t}$中（即更新$c_{t}$）。每次在更新$c_{t}$之前，首先預測兩個整數 $d_f$: 表示歷史信息$h_{t−1}$的層級; $d_i$: 表示當前輸入$x_{t}$的層級：

至於$F_1$, $F_2$的具體結構，我們後面再補充，先把核心思路講清楚。這便是我不滿原論文寫作的原因，一上來就定義cumax，事前事後都沒能把思想講清楚。

現在我們有了$d_f$和$d_i$, 那麼作者根據其大小關係, 有兩種進行細胞狀態的更新的策略.

至此，我們能夠理解ON-LSTM的基本原理了: 它將神經元排序之後，通過位置的前後來表示信息層級的高低，然後在更新神經元時，先分別預測歷史的層級$d_f$和輸入的層級$d_i$，通過這兩個層級來對神經元實行分區間更新。

這樣一來，高層信息就可能保留相當長的距離（因爲高層直接複製歷史信息，導致歷史信息可能不斷被複制而不改變），而低層信息在每一步輸入時都可能被更新（因爲低層直接複製輸入，而輸入是不斷改變的），所以就通過信息分級來嵌入了層級結構。更通俗地說就是分組更新，更高的組信息傳得更遠（跨度更大），更低的組跨度更小，這些不同的跨度就形成了輸入序列的層級結構。

（請反覆閱讀這段話，必要時對照上圖，直接完全理解爲止，這段話稱得上是ON-LSTM的設計總綱。）

3.2 設計：分段軟化

現在要解決的問題就是，這兩個層級怎麼預測，即$F_1$, $F_2$怎麼構建。用一個模型來輸出一個整數不難，但是這樣的模型通常都是不可導的，無法很好地整合到整個模型進行反向傳播，所以，更好的方案是進行“軟化”，即尋求一些光滑近似。

舉個例子: $1_{d_{f}}$ = [0, 0, 0, 1, 0], $1_{d_{i}}$ = [0, 0, 0, 0, 1]. 則根據(6)式, 有$\tilde{f_{t}}$ = [0, 0, 0, 1, 1], $\tilde{i_{t}}$ = [1, 1, 1, 1, 1], 所以$w_{t} = \tilde{f_{t}} ∘ \tilde{i_{t}}$ = [0, 0, 0, 1, 1]. 即對這種情況(公式(4), $d_f$ <= $d_i$, 表示信息有融合), 對4, 5位置處的信息進行融合, 1, 2, 3位置處變爲$\hat{c_{t}}$.

注:

• 1、論文中將$\overrightarrow{cs}(softmax(x))$簡記爲cumax(x)，這只是記號上的轉換而已；
• 2、作爲數列來看，$\tilde{f_{t}}$是一個單調遞增的數列，而$\tilde{i_{t}}$是一個單調遞減的數列；
• 3、$1_{d_{f}}$和$1_{d_{i}}$的dimension是 (N, unit), 其中N爲Batch size, unit是本層的單元數. 同樣的, 隱藏狀態hidden state $h_t$和細胞狀態$c_t$的dimension也都是 (N, unit). 而輸入$x_t$的dimension是 (N, input_dim).
• 4、由於3和1, 需要通過2個全連接層來估計參數, 這會加大計算量, 尤其是當unit和input_dim都被設置的很大的時候. 後面可能會有更好的低複雜度解決方案.

4. 試驗

下面簡單彙總一下ON-LSTM的實驗，其中包括原作者的實現（PyTorch）以及蘇劍林大神自己的復現（Keras）

原作者實現：https://github.com/yikangshen/Ordered-Neurons

蘇神實現：https://github.com/bojone/on-lstm

4.1 分組層級

Notice: 這裏減少的是下式中的(全連接層, 激活函數爲softmax)的輸出的單元個數, 因此減少了一定的計算量.

4.2 語言模型

4.3 無監督語法

如果僅僅是在常規的一些語言任務中超過普通LSTM，那麼ON-LSTM也算不上什麼突破，但ON-LSTM的一個令人興奮的特性是它能夠無監督地從訓練好的模型（比如語言模型）中提取輸入序列的層級樹結構。提取的思路如下：

首先我們考慮：

它是$\tilde{f_{t}}$在$\overrightarrow{cs}$之前的結果，根據我們前面的推導，它就是歷史信息的層級$d_{f}$的一個軟化版本，那麼我們可以寫出：


算法的大概意思是從最高層級處斷開（這意味着當此處包含的歷史信息最少，與前面所有內容的聯繫最爲薄弱，最有可能是一個新的子結構的開始），然後遞歸處理，從而逐漸得到輸入序列隱含的嵌套結構。作者是用三層的ON-LSTM訓練了一個語言模型，然後用中間那層ON-LSTM的$\tilde{f_{t}}$來計算層級，然後跟標註的語法結構對比，發現準確率頗高。我自己也在中文語料下嘗試了一下：https://github.com/bojone/on-lstm/blob/master/lm_model.py

至於效果，因爲我沒做過也不瞭解語法分析，我也不知道怎麼評價，反正好像看着是那麼一回事，但是又好像不大對一樣，所以各位讀者自己評價好了～近一兩年，無監督語法分析其實還有不少研究工作，可能要都讀一讀才能更深刻地理解ON-LSTM。

輸入：蘋果的顏色是什麼
輸出：
[
  [
    [
      ‘蘋果’,
      ‘的’
    ],
    [
      ‘顏色’,
      ‘是’
    ]
  ],
  ‘什麼’
]

5. 思考&發散

文章最後，我們來一起思考幾個問題。

• RNN還有研究價值？

首先，有讀者可能會困惑，都9102年了，居然還有人研究RNN類模型，還有研究價值嗎？近年來，BERT、GPT等基於Attention和語言模型的預訓練模型，在NLP的諸多任務上都提升了效果，甚至有文章直接說“RNN已死”之類的。事實上真的如此嗎？我認爲，RNN活得好好的，並且在將來的相當長時間內都不會死，原因至少包含下面幾個：

第一，BERT之類的模型，以增加好幾個數量級的算力爲代價，在一些任務上提升了也就一兩個百分點的效果，這樣的性價比只有在學術研究和比賽刷榜才有價值，在工程上幾乎沒什麼用（至少沒法直接用）；
第二，RNN類的模型本身具有一些無可比擬的優勢，比如它能輕鬆模擬一個計數函數，在很多序列分析的場景，RNN效果好得很；第三，幾乎所有seq2seq模型（哪怕是BERT中）decoder都是一種RNN，因爲它們基本都是遞歸解碼的，RNN哪會消失？

• 單向ON-LSTM就夠了？

然後，讀者可能會有疑惑：你要析出層級結構，但是隻用了單向的ON-LSTM，這意味着當前的層級分析還不依賴於將來的輸入，這顯然是不大符合事實的。這個筆者也有同樣的困惑，但是作者的實驗表明這樣做效果已經夠好了，可能自然語言的整體結構都傾向於是局部的、單向的（從左往右），所以對於自然語言來說單向也就夠了。

如果一般情況下是否用雙向比較好呢？雙向的話是不是要像BERT那樣用masked language model的方式來訓練呢？雙向的話又怎麼計算層級序列呢？這一切都還沒有完整的答案。至於無監督析出的結構是不是一定就符合人類自身理解的層級結構呢？這個也說不準，因爲比較沒有什麼監督指引，神經網絡就“按照自己的方式去理解”了，而幸運的是，神經網絡的“自己的方式”，似乎跟人類自身的方式有不少重疊之處。

• 爲什麼析出層級考慮的是$d_f$而不是$d_i$？

讀者可能會困惑，明明有兩個master gate，爲什麼析出層級用$d_f$而不是$d_i$？要回答這個問題，我們要理解$d_f$的含義。我們說$d_f$是歷史信息的層級，換言之，它告訴我們做出當前決策還要用多少歷史信息。如果$d_f$很大，意味着當前決策幾乎用不着歷史信息了，這意味着從當前開始就是一個新層級的開始，與歷史輸入幾乎割斷了聯繫。也就是從這種割斷和聯繫中析出了層級結構，所以只能用$d_f$。

• 能否用到CNN或者Attention？

最後，可能想到的一個困惑是，這種設計能不能用到CNN、Attention之中呢？換句話說能不能將CNN、Attention的神經元也排個序，融入層級結構信息呢？個人感覺是有可能的，但需要重新設計，因爲層級結構被假設爲連續嵌套的，RNN的遞歸性正好可以描述了這種連續性，而CNN、Attention的非遞歸性導致我們很難直接去表達這種連續嵌套結構。

不管怎樣，我覺得這是個值得思考的主題，有進一步的思考結果我會和大家分享，當然也歡迎讀者們和我分享你的思考。

6. 文章總結

本文梳理了LSTM的一個新變種ON-LSTM的來龍去脈，主要突出了它在表達層級結構上的設計原理。個人感覺整體的設計還是比較巧妙和有趣的，值得細細思考一番。

最後，學習和研究都關鍵是有自己的判斷能力，不要人云亦云，更不能輕信媒體的“標題黨”。BERT的Transformer固然有它的優勢，但是LSTM等RNN模型的魅力依然不可小覷。我甚至覺得，諸如LSTM之類的RNN模型，會在將來的某天，煥發出更強烈的光彩，transformer與之相比將會相當遜色。

讓我們拭目以待好了。

7. 參考資料

[1] ON-LSTM：用有序神經元表達層次結構
[2] 理解 LSTM(Long Short-Term Memory, LSTM) 網絡
[3] 《最小熵原理（三）：“飛象過河”之句模版和語言結構》