歡迎光臨
每天分享高質量文章

可逆ResNet:極致的暴力美學

作者丨蘇劍林

單位丨廣州火焰信息科技有限公司

研究方向丨NLP,神經網絡

個人主頁丨kexue.fm

今天我們來介紹一個非常“暴力”的模型:可逆 ResNet。

 

為什麼一個模型可以可以用“暴力”來形容呢?當然是因為它確實非常暴力:它綜合了很多數學技巧,活生生地(在一定約束下)把常規的 ResNet 模型搞成了可逆的!

 

 標準ResNet與可逆ResNet對比圖。可逆ResNet允許信息無損可逆流動,而標準ResNet在某處則存在“坍縮”現象。

 

模型出自 Invertible Residual Networks,之前在機器之心也報導過。在這篇文章中,我們來簡單欣賞一下它的原理和內容。

 

 

可逆模型的點滴

 

為什麼要研究可逆 ResNet 模型?它有什麼好處?以前沒有人研究過嗎? 

 

可逆的好處

 

可逆意味著什麼? 

 

意味著它是信息無損的,意味著它或許可以用來做更好的分類網絡,意味著可以直接用最大似然來做生成模型,而且得益於 ResNet 強大的能力,意味著它可能有著比之前的 Glow 模型更好的表現。

 

總而言之,如果一個模型是可逆的,可逆的成本不高而且擬合能力強,那麼它就有很廣的用途(分類、密度估計和生成任務,等等)。 

 

而本文要介紹的可逆 ResNet 基本上滿足這幾點要求,它可逆起來比較簡單,而且基本上不改變 ResNet 的擬合能力。因此,我認為它稱得上是“美”的模型。 

 

舊模型的局限

 

可逆模型的研究由來已久,它們被稱為“流模型(flow-based model)”,代表模型有 NICERealNVPGlow,筆者曾撰文介紹過它們,另外還有一些自回歸流模型。除了用來做生成模型,用可逆模型來做分類任務的也有研究,代表作是 RevNet [1] 和 i-RevNet [2]。

 

這些流模型的設計思路基本上都是一樣的:通過比較巧妙的設計,使得模型每一層的逆變換比較簡單,而且雅可比矩陣是一個三角陣,從而雅可比行列式很容易計算。

 

這樣的模型在理論上很優雅漂亮,但是有一個很嚴重的問題:由於必須保證逆變換簡單和雅可比行列式容易計算,那麼每一層的非線性變換能力都很弱。事實上像 Glow 這樣的模型,每一層只有一半的變數被變換,所以為了保證充分的擬合能力,模型就必須堆得非常深(比如 256 的人臉生成,Glow 模型堆了大概 600 個捲積層,兩億引數量),計算量非常大。 

 

硬“杠”殘差模型

 

而這一次的可逆 ResNet 跟以往的流模型不一樣,它就是在普通的 ResNet 結構基礎上加了一些約束,就使得模型成為可逆的,實際上依然保留了 ResNet 的基本結構和大部分的擬合能力。這樣一來,以往我們在 ResNet 的設計經驗基本上還可以用,而且模型不再需要像 Glow 那樣拼命堆捲積了。 

 

當然,這樣做是有代價的,因為沒有特別的設計,所以我們需要比較暴力的方法才能獲得它的逆函式和雅可比行列式。所以,這次的可逆 ResNet,很美,但也很暴力,稱得上是“極致的暴力美學”。

 

可逆“三要素” 

 

ResNet 模型的基本模塊就是:

 

也就是說本來想用一個神經網絡擬合 y 的,現在變成了用神經網絡擬合 y−x 了,其中 x,y 都是向量(張量)。這樣做的好處是梯度不容易消失,能訓練深層網絡,讓信息多通道傳輸,等等。可逆的意思就是 x+g(x) 是一個一一映射,也就是說每個 x 只有一個 y 與之對應,反過來也是,換言之我們理論中可以從中解出反函式 x=h(y) 來。

 

背景就不多說了,但是要說明一點,我們在分類問題中用的廣義上的 ResNet,是允許通過 1×1 捲積改變維度的,但這裡的 ResNet 指的是不改變維度的 ResNet,也就是說 x,y 的大小保持一樣

 

對於一個號稱“可逆”的模型,必須要回答三個問題:

 

  • 什麼時候可逆?

  • 逆函式是什麼?

  • 雅可比行列式怎麼算?

 

從難度上來看,這三個問題是層層遞進的。當然,如果你只關心做分類問題,那麼事實上保證第一點就行了;如果你關心重構圖片,那麼就需要第二點;如果你還想像 Glow 那樣用最大似然來訓練生成模型,那麼就需要第三點。 

 

下麵按照原論文的思路,逐一解決這三個問題(三道“硬菜”),來一場暴力盛宴。 

 

什麼時候可逆?

 

第一道硬菜是三道硬菜中相對容易啃的,當然只是“相對”,事實上也都用到了泛函分析的一些基礎知識了。 

 

因為 (1) 是 ResNet 的基本模塊,所以我們只需要保證每個模塊都可逆就行了。而 (1) 可逆的一個充分條件是:

 

其中:

 

 

是函式 g 的 Lipschitz 範數。也就是說,g 的 Lipschitz 範數小於 1,就能保證 (1) 可逆了。

 

那什麼時候 g 的 Lipschitz 範數才會小於 1 呢?因為 g 是神經網絡,捲積和全連接都無妨,神經網絡是由矩陣運算和激活函式組合而成的,即代表結構是:

 

 

那麼由鏈式法則,“g 的 Lipschitz 範數小於 1”的一個充分條件是“Activation 的 Lipschitz 範數不超過 1”且“Wx+b 的 Lipschitz 範數小於 1”。而 Activation 只是個標量函式,“Lipschitz 範數不超過 1”意味著導數不超過 1 就行了,目前我們常用的激活函式(sigmoid、tanh、relu、elu、swish 等)都滿足,所以這一點不用管它;而“Wx+b 的 Lipschitz 範數小於 1”,意味著矩陣 W 的 Lipschitz 範數小於 1。

 

矩陣 W 的 Lipschitz 範數其實也就是“譜範數”,記為 Lip(W) 或都行。那啥時候出現過矩陣的譜範數?在《深度學習中的 Lipschitz 約束:泛化與生成模型》一文中我們討論 Lipschitz 約束的時候就出現過,兩者結合起來,結論就是:

 

對模型 g 的所有核權重 W 做譜歸一化,然後乘上一個大於 0 小於 1 的繫數 c(即),就可以使得 x+g(x) 可逆了。

逆函式是什麼?

 

為什麼這樣就可逆了?證明過程可以直接回答第二個問題,也就是說,我們直接把逆函式求出來,然後就知道什麼條件下可逆了。 

 

假如 y=x+g(x) 是可逆的,那麼我們要想辦法求出逆函式 x=h(y),這其實就是解一個非線性方程組。簡單起見,我們考慮迭代:

 

 

顯然,迭代序列 {Xn} 是跟 y 有關的,而一旦 {Xn} 收斂到某個固定的函式:

 

 

那麼我們就有,這意味著就是我們希望求的 x=h(y)。

 

換句話說,如果迭代 (5) 收斂,那麼收斂的結果就是 x+g(x) 的逆函式。所以我們只需要搞清楚 (5) 什麼時候收斂。這時候前面的條件 Lip(g)<1 就能用上了,我們有:

 

 

所以:

 

 

可以看到,的充分條件是 Lip(g)<1。

 

附:單純指出並不能說明序列 {Xn} 收斂,比如 {lnn} 這個序列也滿足這個條件,但發散。所以,為了證明 Lip(g)<1 時 {Xn} 收斂,我們還需要多做一些工作。不過這是比較數學的部分了,考慮到部分讀者可能不感興趣,因此作為附註。

 

對於任意正整數 k,我們繼續考慮

 

 

可以看到我們得到了的一個上界,它只與 n 有關,且可以任意小。也就是說,對於任意 ε>0,我們可以找到一個 n,使得對於任意的正整數 k 都有。這樣的數列我們稱為 Cauchy 列,它是必然收斂的。

 

至此,我們終於證明瞭 {Xn} 的收斂性。順便一提的是,在 (9) 中,取 k→∞,我們得到:

 

 

也就是說,這個迭代演算法的收斂速度跟正比於,那麼自然是 Lip(g) 越小收斂越快,但是 Lip(g) 越小模型的擬合能力越弱,原論文中它的範圍是 0.5~0.9。

 

說宏大一點,其實這就是泛函分析中的“巴拿赫不動點定理”,又稱“壓縮映射定理”(因為 Lip(g) 小於 1,所以 g 被稱為一個壓縮映射)。

 

這樣一來,我們已經回答了為什麼 Lip(g)<1 就有 x+g(x) 可逆了,同時已經給出了求逆函式的方法——就是通過 (5) 迭代到足夠的精度:

 

當做好歸一化操作使得 x+g(x) 可逆後,它的逆函式為的不動點。數值計算時,只需要迭代一定步數,使得滿足精度要求即可。

 

終於,我們啃下了第二道硬菜。

 

雅可比行列式怎麼算?

 

下麵來到三個問題中最“硬核”的一個問題:雅可比行列式怎麼算?為瞭解決它,作者綜合了數學分析、矩陣論、概率論、統計採樣等多方面數學工具,堪稱“暴力之最”、“硬菜之最”。 

 

首先,為什麼要算雅可比行列式?前面已經說了,只有做生成模型時才有這個必要,具體細節請參考筆者最早的對流模型的介紹《細水長flow之NICE:流模型的基本概念與實現》。接著,我們知道雅可比行列式就是雅可比矩陣的行列式,所以要把雅可比矩陣算出來:

 

 

再次提醒,雖然我偷懶沒有加粗,但這裡的 g 輸出是一個向量,x 也是一個向量,∂g/∂x 實際上就是輸入和輸出兩兩配對求偏導數,結果是一個矩陣(雅可比矩陣)。

 

然後,雅可比行列式就是,但事實上,在做生成模型的時候,我們真正要算的,是“雅可比行列式的絕對值的對數”,即:

 

 

最後一個恆等號,是因為 det(I+Jg) 一定是正數,所以可以去掉絕對值。這是可以證明的,但只是細節部分,我們就不糾結了,讀者自行去看作者提供的參考文獻吧。 

 

然後呢?直接按定義來計算雅可比行列式?不行,因為這樣子計算量實在是太大了,而且反向傳播的時候,還要算行列式的導數,那就更複雜了。作者們想出了繁瑣但有效的解決方法,利用恆等式(參考《恆等式 det(exp(A)) = exp(Tr(A)) 賞析》[3]):

 

 

我們得到:

 

如果能求出 ln(I+Jg) 來,然後求跡(trace,對角線元素之和)就行了。怎麼求 ln(I+Jg) 呢?還是參考《恆等式 det(exp(A)) = exp(Tr(A)) 賞析》[3],暴力展開:

 

 

註意這個級數收斂的條件是,這正好意味著 Lip(g)<1,而這正是可逆 ResNet 的約束,前後完全自洽。

 

現在 ln(I+Jg) 變成了一個無窮級數,如果截斷 n 項,那麼誤差也正比於,所以我們要看 Lip(g) 來決定截斷的數目。於是我們可以寫出:

 

 

問題解決了嗎?上式需要我們去計算,註意 Jg 是一個矩陣,我們要算矩陣的 n 次方(想想算兩個矩陣乘法的工作量)。於是作者們想:既然分析的工具用完了,那我們就上概率統計吧。

 

假設 p(u) 是一個多元概率分佈,其均值為 0、協方差為單位矩陣(顯然標準正態分佈符合要求),那麼對於任意矩陣 A,我們有:

 

 

利用“均值為 0、協方差為單位矩陣”這個性質,直接按定義就可以證明上式了,並不困難。然後,作者提出了一個顯得“既無賴又合理”的方法:對於每次迭代,我只從 p(u) 中隨機選兩個向量 u1,u2 出來,然後認為就是 Tr(A),即:

 

讀者可能就有意見了,不是要對所有向量求平均嗎?只隨機挑兩個就行了?其實還真的是可以了,理由如下:

 

1. 我們優化都是基於隨機梯度下降的,本來就帶有誤差,而只隨機挑兩個也有誤差,而每步迭代都重新挑不同的 u1,u2,在一定程度上就能抵消誤差;

 

2. 更重要的原因是,我們要算雅可比行列式的對數,只是用它來做一個額外的 loss,來保證模型不會坍縮,簡單來講,可以將它看稱一個正則項,而既然是一個正則項,有點誤差也無妨。

 

所以,的計算就被作者用這麼粗暴(又有效)的方案解決了。註意到:

 

所以不需要把算出來,而是每步只需要算一個矩陣與一個向量的乘法,並且每步計算可以重覆利用,因此計算量是大大減少的。

 

所以,最終可以總結為:將雅可比矩陣做對數級數展開 (15),然後將行列式計算轉化為跡的計算 (16),並且利用概率採樣的方式 (18),就可以最高效地計算雅可比行列式。

 

縱觀實驗效果

 

其實筆者一開始是就被“可逆 ResNet”這麼美好的構思吸引過來的,但是看到這裡,我發現我也慫了,這絕對對得起“硬杠 ResNet”的評價呀。本來想對照著好歹實現個 mnist 的生成來玩玩,後來確認有這麼多技巧,如此之暴力,我也放棄了。 

 

所以,還是來看看原論文的實驗結果好了。 

 

Toy資料集

 

首先是一個人造的 Toy 資料集,也就是構造一些有規律的隨機點來,然後用生成模型去擬合它的分佈,GAN 中也常做這樣的實驗。

 

從下圖可以看到可逆 ResNet 的效果比 Glow 好得多,歸根結底,我覺得是因為可逆 ResNet 是一個很對稱的、沒有什麼偏置的模型,而 Glow 則是有偏的,它需要我們以某種方式打亂輸入,然後對半切分,對半之後兩部分的運算是不一樣的,這就帶來了不對稱。

 

 可逆ResNet實驗:Toy資料集

 

類任務實驗 

 

一開始我們就說了,既然是 ResNet 的一種,最基本的用途就是分類了。下表表明用可逆 ResNet 做分類任務,效果也是很不錯的,Lipschitz 約束的存在不會明顯影響分類結果(表中的 c 就是 Lip(g))。

 

 可逆ResNet實驗:分類效果

生成模型實驗

 

其實流模型系列在生成複雜圖片方面還遠比不上 GAN,但相互之間的定量對比倒沒有問題。下圖也表明可逆 ResNet 作為一個基於流的生成模型方面也是很優秀的。

 

 可逆ResNet實驗:生成模型效果


終於可以收工了

這樣活生生地取硬杠 ResNet 可真是一件苦力活,我就單純的解讀就已經這麼累了,真佩服作者的數學功底。當然,最後作者終究還是成功了,想必成功的喜悅也是很豐盛的。總的來說,整個工作很暴力,但細細品味之下並沒有什麼違和感,反倒是有一種渾然的美感在裡邊,並非簡單的堆砌數學公式。

 

唯一的問題是,整個“硬杠”的流程還是挺複雜的,因此要推廣使用還是得有好的封裝,而這往往就讓很多人望而卻步了。還有一個問題,就是流模型系列為了保證可逆,自然是不能降維的,但不降維必然就導致計算量大,這是個矛盾的地方。

 

一個有趣的想法是,對於降維的情形,能不能搞個類似矩陣的“偽逆”那樣的模型,來達到類似可逆ResNet的效果呢?非方陣也可以搞個行列式出來的(比如《再談非方陣的行列式》[4]),因此降維變換也應該能夠搞個雅可比行列式的對數出來。貌似不少地方是可以推廣過來的。

 

流模型後面的方向究竟如何呢?讓我們拭目以待。

 

參考文獻

[1] Aidan N. Gomez, Mengye Ren, Raquel Urtasun, Roger B. Grosse. The Reversible Residual Network: Backpropagation Without Storing Activations. arXiv preprint arXiv:1707.04585, 2017. 7

[2] Jörn-Henrik Jacobsen, Arnold Smeulders, Edouard Oyallon. i-RevNet: Deep Invertible Networks. In Proceedings of the International Conference on Learning Representations (ICLR), 2018.

[3] https://kexue.fm/archives/6377

[4] https://kexue.fm/archives/6096

赞(0)

分享創造快樂