機器之心報道

編輯：冷貓、+0

大家都知道，圖像生成和去噪擴散模型是密不可分的。高質量的圖像生成都通過擴散模型實現。

但有沒有一種可能，「去噪擴散模型」實際上并沒有做到「去噪」？

ResNet 之父，超70萬引用的 AI 大神何愷明的新論文敏銳地捕捉了這一現象。

當擴散生成模型最初被提出時 ，核心思想本應是「去噪」，即從被破壞的圖像中預測出干凈圖像。然而，在擴散模型的發展過程中出現的兩個重要里程碑，卻偏離了直接預測干凈圖像這一目標。首先，預測噪聲本身（稱為 「?-prediction」）在生成質量上帶來了關鍵改變，并極大推動了這些模型的普及。隨后，擴散模型被連接到基于流的生成方法，通過預測流速度（「-prediction」 ），這一量結合了干凈數據與噪聲。

流形假設（Manifold Assumption）認為，自然圖像在高維像素空間中位于一個低維流形上。干凈圖像 可以建模為處于流形上（on-manifold），而噪聲 ? 或流速度 （例如 = ? ?）則本質上處于流形之外（off-manifold）。因此，讓神經網絡預測干凈圖像（即 -prediction）在本質上不同于讓其預測噪聲或帶噪的量（即 ?/-prediction）。

當今的去噪擴散模型并不按照經典意義上的「去噪」方式工作，這些神經網絡預測的是噪聲或帶噪的量。

預測干凈數據與預測帶噪量在本質上是不同的。根據流形假設（manifold assumption），自然數據應當位于一個低維流形上，而帶噪的量則不在其上。

基于這一假設，何愷明新論文主張使用直接預測干凈數據的模型，這使得看似容量不足的網絡也能夠在極高維空間中有效運行。

這一研究使去噪擴散模型回歸本源，探索一種在原始自然數據上構建基于 Transformer 的擴散模型的自洽范式。

論文標題：Back to Basics: Let Denoising Generative Models Denoise論文鏈接：https://arxiv.org/abs/2511.13720v1Github 鏈接：https://github.com/LTH14/JiT

論文提出「Just image Transformers（JiT）」架構，僅使用簡單的大 patch 像素級 Transformer 就能成為強大的生成模型，無需 tokenizer、無需預訓練，也無需額外損失項。

該架構在 ImageNet 上實現了有競爭力的像素空間圖像生成，在256×256分辨率下達到1.82 FID，在512×512分辨率下達到1.78 FID，且無需外部組件。

本文的工作朝著一種基于原生數據的自洽 「Diffusion + Transformer」 原則 更進一步。超越計算機視覺領域，這種原則在其他涉及自然數據的領域（如蛋白質、分子或天氣）中也同樣具有潛在價值，因為這些領域的 tokenizer 通常難以設計。通過盡量減少特定領域的建模設計，作者希望源自計算機視覺的通用 「Diffusion + Transformer」 范式能夠在更廣泛的領域中發揮作用。

用于擴散的「Just image Transformers」

Just image Transformers

ViT 的核心理念是 基于 Patch 的 Transformer（ToP），本研究的架構設計遵循這一哲學。

按照標準做法，該架構以時間 t 和給定的類別標簽為條件。研究者使用adaLN-Zero進行條件化處理。從概念上講，這種架構相當于直接應用于像素 patch 的 Diffusion Transformer (DiT)。

整體架構是 Just image Transformer（純圖像 Transformer），將其簡稱為JiT。例如，在 256×256 的圖像上研究 JiT/16（即 patch 大小 p=16），在 512×512 的圖像上研究 JiT/32（p=32）。這些設置分別導致每個 patch 的維度為 768 (16×16×3) 和 3072 (32×32×3)。如此高維的 patch 可以通過 -prediction 來處理。

網絡應該預測什么？

表 1 中總結了損失空間和預測空間的九種可能組合。針對每一種組合，研究者訓練了一個 base 模型（JiT-B），其每個 token 的隱藏層大小為 768 維。

研究者在表 2(a) 中研究了 256×256 分辨率下的 JiT-B/16。作為參考，他們在表 2(b) 中檢查了 64×64 辨率下的 JiT-B/4（即 p=4）。在這兩種設置中，序列長度是相同的 (16×16)。

觀察結果如下：

-prediction 至關重要

在表 2(a) 的 JiT-B/16 中，只有 -prediction 表現良好，并且它在所有三種損失函數下均有效。在這里，一個 patch 是 768 維 (16×16×3)，這與 JiT-B 中 768 的隱藏層大小相吻合。雖然這看起來「剛好足夠」，但在實踐中，模型可能需要額外的容量，例如用于處理位置嵌入。對于 ?-prediction 或 -prediction，模型沒有足夠的容量來分離并保留噪聲量。這些觀察結果與 Toy 案例（圖 2）中的觀察結果相似。

作為對比，研究者檢查了 64×64 分辨率下的 JiT-B/4（表 2(b)）。在這里，所有情況都表現得相當不錯：九種組合之間的準確率差距是邊緣性的，并不具有決定性。每個 patch 的維度是 48 (4×4×3)，遠低于 JiT-B 中 768 的隱藏層大小，這解釋了為什么所有組合都工作得相當好。研究者注意到，許多之前的潛在擴散模型具有類似的小輸入維度，因此未曾暴露在此處討論的問題中。

僅靠損失加權是不夠的

該工作并不是第一個枚舉相關因素組合的，之前的研究探索了損失加權和網絡預測的組合。他們的實驗是在低維的 CIFAR-10 數據集上使用 U-net 完成的。他們的觀察結果與研究者在 ImageNet 64×64 上的結果更為接近。

然而，表 2(a) 在 ImageNet 256×256 上的結果表明，損失加權并非全部。一方面，無論損失空間如何，?-prediction 或 -prediction 在表 2(a) 中都遭受了災難性的失敗，這對應于不同損失空間中的不同有效加權（如前所述）。另一方面，-prediction 在所有三種損失空間中都有效：雖然由 -loss 引起的損失加權更可取，但并非至關重要。

僅靠噪聲水平偏移是不夠的

先前的工作建議增加噪聲水平對于高分辨率的基于像素的擴散是有用的。研究者在表 3 中使用 JiT-B/16 對此進行了檢查。由于研究者使用logit-normal 分布來采樣 t，可以通過改變該分布的參數 μ 來偏移噪聲水平：直觀地說，將 μ 向負側偏移會導致更小的 t，從而增加噪聲水平（公式 (1)）。

表 3 顯示，當模型已經表現不錯時（此處為 x-pred），適當的高噪聲是有益的，這與先前的觀察結果一致。然而，僅調整噪聲水平無法補救 ?-prediction 或 -prediction：它們的失敗本質上源于無法傳播高維信息。

根據表 3，研究者在 ImageNet 256×256 的其他實驗中將 μ 設置為 -0.8。

增加隱藏單元并非必要

既然容量可能受到網絡寬度（即隱藏單元數量）的限制，一個自然的想法是增加它。然而，當觀測維度非常高時，這種補救措施既不符合原則也不可行。研究者表明，在 -prediction 的情況下，這是沒有必要的。

在表 5 和表 6 中，研究者展示了 JiT/32 在分辨率 512 和 JiT/64 在分辨率 1024 下的結果，使用了成比例的大 patch 尺寸 p=32 或 p=64。這相當于每個 patch 為 3072 維（即 32×32×3）或 12288 維，遠大于 B、L 和 H 模型的隱藏層大小。盡管如此，-prediction 仍然運作良好；事實上，除了按比例縮放噪聲（例如，在分辨率 512 和 1024 時分別縮放 2× 和 4×）之外，它無需任何修改即可工作。

這一證據表明，網絡設計可以在很大程度上與觀測維度解耦，就像許多其他神經網絡應用中的情況一樣。增加隱藏單元的數量可能是有益的（如在深度學習中廣泛觀察到的那樣），但這并非決定性因素。

瓶頸結構可能有益

更令人驚訝的是，研究者發現，反過來說，在網絡中引入降低維度的瓶頸結構可能是有益的。

具體來說，研究者將線性 patch 嵌入層轉換為低秩線性層，方法是用一對瓶頸（但仍是線性）層替換它。第一層將維度降低到 d'，第二層將其擴展到 Transformer 的隱藏層大小。這兩層都是線性的，并充當低秩重參數化

圖 4 繪制了使用 JiT-B/16（原始 patch 為 768 維）時 FID 與瓶頸維度 d' 的關系。減小瓶頸維度，即使小到 16 維，也不會導致災難性的失敗。事實上，在大范圍（32 到 512）內的瓶頸維度可以提高質量，FID 的提升幅度高達約 1.3。

從表征學習的更廣闊視角來看，這一觀察并非完全出乎意料。引入瓶頸設計通常是為了鼓勵學習內在的低維表征。

算法

最終算法采用 -prediction 和 -loss，這對應于表 1(3)(a)。形式上，研究者優化：

算法 1 展示了訓練步驟的偽代碼，算法 2 是采樣步驟的偽代碼（歐拉求解器；可擴展到 Heun 或其他求解器）。為簡潔起見，類別條件化和 CFG 被省略了，但兩者都遵循標準做法。為了防止在計算 1/(1-t) 時出現除以零的情況，我們在計算此除法時會截斷其分母（默認值為 0.05）。

更高級的 Transformer

通用 Transformer 的優勢部分在于，當其設計與具體任務解耦時，它可以從其他應用中開發的架構進步中受益。這一特性通過使用任務無關的 Transformer 來構建擴散模型，鞏固了其優勢。

研究者整合了流行的通用改進：SwiGLU、RMSNorm、RoPE、qk-norm，所有這些最初都是為語言模型開發的。

研究者還探索了上下文內類別條件化：但不同于原始 ViT 僅向序列追加一個類別 token，他們追加了多個此類 token（默認為 32 個）。表 4 報告了這些組件的效果。

實驗與對比

圖 5：定性結果。使用 JiT-H/32 在 ImageNet 512×512 上生成的部分示例。

像素空間中的高分辨率生成

在表 5 中，研究者進一步匯報了基礎規模模型（JiT-B）在 ImageNet 上分辨率為 512 以及甚至 1024 時的結果。研究者使用與圖像尺寸成比例的 patch 大小，因此在不同分辨率下序列長度保持不變。單個 patch 的維度可以高達 3072 或 12288，而常見模型都沒有足夠多的隱藏單元來匹配這些維度。

表 5 顯示，模型在不同分辨率下都取得了不錯的性能。所有模型的參數量和計算成本相近，唯一的差別來自輸入 / 輸出 patch 的 embedding。本方法不會遭遇觀測維度災難。

可擴展性

將 Transformer 的結構與任務本身解耦的核心目標之一，是利用其潛在的可擴展能力。表 6 給出了四種模型規模在 ImageNet 256 與 512 分辨率下的結果（注意在 512 分辨率下，這些模型都沒有隱藏單元數量超過 patch 維度）。模型規模和 FLOPs 在表 7 與 8 中列出：本文的模型在 256 分辨率下的成本與在 512 分辨率下的對應模型相似。

有趣的是，隨著模型規模變大，256 與 512 分辨率之間的 FID 差距變得更小。對于 JiT-G，512 分辨率下的 FID 甚至更低。對于在 ImageNet 上的非常大的模型來說，FID 表現在很大程度上取決于過擬合，而 512 分辨率的去噪任務更具挑戰性，使其不太容易過擬合。

以往工作的參考結果

作為參考，研究者們在表 7 與表 8 中與過往結果進行了比較，并標注了各方法所涉及的預訓練組件。與其他像素空間方法相比，本文的方法完全依賴于樸素的、通用型的 Transformer。本文的模型計算友好，并避免了分辨率加倍時常見的二次計算成本增長（參見表 8 中的 FLOPs）。

表7：在 ImageNet 256×256 上的參考結果。評估指標為 5 萬張樣本的 FID 和 IS 。「pre-training」 列列出了獲得這些結果所需的外部預訓練模型。參數量統計包括生成器和 tokenizer 的解碼器（推理階段使用），但不包括其它預訓練組件。Giga-flops 的測量基于一次前向傳播（不包含 tokenizer），并且大致與訓練和推理過程中每次迭代的計算成本成正比。

表 8： 在 ImageNet 512×512 上的參考結果。JiT 采用了更激進的 patch 大小，并能夠在較小的計算量下取得強勁的結果。符號說明與表 7 類似。

本文的方法沒有使用額外損失或預訓練，這些可能會進一步帶來性能提升。

總結

主要集中在概率建模的形式上，而對所使用的神經網絡的能力（以及局限性）關注較少。然而，神經網絡并非具有無限能力，它們更擅長利用自身容量去建模數據，而不是建模噪聲。從這一視角來看，關于-prediction的發現是一個自然的結果。

論文的工作采用了一種極簡且自洽的設計。通過減少特定領域的歸納偏置，研究者希望這種方法能夠推廣到那些難以獲取 tokenizer 的其他領域。這一性質在涉及原始、高維自然數據的科學應用中尤為重要。

研究者設想，這種通用的 「Diffusion + Transformer」 范式，將有潛力成為其他領域的基礎方法之一。

更多信息請參閱原論文。