何愷明帶隊新作「分形生成模型」:逐畫素建模高解析度影像、效率提升4000倍

机器之心發表於2025-02-26
這才過幾天,大神何愷明又放出一篇新論文!

這次構建了一種全新的生成模型。類似於數學中的分形,研究者推出了一種被稱為分形生成模型(Fractal Generative Models)的自相似分形架構。

在電腦科學領域,它的核心是模組化概念,比如深度神經網路由作為模組化單元的原子「層」構建而成。同樣地,現代生成模型(如擴散模型和自迴歸模型)由原子「生成步驟」構建而成,每個步驟都由深度神經網路實現。

透過將複雜函式抽象為這些原子構建塊,模組化使得可以透過組合這些模組來建立更復雜的系統。基於這一概念,研究者提出將生成模型本身抽象為一個模組,以開發更高階的生成模型。一作 Tianhong Li 為 MIT 博士後研究員、二作 Qinyi Sun 為 MIT 本科生(大三)。
image.png
  • 論文標題:Fractal Generative Models

  • 論文地址:https://arxiv.org/pdf/2502.17437v1

  • GitHub 地址:https://github.com/LTH14/fractalgen

具體來講,研究者提出的分形生成模型透過在其內部遞迴呼叫同類生成模型來構建。這種遞迴策略產生了一個生成框架,在下圖 1 中展示了其跨不同模組級別的具有自相似性的複雜架構。
image.png
如前文所述,本文分形生成模型類似於數學中的分形概念。分形是使用被稱為「生成器」的遞迴規則構建的自相似模式。同樣地,本文框架也是透過在生成模型中呼叫生成模型的遞迴過程構建的,並在不同層次上表現出自相似性。因此,研究者將其命名為「分形生成模型」。

本文的分形生成模型的靈感來自於生物神經網路和自然資料中觀察到的分形特性。與自然的分形結構類似,研究者設計的關鍵元件是定義遞迴生成規則的生成器,比如這樣的生成器可以是自迴歸模型,如圖 1 所示。在此例項中,每個自迴歸模型都由本身就是自迴歸模型的模組組成。

具體而言,每個父自迴歸塊都會生成多個子自迴歸塊,每個子塊都會進一步生成更多自迴歸塊。由此產生的架構在不同級別上表現出類似分形的自相似模式。

在實驗環節,研究者在一個具有挑戰性的測試平臺上(逐畫素影像生成)檢驗了這個分形例項。結果顯示,本文的分形框架在這一具有挑戰性的重要任務上表現出色,它不僅可以逐畫素生成原始影像,同時實現了準確的似然估計和高生成質量,效果如下圖 2 所示。
2025-02-25_151626.png
研究者希望這一充滿潛力的的結果能夠激勵大家進一步研究分形生成模型的設計和應用,最終在生成建模中建立一種全新的正規化。

有人評論道,「分形生成模型代表了AI領域一個令人興奮的前沿。自迴歸模型的遞迴性質反映了學習如何反映自然模式。這不僅僅是理論,它是通往更豐富、適應性更強的AI系統的途徑。」
image.png
圖源:https://x.com/abhivendra/status/1894421316012577231

分形生成模型詳解

研究者表示,分形生成模型的關鍵思路是「從現有的原子生成模組中遞迴地構建更高階的生成模型。」

具體來講,該分形生成模型將一個原子生成模組用作了引數分形生成器。這樣一來,神經網路就可以直接從資料中「學習」遞迴規則。透過將指數增長的分形輸出與神經生成模組相結合,分形框架可以對高維非序列資料進行建模。

接下來,研究者展示瞭如何透過將自迴歸模型用作分形生成器來構建分形生成模型。他們將自迴歸模型用作了說明性原子模組,以演示分形生成模型的例項化,並用來對高緯資料分佈進行建模。

假設每個自迴歸模型中的序列長度是一個可管理的常數 k,並使隨機變數的總數為 N = k^n,其中 n = log_k (N) 表示分形框架中的遞迴級別數。然後,分形框架的第一個自迴歸級別將聯合分佈劃分為 k 個子集,每個子集包含 k^n−1 個變數。

在形式上,研究者進行了如下解耦:
微信圖片_20250225150348_副本.jpg
接著每個具有 k^n−1 個變數的條件分佈 p (・・・|・・・) 由第二個遞迴級別的自迴歸模型建模,並依此類推。

研究者表示,透過遞迴地呼叫這種分而治之(divide-and-conquer)的過程,分形框架可以使用 n 級自迴歸模型高效地處理 k^n 個變數的聯合分佈,並且每個模型都對可管理的序列長度 k 進行操作。

這種遞迴過程代表了一種標準的分而治之策略。透過遞迴地解耦聯合分佈,本文分形自迴歸架構不僅相較於單個大型自迴歸模型顯著降低了計算成本,而且還捕獲了資料中的內在層次結構。從概念上講,只要資料表現出可以分而治之的組織結構,就可以在該分形框架內自然地對其進行建模。

實現:影像生成例項化

研究者展示了分形自迴歸架構如何用於解決具有挑戰性的逐畫素影像生成任務。

架構概覽

如下圖 3 所示,每個自迴歸模型將上一級的生成器的輸出作為其輸入,併為下一級生成器生成了多個輸出。該模型還獲取一張影像(也可以是原始影像的 patch),將其分割成 patch,並將它們嵌入以形成一個 transformer 模型的輸入序列。這些 patch 也被饋送到相應的下一級生成器。

接下來,transformer 模型將上一個生成器的輸出作為單獨的 token,放在影像 token 的前面。基於此組合序列,transformer 為下一級生成器生成多個輸出。

研究者將第一級生成器 g_0 的序列長度設定為 256,將原始影像分成 16 × 16 個 patch。然後,第二級生成器對每個 patch 進行建模,並進一步將它們細分為更小的 patch,並繼續遞迴執行此過程。為了管理計算成本,他們逐步減少較小 patch 的寬度和 transformer 塊的數量,這樣做是因為對較小 patch 進行建模通常比對較大 patch 更容易。

在最後一級,研究者使用一個非常輕量級的 transformer 來自迴歸地建模每個畫素的 RGB 通道,並在預測中應用 256 路交叉熵損失。
2025-02-25_150955.png
不同遞迴級別和解析度下,每個 transformer 的精確配置和計算成本如下表 1 所示。值得注意的是,透過本文的分形設計,建模解析度為 256×256 影像的計算成本僅為建模解析度為 64×64 影像的兩倍。
2025-02-25_151118.png

本文方法支援不同的自迴歸設計。研究者主要考慮了兩種變體:光柵順序、類 GPT 的因果 transformer (AR) 和隨機順序、類 BERT 的雙向 transformer (MAR),具體如下圖 6 所示。
image.png
尺度空間自迴歸模型

最近,一些模型已經提出為自迴歸影像生成執行下一尺度(next-scale)預測。這些尺度空間自迴歸模型與本文方法的一個主要區別是:它們使用單個自迴歸模型來逐尺度地預測 token。

相比之下,本文分形框架採用分而治之的策略,使用生成式子模組對原始畫素進行遞迴建模。另一個關鍵區別在於計算複雜性:尺度空間自迴歸模型在生成下一尺度 token 的整個序列時需要執行完全注意力操作,這會導致計算複雜性大大增加。

舉例而言,在生成解析度為 256×256 的影像時,在最後一個尺度上,尺度空間自迴歸模型每個注意力塊中的注意力矩陣大小為 (256 ×256)^2 即 4,294,967,296。相比之下,本文方法在對畫素 (4×4) 相互依賴性進行建模時對非常小的 patch 執行注意力,其中每個 patch 的注意力矩陣只有 (4 × 4)^2 = 256,導致總注意力矩陣大小為 (64 × 64) × (4 × 4)^2 = 1,048,576 次操作。

這種減少使得本文方法在最精細解析度下的計算效率提高了 4000 倍,從而首次能夠逐畫素建模高解析度影像。

長序列建模

之前大多數關於逐畫素生成的研究都將問題表述為長序列建模,並利用語言建模的方法來解決。與這些方法不同,研究者將此類資料視為由多個元素組成的集合(而不是序列),並採用分而治之的策略以遞迴方式對具有較少元素的較小子集進行建模。

這種方法的動機是觀察到大部分資料都呈現出了近乎分形的結構。影像由子影像組成,分子由子分子組成,生物神經網路由子網路組成。因此,設計用於處理此類資料的生成模型應該由本身就是生成模型的子模組組成。

實驗結果

本文在 ImageNet 資料集上進行了實驗,影像解析度分別為 64×64 和 256×256。評估包括無條件和類條件影像生成,涵蓋模型的各個方面,如似然估計、保真度、多樣性和生成質量。

因此,本文報告了負對數似然(NLL)、Frechet Inception Distance(FID)、Inception Score(IS)、精度(Precision)和調回率(Recall)以及視覺化結果,以全面評估分形框架。

似然估計。本文首先在無條件 ImageNet 64×64 生成任務上進行了評估,以檢驗其似然估計能力。為了驗證分形框架的有效性,本文比較了不同分形層級數量下框架的似然估計效能,如表 2 所示。
image.png
再來看生成質量評估。研究者在解析度為 256×256 的類條件影像生成這一挑戰性任務上,使用四個分形級別對 FractalMAR 進行了評估。指標包括了 FID、Inception Score、精度和召回率,具體如下表 4 所示。
image.png
值得注意的是,本文方法實現了強大的 Inception Score 和精度,表明它能夠生成具有高保真度和細粒度細節的影像,如下圖 4 所示。
2025-02-26_091306.png
最後是條件逐畫素預測評估。

研究者進一步使用影像編輯中的常規任務來檢驗本文方法的條件逐畫素預測效能。下圖 5 提供了幾個具體示例,包括修復、去除修復、取消裁剪和類條件編輯。

結果顯示,本文方法可以根據未遮蔽區域來準確預測被遮蔽的畫素,還可以有效地從類標籤中捕獲高階語義並將其反映在預測畫素中。
2025-02-26_091614.png
更多實驗結果參閱原論文。

相關文章