標準化流(Normalizing Flow)(二):現(xiàn)代標準化流技術

翻譯自https://blog.evjang.com/2018/01/nf2.html
原作者:Eric Jang
譯者:尹肖貽

0. 交代故事

我在下面的教程里教你干一件很酷的事兒:


你將學到使用MAF、IAF和Real-NVP等標準化流來變形二維高斯分布,而形成復雜的圖形“SIGGRAPH”。就像……拉扯橡皮泥

在教程的上半部分,你學到使用標準化流,把高斯分布這類地攤貨分布,“變形”為高大上的概率分布。你還親自使用PReLU搭建了小型可逆神經(jīng)網(wǎng)絡,實現(xiàn)了一個簡單的鏈式二維仿射雙射函數(shù)。
不過,上次的全連接網(wǎng)絡只用了兩層隱含層,搭建的流弱爆了。更糟的是,非線性激活函數(shù)是單調(diào)的、分段線性的。這樣一來,網(wǎng)絡只能在原點附近,稍微地扭捏數(shù)據(jù)分布的形狀【譯者按:1)原文用的是manifold流形,對于非從業(yè)者來說,只需知道數(shù)據(jù)在空間中形成一定形狀就可以了2)因為激活函數(shù)在原點附近是非線性的】。這個流想要實現(xiàn)更高級的變換,就完全不給力。比如把各向同性的高斯變成下面雙模式的“雙月”數(shù)據(jù)。


幸好最近的研究發(fā)明了幾個更強大的標準化流。下面我們就去探索其中的幾個技術。

1. 標準化流之 自回歸(autoregressive)模型和MAF

自回歸概率密度估計技術,比如 WaveNetPixelRNN,是用來學習聯(lián)合概率密度p(x_{1:D})的。該技術將聯(lián)合概率密度,分解為以為條件概率密度的乘積,這里的x_i依賴前i-1個數(shù)據(jù)的取值:
p(x) = \prod_i{p(x_i \,|\, x_{1:i-1})}

1.1 學習階段

上式中的條件概率密度大都有可學習的參數(shù)。例如,常見的選擇是單變量高斯分布,自回歸的概率密度p(x_{1:D})就有兩個參數(shù),均值和標準差。這兩個參數(shù)取決于先前的變量x_{1:i-1}。
p(x_i \,|\, x_{1:i-1}) = \mathcal{N}(x_i \,|\,\mu_i, (\exp\alpha_i)^2)
\mu_i = f_{\mu_i}(x_{1:i-1})
\alpha_i = f_{\alpha_i}(x_{1:i-1})
這種方法基于一個簡單粗暴的假設:每個變量都依賴于先于其出現(xiàn)的變量,而不是后來的變量。拍腦門就能想明白這不(全)是真的:頂上的像素和底下的像素能有啥因果關系?不過(令許多研究者驚訝的是),這東西產(chǎn)生出的圖像效果居然還不錯!

1.2 采樣階段

采樣階段,我們從標準高斯N(0,1)采D個"噪聲變量"u_{1:D},進而依次地生成所有樣本:
x_i = u_i\exp{\alpha_i} + \mu_i
u_i \sim \mathcal{N}(0, 1)
自回歸采樣技術可以這么理解,把(從標準高斯分布\mathcal{N}(0, \mathbb{I}))采樣得到的噪聲變量轉換為一個新的分布形式,或者說新的分布形式是標準高斯分布的轉換分布(TransformedDistribution)。

1.3 搭建一個流

有了這個認識,我們可以堆疊多個自回歸模型,搭建一個標準化流。這樣做的好處是,你可以更改流中任意的一個雙射函數(shù)的輸入順序x_1,...x_D,這樣就可以在某一環(huán)節(jié)(因為不合時宜的排序)效果不佳,而在另一個環(huán)節(jié)彌補。
掩模自回歸流(MAF)實現(xiàn)了一個條件高斯自回歸模型。這里對于任何一個分布x_i:


灰色的單元是當前待計算的單元,藍色的單元表示其依賴的單元。和通過計算傳遞給網(wǎng)絡(分別是涂上了洋紅色、橘色的圓圈)。即使變換只有尺度變換和平移變換,這些變換也錯綜復雜地依賴先前的變量。對于第一個單元,其和不依賴于x和u,而設置為可學習的變量。
注意,之所以如此這般地設計變換函數(shù),是因為計算的時候,不必要再計算和的逆。由于變換僅有尺度變換和平移變換,我們只需要逆運算這兩個變換:。雙射函數(shù)的前傳和反傳,僅僅取決于和,這樣我們就可以在和的網(wǎng)絡內(nèi)部使用非線性激活函數(shù)(如ReLU)和非方形矩陣操作。
MAF模型的梯度反傳可以這樣來評估密度函數(shù):
distribution.log_prob(bijector.inverse(x)) + bijector.inverse_log_det_jacobian(x))

2. 時間復雜度和MADE

自回歸模型和MAF訓練較快,因為計算下面D個似然概率
p(x_1), p(x_2\,|\, x_1), ... p(x_D\,|\, x_{1:D-1})
可以利用GPU的并行技術,一次性地用D個線程計算。這里假設計算任務,如在計算SIMD向量化用CPU或GPU并行計算時那樣,資源沒有上限。
從另一方面看,自回歸采樣較慢,因為你必須等到x_{1:i-1}都算出來,才能計算x_i。如此只能利用一條線程采樣D個序列,而無法發(fā)揮并行計算的優(yōu)勢。
另一個問題是,在并行的反傳計算時,我們要不要使用兩個不同的(具有不同的輸入長度的)網(wǎng)絡,來計算\alpha_i\mu_i?這樣做不夠高效,尤其考慮到D個網(wǎng)絡學到的表示(learned representation)需要共享(只要不違背自回歸的依賴關系)。在“密度估計的掩模自編碼器”(MADE)一文中,作者給這個問題提供了一個不錯的解決方案:使用一個網(wǎng)絡,把所有的\alpha\mu同時計算出來,而用來掩模的權重保證了自回歸的有效性。
這一技巧也大大簡化了反傳的計算,從所有的x,只需一步即可計算出所有的u(D個輸入,D個輸出)。這就比處理D個不同的網(wǎng)絡容易的多((D+1)*D/2個輸入,D個輸出)

3.Inverse Autoregressive Flow (IAF)

可逆自回歸流(IAF)的非線性平移和尺度變換操作,其輸入是前一時刻的噪聲變量,而不是輸入前一時刻的數(shù)據(jù)采樣。
x_i = u_i\exp{\alpha_i} + \mu_i
\mu_i = f_{\mu_i}(u_{1:i-1})
\alpha_i = f_{\alpha_i}(u_{1:i-1})


前傳(采樣)速度很快:對于所有的,并行D個線程可以一次算出來。IAF也可以使用MADE框架高效地實現(xiàn)并行。
不過,要是你想估計某一新數(shù)據(jù)附近的概率密度,就需要恢復所有的u,過程很長:首先計算,然后逐次計算。好在計算這一趟以后,IAF生成的數(shù)據(jù)的(log)概率直接就得到了。因為我們已知u的情況下,不需要再通過反向計算x得到u。
細心的你可能已經(jīng)察覺到了,要是把底下標注,頂層標注,IAF和MAF的反傳計算沒有差別。對應而言,IAF的反傳計算恰好也就是MAF(把x和u交換一下)。所以在tensorflow實現(xiàn)的時候,雙射函數(shù)的類沒有區(qū)別,而且逆轉特征來交換正反傳的操作很容易:
iaf_bijector = tfb.Invert(maf_bijector)
IAF和MAF有互補的計算考慮——MAF訓練快采樣慢,IFA訓練慢采樣快。對于訓練網(wǎng)絡,需要做的密度估計操作更多,采樣操作較少,所以通常在學習密度的任務中選擇MAF更為合適。

4.平行聲波網(wǎng)絡(Parallel Wavenet)

顯而易見的問題是,IAF和MAF能否合二為一,達到最優(yōu)性能?比如快速的訓練和采樣。
答案是:當然能!DeepMind實驗室公布的平行聲波網(wǎng)絡正是這么做的:MAF快速訓練一個產(chǎn)生式網(wǎng)絡,IAF在這個網(wǎng)絡的“啟發(fā)”下最大化采樣的似然概率。回顧一下IAF,計算外來數(shù)據(jù)的概率密度很慢(比如訓練集中的數(shù)據(jù)),但是計算采樣數(shù)據(jù)u_{1:D}的密度很容易,這樣就省了反傳計算。只要約束“學生”IAF和“老師”MAF的概率分布差異,就可以完成訓練。


這項研究在標準化流的領域里非常重要——最終實現(xiàn)的效果是,實時的音頻合成的速度比采樣的方法快20倍,而且已經(jīng)在谷歌助手之類的產(chǎn)品上線了。

5.NICE and Real-NVP

最后請看Real-NVP,你可以認為是IAF雙射函數(shù)的特殊情況。
訓練時,在NVP的“配對層”(coupling layer),依次地調(diào)整標號d=1,2,...,D-1。像IAF一樣,x_{d+1}的尺度變換和平移變換取決于u_d的值。所不同的是,x_{d+2}, x_{d+3}, … x_{D}都只依賴于u_u0z1t8os, 所以單次的網(wǎng)絡前傳就可以得到\alpha_{d+1:D}\mu_{d+1:D}?!咀g者按:這里可能的意思是,算出來x以后,求個逆就得到\alpha\mu
對于x_{1:d},這些單元“放行”了所有信號,并被設置為u_{1:d}。所以可以認為Real-NVP是一種特殊的MAF雙射(當\alpha(u_{1:d}) = \alpha(x_{1:d}))。


因為尺度平移變換后,整個層的統(tǒng)計量可以從或者一次前傳得到,NVP可以在一次前傳和后傳完成所有的計算(也就是說,采樣和估計都很快)。MADE架構就不再需要了。
從實驗效果來看,Real-NVP比MAF或IAF都要差,使用同樣多的層數(shù)NVP在我的例子中(比如SIGGRAPH形)也更差。Real-NVP和IAF在二維圖形的情況下幾乎是等價的。唯一的區(qū)別是在第一個單元,IAF通過尺度平移變換得到,不依賴于,而Real-NVP對第一個單元不作處理。
Real-NVP是NICE雙射的后續(xù)工作。NICE只有平移變換并假定。因為NICE不做尺度變換,ILDJ始終是個常數(shù)!

6.批正則化的雙射函數(shù)(batch normalization bijector)

Real-NVP 的論文中提到許多新奇的分布,其中之一是批正則化的雙射函數(shù),用來穩(wěn)定訓練過程。在傳統(tǒng)的觀念中,批正則化應用在訓練神經(jīng)網(wǎng)絡的場景中,前傳數(shù)據(jù)服從統(tǒng)計意義上的中心集中、方差是對角陣的高斯分布,批數(shù)據(jù)是正態(tài)分布的統(tǒng)計特性(running mean,running variance)通過指數(shù)移動平均值積累。測試時,積累的統(tǒng)計數(shù)據(jù)用來做標準化處理。
在標準化流中,訓練時批正則化用在雙射函數(shù)的逆的計算中bijector.inverse,在測試時積累的統(tǒng)計數(shù)據(jù)要去標準化(bijector.forward)。具體而言,批正則化雙射函數(shù)這樣實現(xiàn):
反傳
1.計算現(xiàn)有變量x的均值和方差
2.更新當前的均值和方差
3.用當前的均值方差批正則化當前數(shù)據(jù)。
前傳
1.用當前的均值方差去標準化,得到數(shù)據(jù)的分布。
歸功于TF雙射函數(shù),這一過程可用這些代碼來實現(xiàn):

class BatchNorm(tfb.Bijector):
    def __init__(self, eps=1e-5, decay=0.95, validate_args=False, name="batch_norm"):
        super(BatchNorm, self).__init__(
            event_ndims=1, validate_args=validate_args, name=name)
        self._vars_created = False
        self.eps = eps
        self.decay = decay

    def _create_vars(self, x):
        n = x.get_shape().as_list()[1]
        with tf.variable_scope(self.name):
            self.beta = tf.get_variable('beta', [1, n], dtype=DTYPE)
            self.gamma = tf.get_variable('gamma', [1, n], dtype=DTYPE)
            self.train_m = tf.get_variable(
                'mean', [1, n], dtype=DTYPE, trainable=False)
            self.train_v = tf.get_variable(
                'var', [1, n], dtype=DTYPE, initializer=tf.ones_initializer, trainable=False)
        self._vars_created = True

    def _forward(self, u):
        if not self._vars_created:
            self._create_vars(u)
        return (u - self.beta) * tf.exp(-self.gamma) * tf.sqrt(self.train_v + self.eps) + self.train_m

    def _inverse(self, x):
        # Eq 22. Called during training of a normalizing flow.
        if not self._vars_created:
            self._create_vars(x)
        # statistics of current minibatch
        m, v = tf.nn.moments(x, axes=[0], keep_dims=True)
        # update train statistics via exponential moving average
        update_train_m = tf.assign_sub(
            self.train_m, self.decay * (self.train_m - m))
        update_train_v = tf.assign_sub(
            self.train_v, self.decay * (self.train_v - v))
        # normalize using current minibatch statistics, followed by BN scale and shift
        with tf.control_dependencies([update_train_m, update_train_v]):
            return (x - m) * 1. / tf.sqrt(v + self.eps) * tf.exp(self.gamma) + self.beta

    def _inverse_log_det_jacobian(self, x):
        # at training time, the log_det_jacobian is computed from statistics of the
        # current minibatch.
        if not self._vars_created:
            self._create_vars(x)
        _, v = tf.nn.moments(x, axes=[0], keep_dims=True)
        abs_log_det_J_inv = tf.reduce_sum(
            self.gamma - .5 * tf.log(v + self.eps))
        return abs_log_det_J_inv

ILDJ的數(shù)學形式可方便地從逆函數(shù)的對數(shù)形式導出(參考單變量的例子)。

7.代碼在此

歸功于JoshDillon和谷歌貝葉斯流研究組的努力,在MADE架構下的Masked Autoregressive Flow 來實現(xiàn)u的快速訓練的代碼已經(jīng)寫好了。
我創(chuàng)建了一個復雜的二維分布,是利用這個混合腳本點云“SIGGRAPH”的形狀。我們搭建了數(shù)據(jù)庫、雙射函數(shù)、轉換分布,方法和教程一中非常類似,所以就不重復了,你可以參看這個Jupyter notebook 文件。這個代碼可以通過MAF、IAF、Real-NVP實現(xiàn),使用或不使用批正則化,重構雙月形和“SIGGRAPH”形。
一個細節(jié)很容易被忽略,如果你忘記重置變量順序,這個代碼就不能正常工作。這是因為不重置順序的話,所有的自回歸分解就只能學到p(x_1 | x_2)。幸好,Tensorflow有一個換序雙射函數(shù),可以專門實現(xiàn)這一功能。
這里可視化了流的學習過程,從一而終。看到它,我想起一臺太妃糖攪拌機


8.教程總結

Tensorflow讓標準化流的實現(xiàn)變得簡單,自動收集所有雅各比矩陣行列式的操作讓代碼變得簡單易讀。如果你要選擇一個實現(xiàn)標準化流的方案,要同時考慮前傳和后傳的快捷性,以及平衡流的表達能力和ILJD的計算速度。
在該教程的第一部分,我介紹了標準化流的動機:我們需要更強大的分布函數(shù),用于強化學習和產(chǎn)生式模型。從大背景來看,在變分推斷和隱式密度估計模型成功應用的今天,能夠追蹤變換的標準化流模型是否是人工智能應用(如機器人、架構估計)中最適合的模型,尚不明確。即便如此,標準化流仍是你工具箱中非常好用的工具,它們已經(jīng)在現(xiàn)實應用中發(fā)揮作用,比如谷歌助手團隊開發(fā)的實時的產(chǎn)生音頻的模型。
即便像標準化流這樣的顯式的密度估計模型能方便地通過最大似然法訓練,這些模型的作用遠不止于此。它們和VAE、GAN這類模型是互補的。任何模型中的高斯分布直接就換上一個標準化流,比如VAE的先驗分布和GAN的隱變量。舉個例子,這篇論文用標準化流靈活選擇變分先驗,Tensorflow分布的論文提到VAE中標準化流做PixelCNN的解碼器。平行聲波網(wǎng)絡通過最小化KL散度訓練一個IAF“學生”模型。
標準化流最具啟發(fā)的性質(zhì)之一是每一步計算都是可逆的(即具有函數(shù)的定義清楚的逆)。這就意味著,如果我們想做一次反傳計算時,不必在前傳時預先存儲激活值,而是重新計算一遍前傳激活值。(存儲這些值有可能代價很大。)在任務分配的過程很長的情況下,我們可以使用計算的可逆性來“恢復”過去的每一個選擇狀態(tài),從而限制了計算時的內(nèi)存消耗。這個主意在論文 RevNets中有所體現(xiàn),正是受NICE雙射函數(shù)運算可逆的啟發(fā)。我回憶起電影Memento中的情節(jié),主角不能存貯記憶了,于是通過可逆計算來記住事情。
謝謝閱讀!


Code on Github

9.致謝

非常感謝 Dustin Tran, Luke Metz, Jonathan Shen, Katherine Lee, Samy Bengio 等人預先檢查閱讀了該教程。

10.引用文獻和推薦閱讀

11.譯者補充與提問解答

我很喜歡這個教程,清晰易懂,所以把它全文翻譯成中文??紤]到中英文表達習慣的差異,譯文并不完全是原文的原句(原詞)對應,但意思(大概)還原了原文的意思。
這里補充幾個問題與其解答,也歡迎讀者向我(或原作者)提問!

問題1:這篇論文Parallel WaveNet:Fast High-Fidelity Speech Synthesis涉及到了并行化,MAF作為teacher模型,IVF作為student模型,我始終不是很理解這個方法的原理。

回答1:
這個問題拆解為兩部分。第一,為什么要讓MAF做教師,IAF做學生;第二,并行化的訓練是怎樣完成的。
============================================
第一,原文提到了MAF訓練快采樣慢,IAF訓練慢采樣快。所謂訓練就是從數(shù)據(jù)算出模型參數(shù)的過程,采樣就是從模型參數(shù)產(chǎn)生新數(shù)據(jù)的過程。
+++++++++++++++++++++++
對于MAF,訓練的時候,把所有的數(shù)據(jù)送進兩個網(wǎng)絡,f_\alphaf_\mu。在允許并行的情況下,一步就可以算出來所有的\alpha_{2,...,n}\mu_{2,...,n},結合采樣得到的所有的u_{1...n},就可以算出\hat{x}_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}。(注意這里\alpha_1\mu_1是超參,不是從x_1計算出來的)然后就可以去訓練f_\alphaf_\mu的參數(shù),這個過程非???。
采樣的時候,我們手里拿著所有f_\alphaf_\mu的參數(shù),并\alpha_1\mu_1,從x_1=u_1exp(\alpha_1)+\mu_1,\alpha_2=f_\alpha(x_1),\mu_2=f_\mu(x_1),開始算,x_2=u_2exp(\alpha_2)+\mu_2,\alpha_3=f_\alpha(x_1,x_2),\mu_3=f_\mu(x_1,x_2)...這樣一直算到x_n算完,這個過程是串行的而不能做到并行。
注意,采樣階段,u_i = (x_i-f_\mu(x_i))/\exp(f_\alpha(x_i))這個式子里,u_i其實就是p(x_i)。
+++++++++++++++++++++++
對于IAF,訓練的時候,從超參\alpha_1\mu_1開始計算得到u_1 = (x-\mu_1) * \exp(-\alpha_1),把它們送進兩個網(wǎng)絡,f_\alphaf_\mu\mu_2 = f_{\mu}(u_{1}),\alpha_2 = f_{\alpha}(u_{1})。繼續(xù)計算u_2 = (x-\mu_2) * \exp(-\alpha_2)。其余的以此類推u_i = (x-\mu(u_{1:i-1})) * \exp(-\alpha(u_{1:i-1}))。用所有的u_{1...n}\alpha_{1...n}\mu_{1...n},就可以算出\hat{x}_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}。然后就可以去訓練f_\alphaf_\mu的參數(shù)。這個過程不能并行。
采樣的時候,我們手里拿著所有f_\alphaf_\mu的參數(shù),采樣基礎分部產(chǎn)生u_{1...n},然后一步就可以算出所有的\alpha_{1,...,n}\mu_{1,...,n},于是x_{1,...,n}=u_{1...n}exp(\alpha_{1,...,n})+\mu_{1,...,n}。這個過程是非???。
+++++++++++++++++++++++
============================================
Parallel WaveNet訓練過程是這樣的:
1.訓練一個MAF,得到合理的f_{\alpha MAF}、f_{\mu MAF}。詳細步驟見上文。
2.采樣基礎網(wǎng)絡,得到的隨機噪聲u_i,送進IAF里得到x^i=u_iexp\alpha_i+\mu_i(注意這個IAF沒有訓練好),和log p_s(x^i) = log(x^i - f_{\mu IAF}(x^i)) -f_{\alpha IAF}(x^i)。
3.給定測試數(shù)據(jù)x^i,利用MAF的反向傳遞,計算(log)似然概率log p_t(x^i) = log(x^i -f_{\mu MAF}(x^i)) - f_{\alpha MAF}(x^i)。
4.計算KL距離,作為損失。
============================================
要徹底理解這個框架,要對最大似然法和EM算法有個大概了解。我貼一張圖,從知乎拷貝的。侵權就刪:

EM算法

最后編輯于
?著作權歸作者所有,轉載或內(nèi)容合作請聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成,瀏覽時請結合常識與多方信息審慎甄別。
平臺聲明:文章內(nèi)容(如有圖片或視頻亦包括在內(nèi))由作者上傳并發(fā)布,文章內(nèi)容僅代表作者本人觀點,簡書系信息發(fā)布平臺,僅提供信息存儲服務。

友情鏈接更多精彩內(nèi)容