引言：高分辨率圖像合成的挑戰(zhàn)與進展
在人工智能領域应闯，高分辨率圖像的合成一直是一個極具挑戰(zhàn)性的課題。隨著技術(shù)的不斷進步，我們已經(jīng)能夠通過各種模型纬霞，如穩(wěn)定擴散模型（Stable Diffusion）募逞、Imgen蛋铆、eDiff-I、DALL·E 和 Emu放接，從簡單的文本提示中生成引人入勝的圖像和設計刺啦。這些模型在多樣化的開放網(wǎng)絡數(shù)據(jù)上進行訓練，展現(xiàn)出了令人驚嘆的能力纠脾。然而玛瘸，當涉及到生成高分辨率圖像時，巨大的計算成本成為了一個障礙苟蹈，尤其是在互動應用中糊渊，高分辨率導致的延遲問題尤為突出。
為了解決這一問題慧脱，研究人員提出了各種加速擴散模型推理的方法渺绒，包括減少采樣步驟和優(yōu)化神經(jīng)網(wǎng)絡推理。盡管計算資源的快速增長使得利用多個GPU加速推理變得可行，但在生成單個圖像時芒篷，通常只涉及到一個GPU的使用搜变，這限制了加速的潛力。因此针炉，開發(fā)一種能夠利用多個GPU加速單圖像生成的方法變得尤為重要挠他。
在這篇論文中，我們提出了一種名為DistriFusion的方法篡帕，它通過在多個GPU上并行運行擴散模型來減少單樣本生成的延遲殖侵，而不會損害圖像質(zhì)量。DistriFusion的核心思想是利用相鄰擴散步驟之間的輸入相似性镰烧，通過重用前一時間步的預計算特征圖來為當前步驟提供上下文拢军。這種方法支持異步通信，可以通過計算來流水線化怔鳖。廣泛的實驗表明茉唉，我們的方法可以應用于最新的Stable Diffusion XL，無需降低質(zhì)量结执，且在八個NVIDIA A100 GPU上相比單個GPU可實現(xiàn)高達6.1倍的加速度陆。我們的代碼已在GitHub上公開。
論文標題: DistriFusion: Distributed Parallel Inference for High-Resolution Diffusion Models
項目地址: https://github.com/mit-han-lab/distrifuser
公眾號【AI論文解讀】后臺回復“論文解讀” 獲取論文PDF!
DistriFusion方法概述：分布式并行推理的動機與基本原理
隨著AI生成內(nèi)容（AIGC）的興起献幔，如Adobe Firefly懂傀、Midjourney和Sora等工具展示了從簡單文本提示生成引人入勝的圖像和設計的驚人能力。這些成就得益于擴散模型的進步蜡感，這些模型通過迭代去噪過程合成內(nèi)容蹬蚁。然而，生成高分辨率圖像仍然具有挑戰(zhàn)性郑兴，因為它需要巨大的計算成本犀斋，導致交互式應用程序的延遲過高。為了解決這個問題杈笔，我們提出了DistriFusion方法，通過利用多個GPU上的并行性來降低單樣本生成的延遲蒙具，同時不損害圖像質(zhì)量球榆。DistriFusion的核心思想是將模型輸入分割成多個補丁持钉，并將每個補丁分配給一個GPU。然而篱昔，簡單地實現(xiàn)這樣的算法會破壞補丁之間的交互并且失去保真度始腾，而引入這種交互將會產(chǎn)生巨大的通信開銷空执。為了克服這個困境，我們觀察到相鄰擴散步驟之間的輸入具有高度相似性辨绊，并提出了位移補丁并行主義奶栖，它利用擴散過程的順序性質(zhì)，通過重用前一個時間步的預計算特征圖為當前步驟提供上下文门坷。因此宣鄙，我們的方法支持異步通信，可以通過計算進行流水線處理默蚌。廣泛的實驗表明冻晤，我們的方法可以應用于最新的Stable Diffusion XL，無質(zhì)量退化绸吸，并在八個NVIDIA A100s上與一個相比實現(xiàn)高達6.1倍的加速鼻弧。
實現(xiàn)高效并行的關(guān)鍵：位移補丁并行主義
1. 位移補丁并行主義的定義與工作原理
位移補丁并行主義是一種新的并行化機會，它基于對擴散模型順序特性的觀察锦茁，通過重用前一步驟的特征作為當前步驟的上下文來重疊通信和計算温数。在DistriFusion中，我們首先將輸入圖像分割成多個補丁蜻势，每個補丁分配給不同的設備。對于每個設備鹉胖，我們使用從上一步驟預計算的激活來提供全局上下文和補丁交互握玛。我們進一步共同設計了一個推理框架來實現(xiàn)我們的算法。具體來說甫菠，我們的框架通過異步通信有效地隱藏了通信開銷挠铲，并且僅在分配的區(qū)域上稀疏地運行卷積和注意力層，從而按比例減少了每個設備的計算量寂诱。
2. 通過異步通信優(yōu)化計算與通信的重疊
在DistriFusion中拂苹，每個設備都有一個模型副本，并將獨立地痰洒、并行地處理單個補丁瓢棒。每個設備上的補丁首先被分散到上一步驟的陳舊激活中，然后選擇性地應用層操作（線性丘喻、卷積或注意力層）到這些新鮮區(qū)域脯宿，從而生成相應區(qū)域的輸出。這個過程對每一層重復泉粉。最后连霉，所有層的輸出被同步在一起榴芳，以近似?θ(xt)。通過這種方法跺撼，每個設備只負責總計算量的1/N窟感，實現(xiàn)了高效的并行化。
為了獲得上一步驟的陳舊激活歉井，每個設備在每個時間步驟t獲取Al,(i)柿祈，然后將激活廣播到所有其他設備并執(zhí)行AllGather操作〉玻現(xiàn)代GPU通常支持異步通信和計算憨募，這意味著AllGather過程不會阻塞正在進行的計算。到達下一個時間步驟的層l時菜谣，每個設備應該已經(jīng)接收到Al t的副本尾膊。這種方法有效地將通信開銷隱藏在計算階段內(nèi)。然而待笑，有一個例外：第一步（即xT）暮蹂。在這種情況下癌压，每個設備簡單地執(zhí)行標準的同步通信，并緩存中間激活以備下一步使用集侯。
稀疏操作與校正異步GroupNorm
1. 稀疏操作在減少每個設備計算量中的作用
稀疏操作是DistriFusion方法中的一個關(guān)鍵技術(shù)棠枉，它通過在每個設備上僅對新鮮區(qū)域執(zhí)行計算來減少每個設備的計算量术健。在這種設置中粘衬，如果Fl是卷積、線性或跨注意力層勘伺，操作符僅應用于新鮮區(qū)域，而不是完整的特征圖冲茸。對于自注意力層轴术，它被轉(zhuǎn)換為跨注意力層钦无，只保留來自新鮮區(qū)域的查詢令牌失暂，而鍵和值令牌則覆蓋整個特征圖。這樣凭峡，F(xiàn)l的計算成本與新鮮區(qū)域的大小成正比摧冀，從而實現(xiàn)了稀疏操作系宫。
2. 校正異步GroupNorm以保持圖像質(zhì)量
在DistriFusion中笙瑟，為了保持圖像質(zhì)量往枷，采用了校正異步GroupNorm技術(shù)凄杯。由于擴散模型通常采用組歸一化（GN）層戒突，需要聚合激活以恢復其完整的空間形狀。研究發(fā)現(xiàn)导而，僅歸一化新鮮補丁或重用陳舊特征會降低圖像質(zhì)量。為了解決這一問題韵丑，引入了一個校正項來近似全局均值和方差撵彻。這種方法通過異步通信計算精確的統(tǒng)計數(shù)據(jù)实牡，從而將通信成本管道化到計算中。在極少數(shù)情況下碗短，如果近似方差為負豪椿，則回退使用新鮮補丁的局部方差携栋。
實驗設置與基線比較
1. 實驗模型與數(shù)據(jù)集介紹
實驗中主要使用了最先進的公開文本到圖像模型Stable Diffusion XL（SDXL）。SDXL首先使用預訓練的自動編碼器將圖像壓縮為原始大小的1/8鸯隅，然后在這個潛在空間中應用擴散模型蝌以。此外，SDXL還采用了多個跨注意力層以促進文本條件跟畅。與原始的Stable Diffusion相比徊件，SDXL采用了更多的注意力層蒜危，導致模型計算更加密集辐赞。實驗使用了COCO Captions 2014數(shù)據(jù)集的HuggingFace版本進行基準測試。
2. 與現(xiàn)有方法的比較基準
DistriFusion與以下基線進行了比較：Na?ve Patch方法將輸入圖像分成多個補丁新思，每個補丁獨立處理，然后拼接在一起窄绒，但由于缺乏補丁間的交互崔兴，導致明顯的邊界縫隙。ParaDiGMS技術(shù)通過并行降噪多個步驟來加速預訓練的擴散模型位谋，使用Picard迭代來猜測未來步驟的解掏父，并迭代細化直至收斂赊淑。實驗使用了PSNR、LPIPS和FID等標準指標來評估圖像質(zhì)量陶缺。DistriFusion在保持與原始模型相當?shù)膱D像質(zhì)量的同時饱岸，實現(xiàn)了與使用設備數(shù)量成比例的延遲減少苫费。
主要結(jié)果：質(zhì)量與效率的評估
1. DistriFusion在圖像質(zhì)量保持上的表現(xiàn)
DistriFusion是一種新穎的并行計算方法双抽，它通過將圖像分割成多個補丁并在多個GPU上并行處理，旨在降低單樣本生成的延遲琅翻，同時不損害圖像質(zhì)量。實驗表明钧嘶，DistriFusion能夠在不降低視覺保真度的前提下琳疏，復制原始Stable Diffusion XL (SDXL)的性能闸拿。在生成1024×1024圖像的設置中新荤，即使使用8個設備台汇，DistriFusion也能夠在PSNR苟呐、LPIPS和FID評分上與原始模型的結(jié)果相媲美牵素。
2. 在不同分辨率下的延遲減少與加速比
DistriFusion在多個NVIDIA A100 GPUs上的性能測試顯示，與單個GPU相比请琳，使用2俄精、4和8個GPU時劣坊，生成單張圖像的延遲分別降低了1.8倍、3.4倍和6.1倍测蘑。當分辨率提高到2048×2048和3840×3840時碳胳，GPU設備的利用率得到了更好的提升挨约，從而進一步減少了延遲产雹。特別是在生成3840×3840圖像時蔓挖，DistriFusion實現(xiàn)了顯著的加速，達到了使用2怨绣、4和8個A100s分別減少1.8倍篮撑、3.4倍和6.1倍的延遲。
消融研究：驗證設計選擇的有效性
1. 通信成本分析與比較
在使用8個A100 GPUs進行的通信成本比較中未蝌，與同步張量并行（Sync. TP）和同步補丁并行（Sync. PP）相比树埠，DistriFusion進一步減少了50%至60%的通信開銷怎憋。這一結(jié)果表明绊袋，DistriFusion通過異步通信和計算重疊铸鹰，有效地隱藏了通信開銷蹋笼。
2. 輸入相似性的量化驗證
DistriFusion的設計基于一個關(guān)鍵觀察：相鄰去噪步驟的輸入具有高度相似性剖毯。通過量化計算使用50步DDIM采樣器的所有連續(xù)步驟的模型輸入差異逊谋，平均差異僅為0.02（輸入范圍在[-4, 4]內(nèi)，約占0.3%）胶滋。這一發(fā)現(xiàn)支持了DistriFusion使用來自前一步驟的稍微過時的激活來促進補丁間交互的假設究恤。
3. 少步采樣與預熱步驟的影響
在使用10步DPM-Solver的配置中部宿，DistriFusion在沒有預熱步驟的情況下難以保持圖像質(zhì)量。然而，通過增加額外的兩步預熱步驟棍掐，顯著提高了性能拷况，同時僅略微增加了延遲赚瘦。這表明預熱步驟對于維持圖像質(zhì)量，尤其是在減少步驟數(shù)的采樣器中鹰服，是非常重要的悲酷。
關(guān)注公眾號 AI論文解讀
獲取最新AI論文解讀
