Open-sourcing del codice di Knowledge Distillation e dei pesi di SD-Small e SD-Tiny

Open-sourcing del codice di Knowledge Distillation e pesi di SD-Small e SD-Tiny

Recentemente, la comunità di intelligenza artificiale ha assistito a un notevole aumento nello sviluppo di modelli di linguaggio sempre più grandi e performanti, come Falcon 40B, LLaMa-2 70B, Falcon 40B, MPT 30B, e nel dominio dell’immagine con modelli come SD2.1 e SDXL. Questi progressi hanno senza dubbio spinto i limiti di quello che l’IA può realizzare, consentendo capacità di generazione di immagini e comprensione del linguaggio altamente versatili e all’avanguardia. Tuttavia, mentre ci meravigliamo del potere e della complessità di questi modelli, è essenziale riconoscere la crescente necessità di rendere i modelli di intelligenza artificiale più piccoli, efficienti e accessibili, in particolare attraverso la loro distribuzione come open-source.

Da Segmind, abbiamo lavorato su come rendere i modelli di intelligenza artificiale generativi più veloci e più economici. L’anno scorso, abbiamo reso open-source la nostra libreria SD-WebUI accelerata chiamata voltaML, che è una libreria di accelerazione dell’inferenza basata su AITemplate/TensorRT che ha fornito un aumento della velocità di inferenza compreso tra 4 e 6 volte. Per continuare a perseguire l’obiettivo di rendere i modelli generativi più veloci, più piccoli e più economici, stiamo rendendo open-source i pesi e il codice di addestramento dei nostri modelli compressi SD; SD-Small e SD-Tiny. I checkpoint preaddestrati sono disponibili su Huggingface 🤗

Distillazione della Conoscenza

• Generazione pratica di asset 3D Una guida passo passo
• Misurare la creatività dell’IA
• Cigna accusata di respingere impropriamente migliaia di richieste di pazienti utilizzando un algoritmo

I nostri nuovi modelli compressi sono stati addestrati utilizzando le tecniche di distillazione della conoscenza (KD) e il lavoro si basa principalmente su questo articolo. Gli autori descrivono un metodo di distillazione della conoscenza con rimozione di blocchi in cui alcuni dei livelli UNet vengono rimossi e vengono addestrati i pesi del modello studente. Utilizzando i metodi di KD descritti nell’articolo, siamo stati in grado di addestrare due modelli compressi utilizzando la libreria 🧨 diffusers; Small e Tiny, che hanno rispettivamente il 35% e il 55% in meno di parametri rispetto al modello di base, pur raggiungendo una fedeltà delle immagini comparabile a quella del modello di base. Abbiamo reso open-source il nostro codice di distillazione in questa repository e i checkpoint preaddestrati su Huggingface 🤗.

L’addestramento della distillazione della conoscenza di una rete neurale è simile a un insegnante che guida uno studente passo dopo passo. Un grande modello insegnante viene preaddestrato su una grande quantità di dati e quindi un modello più piccolo viene addestrato su un dataset più piccolo, per imitare le uscite del modello più grande insieme all’addestramento classico sul dataset.

In questo particolare tipo di distillazione della conoscenza, il modello studente viene addestrato a eseguire il normale compito di diffusione di ripristinare un’immagine da rumore puro, ma nello stesso tempo, il modello viene fatto corrispondere all’output del modello insegnante più grande. Il confronto degli output avviene ad ogni blocco delle U-net, quindi la qualità del modello è per lo più preservata. Quindi, usando la precedente analogia, possiamo dire che durante questo tipo di distillazione, lo studente cercherà di imparare non solo dalle domande e risposte ma anche dalle risposte dell’insegnante, nonché dal metodo passo dopo passo per arrivare alla risposta. Abbiamo 3 componenti nella funzione di loss per raggiungere questo risultato, in primo luogo la loss tradizionale tra i latenti dell’immagine di destinazione e i latenti dell’immagine generata. In secondo luogo, la loss tra i latenti dell’immagine generata dall’insegnante e i latenti dell’immagine generata dallo studente. E infine, la componente più importante, è la loss a livello di caratteristica, che è la loss tra le uscite di ciascuno dei blocchi dell’insegnante e dello studente.

Combina tutto questo per ottenere l’addestramento della distillazione della conoscenza. Di seguito è riportata un’architettura del UNet con blocchi rimossi utilizzata nella KD come descritto nell’articolo.

Immagine tratta dall’articolo “On Architectural Compression of Text-to-Image Diffusion Models” di Shinkook. et. al

Abbiamo preso come base il nostro modello insegnante Realistic-Vision 4.0 e lo abbiamo addestrato sul dataset LAION Art Aesthetic con punteggi delle immagini superiori a 7.5, a causa delle loro descrizioni di immagini di alta qualità. A differenza dell’articolo, abbiamo scelto di addestrare i due modelli su 1M di immagini per 100K passi per la modalità Small e 125K passi per la modalità Tiny rispettivamente. Il codice per l’addestramento della distillazione può essere trovato qui.

Utilizzo del Modello

Il modello può essere utilizzato utilizzando la DiffusionPipeline di 🧨 diffusers

from diffusers import DiffusionPipeline
import torch

pipeline = DiffusionPipeline.from_pretrained("segmind/small-sd", torch_dtype=torch.float16)
prompt = "Ritratto di una ragazza carina"
negative_prompt = "(iride deformata, pupille deformate, semi-realistiche, cgi, 3d, render, sketch, cartoon, disegno, anime:1.4), testo, primo piano, ritagliato, fuori dal frame, qualità peggiore, qualità bassa, artefatti jpeg, brutto, duplicato, morboso, mutilato, dita extra, mani mutate, mani mal disegnate, volto mutato, deforme, sfocato, disidratato, anatomia sbagliata, proporzioni sbagliate, arti extra, volto clonato, proporzioni brutte, arti malformati, braccia mancanti, gambe mancanti, braccia extra, gambe extra, dita fuse, troppe dita, collo lungo"
image = pipeline(prompt, negative_prompt = negative_prompt).images[0]
image.save("mia_immagine.png")

Velocità in termini di latenza di inferenza

Abbiamo osservato che i modelli distillati sono fino al 100% più veloci dei modelli di base originali. Il codice di benchmarking può essere trovato qui.

Limitazioni potenziali

I modelli distillati sono ancora nella fase iniziale e le loro uscite potrebbero non essere ancora di qualità produttiva. Questi modelli potrebbero non essere i migliori modelli generali. Sono meglio utilizzati come modelli sintonizzati o LoRA addestrati su concetti/stili specifici. I modelli distillati non sono ancora molto adatti alla composizione o ai multiconcetti.

Sintonizzazione fine del modello SD-tiny sul dataset di ritratti

Abbiamo sintonizzato il nostro modello sd-tiny sulle immagini di ritratti generate con il modello Realistic Vision v4.0. Di seguito sono riportati i parametri di sintonizzazione fine utilizzati.

• Passaggi: 131000
• Tasso di apprendimento: 1e-4
• Dimensione del batch: 32
• Passaggi di accumulo del gradiente: 4
• Risoluzione dell’immagine: 768
• Dimensione del dataset – 7k immagini
• Precisione mista: fp16

Siamo stati in grado di produrre una qualità delle immagini simile alle immagini prodotte dal modello originale, con quasi il 40% in meno di parametri e i risultati campione qui sotto parlano da soli:

Il codice per la sintonizzazione fine dei modelli di base può essere trovato qui.

Addestramento LoRA

Uno dei vantaggi dell’addestramento LoRA su un modello distillato è l’addestramento più veloce. Di seguito ci sono alcune immagini del primo LoRA che abbiamo addestrato sul modello distillato su alcuni concetti astratti. Il codice per l’addestramento LoRA può essere trovato qui.

Conclusione

Invitiamo la comunità open-source ad aiutarci a migliorare e ottenere una più ampia adozione di questi modelli SD distillati. Gli utenti possono unirsi al nostro server Discord, dove annunceremo gli ultimi aggiornamenti su questi modelli, rilasceremo più checkpoint e nuovi entusiasmanti LoRA. E se ti piace il nostro lavoro, per favore dacci una stella su Github.