最近，人工智能社區(qū)在開發(fā)更大、更高性能的語(yǔ)言模型方面取得了顯著的進(jìn)展，例如 Falcon 40B、LLaMa-2 70B、Falcon 40B、MPT 30B; 以及在圖像領(lǐng)域的模型，如 SD2.1 和 SDXL 。這些進(jìn)步無(wú)疑推動(dòng)了人工智能的發(fā)展，使其具有高度多功能和最先進(jìn)的圖像生成和語(yǔ)言理解能力。然而，在我們驚嘆于這些模型的強(qiáng)大和復(fù)雜性之余，必須認(rèn)識(shí)到一個(gè)日益增長(zhǎng)的需求: 使人工智能模型體量更小、運(yùn)行更高效、更易于訪問(wèn)，特別是通過(guò)開源它們來(lái)共建生態(tài)。

在 Segmind，我們一直致力于如何使生成式 AI 更快、更便宜。去年，我們開源了我們加速的 SD-WebUI 庫(kù) voltaML，它是一個(gè)基于 AITemplate/TensorRT 的推理加速庫(kù)，推理速度提高了 4-6 倍。為了繼續(xù)實(shí)現(xiàn)使生成模型更快、更小、更便宜的目標(biāo)，我們正在開源我們壓縮的?SD 模型：SD-Small 和 SD-Tiny?的權(quán)重和訓(xùn)練代碼。預(yù)訓(xùn)練的檢查點(diǎn)可在 Hugging Face ?? 上獲取。

知識(shí)蒸餾

我們的新壓縮模型已經(jīng)經(jīng)過(guò)知識(shí)蒸餾 (KD) 技術(shù)的訓(xùn)練，這項(xiàng)工作主要基于 這篇論文。作者描述了一種塊移除知識(shí)蒸餾方法，其中一些 UNet 層被移除，學(xué)生模型權(quán)重被訓(xùn)練。使用論文中描述的 KD 方法，我們能夠使用 ?? diffusers 庫(kù)訓(xùn)練兩個(gè)壓縮模型;?Small (微小版本)?和?Tiny (極小版本)，分別比基礎(chǔ)模型少 35% 和 55% 的參數(shù)，同時(shí)實(shí)現(xiàn)與基礎(chǔ)模型相當(dāng)?shù)膱D像保真度。我們已經(jīng)在這個(gè) repo 中開源了我們的蒸餾代碼，并將預(yù)訓(xùn)練檢查點(diǎn)上傳到了 Hugging Face ??。

知識(shí)蒸餾訓(xùn)練神經(jīng)網(wǎng)絡(luò)類似于老師一步一步指導(dǎo)學(xué)生。一個(gè)大的老師模型 (teacher model) 預(yù)先在大量數(shù)據(jù)上訓(xùn)練，然后一個(gè)較小的模型在較小的數(shù)據(jù)集上訓(xùn)練，以模仿大模型的輸出并在數(shù)據(jù)集上進(jìn)行經(jīng)典訓(xùn)練。

在這種特殊類型的知識(shí)蒸餾中，學(xué)生模型被訓(xùn)練來(lái)完成從純?cè)肼暬謴?fù)圖像的正常擴(kuò)散任務(wù)，但同時(shí)，模型被迫與更大的老師模型的輸出匹配。輸出匹配發(fā)生在 U-nets 的每個(gè)塊，因此模型質(zhì)量基本保持不變。所以，使用前面的類比，我們可以說(shuō)，在這種蒸餾過(guò)程中，學(xué)生不僅會(huì)試圖從問(wèn)題和答案中學(xué)習(xí)，還會(huì)從老師的答案以及逐步得到答案的方法中學(xué)習(xí)。我們?cè)趽p失函數(shù)中有 3 個(gè)組成部分來(lái)實(shí)現(xiàn)這一點(diǎn)，首先是目標(biāo)圖像隱變量和生成圖像隱變量之間的傳統(tǒng)損失。其次是老師生成的圖像隱變量和學(xué)生生成的圖像隱變量之間的損失。最后，也是最重要的組成部分，是特征級(jí)損失，即老師和學(xué)生每個(gè)塊輸出之間的損失。

結(jié)合所有這些構(gòu)成了知識(shí)蒸餾訓(xùn)練。下面是論文中描述的用于 KD 的塊移除 UNet 架構(gòu)。

圖片來(lái)自 Shinkook 等人的 論文 “On Architectural Compression of Text-to-Image Diffusion Models”。

我們以 Realistic-Vision 4.0 為基礎(chǔ)老師模型，并在LAION Art Aesthetic 數(shù)據(jù)集 上訓(xùn)練，圖像分?jǐn)?shù)高于 7.5，因?yàn)樗鼈兙哂懈哔|(zhì)量的圖像描述。與論文不同，我們選擇分別為 Small 和 Tiny 模式訓(xùn)練兩個(gè)模型，分別在 1M 張圖像上進(jìn)行 100K 步和 125K 步的訓(xùn)練。蒸餾訓(xùn)練的代碼可以在 這里 找到。

模型使用

模型可以通過(guò) ?? diffusers 中的 DiffusionPipeline 來(lái)使用。

推理延遲方面的速度表現(xiàn)

我們觀察到，蒸餾模型比原始基礎(chǔ)模型快了一倍?；鶞?zhǔn)測(cè)試代碼可以在 這里 找到。

潛在的局限性

蒸餾模型處于早期階段，輸出可能還不具備生產(chǎn)水平的質(zhì)量。這些模型可能不是最好的通用模型，它們最好用作針對(duì)特定概念/風(fēng)格進(jìn)行微調(diào)或 LoRA 訓(xùn)練。蒸餾模型目前還不太擅長(zhǎng)組合性或多概念。

在人像數(shù)據(jù)集上微調(diào) SD-tiny 模型

我們已經(jīng)在 Realistic Vision v4.0 模型生成的人像圖像上微調(diào)了我們的 sd-tiny 模型。下面是使用的微調(diào)參數(shù)。

原版參數(shù)中文釋義Steps: 131000步數(shù): 131000Learning rate: 1e-4學(xué)習(xí)率: 1e-4Batch size: 32批量大小: 32Gradient accumulation steps: 4梯度累積步數(shù): 4Image resolution: 768圖像分辨率: 768Dataset size: 7k images數(shù)據(jù)集大小: 7 千張圖像Mixed precision: fp16混合精度: fp16

我們能夠產(chǎn)生接近原始模型產(chǎn)生的圖像質(zhì)量，參數(shù)減少了近 40%，下面的樣本結(jié)果不言自明:

微調(diào)基礎(chǔ)模型的代碼可以在 這里 找到。

LoRA 訓(xùn)練

在蒸餾模型上進(jìn)行 LoRA 訓(xùn)練的一個(gè)優(yōu)點(diǎn)是訓(xùn)練更快。下面是我們?cè)谡麴s模型上對(duì)一些抽象概念進(jìn)行的第一個(gè) LoRA 訓(xùn)練的一些圖像。LoRA 訓(xùn)練的代碼可以在 這里 找到。

結(jié)論

我們邀請(qǐng)開源社區(qū)幫助我們改進(jìn)并實(shí)現(xiàn)這些蒸餾 SD 模型的更廣泛采用。用戶可以加入我們的 Discord 服務(wù)器，在那里我們將宣布這些模型的最新更新，發(fā)布更多的檢查點(diǎn)和一些令人興奮的新 LoRAs。如果你喜歡我們的工作，請(qǐng)?jiān)谖覀兊?Github 上點(diǎn)一下 star。

英文原文:?https://hf.co/blog/sd_distillation

原文作者: Yatharth Gupta

譯者: innovation64

審校/排版: zhongdongy (阿東)