何愷明ResNet級神作,分形生成模型計算效率狂飆4000倍!清華校友一作


新智元報道  

編輯:桃子 KingHZ
【新智元導讀】大自然的分形之美,蘊藏著宇宙的設計規則。剛剛,何愷明團隊祭出「分形生成模型」,首次實現高解析度逐畫素建模,讓計算效率飆升4000倍,開闢AI影像生成新正規化。
影像生成建模全新正規化來了。
你是否曾凝視過雪花的精緻對稱,或驚歎於樹枝的無窮分支?
這些都是大自然中的「分形」。早在1983年,數學家Mandelbrot就揭示了這一現象。
而如今,何愷明團隊將這一概念注入AI,重磅提出「分形生成模型」(fractal generative models),將GenAI模型的模組化層次提升到全新的高度。
論文連結:https://arxiv.org/abs/2502.17437
類似於數學中的分形,它採用了「遞迴結構」,遞迴呼叫原子生成模組,構建了新型的生成模型,形成了自相似的分形架構。
具體來說,每個生成模組內部包含了更小的生成模組,而這些小模組內又巢狀著更小的模組。
這也並非憑空想象,科學研究早已證明,大腦的神經網路正是分形的傑作。人類大腦同樣是透過模組化遞迴,將微型神經網路組合成更大的網路。
在畫素級影像生成上,研究團隊驗證了新方法的強大——
「分形生成模型」首次將逐畫素建模的精細解析度的計算效率,提升了4000倍。
分形生成模不僅是一種新模型,更是生成建模領域的全新正規化。
它將AI設計與自然界奧秘合二為一,或許通往真正智慧道路,就是更深入理解、模擬自然界已有的設計模式。
這篇神作一齣世,便有網友表示,何愷明的ResNet 2?
還有大佬稱,「分形生成模型代表了AI領域一個激動人心的新前沿。自迴歸模型的遞迴特性,就是在學習模仿大自然的模式。
這不僅僅是理論,而是一條通往更豐富、更具適應性AI系統的道路」。
自然界終極設計模式,「分形」無處不在
計算機科學的核心概念之一是模組化。
現代生成模型(如擴散模型和自迴歸模型)是由基本的「生成步驟」組成的,而每個步驟本身都是由深度神經網路實現的。
將複雜的功能抽象成基本模組,透過組合這些模組來構建更復雜的系統。這就是模組化方法。
基於這一理念,研究團隊提出將生成模型本身作為一個模組,從而開發更高階的生成模型。
新方法受到了生物神經網路和自然資料中觀察到的分形特性的啟發。
與自然分形結構類似,設計的關鍵元件是定義遞迴生成規則的「生成器」。
例如,生成器可以是一個自迴歸模型,如圖1所示。在這種例項化中,每個自迴歸模型由本身也是自迴歸模型的模組組成。具體來說,每個父自迴歸模組生成多個子自迴歸模組,而每個子模組進一步生成更多的自迴歸模組。
最終的架構在不同層次上展現出類似分形的自相似模式,如圖1所示。
圖1:分形生成模型
在這個例項中,使用自迴歸模型作為分形生成器。在自迴歸模型中,遞迴呼叫自迴歸模型,構建了一個具有不同層級之間自相似性的類似分形的框架。

動機和直覺

從簡單的遞迴規則中,分形可以產生複雜的模式。
這也是分形生成模型的核心思想:利用現有的原子生成模組,遞迴地構建成更高階的生成模型。
在分形幾何中,這些規則通常被稱為「生成器」。
透過不同的生成器,分形方法可以構建許多自然模式,如雲、山脈、雪花和樹枝,並且和更復雜的系統有關,如生物神經網路的結構、非線性動力學和混沌系統。
Mathworld中不同的分形模式
形式上,分形生成器g_i,指定了如何基於上一級生成器的輸出xi,生成下一級生成器的新資料集

例如,如圖1所示,生成器可以透過在每個灰色框內遞迴呼叫類似的生成器來構建分形。
由於每個生成器層級可以從單個輸入生成多個輸出,在僅需要線性遞迴層級的情況下,分形框架可以實現生成輸出的指數級增長。
這特別適合用相對較少的生成器層級,來建模高維資料。
「分形生成模型」核心架構

分而治之

在理論上,生成模型就是建模多個高維隨機變數的聯合分佈,但直接用單一的自迴歸模型建模,在計算上是不可行的。
為了解決這個問題,採用了分而治之的策略。
關鍵模組化是將自迴歸模型抽象為一個模組化單元,用於建模機率分佈p(x∣c)。
透過這種模組化,可以在多個下一級自迴歸模型的基礎上構建一個更強大的自迴歸模型。
假設每個自迴歸模型中的序列長度是一個可管理的常數k,並且總隨機變數數N=k^n,其中n=log⁡_k(N)表示框架中的遞迴層級數。
分形框架的第一層自迴歸模型將聯合分佈劃分為k個子集,每個子集包含k^{n−1}個變數。形式上,我們將聯合分佈分解為:
每個包含k^{n−1}個變數的條件分佈p(⋯∣⋯ ),由第二層遞迴的自迴歸模型建模,以此類推。
透過遞迴呼叫這種分而治之的過程,分形框架可以使用n層自迴歸模型,高效地處理k^n個變數的聯合分佈。

架構例項

如圖3所示,每個自迴歸模型將上一層生成器的輸出作為輸入,併為下一層生成器生成多個輸出。
它還接受一張影像(可以是原始影像的一部分),將其切分成多個patch,並將這些patch嵌入以形成Transformer模型的輸入序列。這些patch也會被傳遞給相應的下一層生成器。
然後,Transformer將上一層生成器的輸出作為一個獨立的token,放置在影像token之前。
基於合併的序列,Transformer生成多個輸出,供下一層生成器使用。
遵循領域內模型中的常見做法,將第一個生成器g_0的序列長度設定為256,將原始影像劃分為16×16的patch。
然後,第二層生成器對每個patch進行建模,並進一步將這些patch細分為更小的patch,遞迴地繼續這一過程。
為了管理計算成本,逐漸減少較小patch的Transformer寬度和Transformer塊的數量,因為對較小的patch建模,通常比較大的patch更容易。
在最後一級,使用非常輕量的Transformer,以自迴歸方式建模每個畫素的RGB通道,並對預測應用256路交叉熵損失。
值得注意的是,分形設計建模256×256影像的計算,成本僅為建模64×64影像的兩倍。

實現

採用寬度優先的方式,端到端訓練原始影像畫素。
在訓練過程中,每個自迴歸模型從上一層的自迴歸模型接收輸入,併為下一層自迴歸模型生成一組輸出作為輸入。這個過程一直持續到最終層級,在那裡影像被表示為畫素序列。
最後的模型使用每個畫素的輸出,以自迴歸的方式預測RGB通道。
對預測的logits計算交叉熵損失(將RGB值視為從0到255的離散整數),並透過所有層級的自迴歸模型,進行反向傳播,從而端到端地訓練整個分形框架。
分形模型以逐畫素的方式生成影像,按照深度優先的順序遍歷分形架構,如圖2所示。
這裡,以下文中的MAR的隨機順生成方案為例。
論文連結:https://arxiv.org/abs/2406.11838
第一層:自迴歸模型捕捉16×16影像patch之間的相互依賴關係,並在每一步根據已知的影像patch生成下一層的輸出。
第二層:模型利用這些輸出,對每個16×16影像patch內4×4影像patch之間的相互依賴關係建模。
類似地,第三層自迴歸模型建模每個4×4影像patch內的畫素之間的相互依賴關係。
最後,從自迴歸預測的RGB logits中,最頂層的自迴歸模型取樣出實際的RGB值。

與尺度空間自迴歸模型的關係

尺度空間自迴歸模型(Scale-space Autoregressive Models),與分形方法之間的一個主要區別在於:它們使用單一的自迴歸模型,按尺度逐步預測token。
相較之下,分形框架採用了分而治之的策略,透過生成子模組遞迴地建模原始畫素。
另一個關鍵區別在於計算複雜度:尺度空間自迴歸模型在生成下一個尺度的token時,需要對整個序列執行全注意力操作,這導致計算複雜度顯著更高。
舉個栗子,當生成256×256解析度的影像時,在最後一個尺度下,尺度空間自迴歸模型中每個注意力patch的注意力矩陣大小為(256 × 256)² = 4,294,967,296。
而新方法在建模畫素間依賴關係時,對非常小的圖patch(4×4)進行注意力操作,其中每個圖patch的注意力矩陣僅為(4 × 4)² = 256,從而使得總的注意力矩陣大小為(64 × 64) × (4 × 4)² = 1,048,576次操作。
這種縮減使得分形方法在最精細的解析度下,比傳統方法計算效率提高了4000倍,從而首次實現了畫素逐畫素建模高解析度影像。

與長序列建模的關係

大多數關於逐畫素生成的前期研究,將問題表述為長序列建模,並利用語言建模中的方法來解決這個問題。
然而,許多資料型別的內在結構,包括但不限於影像,超出了單一維度的序列。
與這些方法不同,研究團隊將這類資料視為由多個元素組成的集合(而非序列),並採用分治策略遞迴地對包含較少元素的子集進行建模。
這一方法的動機來源於對這些資料的觀察——大量資料展現出近似分形結構:
影像由子影像組成,
分子由子分子組成,
生物神經網路由子網路組成。
因此,旨在處理這類資料的生成模型應當由子模組組成,而這些子模組本身也是生成模型。
實驗結果
研究人員在ImageNet資料集上,對「分形生成模型」進行了廣泛的實驗,分別使用64×64和256×256的解析度。
評估包括無條件和類別條件影像生成,涵蓋了模型的多個方面,如似然估計、保真度、多樣性和生成質量。

似然估計

研究人員首先在ImageNet 64×64無條件生成任務上展開評估,以檢驗其似然估計的能力。
為了考察分形框架的有效性,他們比較了不同分形層級數量下框架的似然估計效能,如下表2所示。
使用單個自迴歸模型,對整個64×64×3=12,288畫素序列建模會產生過高的計算成本,讓訓練變得不可行。
此外,先對整個畫素序列然後對RGB通道建模的兩級分形框架,需要的計算量是三級分形模型的十多倍。
在表5中,進一步將分形框架與其他基於似然的模型進行了比較。
分形生成模型,同時使用因果自迴歸和掩碼自迴歸分形生成器實現,取得了強大的似然效能。
特別是,它實現了每維3.14位元的負對數似然,顯著優於此前最佳的自迴歸模型(每維3.40位元),並且與SOTA模型相比有著強大得競爭力。
這些發現證明了,分形框架在具有挑戰性的逐畫素影像生成任務上的有效性,突顯了其在建模高維非序列資料分佈方面的潛力。

生成質量

此外,研究人員還使用「四級分形結構」評估了FractalMAR在256×256解析度下,在具有挑戰性的類別條件影像生成任務上的表現。
如下表4中,FractalMAR-H實現了6.15的FID,以及348.9的IS。
在單個Nvidia H100 PCIe GPU上以1024的批大小評估時,平均每張影像的生成時間為1.29秒。
值得注意的是,新方法在IS和精確率上,具備了顯著優勢,表明其能夠生成具有高保真度和精細細節的影像,這一點在圖4中也得到了展示。
然而,其FID、召回率相對較弱,與其他方法相比,生成的樣本多樣性較低。
研究人員推測,這是由於以逐畫素方式對近200,000個畫素進行建模的巨大挑戰所致。
此外,研究人員進一步觀察到一個有前景的Scaling趨勢:
將模型規模從1.86億引數增加到8.48億引數,顯著改善了FID(從11.80降至6.15)和召回率(從0.29提升至0.46)。
他們預期,進一步增加引數規模,可能會進一步縮小FID和召回率的差距。

條件逐畫素預測

進一步地,作者透過影像編輯的常規任務,檢驗了分形方法的條件逐畫素預測效能。
如下圖5中的幾個示例,包括修復(inpainting)、外延(outpainting)、基於掩碼外延(uncropping)和類別條件編輯。
結果顯示,分享方法能夠基於未掩碼區域,準確預測被掩碼的畫素。
此外,它能夠有效地從類別標籤中捕捉高層語義資訊,並反映在預測的畫素中。
這一點在類別條件編輯示例中得到了展示,其中模型透過條件化狗的類別標籤,將貓的臉替換成了狗的臉。這些結果證明了,新方法在給定已知條件下預測未知資料的有效性。
更廣泛地說,透過逐畫素生成資料,新方法提供了一個相比擴散模型或在潛空間運作的生成模型,更易於人類理解的生成過程。
這種可解釋的生成過程不僅讓我們能夠更好地理解資料是如何生成的,還提供了一種控制和互動生成的方式。
未來,這些能力在視覺內容創作、建築設計和藥物發現等應用中,尤為重要。
作者介紹
Tianhong Li(黎天鴻)
黎天鴻目前是MIT CSAIL(麻省理工學院計算機科學與人工智慧實驗室)的博士後研究員,師從何愷明教授。
在此之前,他獲得了MIT博士、碩士學位,導師是Dina Katabi教授。他本科畢業於清華大學姚班計算機科學專業。
黎天鴻的研究興趣主要集中在表示學習、生成模型,以及這兩者之間的協同效應。他的目標是構建能夠理解和建模,超越人類感知的智慧視覺系統。
他也非常喜歡烹飪,這種熱愛程度幾乎和做研究一樣。
有趣的是,在個人主頁裡,他列出了一些自己最喜歡的菜譜。
Qinyi Sun
Qinyi Sun目前是麻省理工學院(MIT)電氣工程與計算機科學系(EECS)本科生,師從何愷明教授。
Lijie Fan
Lijie Fan目前是谷歌DeepMind研究科學家。
他於2024年獲得了MIT計算機科學專業博士學位,於2018年獲得了清華大學計算機科學學士學位。
他的個人研究重點在生成模型和合成資料。
何愷明
何愷明目前是麻省理工學院(MIT)電氣工程與計算機科學系(EECS)的副教授,於2024年2月加入。
他本人的研究重點是,構建能夠從複雜世界中學習表示並發展智慧的計算機模型,研究的長期目標是用更強大的AI來增強人類智慧。
何愷明最為人熟知的研究是深度殘差網路(ResNets)。ResNets的影響力不僅限於計算機視覺領域,它的設計思想被廣泛應用於現代深度學習模型中。
無論是自然語言處理中的Transformer(如GPT、ChatGPT),還是強化學習中的AlphaGo Zero,甚至是蛋白質結構預測的AlphaFold,殘差連線都成為了這些模型的核心元件之一。
除了ResNets,何愷明在計算機視覺領域的貢獻同樣令人矚目。他提出的Faster R-CNN和Mask R-CNN,極大地推動了目標檢測和影像分割技術的發展。
在加入MIT之前,他於2016年-2024年在Facebook AI研究院(FAIR)擔任研究科學家,2011年-2016年在微軟亞洲研究院(MSRA)擔任研究員。
此前,他於2011年在香港中文大學獲得博士學位,2007年在清華大學獲得學士學位。
參考資料:
https://arxiv.org/abs/2502.17437

相關文章