開源 SD-Small 和 SD-Tiny 的知識蒸餾程式碼和權重

釋出日期：2023年8月1日

訪客

近期，AI 社群見證了更大、效能更強的語言模型（如 Falcon 40B、LLaMa-2 70B、Falcon 40B、MPT 30B）以及影像領域模型（如 SD2.1 和 SDXL）的顯著發展。這些進步無疑推動了 AI 能力的邊界，實現了高度通用和最先進的影像生成和語言理解能力。然而，當我們驚歎於這些模型的強大和複雜性時，也必須認識到對 AI 模型進行小型化、高效化和更易於訪問（尤其是透過開源）的需求日益增長。

在 Segmind，我們一直在致力於如何使生成式 AI 模型更快、更便宜。去年，我們開源了加速 SD-WebUI 庫，名為 voltaML，它是一個基於 AITemplate/TensorRT 的推理加速庫，已將推理速度提高了 4-6 倍。為了繼續實現使生成式模型更快、更小、更便宜的目標，我們正在開源我們壓縮的 SD 模型 SD-Small 和 SD-Tiny 的權重和訓練程式碼。預訓練檢查點可在 Huggingface 🤗 上獲取。

知識蒸餾

我們的新壓縮模型已採用知識蒸餾（KD）技術進行訓練，這項工作主要基於這篇論文。作者描述了一種塊移除知識蒸餾方法，其中一些 UNet 層被移除，然後訓練學生模型權重。利用論文中描述的 KD 方法，我們能夠使用 🧨 diffusers 庫訓練出兩個壓縮模型：Small 和 Tiny，它們比基礎模型分別減少了 35% 和 55% 的引數，同時實現了與基礎模型相當的影像保真度。我們已將蒸餾程式碼開源到此儲存庫，並將預訓練檢查點發布到 Huggingface 🤗。

訓練神經網路的知識蒸餾類似於老師一步步引導學生。一個大型教師模型在大資料上進行預訓練，然後一個較小的模型在較小的資料集上進行訓練，以模仿大型模型的輸出，並結合資料集上的傳統訓練。

在這種特定的知識蒸餾型別中，學生模型被訓練來執行從純噪聲中恢復影像的正常擴散任務，但同時，模型被要求匹配大型教師模型的輸出。輸出匹配發生在 U-net 的每個塊中，因此模型質量得以大部分保留。因此，使用之前的類比，我們可以說，在這種蒸餾過程中，學生不僅會嘗試從問題和答案中學習，還會從教師的答案以及獲取答案的循序漸進的方法中學習。我們損失函式中有 3 個組成部分來實現這一點：首先是目標影像的潛在變數與生成影像的潛在變數之間的傳統損失。其次是教師生成的影像的潛在變數與學生生成的影像的潛在變數之間的損失。最後，也是最重要的組成部分是特徵級別損失，即教師和學生每個塊的輸出之間的損失。

所有這些結合起來構成了知識蒸餾訓練。下面是論文中描述的 KD 中使用的塊移除 UNet 的架構。

圖片摘自 Shinkook 等人的論文《文字到影像擴散模型的架構壓縮》（“On Architectural Compression of Text-to-Image Diffusion Models”）

我們以 Realistic-Vision 4.0 作為我們的基礎教師模型，並在 LAION Art Aesthetic 資料集（影像評分高於 7.5 分，因其高質量影像描述）上進行訓練。與論文不同，我們選擇在 100 萬張影像上分別對 Small 和 Tiny 模型訓練 10 萬步和 12.5 萬步。蒸餾訓練程式碼可以在這裡找到。

模型使用

該模型可使用 🧨 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")