返璞歸真！MIT何愷明再次顛覆生成模型：拋棄“去噪”，Diffusion模型直接預測干凈圖像！
文章鏈接：https://arxiv.org/pdf/2511.13720
亮點直擊
解決的問題
提出的方案
達到的效果
- 優秀的可擴展性: JiT模型展現了卓越的伸縮性。
- 分辨率擴展: 通過成比例地增大patch大小，JiT-B模型可以輕松地從256x256擴展到512x512甚至1024x1024分辨率，而模型參數量和計算量幾乎保持不變，FID分數也保持在很低的水平。這證明該方法不受觀測維度急劇增加的影響。
- 模型規模擴展: 從JiT-B（Base）到JiT-G（Giant）模型，性能隨著模型尺寸的增大而持續穩定提升，在200-epoch和600-epoch的訓練設置下都展示了清晰的擴展趨勢，在512x512分辨率上取得了低至1.78的FID分數。
擴散模型的預測輸出
擴散模型可以在 x、? 或 v 空間中建立。空間的選擇不僅決定了損失的定義，還決定了網絡的預測結果。重要的是，損失空間和網絡輸出空間不必相同。這一選擇會產生關鍵性的差異。
背景：擴散模型與流模型
擴散模型可以從常微分方程（ODE）的角度進行公式化。我們從基于流的框架出發進行建模，即首先在 v 空間中展開，作為更簡單的起點，然后再討論其他空間。
預測空間與損失空間
玩具實驗
圖 1.流形假設假設自然圖像位于高維流形中的低維流形上。 圖像位于高維像素空間內的低維流形上。 維像素空間內的低維流形上。干凈的圖像 x 可以建模為 在流形上，而噪聲 ? 或流速 v（例如，v = x - ?）則完全不在流形上。 此時為非褶皺狀態。訓練神經網絡來預測干凈的 即 x 預測）與訓練神經網絡預測噪聲或有噪點圖像有著本質區別。 訓練神經網絡來預測噪聲或噪聲量（即 ?/v 預測）有著本質的區別
圖 2.玩具實驗：D 維（d = 2）底層 數據 "埋藏 "在 D 維空間中，由一個固定的隨機 列正交投影矩陣。在 D 維空間中，我們訓練一個簡單的生成模型（5 層 ReLU MLP，帶 256 維的 隱藏單元）。投影矩陣對模型來說是未知的。 我們僅將其用于可視化輸出。在這個玩具實驗中、 隨著觀測維度 D 的增加，只有 x 預測能產生合理的結果。
JiT: Just image Transformers (純粹的圖像Transformer)
基于流形假設的啟發，本文提倡使用 x-prediction。將介紹一個用于像素空間擴散的、概念上簡單的方法，稱之為JiT。
架構
圖 3.圖像Transformer"（JiT）架構： 在像素斑塊上簡單地使用普通 ViT [13] 進行x-prediction
JiT的骨干網絡是一個標準的Vision Transformer (ViT) 。網絡接收一個帶噪圖像Zt作為輸入，將其劃分為不重疊的圖像塊（patches）。每個圖像塊被一個線性層獨立地嵌入。這些嵌入后的向量，連同一個從時間步t派生的嵌入向量，共同形成一個序列，輸入到Transformer編碼器中。在經過Transformer編碼器處理后，輸出序列中的圖像塊 token 會被另一個線性層獨立地投影回原始像素空間，以重構出預測的清晰圖像。
在ViT中，一個核心設計是使用大尺寸的圖像塊。例如，在256×256分辨率下使用16×16的patch大小，每個patch是768維（16×16×3）的向量。這是一種非常高的維度，為驗證我們的假設提供了一個很好的平臺。
主要實驗
本文在ImageNet 256×256上進行了主要實驗，使用JiT-B/16（ViT-Base架構，patch大小為16）。表2(a)展示了所有九種損失空間和預測空間組合的結果。
這個證據表明，網絡設計可以在很大程度上與觀測維度解耦，這與許多其他神經網絡應用中的情況一致。增加隱藏單元的數量可能是有益的（正如在深度學習中廣泛觀察到的那樣），但它不是決定性因素。
瓶頸結構可能是有益的。更令人驚訝的是，本文發現，反過來，在網絡中引入一個降低維度的瓶頸結構可能是有益的。
從表示學習的更廣闊視角來看，這一觀察并非完全出乎意料。引入瓶頸設計的目的通常是為了鼓勵學習內在的低維表示，如文獻[64, 48, 41, 2]。
本文算法
“Just Advanced” Transformers (純粹的先進Transformer)
通用Transformer的強大之處部分在于，當其設計與特定任務解耦時，它可以從其他應用中發展出來的架構進步中受益。這一特性支撐了用任務無關的Transformer來構建擴散模型的優勢。
遵循 [73]的做法，本文整合了流行的通用改進措施：SwiGLU、RMSNorm 、RoPE、qk-norm，所有這些最初都是為語言模型開發的。本文還探索了上下文類別條件化（in-context class conditioning）：但與原始ViT [13] 在序列后附加一個類別token不同，本文附加了多個此類token（默認為32個；見附錄），這一做法遵循了[35]。表4報告了這些組件的效果。
對比
高分辨率像素生成。在表5中，進一步報告了基礎尺寸模型（JiT-B）在ImageNet上512甚至1024分辨率的結果。使用的patch大小與圖像大小成正比，因此不同分辨率下的序列長度保持不變。每個patch的維度可以高達3072或12288，沒有任何一個常見模型會有足夠多的隱藏單元來匹配。
表5顯示，本文的模型在不同分辨率下都表現得相當不錯。所有模型的參數數量和計算成本都相似，僅在輸入/輸出的patch嵌入層有所不同。本文的方法沒有受到觀測維度災難的影響。
可擴展性。將Transformer設計與任務解耦的一個核心目標是利用其可擴展的潛力。表6提供了四種模型尺寸在ImageNet 256和512分辨率下的結果（請注意，在512分辨率下，這些模型沒有一個的隱藏單元數比patch維度更多）。模型尺寸和浮點運算次數（FLOPs）顯示在表7和表8中：本文的模型在256分辨率下的成本與其在512分辨率下的對應模型相似。
表 7.ImageNet 256×256 的參考結果。對 50K 個樣本的 FID 和 IS 的 50K 個樣本進行了評估。預訓練 "列 列列出了獲得結果所需的外部模型（注意 感知損失使用的是預先訓練好的 VGG 分類器）。參數包括 參數包括生成器和標記解碼器（在推理時使用 推理時使用），但不包括其他預訓練組件。參數 Giga-flops 是針對單個前向傳遞（不計 標記符），與訓練和推理過程中一次迭代的計算成本大致成正比。 在訓練和推理過程中的迭代成本成正比（對于 多尺度方法，我們測量的是最精細的層次）。
本文的方法受益于模型規模的擴展。有趣的是，隨著模型變大，256和512分辨率之間的FID差距變小了。對于JiT-G模型，512分辨率下的FID甚至更低。對于ImageNet上的超大型模型，FID性能在很大程度上取決于過擬合，而512分辨率下的去噪任務更具挑戰性，使其不易受到過擬合的影響。
結論
本文通過回歸擴散模型中“去噪”的核心理念，展示了一個在原始像素上操作的、簡單而強大的“Diffusion + Transformer”范式是可行的。研究表明，直接預測清晰數據（即 x-prediction），而不是噪聲或一個帶噪的量，是成功的關鍵。這一選擇源于一個基礎的機器學習原則——流形假設——該假設認為自然數據位于一個低維流形上。本文的JiT模型，一個純粹的圖像Transformer，在ImageNet上取得了強大的結果，并且可以優雅地擴展到1024分辨率，而無需依賴分詞器、預訓練或輔助損失。這項工作朝著在原生數據上實現一個自包含的“Diffusion + Transformer”哲學邁出了一步，并有望在其他難以設計分詞器的科學領域找到更廣泛的應用。
本文轉自AI生成未來 ，作者：AI生成未來
