開源 SD-Small 和 SD-Tiny 知識(shí)蒸餾代碼與權(quán)重
最近，人工智能社區(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)行更高效、更易于訪問，特別是通過開源它們來共建生態(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)過知識(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)練來完成從純?cè)肼暬謴?fù)圖像的正常擴(kuò)散任務(wù)，但同時(shí)，模型被迫與更大的老師模型的輸出匹配。輸出匹配發(fā)生在 U-nets 的每個(gè)塊，因此模型質(zhì)量基本保持不變。所以，使用前面的類比，我們可以說，在這種蒸餾過程中，學(xué)生不僅會(huì)試圖從問題和答案中學(xué)習(xí)，還會(huì)從老師的答案以及逐步得到答案的方法中學(xué)習(xí)。我們?cè)趽p失函數(shù)中有 3 個(gè)組成部分來實(shí)現(xiàn)這一點(diǎn)，首先是目標(biāo)圖像隱變量和生成圖像隱變量之間的傳統(tǒng)損失。其次是老師生成的圖像隱變量和學(xué)生生成的圖像隱變量之間的損失。最后，也是最重要的組成部分，是特征級(jí)損失，即老師和學(xué)生每個(gè)塊輸出之間的損失。
結(jié)合所有這些構(gòu)成了知識(shí)蒸餾訓(xùn)練。下面是論文中描述的用于 KD 的塊移除 UNet 架構(gòu)。
圖片來自 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)練的代碼可以在 這里 找到。
模型使用
模型可以通過 ?? diffusers 中的 DiffusionPipeline 來使用。
from diffusers import DiffusionPipeline
import torch
pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float16)
prompt = "Portrait of a pretty girl"
negative_prompt = "(deformed iris, deformed pupils, semi-realistic, cgi, 3d, render, sketch, cartoon, drawing, anime:1.4), text, close up, cropped, out of frame, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck"
image = pipeline(prompt, negative_prompt = negative_prompt).images[0]
image.save("my_image.png")
推理延遲方面的速度表現(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ù): 131000 |
| Learning rate: 1e-4 | 學(xué)習(xí)率: 1e-4 |
| Batch size: 32 | 批量大小: 32 |
| Gradient accumulation steps: 4 | 梯度累積步數(shù): 4 |
| Image resolution: 768 | 圖像分辨率: 768 |
| Dataset 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 (阿東)