細水長flow之TARFLOW：流模型滿血歸來？

©PaperWeekly 原創 · 作者 | 蘇劍林

單位 | 科學空間

研究方向 | NLP、神經網路

不知道還有沒有讀者對這個系列有印象？這個系列取名“細水長 flow”，主要介紹 flow 模型的相關工作，起因是當年（2018 年）OpenAI 釋出了一個新的流模型 Glow，在以 GAN 為主流的當時來說著實讓人驚豔了一番。

但驚豔歸驚豔，事實上在相當長的時間內，Glow 及後期的一些改進在生成效果方面都是比不上 GAN 的，更不用說現在主流的擴散模型了。

不過局面可能要改變了，前段時間的論文《Normalizing Flows are Capable Generative Models》[1] 提出了新的流模型 TARFLOW，它在幾乎在所有的生成任務效果上都逼近了當前 SOTA，可謂是流模型的“滿血”迴歸。

寫在前面

這裡的流模型，特指 Normalizing Flow，是指模型架構具有可逆特點、以最大似然為訓練目標、能實現一步生成的相關工作，當前擴散模型的分支 Flow Matching 不歸入此列。

自從 Glow 閃耀登場之後，流模型的後續進展可謂“乏善可陳”，簡單來說就是讓它生成沒有明顯瑕疵的 CelebA 人臉都難，更不用說更復雜的 ImageNet 了，所以“細水長 flow” 系列也止步於 2019 年的《細水長 flow 之可逆 ResNet：極致的暴力美學》。

不過，TARFLOW 的出現，證明了流模型“尚能一戰”，這一次它的生成畫風是這樣的：

▲ TARFLOW 的生成效果

相比之下，此前 Glow 的生成畫風是這樣的：

▲ Glow 的生成效果

Glow 演示的還只是相對簡單的人臉生成，但瑕疵已經很明顯了，更不用說更復雜的自然影像生成了，由此可見 TARFLOW 的進步並不只是一星半點。從資料上看，它的表現也逼近模型模型的最佳表現，超過了 GAN 的 SOTA 代表 BigGAN：

▲ TARFLOW 與其他模型的定量對比

要知道，流模型天然就是一步生成模型，並且不像 GAN 那樣對抗訓練，它也是單個損失函式訓練到底，某種程度上它的訓練比擴散模型還簡單。所以，TARFLOW 把流模型的效果提升了上來，意味著它同時具備了 GAN 和擴散模型的優點，同時還有自己獨特的優勢（可逆、可以用來估計對數似然等）。

模型回顧

言歸正傳，我們來看看 TARFLOW 用了什麼“靈丹妙藥”來讓流模型重新煥發活力。不過在此之前，我們簡要回顧一下流模型的理論基礎，更詳細的歷史溯源，可以參考《細水長 flow之NICE：流模型的基本概念與實現》和《細水長 flow 之 RealNVP 與 Glow：流模型的傳承與昇華》。

從最終目標來看，流模型和 GAN 都是希望得到一個確定性函式，將隨機噪聲對映到目標分佈的圖片，用機率分佈的語言說，就是用如下形式的分佈來建模目標分佈：

其中是狄拉克 δ 函式 [2]。訓練機率模型的理想目標是最大似然，即以為損失函式。但目前的帶有積分，只有形式意義，無法用於訓練。

此時流模型和 GAN 就“分道揚鑣”了：GAN 大致上是用另外一個模型（判別器）去近似，這導致了交替訓練；流模型則透過設計適當的，讓積分（1）可以直接算出來。把積分（1）算出來需要什麼條件呢？設，其逆函式為，那麼

以及

因此

這表明，將積分（1）算出來需要兩個條件：一、需要知道的逆函式；二、需要計算雅可比矩陣的行列式。

仿射耦合

為此，流模型提出了一個關鍵設計——“仿射耦合層”：

其中，、是以為輸入、輸出形狀跟一致的模型，是 Hadamard 積。上式說的是，將分成兩部份，分法隨意，也不要求等分，將其中一份原封不動輸出，另一邊按照指定規則運算輸出。注意仿射耦合層是可逆的，其逆為

這就滿足了第一個條件可逆性。另一方面，仿射耦合層的雅可比矩陣是一個下三角陣：

三角矩陣的行列式，等於對角線元素之積，所以

即雅可比矩陣行列式的絕對值對數等於的各分量之和，這就滿足了第二個條件雅可比矩陣行列式的可計算性。

仿射耦合層首次提出自 RealNVP [3]，NVP 的含義是 “Non-Volume Preserving”，也就是“不保體積”，這個名字對標的是恆等於零時的特殊情形，此時稱為“加性耦合層”，提出自 NICE [4]，特點是其雅可比行列式等於 1，即加性耦合層是“保體積”的（行列式的幾何意義是體積）。

注意，如果直接堆疊多個仿射耦合層，那麼將一直保持不變，這並不是我們想要的，我們想要做的是將整個對映到標準正態分佈。為了解決這個問題，我們每次應用仿射耦合層前，都要輸入的分量以某種方式“打亂”，這樣就不至於出現始終不變的分量。“打亂”運算對應於置換矩陣變換，行列式絕對值始終為 1。

核心改進

到目前為止，這些內容都還只是流模型的基礎內容，接下來才正式進入到 TARFLOW 的貢獻點。

首先，TARFLOW 留意到仿射耦合層（5）可以推廣到多塊劃分，即將分成更多份，然後按照類似的規則運算

其中，，其逆運算為

類似地，推廣版本的仿射耦合層的雅可比行列式絕對值對數為的所有分量之和，所以流模型要求的兩個條件都能滿足。這個推廣的雛形最早在 2016 年的 IAF 就提出了，比 Glow 還早。

可為什麼後來的工作鮮往這個方向深入呢？這大體是歷史原因。早年 CV 模型的主要架構是 CNN，用 CNN 的前提是特徵滿足區域性相關性，這就導致了將分塊時，往往只考慮在通道（channel）維度劃分。

因為每一層還有一個必要的打亂運算，一旦選擇在長、寬兩個空間維度劃分，那麼隨機打亂後特徵就失去了區域性相關性了，從而沒法用 CNN。而如果在通道維度劃分多份，那麼多個通道特徵圖比較難高效地互動。

然而，到了 Transformer 時代，情況截然不同了。Transformer 的輸入本質上是一個無序的向量集合，換言之不依賴區域性相關性，因此以 Transformer 為主架構，我們就可以選擇在空間維度劃分，這就是 Patchify。

此外，式（9）中形式，意味著這是一個 Causal 模型，這也正好可以用 Transformer 高效實現。

除了形式上的契合外，在空間維度劃分有什麼本質好處呢？

這就要回到流模型的目標了：將噪聲變成圖片，逆模型則將圖片變成噪聲，噪聲的特點是隨機，說白了就是亂，圖片的顯著特點是區域性相關性，所以圖片變噪聲的關鍵之一就是要打亂這種區域性相關性，直接在空間維度 Patchify 配合耦合層自帶的打亂操作，無疑是最高效的選擇。

所以，式（9）跟 Transformer 可謂是“一拍即合”、“相得益彰”，這就是 TARFLOW 前三個字母 TAR 的含義（Transformer AutoRegressive Flow），也是它的核心改進。

加噪去噪

流模型常用的一個訓練技巧是加噪，也就是往圖片加入微量噪聲後再送入模型進行訓練。雖然我們將圖片視為連續向量，但它實際是以離散的格式儲存的，加噪可以進一步平滑這種不連續性，使得圖片更接近連續向量。噪聲的加入還可以防止模型過度依賴訓練資料中的特定細節，從而減少過擬合的風險。

加噪是流模型的基本操作，並不是 TARFLOW 首先提出的，TARFLOW 提出的是去噪。理論上來說，圖片加噪後訓練出來的流模型，生成結果也是帶有噪聲的，只不過以往的流模型生成效果也沒多好，所以這點噪聲也無所謂了。但 TARFLOW 把流模型的能力提上去後，去噪就“勢在必行”了，不然噪聲就稱為影響效果的主要因素了。

怎麼去噪呢？另外訓練一個去噪模型？沒這個必要。我們在《從去噪自編碼器到生成模型》已經證明了，如果是加噪訓練後的機率密度函式，那麼

就是去噪模型的理論最優解。所以有了後，我們就不用另外訓練去噪模型了，直接按照上式計算就可以去噪，這也是流模型的優勢之一。而正因為加入了去噪步驟，所以 TARFLOW 的輸入的噪聲改為高斯分佈，並適當增大了噪聲的方差，這也是它效能更好的原因之一。

綜上所述，TARFLOW 完整的取樣流程是

延伸思考

至此，TARFLOW 相比以往流模型的一些關鍵變化已經介紹完畢，剩下的一些模型細節，大家自行讀原論文就好，如果還有不懂的也可以參考官方開源的程式碼。

Github:

https://github.com/apple/ml-tarflow

下面主要談談筆者對 TARFLOW 的一些思考。

首先，需要指出的是，TARFLOW 雖然效果上達到了 SOTA，但它取樣速度實際上不如我們期望，原論文附錄提到，在 A100 上取樣 32 張 ImageNet64 圖片大概需要 2 分鐘。

為什麼會這麼慢呢？我們仔細觀察耦合層的逆（10）就會發現，它實際上是一個非線性 RNN！非線性 RNN 只能序列計算，這就是它慢的根本原因。

換句話說，TARFLOW 實際上是一個訓練快、取樣慢的模型，當然如果我們願意也可以改為訓練慢、取樣快，總之正向和逆向不可避免會有一側慢，這是分多塊的仿射耦合層的缺點，也是 TARFLOW 想要進一步推廣的主要改進方向。

其次，TARFLOW 中的 AR 一詞，容易讓人聯想到現在主流的自迴歸式 LLM，那麼它們倆是否可以整合在一起做多模態生成？

說實話很難。因為 TARFLOW 的 AR 純粹是仿射耦合層的要求，而耦合層之前還要打亂，所以它並非一個真正的 Causal 模型，反而是徹頭徹尾的 Bi-Directional 模型，所以它並不好跟文字的 AR 強行整合在一起。

總的來說，如果 TARFLOW 進一步將取樣速度也提上去，那麼它將會是一個非常有競爭力的純視覺生成模型。

因為除了訓練簡單和效果優異，流模型的可逆性還有另外一個優點，就是《The Reversible Residual Network: Backpropagation Without Storing Activations》[5] 所提的反向傳播可以完全不用存啟用值，並且重計算的成本比普通模型低得多。

至於它有沒有可能成為多模態 LLM 的統一架構，只能說現在還不大明朗。

文藝復興

最後，談談深度學習模型的“文藝復興”。

近年來，已經有不少工作嘗試結合當前最新的認知，來反思和改進一些看起來已經落後的模型，並得到了一些新的結果。除了 TARFLOW 試圖讓流模型重新煥發活力外，最近還有《The GAN is dead; long live the GAN! A Modern GAN Baseline》[6] 對 GAN 的各種排列組合去蕪存菁，得到了同樣有競爭力的結果。

更早一些，還有《Improved Residual Networks for Image and Video Recognition》[7]、《Revisiting ResNets: Improved Training and Scaling Strategies》[8] 等工作讓 ResNet 更上一層樓，甚至還有《RepVGG: Making VGG-style ConvNets Great Again》[9] 讓 VGG 經典再現。

當然，SSM、線性 Attention 等工作也不能不提，它們代表著 RNN 的“文藝復興”。

期待這種百花齊放的“復興”潮能更熱烈一些，它能讓我們獲得對模型更全面和準確的認知。