一臺16GB的4090筆記本，僅需0.37秒，直接吐出1024×1024像素圖片。

如此神速AI生圖工具，竟是出自英偉達(dá)MIT清華全華人團(tuán)隊(duì)之筆！

正如其名字一樣，Sana能以驚人速度合成高分辨率、高質(zhì)量，且具有強(qiáng)文本-圖像對其能力的模型。

而且，它還能高效生成高達(dá)4096×4096像素的圖像。

項(xiàng)目主頁：https://nvlabs.github.io/Sana/

論文地址：https://arxiv.org/abs/2410.10629

Sana的核心設(shè)計(jì)包含了以下幾個(gè)要素：

• 深度壓縮自編碼器（AE）：傳統(tǒng)自編碼器只能將圖像壓縮8倍，全新AE可將圖像壓縮32倍，有效減少了潛在token的數(shù)量。
• 線性DiT（Diffusion Transformer）：用「線性注意力」替換了DiT中所有的普通注意力，在高分辨率下更加高效，且不會(huì)犧牲質(zhì)量。
• 基于僅解碼器模型的文本編碼器：用現(xiàn)代的僅解碼器SLM替換T5作為文本編碼器，并設(shè)計(jì)了復(fù)雜的人類指令，通過上下文學(xué)習(xí)來增強(qiáng)圖像-文本對齊。
• 高效的訓(xùn)練和采樣：提出Flow-DPM-Solver來減少采樣步驟，并通過高效的標(biāo)題標(biāo)注和選擇來加速收斂。

基于以上的算法創(chuàng)新，相較于領(lǐng)先擴(kuò)散模型Flux-12B，Sana-0.6B不僅參數(shù)小12倍，重要的是吞吐量飆升100倍。

以后，低成本的內(nèi)容創(chuàng)作，Sana才堪稱這一領(lǐng)域的王者。

效果一覽

一只賽博貓，和一個(gè)帶有「SANA」字樣的霓虹燈牌。

一位站在山頂上的巫師，在夜空中施展魔法，形成了由彩色能量組成的「NV」字樣。

在人物的生成方面，Sana對小女孩面部的描繪可以說是非常地細(xì)致了。

下面來看個(gè)更復(fù)雜的：

一艘海盜船被困在宇宙漩渦星云中，通過模擬宇宙海灘旋渦的特效引擎渲染，呈現(xiàn)出令人驚嘆的立體光效。場景中彌漫著壯麗的環(huán)境光和光污染，營造出電影般的氛圍。整幅作品采用新藝術(shù)風(fēng)格，由藝術(shù)家SenseiJaye創(chuàng)作的插畫藝術(shù)，充滿精致細(xì)節(jié)。

甚至，像下面這種超級復(fù)雜的提示，Sana也能get到其中的關(guān)鍵信息，并生成相應(yīng)的元素和風(fēng)格。

Prompt：a stunning and luxurious bedroom carved into a rocky mountainside seamlessly blending nature with modern design with a plush earth-toned bed textured stone walls circular fireplace massive uniquely shaped window framing snow-capped mountains dense forests, tranquil mountain retreat offering breathtaking views of alpine landscape wooden floors soft rugs rustic sophisticated charm, cozy tranquil peaceful relaxing perfect escape unwind connect with nature, soothing intimate elegance modern design raw beauty of nature harmonious blend captivating view enchanting inviting space, soft ambient lighting warm hues indirect lighting natural daylight balanced inviting glow

順便，團(tuán)隊(duì)還給經(jīng)典梗圖，生成了一個(gè)卡通版變體（右）。

設(shè)計(jì)細(xì)節(jié)

Sana的核心組件，已在開頭簡要給出介紹。接下來，將更進(jìn)一步展開它們實(shí)現(xiàn)的細(xì)節(jié)。

模型架構(gòu)的細(xì)節(jié)，如下表所示。

- 深度壓縮自編碼器

研究人員引入的全新自編碼器（AE），大幅將縮放因子提高至32倍。

過去，主流的AE將圖像的長度和寬度，只能壓縮8倍（AE-F8）。

與AE-F8相比，AE-F32輸出的潛在token數(shù)量減少了16倍，這對于高效訓(xùn)練和生成超高分辨率圖像（如4K分辨率）至關(guān)重要。

- 高效線性DiT（Diffusion Transformer）

原始DiT的自注意力計(jì)算復(fù)雜度為O(N2)，在處理高分辨率圖像時(shí)呈二次增長。

線性DiT在此替換了傳統(tǒng)的二次注意力機(jī)制，將計(jì)算復(fù)雜度從O(N2)降低到O(N)。

與此同時(shí)，研究人員還提出了Mix-FFN，可以在多層感知器（MLP）中使用3×3深度卷積，增強(qiáng)了token的局部信息。

實(shí)驗(yàn)結(jié)果顯示，線性注意力達(dá)到了與傳統(tǒng)注意力相當(dāng)?shù)慕Y(jié)果，在4K圖像生成方面將延遲縮短了1.7倍。

此外，Mix-FFN無需位置編碼（NoPE）就能保持生成質(zhì)量，成為首個(gè)不使用位置嵌入的DiT。

- 基于僅解碼器「小語言模型」的文本編碼器

這里，研究人員使用了Gemma（僅解碼器LLM）作為文本編碼器，以增強(qiáng)對提示詞的理解和推理能力。

盡管T2I生成模型多年來取得了顯著進(jìn)展，但大多數(shù)現(xiàn)有模型仍依賴CLIP或T5進(jìn)行文本編碼，這些模型往往缺乏強(qiáng)大的文本理解和指令跟隨能力。

與CLIP或T5不同，Gemma提供了更優(yōu)的文本理解和指令跟隨能力，由此解訓(xùn)練了不穩(wěn)定的問題。

他們還設(shè)計(jì)了復(fù)雜人類指令（CHI），來利用Gemma強(qiáng)大指令跟隨、上下文學(xué)習(xí)和推理能力，改善了圖像-文本對齊。

在速度相近的情況下，Gemma-2B模型比T5-large性能更好，與更大更慢的T5-XXL性能相當(dāng)。

- 高效訓(xùn)練和推理策略

另外，研究人員還提出了一套自動(dòng)標(biāo)注和訓(xùn)練策略，以提高文本和圖像之間的一致性。

首先，對于每張圖像，利用多個(gè)視覺語言模型（VLM）生成重新描述。盡管這些VLM的能力各不相同，但它們的互補(bǔ)優(yōu)勢提高了描述的多樣性。

此外，他們還提出了一種基于clipscore的訓(xùn)練策略，根據(jù)概率動(dòng)態(tài)選擇與圖像對應(yīng)的多個(gè)描述中具有高clip分?jǐn)?shù)的描述。

實(shí)驗(yàn)表明，這種方法改善了訓(xùn)練收斂和文本-圖像對齊能力。

此外，與廣泛使用的Flow-Euler-Solver相比，團(tuán)隊(duì)提出的Flow-DPM-Solver將推理采樣步驟從28-50步顯著減少到14-20步，同時(shí)還能獲得更優(yōu)的結(jié)果。

整體性能

如下表1中，將Sana與當(dāng)前最先進(jìn)的文本生成圖像擴(kuò)散模型進(jìn)行了比較。

對于512×512分辨率：- Sana-0.6的吞吐量比具有相似模型大小的PixArt-Σ快5倍- 在FID、Clip Score、GenEval和DPG-Bench等方面，Sana-0.6顯著優(yōu)于PixArt-Σ

對于1024×1024分辨率：- Sana比大多數(shù)參數(shù)量少于3B的模型性能強(qiáng)得多- 在推理延遲方面表現(xiàn)尤為出色

與最先進(jìn)的大型模型FLUX-dev的比較：- 在DPG-Bench上，準(zhǔn)確率相當(dāng)- 在GenEval上，性能略低- 然而，Sana-0.6B的吞吐量快39倍，Sana-1.6B快23倍

Sana-0.6吞吐量，要比當(dāng)前最先進(jìn)4096x4096圖像生成方法Flux，快100倍。

而在1024×1024分辨率下，Sana的吞吐量要快40倍。

如下是，Sana-1.6B與其他模型可視化性能比較。很顯然，Sana模型生成速度更快，質(zhì)量更高。

終端設(shè)備部署

為了增強(qiáng)邊緣部署，研究人員使用8位整數(shù)對模型進(jìn)行量化。

而且，他們還在CUDA C++中實(shí)現(xiàn)了W8A8 GEMM內(nèi)核，并采用內(nèi)核融合技術(shù)來減少不必要的激活加載和存儲帶來的開銷，從而提高整體性能。

如下表5所示，研究人員在消費(fèi)級4090上部署優(yōu)化前后模型的結(jié)果比較。

在生成1024x1024圖像方面，優(yōu)化后模型實(shí)現(xiàn)了2.4倍加速，僅用0.37秒就生成了同等高質(zhì)量圖像。

作者介紹

Enze Xie（謝恩澤）

共同一作Enze Xie是NVIDIA Research的高級研究科學(xué)家，隸屬于由麻省理工學(xué)院的Song Han教授領(lǐng)導(dǎo)的高效AI團(tuán)隊(duì)。此前，曾在華為諾亞方舟實(shí)驗(yàn)室（香港）AI理論實(shí)驗(yàn)室擔(dān)任高級研究員和生成式AI研究主管。

他于2022年在香港大學(xué)計(jì)算機(jī)科學(xué)系獲得博士學(xué)位，導(dǎo)師是Ping Luo教授，聯(lián)合導(dǎo)師是Wenping Wang教授。并于朋友Wenhai Wang密切合作。

在攻讀博士學(xué)習(xí)期間，他與阿德萊德大學(xué)的Chunhua Shen教授、加州理工學(xué)院的Anima Anandkumar教授以及多倫多大學(xué)的Sanja Fidler教授共事。同時(shí)，還與Facebook和NVIDIA等業(yè)界的多位研究人員進(jìn)行了合作。

他的研究方向是高效的AIGC/LLM/VLM，并在實(shí)例級檢測和自監(jiān)督/半監(jiān)督/弱監(jiān)督學(xué)習(xí)領(lǐng)域做了一些工作——開發(fā)了多個(gè)CV領(lǐng)域非常知名的算法，以及一個(gè)2000多星的自監(jiān)督學(xué)習(xí)框架OpenSelfSup（現(xiàn)名為mmselfsup）。

- PolarMask（CVPR 2020十大影響力論文排名第十）

- PVT（ICCV 2021十大影響力論文排名第二）

- SegFormer（NeurIPS 2021十大影響力論文排名第三）

- BEVFormer（ECCV 2022十大影響力論文排名第六）

Junsong Chen

共同一作Junsong Chen是NVIDIA Research的研究實(shí)習(xí)生，由Enze Xie博士和Song Han教授指導(dǎo)。同時(shí)，他也是大連理工大學(xué)IIAU實(shí)驗(yàn)室的博士生，導(dǎo)師是Huchuan Lu教授。

他的研究領(lǐng)域是生成式AI和機(jī)器學(xué)習(xí)的交叉，特別是深度學(xué)習(xí)及其應(yīng)用的算法與系統(tǒng)協(xié)同設(shè)計(jì)。

此前，他曾在香港大學(xué)擔(dān)任研究助理，由Ping Luo教授的指導(dǎo)。

Song Han（韓松）

Song Han是MIT電氣工程與計(jì)算機(jī)科學(xué)系的副教授。此前，他在斯坦福大學(xué)獲得博士學(xué)位。

他提出了包括剪枝和量化在內(nèi)廣泛用于高效AI計(jì)算的「深度壓縮」技術(shù)，以及首次將權(quán)重稀疏性引入現(xiàn)代AI芯片的「高效推理引擎」——ISCA 50年歷史上引用次數(shù)最多的前五篇論文之一。

他開創(chuàng)了TinyML研究，將深度學(xué)習(xí)引入物聯(lián)網(wǎng)設(shè)備，實(shí)現(xiàn)邊緣學(xué)習(xí)。

他的團(tuán)隊(duì)在硬件感知神經(jīng)架構(gòu)搜索方面的工作使用戶能夠設(shè)計(jì)、優(yōu)化、縮小和部署 AI 模型到資源受限的硬件設(shè)備，在多個(gè)AI頂會(huì)的低功耗計(jì)算機(jī)視覺比賽中獲得第一名。

最近，團(tuán)隊(duì)在大語言模型量化/加速（SmoothQuant、AWQ、StreamingLLM）方面的工作，有效提高了LLM推理的效率，并被NVIDIA TensorRT-LLM采用。

Song Han憑借著在「深度壓縮」方面的貢獻(xiàn)獲得了ICLR和FPGA的最佳論文獎(jiǎng)，并被MIT Technology Review評選為「35歲以下科技創(chuàng)新35人」。與此同時(shí)，他在「加速機(jī)器學(xué)習(xí)的高效算法和硬件」方面的研究，則獲得了NSF CAREER獎(jiǎng)、IEEE「AIs 10 to Watch: The Future of AI」獎(jiǎng)和斯隆研究獎(jiǎng)學(xué)金。

他是DeePhi（被AMD收購）的聯(lián)合創(chuàng)始人，也是OmniML（被NVIDIA收購）的聯(lián)合創(chuàng)始人。