FrEIA：Framework for Easily Invertible Architectures

FrEIA 是實(shí)現(xiàn) INN 的基礎(chǔ)翰苫，可以理解為實(shí)現(xiàn) INN 的最重要的工具包一罩。其分為兩個(gè)模塊，每個(gè)模塊又有很多重要的類粱腻，以下是對(duì)這些類及其方法的描述。

一、FrEIA.framework

framework 模塊包含用于構(gòu)建網(wǎng)絡(luò)模型和推斷節(jié)點(diǎn)在向前和向后方向上執(zhí)行的順序的邏輯矛双。該模塊包括四個(gè)類敬辣。

<1> Node類
一個(gè) Node 即為一個(gè) INN 基礎(chǔ)構(gòu)建塊雪标。

build_modules(verbose=True)
該方法返回該結(jié)點(diǎn)的輸出結(jié)點(diǎn)的維度列表。會(huì)遞歸調(diào)用其輸入結(jié)點(diǎn)的 build_modules 方法溉跃。使用這些信息來初始化該結(jié)點(diǎn)的 pytorch.nn.Module村刨，即構(gòu)建起基礎(chǔ)塊。

run_forward(op_list)
調(diào)用該方法來確定執(zhí)行到該結(jié)點(diǎn)前的操作順序喊积。會(huì)遞歸調(diào)用其輸入結(jié)點(diǎn)的 run_forward 方法烹困。每個(gè)操作被追加到全局列表 op_list 中，操作的表示形式為 (node ID, input variable IDs, output variable IDs)乾吻。

run_backward(op_list)
該方法和 run_forward 相似髓梅，只是用來確定逆向執(zhí)行到該結(jié)點(diǎn)前的操作順序拟蜻。調(diào)用該方法前須先調(diào)用 run_forward。

<2> InputNode類
InputNode 類是 Node 類的子類枯饿，是一種特殊的結(jié)點(diǎn)——輸入結(jié)點(diǎn)酝锅，表示整個(gè) INN 網(wǎng)絡(luò)的輸入數(shù)據(jù)（及逆向過程的輸出數(shù)據(jù)）。

<3>OutputNode類
OutputNode 類也是 Node 類的子類奢方，是一種特殊的結(jié)點(diǎn)——輸出結(jié)點(diǎn)搔扁，表示整個(gè) INN 網(wǎng)絡(luò)的輸出數(shù)據(jù)（及逆向過程的輸入數(shù)據(jù)）。

<4>ReversibleGraphNet類
是 torch.nn.modules.module.Module 類的子類蟋字。這個(gè)類表示了 INN 本身稿蹲。這個(gè)類的構(gòu)造函數(shù)會(huì)確定 node list 中的輸入結(jié)點(diǎn)和輸出結(jié)點(diǎn)，并調(diào)用輸出結(jié)點(diǎn)的 build_modules 方法以獲得結(jié)點(diǎn)間的連接關(guān)系及維度鹊奖，調(diào)用輸出結(jié)點(diǎn)的 run_forward 方法獲得正向過程的所有操作及涉及的變量苛聘；再調(diào)用輸入結(jié)點(diǎn)的 run_backward 方法獲得逆向過程的所有操作及涉及的變量。

二忠聚、FrEIA.modules

modules 模塊中都是 torch.nn.Module 的子類设哗，都是可逆的，可用于 ReversibleGraphNet 中 Node 的構(gòu)建两蟀。

1. Coefficient functions

class FrEIA.modules.coeff_functs.F_conv(in_channels, channels, channels_hidden=None, stride=None, kernel_size=3, leaky_slope=0.1, batch_norm=False)
F_conv 類使用多層卷積網(wǎng)絡(luò)實(shí)現(xiàn) INN 基礎(chǔ)構(gòu)建塊中的 s网梢、t 轉(zhuǎn)換。其本身是不可逆的赂毯。

class FrEIA.modules.coeff_functs.F_fully_connected(size_in, size, internal_size=None, dropout=0.0)
F_fully_connected 類使用包含四層全連接層的神經(jīng)網(wǎng)絡(luò)實(shí)現(xiàn) INN 基礎(chǔ)構(gòu)建塊中的 s战虏、t 轉(zhuǎn)換。其本身是不可逆的欢瞪。

2. Coupling layers

F_* 類本身都是不可逆的活烙，而 coupling_layers 層中的類就是賦予 INN 可逆能力的。

class FrEIA.modules.coupling_layers.rev_layer(dims_in, F_class=<class 'FrEIA.modules.coeff_functs.F_conv'>, F_args={})
當(dāng)轉(zhuǎn)換是 F_conv 時(shí)使用 rev_layer 類實(shí)現(xiàn)可逆遣鼓。

class FrEIA.modules.coupling_layers.rev_multiplicative_layer(dims_in, F_class=<class 'FrEIA.modules.coeff_functs.F_fully_connected'>, F_args={}, clamp=5.0)
當(dāng)轉(zhuǎn)換是 F_fully_connected 時(shí)使用 rev_multiplicative_layer 實(shí)現(xiàn)可逆啸盏，思想是基于 RealNVP 基礎(chǔ)構(gòu)建塊的雙向轉(zhuǎn)換。

class FrEIA.modules.coupling_layers.glow_coupling_layer(dims_in, F_class=<class 'FrEIA.modules.coeff_functs.F_fully_connected'>, F_args={}, clamp=5.0)
和 rev_multiplicative_layer 相似骑祟，只是思想是基于 Glow 基礎(chǔ)構(gòu)建塊的雙向轉(zhuǎn)換回懦。

3. Fixed transforms

這個(gè)模塊的類也都是 torch.nn.Module 的子類，用于實(shí)現(xiàn)一些固定轉(zhuǎn)換次企。在構(gòu)建塊之間插入置換層怯晕，用于隨機(jī)打亂下一層的輸入元素，這使得 u = [u1缸棵，u2] 的分割在每一層都不同舟茶，也因此促進(jìn)了獨(dú)立變量間的交互。

class FrEIA.modules.fixed_transforms.linear_transform(dims_in, M, b)
根據(jù) y=Mx+b 做固定轉(zhuǎn)換，其中 M 是可逆矩陣吧凉。其實(shí)現(xiàn)為：

class FrEIA.modules.fixed_transforms.permute_layer(dims_in, seed)
隨機(jī)轉(zhuǎn)換輸入向量列順序隧出。其實(shí)現(xiàn)為：

4. Graph topology

用于對(duì) INN 網(wǎng)絡(luò)結(jié)構(gòu)進(jìn)行調(diào)整，如構(gòu)建殘差 INN 等阀捅。

class FrEIA.modules.graph_topology.cat_layer(dims_in, dim)
沿給定維度合并多個(gè)張量胀瞪。

class FrEIA.modules.graph_topology.channel_merge_layer(dims_in)
沿著通道從兩個(gè)獨(dú)立的輸入合并到一個(gè)輸出（用于跳過連接等）。

class FrEIA.modules.graph_topology.channel_split_layer(dims_in)
沿通道拆分以產(chǎn)生兩個(gè)單獨(dú)的輸出（用于跳過連接等）饲鄙。

class FrEIA.modules.graph_topology.split_layer(dims_in, split_size_or_sections, dim)
沿給定維度拆分以生成具有給定大小的單獨(dú)輸出的列表凄诞。

5.Reshapes

class FrEIA.modules.reshapes.flattening_layer(dims_in)
正向過程是將 N 維張量拉平為 1 維的張量。

class FrEIA.modules.reshapes.haar_multiplex_layer(dims_in, order_by_wavelet=False)
使用 Haar wavelets 將通道按寬度和高度對(duì)半分為 4 個(gè)通道忍级。

class FrEIA.modules.reshapes.haar_restore_layer(dims_in)
使用 Haar wavelets 將 4 個(gè)通道合并為一個(gè)通道帆谍。

class FrEIA.modules.reshapes.i_revnet_downsampling(dims_in)
i-RevNet 中使用的可逆空間 downsampling。

class FrEIA.modules.reshapes.i_revnet_upsampling(dims_in)
與 i_revnet_downsampling 相反的操作颤练。

loss function 深入解析

INN 的 loss 分為兩部分：前向訓(xùn)練的損失和后向訓(xùn)練的損失既忆。前后向訓(xùn)練的簡單過程如下：

對(duì)應(yīng)的損失函數(shù)如下：

其中驱负，L2 是一種求平方誤差的損失函數(shù)嗦玖；MMD 則用于求兩個(gè)分布之間的差異。

• 在 lossf 中跃脊，前半部分是用于求正向網(wǎng)絡(luò)預(yù)測值 ygen 與通過已知正向模型計(jì)算而得的 ytrue 之間的平方誤差宇挫，降低這部分損失是為了使 INN 正向模型更能擬合真實(shí)模型；后半部分損失則是計(jì)算真實(shí)隱藏信息的分布 ztrue 和服從標(biāo)準(zhǔn)正態(tài)分布的 z 之間的分布差異酪术，降低這部分損失是為了將在正向過程中丟失的信息 z' 映射到服從標(biāo)準(zhǔn)正態(tài)分布的 latent-z 空間器瘪，這樣我們通過對(duì)標(biāo)準(zhǔn)正態(tài)分布采樣，加上觀測值 y绘雁，即可得到一個(gè) x橡疼。另外，在反向傳播時(shí)庐舟，忽略 MMD 損失上的 y 的梯度欣除，以便訓(xùn)練神經(jīng)元學(xué)習(xí)從真實(shí)潛在分布到正態(tài)分布的映射而不妨礙正向模型 x→y 的訓(xùn)練，因此這種 MMD 損失的收斂確保了 z 與 y 的獨(dú)立性挪略。
• 在 lossb 中历帚，前半部分用于求逆向網(wǎng)絡(luò)根據(jù)真實(shí)潛在信息 ztrue 和 ytrue 得到的預(yù)測值 xgen1 與 xtrue 的平方誤差，降低這部分損失是為了讓 INN 逆向模型更準(zhǔn)確杠娱；后半部分損失是計(jì)算逆向網(wǎng)絡(luò)根據(jù)標(biāo)準(zhǔn)正態(tài)分布 z 和 ytrue 得到的預(yù)測值 xgen2 與 xtrue 的分布差異挽牢，降低這部分 MMD 損失是為了確保根據(jù)標(biāo)準(zhǔn)正態(tài)分布中的采樣點(diǎn)和 ytrue 得到的 x 的分布與其 xtrue 的真實(shí)分布看起來相似。而之所以給 MMD 損失增加了隨訓(xùn)練輪數(shù) n 逐漸增加的權(quán)重 ξ(n)摊求，是為了避免最初訓(xùn)練時(shí) MMD 的大梯度使網(wǎng)絡(luò)遠(yuǎn)離正確的解決方案禽拔。

重讀 Analyzing Inverse Problems with INN

這篇文章前后讀了很多次，這次重讀，又有許多新收獲睹栖。

INN 的基礎(chǔ)構(gòu)建塊是 RealNVP 中的仿射耦合層寥闪；如果要構(gòu)建深度可逆網(wǎng)絡(luò)，可將基礎(chǔ)構(gòu)建塊和殘差訓(xùn)練思想結(jié)合磨淌，i-RevNet 就是深度可逆網(wǎng)絡(luò)疲憋。基礎(chǔ)構(gòu)建塊本身是可逆結(jié)構(gòu)梁只，因此其中的 s缚柳、t 轉(zhuǎn)換不需要是可逆的，其可以是任意復(fù)雜的網(wǎng)絡(luò)結(jié)構(gòu)搪锣；更重要的是秋忙，構(gòu)建塊的可逆性允許我們同時(shí)為正向訓(xùn)練和反向訓(xùn)練應(yīng)用損失函數(shù)，并可以從任一方向計(jì)算 s 和 t 的梯度构舟。

至于如何將輸入分割為兩部分灰追，目前采用的方法是隨機(jī)（但為了可逆，必須是固定的）分割狗超。但為了實(shí)現(xiàn)更好的訓(xùn)練結(jié)果弹澎，往往在基礎(chǔ)構(gòu)建塊之間加入置換層，用于 shuffle 輸入努咐，這樣是為了使得每次分割得到的 u1苦蒿、u2 和上一構(gòu)建塊分割得到的不同。如果沒有置換層渗稍，那么在訓(xùn)練過程中佩迟，u1、u2 都是相互獨(dú)立的竿屹。

正向過程的信息丟失使得逆向過程存在模糊性报强，如果使用貝葉斯方法，用 p(x|y) 表示這種模糊性拱燃，也是可行的秉溉，但非常復(fù)雜。實(shí)際上我們希望網(wǎng)絡(luò)學(xué)習(xí)完整的后驗(yàn)分布 p(x|y)扼雏，我們可以嘗試預(yù)測簡單分布的擬合參數(shù)（如均值和方差）坚嗜，或者使用分布（而非固定值）來表示網(wǎng)絡(luò)權(quán)重，甚至兩者結(jié)合诗充，但無論如何苍蔬，這將限制我們預(yù)先選擇一種固定的分布或分布族。 也可以使用 cGAN蝴蜓，但它難以訓(xùn)練的碟绑，并且經(jīng)常遭遇難以檢測的模式崩潰俺猿。

因此，INN 中引入 latent variable z格仲，其捕獲了從 x 到 y 的正向過程中丟失的信息押袍，即 x 與 [y,z] 形成了雙射。在正向訓(xùn)練過程中凯肋，我們可以得到真實(shí)潛在分布 ztrue谊惭，從 ztrue 中采樣，再結(jié)合 y 值侮东，即可得到一個(gè)固定的 x 值圈盔。即，p(x|y) 轉(zhuǎn)換成一個(gè)固定函數(shù) x=g(y,z)悄雅。不過驱敲，為了易于采樣，我們需要將 ztrue 調(diào)整為一個(gè)簡單的分布（如標(biāo)準(zhǔn)正態(tài)分布）宽闲，這樣我們就可以從簡單分布中采樣众眨，結(jié)合 y 值得到 x，這樣的準(zhǔn)確率依然是很高的容诬，因?yàn)楹唵畏植己?ztrue 之間只是一個(gè)簡單映射關(guān)系娩梨。至此看來，INN 和 cGAN 思想是很接近的放案，不過 INN 的可逆性使得其訓(xùn)練過程與 cGAN 不同姚建，也有其特殊的優(yōu)點(diǎn)。

INN 同時(shí)對(duì)正向和逆向過程進(jìn)行訓(xùn)練吱殉，累積網(wǎng)絡(luò)兩端的損失項(xiàng)的梯度。 我們還可以在 x 上添加類似 GAN 的鑒別器損失厘托，但到目前為止友雳，在我們的應(yīng)用程序中，MMD 是足夠的铅匹，因此我們可以避免對(duì)抗性訓(xùn)練的麻煩押赊。在大量訓(xùn)練數(shù)據(jù)的支持下，我們可以訓(xùn)練網(wǎng)絡(luò)擬合正向模型包斑，由于網(wǎng)絡(luò)的可逆構(gòu)造流礁，我們可以免費(fèi)得到逆向模型。