本文作者楊磊，目前在大模型初創(chuàng)公司階躍星辰擔(dān)任后訓(xùn)練算法工程師，其研究領(lǐng)域包括生成模型和語言模型后訓(xùn)練。在這之前，他曾在曠視科技擔(dān)任了六年的計算機視覺算法工程師，從事三維視覺、數(shù)據(jù)合成等方向。他于 2018 年本科畢業(yè)于北京化工大學(xué)。

當(dāng)前，主流的基礎(chǔ)生成模型大概有五大類，分別是 ：Energy-Based Models (Diffusion)、GAN、Autoregressive、VAE 和 Flow-Based Models。

本項工作提出了一種全新的生成模型：離散分布網(wǎng)絡(luò)（Discrete Distribution Networks），簡稱 DDN。相關(guān)論文已發(fā)表于 ICLR 2025。

DDN 采用一種簡潔且獨特的機制來建模目標(biāo)分布：

1.在單次前向傳播中，DDN 會同時生成 K 個輸出（而非單一輸出）。

2.這些輸出共同構(gòu)成一個包含 K 個等權(quán)重（概率均為 1/K）樣本點的離散分布，這也是「離散分布網(wǎng)絡(luò)」名稱的由來。

3.訓(xùn)練目標(biāo)是通過優(yōu)化樣本點的位置，使網(wǎng)絡(luò)輸出的離散分布盡可能逼近訓(xùn)練數(shù)據(jù)的真實分布。

每一類生成模型都有其獨特的性質(zhì)，DDN 也不例外。本文將重點介紹 DDN 的三個特性：

• 零樣本條件生成 (Zero-Shot Conditional Generation, ZSCG)
• 樹狀結(jié)構(gòu)的一維離散潛變量 (Tree-Structured 1D Discrete Latent)
• 完全的端到端可微分 (Fully End-to-End Differentiable)

• 論文標(biāo)題： 《Discrete Distribution Networks》
• 論文鏈接： https://arxiv.org/abs/2401.00036
• 項目鏈接： https://discrete-distribution-networks.github.io/
• 代碼地址： https://github.com/DIYer22/discrete_distribution_networks

離散分布網(wǎng)絡(luò)原理

圖1: DDN 的重建過程示意圖

首先，借助上圖所示的 DDN 重建流程作為切入點來一窺其原理。與 diffusion 和 GAN 不同，它們無法重建數(shù)據(jù)，DDN 能像 VAE 一樣具有數(shù)據(jù)重建能力：先將數(shù)據(jù)映射為 latent ，再由 latent 生成與原始圖像高度相似的重建圖像。

上圖展示了 DDN 重建 target 并獲得其 latent 的過程。一般 DDN 內(nèi)部包含多個層級結(jié)構(gòu)，其層數(shù)為 L，示意圖里 L=3。但先讓我們把目光集中在最左側(cè)的第一層。

離散分布： 正如上文所言，DDN 的核心思想在于讓網(wǎng)絡(luò)同時生成 K 個輸出，從而表示「網(wǎng)絡(luò)輸出了一個離散分布」。因此每一層 DDN 都有 K 個 outputs，即一次性輸出 K 張不同的圖像，示意圖中 K=3。每個 output 都代表了這個離散分布中的一個樣本點，每個樣本點的概率質(zhì)量相等，均為 1/K。

層次化生成： 最終目標(biāo)是讓這個離散分布 (K 個 outputs)，和目標(biāo)分布（訓(xùn)練集）越接近越好，顯然，單靠第一層的 K 個 outputs 無法清晰地刻畫整個 MNIST 數(shù)據(jù)集。第一層獲得的 K 張圖像更像是將 MNIST 聚為 K 類后得到的平均圖像。因此，我們引入「層次化生成」設(shè)計以獲得更加清晰的圖像。

在第一層，橙色 Sampler 根據(jù)  距離從 K 個 outputs 中選出和重建 target 最相似的一張 output。再把被選中的 output 圖輸入回網(wǎng)絡(luò)，作為第二層 DDN 的 condition。這樣，第二層 DDN 就會基于 condition（被選中的圖）生成新的 K 張和 target 更相似的 outputs。

接著，從第二層的 outputs 中繼續(xù)選擇出和 target 最相似的一張作為第三層的 condition，并重復(fù)上述過程。隨著層數(shù)增加，生成的圖像和 target 會越來越相似，最終完成對 target 的重建。

Latent： 這一路選下來，每一層被選中 output 的 index 就組成了 target 的 latent（圖中綠色部分「3-1-2」）。因此 latent 是一個長度為 L, 取值范圍 [1,K] 的整數(shù)數(shù)組。

訓(xùn)練： DDN 的訓(xùn)練過程和重建過程一樣，只需額外在每一層中，對選中的 output 和 target 計算  loss 即可?？偟?loss 就是對每一層  loss 取平均。

生成： 在生成階段，將 Sampler 替換為 random choice 即可：每一層從 K 個 outputs 中隨機抽取一個作為下一層的 condition。由于生成空間包含  個樣本點，復(fù)雜度隨 K 和 L 指數(shù)級增長，隨機采樣的 latent 幾乎不可能與訓(xùn)練集中的 latent 重合，因此可視為模型生成的新樣本。

網(wǎng)絡(luò)結(jié)構(gòu)

將「重建過程示意圖」進一步細化，就有下圖 (a) 的網(wǎng)絡(luò)結(jié)構(gòu)圖：

DDN 網(wǎng)絡(luò)結(jié)構(gòu)示意圖和支持的兩種網(wǎng)絡(luò)結(jié)構(gòu)形式

在圖 (a) 中，把生成相關(guān)的設(shè)計整合為 Discrete Distribution Layer (DDL)， 把僅提供基礎(chǔ)計算的模塊封裝為了 NN Block，并重點展示訓(xùn)練時 DDL 內(nèi)部的數(shù)據(jù)流。主要關(guān)注以下幾點：

• 第一層 DDN 的輸入為 zero tensor，不需要任何 condition；
• DDL 內(nèi)部通過 K 個 conv1x1 來同時生成 K 個 outputs；
• 然后，Guided Sampler 從這些 outputs 中選出和 training image    距離最小的 output；
• 被選中的 output 圖像承擔(dān)兩項任務(wù)：[1]. concat 回 feature 中，作為下一層 DDL 的 condition；[2]. 和 training image 計算  loss。

右側(cè)的 (b)、 (c) 兩圖分別展示了 DDN 支持的兩種網(wǎng)絡(luò)結(jié)構(gòu)形式:

• (b)Single Shot Generator: 類似 GAN 中生成器的 decoder 結(jié)構(gòu)，但需要在網(wǎng)絡(luò)中插入足夠數(shù)量的 DDL 以確保生成空間  足夠大。
• (c)Recurrence Iteration: 各層 DDL 共享相同參數(shù)，類似 diffusion 模型，需要做多次 forward 才能生成樣本。

出于計算效率考慮，DDN 默認采用具有 coarse-to-fine 特性的 single shot generator 形式。

損失函數(shù)

DDN 是由 L 層 DDL 組成，以第 層 DDL  為例，輸入上一層選中的樣本 ，生成 K 個新的樣本 ，并從中找出和當(dāng)前訓(xùn)練樣本 x 最相似的樣本  及其 index。最后，只在選中的樣本  上計算這一層 DDL 的 loss。公式及說明如下：

其中，代表第一層 DDL 的輸入為 zero tensor。DDN 的總 loss 就是每一層的 loss 取平均。

此外，本文還提出了 Split-and-Prune 優(yōu)化算法來使得訓(xùn)練時每個節(jié)點被 GT 匹配上的概率均勻，都是 1/K。

下圖展示了 DDN 做二維概率密度估計的優(yōu)化過程：

左：生成樣本集；右：概率密度GT

實驗與特性展示

隨機采樣效果展示

在人臉數(shù)據(jù)集上的隨機采樣效果

更通用的零樣本條件生成

先描述一下「零樣本條件生成」（Zero-Shot Conditional Generation, ZSCG）這個任務(wù)：

• 首先，Unconditional 地訓(xùn)練一個生成模型，即訓(xùn)練階段，模型只見過圖像，沒有見過任何 condition 信號。
• 在生成階段，用戶會提供 condition，比如 text prompt、低分辨率圖像、黑白圖像。
• 任務(wù)目標(biāo)：讓已經(jīng) unconditional 訓(xùn)練好的生成模型能根據(jù) condition 生成符合對應(yīng) condition 的圖像。
• 因為在訓(xùn)練階段，模型沒見過任何的 condition 信號，所以叫 Zero-Shot Conditional Generation。
  

用 Unconditional DDN 做零樣本條件生成效果：DDN 能在不需要梯度的情況下，使不同模態(tài)的 Condition (比如 text prompt 加 CLIP) 來引導(dǎo) Unconditional trained DDN 做條件生成。黃色框圈起來部分就是用于參考的 GT。SR 代表超分辨率、ST 代表 Style Transfer。

如上圖所示，DDN 支持豐富的零樣本條件生成任務(wù)，其做法和圖 1 中的 DDN 重建過程幾乎一樣。

具體而言，只需把圖 1 中的 target 替換為對應(yīng)的 condition，并且，把采樣邏輯調(diào)整為從每一層的多個 outputs 中選出最符合當(dāng)前 condition 的那一個 output 作為當(dāng)前層的輸出。這樣隨著層數(shù)的增加，生成的 output 越來越符合 condition。整個過程中不需要計算任何梯度，僅靠一個黑盒判別模型就能引導(dǎo)網(wǎng)絡(luò)做零樣本條件生成。DDN 是第一個支持如此特性的生成模型。

換為更專業(yè)的術(shù)語描述便是：

> DDN 是首個支持用純粹判別模型引導(dǎo)采樣過程的生成模型；

> 某種意義上促進了生成模型和判別模型的大一統(tǒng)。

這也意味著用戶能夠通過 DDN 高效地對整個分布空間進行篩選和操作。這個性質(zhì)非常有趣，可玩性很高，個人感覺「零樣本條件生成」將會得到廣泛的應(yīng)用。

Conditional Training

訓(xùn)練 conditional DDN 非常簡單，只需要把 condition 或者 condition 的特征直接輸入網(wǎng)絡(luò)中，網(wǎng)絡(luò)便自動學(xué)會了 P (X|Y)。

此外，conditional DDN 也可以和 ZSCG 結(jié)合以增強生成過程的可控性，下圖的第四 / 五列就展示了以其它圖像為 ZSCG 引導(dǎo)的情況下 conditional DDN 的生成效果。

Conditional-DDNs 做上色和邊緣轉(zhuǎn) RGB 任務(wù)。第四、五列展示了以其它圖像為引導(dǎo)的情況下，零樣本條件生成的效果，生成的圖像會在保證符合 condition 的情況下盡可能靠近 guided 圖像的色調(diào)。

端到端可微分

DDN 生成的樣本對產(chǎn)生該樣本的計算圖完全可微，使用標(biāo)準(zhǔn)鏈?zhǔn)椒▌t就能對所有參數(shù)做端到端優(yōu)化。這種梯度全鏈路暢通的性質(zhì)，體現(xiàn)在了兩個方面：

1.DDN 有個一脈相承的主干 feature，梯度能沿著主干 feature 高效反傳。而 diffusion 在傳遞梯度時，需多次將梯度轉(zhuǎn)換到帶噪聲的樣本空間進行反傳。

2.DDN 的采樣過程不會阻斷梯度，意味著網(wǎng)絡(luò)中間生成的 outputs 也是完全可微的，不需要近似操作，也不會引入噪聲。

理論上，在利用判別模型做 fine-tuning 的場景或著強化學(xué)習(xí)任務(wù)中，使用 DDN 作為生成模型能更高效地 fine-tuning。

獨特的一維離散 latent

DDN 天然具有一維的離散 latent。由于每一層 outputs 都 condition on 前面所有的 results，所以其 latent space 是一個樹狀結(jié)構(gòu)。樹的度為 K，層數(shù)為 L，每一個葉子節(jié)點都對應(yīng)一個 DDN 的采樣結(jié)果。

DDN 的 latent 空間為樹狀結(jié)構(gòu)，綠色路徑展示了圖 1 中的 target 所對應(yīng)的 latent

DDN 具有較強的數(shù)據(jù)壓縮能力（有損壓縮）。DDN 的 latent 是一列整數(shù) (list of ints)，屬于高度壓縮的離散表征。一個 DDN latent 有  個 bits 的信息量，以人臉圖像實驗?zāi)J的 K=512, L=128 為例，一個樣本可以被壓縮到 1152 bits。

Latent 可視化

為了可視化 latent 的結(jié)構(gòu)，我們在 MNIST 上訓(xùn)練了一個 output level 層數(shù) L=3，每一層 output nodes 數(shù)目 K=8 的 DDN，并以遞歸九宮格的形式來展示其 latent 的樹形結(jié)構(gòu)。九宮格的中心格子就是 condition，即上一層被采樣到的 output，相鄰的 8 個格子都代表基于中心格子為 condition 生成的 8 個新 outputs。

Hierarchical Generation Visualization of DDN

未來可能的研究方向

• 通過調(diào)參工作、探索實驗、理論分析以改進 DDN 自身，Scaling up 到 ImageNet 級別，打造出能實際使用、以零樣本條件生成為特色的生成模型。
• 把 DDN 應(yīng)用在生成空間不大的領(lǐng)域，例如圖像上色、圖像去噪。又或者 Robot Learning 領(lǐng)域的 Diffusion Policy。
• 把 DDN 應(yīng)用在非生成類任務(wù)上，比如 DDN 天然支持無監(jiān)督聚類，或者將其特殊的 latent 應(yīng)用在數(shù)據(jù)壓縮、相似性檢索等領(lǐng)域。
• 用 DDN 的設(shè)計思想來改進現(xiàn)有生成模型，或者和其它生成模型相結(jié)合，做到優(yōu)勢互補。
• 將 DDN 應(yīng)用在 LLM 領(lǐng)域，做序列建模任務(wù)。