注意力 (attention) 模型在神經網絡中被廣泛應用，不過注意力機制一般是決定性的而非隨機變量。來自哈佛大學的研究人員提出了將注意力建模成隱變量，應用變分自編碼器（Variational Auto-Encoder，VAE）和梯度策略來訓練模型，在不使用 kl annealing 等訓練技巧的情況下進行訓練，目前在 IWSLT German-English 上取得了非常不錯的成果。
論文連結：https://arxiv.org/abs/1807.03756v1相關代碼：https://github.com/harvardnlp/var-attn
一、背景
近年來很多論文將 VAE 應用到文本生成上，通過引入隱變量對不確定性進行建模。不過這會導致一個常見的 KL collapsing 問題，導致的現象就是直接訓練的 VAE 得到的 KL 接近於 0，也就意味著近似後驗和先驗一樣，隱變量被模型所忽略。
為了解決這個問題，在哈佛研究人員的工作中，注意力被建模成隱變量，由於解碼器 (decoder) 和解碼器 (encoder) 之間的主要信息傳輸通道是通過注意力來傳輸的，如果忽略了這個隱變量，就會因無法得到源文本的信息而得到很大的懲罰 (penalty)（這相比之前的許多工作中直接把隱變量加入到每個解碼步驟不同，因為那樣即使解碼器忽略了隱變量，也可以達到很好的模型表現）。因此通過直接優化目標函數才能使得這個隱變量也不易被忽略，研究人員的實驗完全驗證了這一點。
值得注意的是，研究人員將注意力建模成隱變量並不是單純為了應用 VAE 這個工具，而是因為研究人員認為將注意力建模成隱變量可以為解碼器提供更乾淨 (clean) 的特徵，從而在不增加模型參數的情況下提高模型的表現能力。
二、實驗效果
與傳統的注意力機制進行對比，傳統的注意力機制僅能通過之前生成的單詞確定當前即將生成單詞的注意力（上圖藍色部分，僅能觀測到已經預測的單詞，由於存在多種翻譯方式，因此會注意到和實際翻譯的單詞並不對應的位置），研究人員通過全部的源文本和目標文本去得到更準確的後驗注意力（上圖紅色部分，通過全部信息得到後驗，因此注意力和實際翻譯應該與注意 (attend) 的源單詞對應），並把更好的後驗注意力提供給解碼器，從而使解碼器拿到更為乾淨的特徵，藉此希望得到更好的模型。
三、核心思想
方法：假定 x 是源文本，y 是目標文本，z 是注意力，根據標準的 VAE 方法，研究人員引入推理網絡 (inference network) q(z | x, y) 去近似後驗，那麼 ELBO 可以表達為（為了方便，只考慮目標文本只有一個單詞的情況）：
上面不等式的右側是 ELBO，其中第一項是從 q(z | x, y) 中採樣出注意力，使用採樣出的注意力作為解碼器的輸入計算交叉熵損失，第二項是確保後驗分布接近於先驗分布（注意到此處的先驗和一般的 VAE 不同，這裡的先驗是和模型一起學習的）。此時的 p(z | x) 和 q(z | x, y) 都是分類分布的，因此我們使用梯度策略去優化上面的目標函數。
由於此時的後驗 q 能看到全部的 x 和 y，因此後驗中採樣的注意力可以比先驗 p(z | x) 好，比如以下的例子：
這裡需要把德語（縱向）翻譯成英語（橫向），紅色代表先驗（即只觀測到 x 而不觀測到 y 的 p(z | x)），藍色代表後驗（即觀測到全部信息的 q(z | x, y)）。注意到在第二個單詞「actually」處，紅色的先驗試圖注意到「nun」後面的逗號「,」，從而試圖生成一個「well,」的翻譯結果，然而實際的英語翻譯中並沒有逗號，而是直接是」well actually」，由於後驗 q(z | x, y) 可以看到實際的翻譯，因此藍色的後驗正確地注意到了「tatsachlich」上。注意到訓練目標 ELBO 中我們從後驗 q 中採樣注意力給解碼器，因此通過使用 VAE 的方法，解碼器得到了更準確的注意力信號，從而提高了模型的表達能力。
四、展望
注意力方法是自然語言處理等領域普遍存在的工具，但它們很難用作隱變量模型。這項工作通過具有良好實驗結果的變分注意力機制來探索潛在對齊的替代方法。研究人員同時也表示未來的工作將實驗規模較大的任務和更複雜的模型，如多跳注意力模型 (multi-hop attention models)，變壓器模型 (transformer models) 和結構化模型 (structured models)，以及利用這些潛在變量的可解釋性作為一種方式去將先驗知識進行結合。