何愷明ResNet級神作,分形生成模型計算效率狂飆4000倍!清華校友一作
圖像生成建模全新范式來了。
你是否曾凝視過雪花的精致對稱,或驚嘆于樹枝的無窮分支?
這些都是大自然中的「分形」。早在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)2 = 4,294,967,296。
而新方法在建模像素間依賴關系時,對非常小的圖patch(4×4)進行注意力操作,其中每個圖patch的注意力矩陣僅為(4 × 4)2 = 256,從而使得總的注意力矩陣大小為(64 × 64) × (4 × 4)2 = 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年在清華大學獲得學士學位。