Unified Latents (UL): How to train your latents

Die Arbeit stellt Unified Latents (UL) vor, ein Framework zur gemeinsamen Regularisierung latenter Darstellungen durch einen Diffusionsprior und einen Diffusionsdecoder, das auf ImageNet-512 und Kinetics-600 neue State-of-the-Art-Ergebnisse bei gleichzeitig reduzierter Rechenkomplexität erzielt.

Das Wesentliche

Stell dir vor, du möchtest ein riesiges, detailliertes Ölgemälde (ein Foto) in eine kleine, handliche Schatulle packen, um es später wieder herauszuholen und perfekt zu rekonstruieren. Das ist im Grunde das Problem, das sich KI-Modelle stellen müssen, wenn sie Bilder oder Videos generieren.

Die Forscher von Google DeepMind haben mit ihrer Methode „Unified Latents" (UL) einen neuen, cleveren Weg gefunden, diese Schatulle zu füllen. Hier ist die Erklärung in einfachen Worten, mit ein paar anschaulichen Vergleichen:

1. Das Problem: Die „zu dicke" oder „zu dünne" Schatulle

Bisher gab es zwei Hauptprobleme beim Packen von Bildern in eine KI:

• Die alte Methode (VAE): Man packte das Bild in eine Schatulle, aber die Regeln dafür waren etwas willkürlich. Es war wie ein Spiel, bei dem man nicht genau wusste, wie viele Informationen man weglassen durfte. Oft verlor das Bild dabei wichtige Details (wie feine Textur oder Haare), oder die KI hatte Mühe, das Bild später wiederherzustellen.
• Die neue Herausforderung: Man möchte die Schatulle so klein wie möglich halten (wenig Daten), damit die KI sie leicht verarbeiten kann, aber sie muss trotzdem genug Informationen enthalten, damit das Bild später scharf aussieht.

2. Die Lösung: Der „Doppel-Check" mit dem Diffusions-Modell

Die Autoren schlagen vor, die Schatulle nicht einfach nur zu füllen, sondern sie mit einem zweiten, strengen Prüfer zu versehen.

Stell dir den Prozess so vor:

1. Der Packmeister (Encoder): Er nimmt das Originalbild und packt es in die Schatulle. Aber er macht es nicht perfekt sauber. Er wirft ein wenig „Rauschen" (Störgeräusch) hinein, genau wie bei einem alten Radio, das leicht knistert.
2. Der Prüfer (Diffusion Prior): Hier kommt der Clou. Dieser Prüfer ist wie ein perfekter Restaurator, der gelernt hat, wie ein „sauberer" Schatullen-Inhalt aussehen sollte. Er versucht, das Bild aus dem Rauschen in der Schatulle wiederherzustellen.
   • Wenn der Packmeister zu viel Information weggeworfen hat, scheitert der Prüfer.
   • Wenn er zu viel Information hineingepackt hat, ist die Schatulle unnötig groß.
   • Die Magie: Der Prüfer zwingt den Packmeister, genau die richtige Menge an Information zu hinterlassen. Es ist wie ein Tanz, bei dem beide Partner (Packmeister und Prüfer) gleichzeitig lernen, wie man die Schatulle optimal füllt.

3. Der „Decoder": Der Künstler, der das Bild malt

Sobald die Schatulle perfekt gefüllt ist, kommt der Decoder (der Künstler) ins Spiel.

• In der alten Welt malte der Künstler oft nur grobe Umrisse und ließ die Details im Rauschen.
• Bei der neuen Methode (Unified Latents) bekommt der Künstler die Schatulle mit dem „geprüften" Inhalt. Er ist ein sehr talentierter Maler (ein Diffusions-Modell), der die feinen Details aus dem Inhalt der Schatulle und dem leichten Rauschen herausarbeiten kann.
• Der Trick: Der Künstler darf sich auf die feinen Details konzentrieren, weil der Prüfer (der Prior) schon dafür gesorgt hat, dass die grobe Struktur und die wichtigsten Informationen in der Schatulle sind.

4. Warum ist das besser? (Die Analogie des „Lernens")

Stell dir vor, du lernst für eine Prüfung:

• Früher: Du hast versucht, alles auswendig zu lernen (hohe Datenmenge), aber das war schwer und ineffizient. Oder du hast nur Stichpunkte gelernt (niedrige Datenmenge), aber dann hast du Details vergessen.
• Jetzt (Unified Latents): Du hast einen Tutor (den Prior), der dir sagt: „Du musst nur diese 50 wichtigsten Fakten lernen, und ich garantiere dir, dass du damit die Prüfung bestehst." Der Rest wird vom Lehrer (dem Decoder) ergänzt.
  • Das Ergebnis: Du lernst schneller (weniger Rechenleistung nötig) und bestehst die Prüfung trotzdem mit einer sehr guten Note (hervorragende Bildqualität).

Was haben sie erreicht?

• Bessere Bilder: Auf dem berühmten Datensatz ImageNet haben sie Bilder generiert, die so gut aussehen wie die besten bisherigen Modelle (FID von 1,4 – das ist extrem niedrig, also sehr nah am Original).
• Schnelleres Lernen: Sie brauchen weniger Rechenleistung, um diese Qualität zu erreichen, als andere Modelle.
• Videos: Auch bei Videos (Kinetics-600) haben sie einen neuen Weltrekord aufgestellt.

Zusammenfassung in einem Satz

Unified Latents ist wie ein perfekt abgestimmtes Team aus einem Packmeister, einem strengen Prüfer und einem Künstler, die gemeinsam lernen, wie man ein Bild in eine winzige Schatulle packt, ohne dabei auch nur ein einziges wichtiges Detail zu verlieren – und das alles viel effizienter als bisherige Methoden.

Es ist der erste Schritt zu einer Welt, in der KI nicht nur „schöne Bilder" macht, sondern diese Bilder auch effizient und präzise versteht und speichert.

Technische Zusammenfassung

1. Problemstellung

Diffusionsmodelle haben sich als äußerst erfolgreich für die Generierung von Bildern, Videos und Audio erwiesen. Ein Schlüsselfaktor für diesen Erfolg sind latente Darstellungen (Latents), die kompakte Kodierungen darstellen und es Diffusionsmodellen ermöglichen, effizienter auf höhere Auflösungen zu skalieren.

Trotzdem bleibt die Frage offen, wie solche Latents optimal gelernt werden sollten:

• Herausforderung bei VAEs: Herkömmliche Latent Diffusion Models (wie Stable Diffusion) verwenden oft einen VAE-Ansatz mit einer KL-Strafe (Kullback-Leibler-Divergenz) gegenüber einer Standard-Gauß-Verteilung. Da der Decoder hier keine likelihood-basierte Verlustfunktion hat, muss das Gewicht des KL-Terms manuell gesetzt werden. Dies erschwert die Kontrolle über den Informationsgehalt der Latents.
• Herausforderung bei semantischen Latents: Ansätze, die auf vortrainierten Netzwerken (z. B. DINO) oder stark regularisierten Autoencodern basieren, liefern oft gute Generierungsergebnisse (niedriger FID), verlieren aber jedoch hochfrequente Informationen, was zu schlechteren Rekonstruktionen (niedriger PSNR) und Artefakten führt.
• Der Trade-off: Es besteht ein fundamentaler Zielkonflikt zwischen dem Informationsgehalt des Latents (Bitrate) und der Rekonstruktionsqualität. Ein einfacher zu lernendes Latent führt zu besserer Generierung, während ein informationsreiches Latent eine bessere Rekonstruktion ermöglicht, aber schwerer zu modellieren ist.

Das Paper zielt darauf ab, diesen Trade-off systematisch zu steuern und eine Methode zu entwickeln, die sowohl effizientes Training als auch hohe Rekonstruktionsqualität ermöglicht.

2. Methodik: Unified Latents (UL)

Das vorgestellte Framework Unified Latents trainiert Encoder, Prior und Decoder gemeinsam, wobei alle Komponenten durch Diffusionsmodelle reguliert werden. Die Kernidee besteht darin, den Encoder-Ausgangsrauschen direkt mit dem minimalen Rauschniveau des Priors zu verknüpfen.

Das System besteht aus drei Hauptkomponenten:

A. Kodierung und Prior (Encoder & Diffusion Prior)

• Deterministischer Encoder: Anstatt eine komplexe Verteilung (Mittelwert und Varianz) zu lernen, sagt der Encoder eine deterministische latente Darstellung $z_{clean}$ vorher.
• Feste Rausch-Verknüpfung: Um die Information zu regulieren, wird $z_{clean}$ mit einem festen Betrag an Gauß-Rauschen versehen, um $z_0$ zu erhalten. Dies entspricht einem festen log-SNR-Wert ($\lambda(0) = 5$).
• Diffusion Prior: Ein Diffusionsmodell lernt, diesen leicht verrauschten Latent $z_0$ aus reinem Rauschen ($z_1$) zu rekonstruieren.
• Vorteil: Durch diese Verknüpfung reduziert sich der KL-Term der VAE-Loss-Funktion auf einen einfachen, gewichteten MSE über die Rauschniveaus. Dies liefert eine enge obere Schranke für die Latent-Bitrate und macht den Informationsgehalt interpretierbar.

B. Dekodierung (Diffusion Decoder)

• Der Decoder ist ebenfalls ein Diffusionsmodell, das im Bildraum operiert und auf dem Latent $z_0$ sowie dem verrauschten Bild $x_t$ konditioniert.
• Gewichtung (Reweighted ELBO): Im Gegensatz zum Prior, der eine ungewichtete ELBO verwendet (um den Encoder nicht zu täuschen), nutzt der Decoder eine sigmoid-gewichtete ELBO.
• Loss Factor: Um das Problem des „Posterior Collapse" (wo der Decoder die Latents ignoriert) zu vermeiden, wird der Decoder-Loss mit einem Faktor (Loss Factor, z. B. 1.3–1.7) hochgewichtet. Dies erzwingt, dass der Encoder relevante Informationen liefert.
• Steuerung: Die Hyperparameter (Loss Factor und Sigmoid-Bias) steuern direkt den Kompromiss zwischen Rekonstruktionsqualität und Modellierungskomplexität.

C. Training in zwei Stufen

1. Stufe 1 (Co-Training): Encoder, Prior und Decoder werden gemeinsam trainiert. Der Prior regularisiert den Encoder, und der Decoder rekonstruiert das Bild.
2. Stufe 2 (Base Model Training): Der Encoder und Decoder werden eingefroren. Der Prior wird als „Base Model" neu trainiert, jedoch mit einer sigmoid-gewichteten Loss-Funktion (statt der ungewichteten ELBO aus Stufe 1), um bessere Stichproben zu erzeugen. Dies ermöglicht größere Batch-Größen und Modelle.

3. Schlüsselbeiträge

• Unified Framework: Eine Methode, bei der Encoder, Prior und Decoder alle Diffusionsmodelle sind und gemeinsam optimiert werden.
• Interpretierbare Bitrate: Durch die Verknüpfung von Encoder-Rauschen und Prior-Präzision entsteht eine klare, theoretisch fundierte Obergrenze für die Bits pro Dimension (bpd) im Latent-Raum.
• Einfache Hyperparameter: Der Trade-off zwischen Rekonstruktion und Generierung wird durch einfache Hyperparameter (Loss Factor, Sigmoid-Bias) gesteuert, anstatt durch manuelle KL-Gewichtung.
• Stabilität: Die Verwendung eines deterministischen Encoders mit festem Rauschen vermeidet die Instabilitäten, die bei der Lernung komplexer Encoder-Verteilungen (wie in LSGM) auftreten.

4. Ergebnisse

Die Methode wurde auf ImageNet-512 (Bilder) und Kinetics-600 (Videos) evaluiert.

• Bildgenerierung (ImageNet-512):

  • Unified Latents erreichen einen FID von 1.4 bei hoher Rekonstruktionsqualität (PSNR).
  • Im Vergleich zu Modellen, die auf Stable-Diffusion-Latents basieren, benötigt UL weniger Trainings-FLOPs für vergleichbare oder bessere Ergebnisse.
  • Es zeigt sich, dass kleinere Modelle von niedrigeren Bitraten profitieren, während größere Modelle auch höhere Bitraten effektiv nutzen können.

• Videogenerierung (Kinetics-600):

  • Das Modell setzt einen neuen State-of-the-Art (SOTA) FVD-Wert von 1.3 (für das mittlere Modell).
  • Auch hier übertrifft UL den Trade-off zwischen Trainingskosten und Generierungsqualität (FVD) bestehende Ansätze wie Video Diffusion, MAGVIT oder W.A.L.T.

• Text-zu-Bild:

  • Auf großen Text-zu-Bild-Datensätzen zeigen die Modelle eine überlegene Wahrnehmungsqualität (gFID) und eine leicht bessere Textausrichtung (CLIP-Score) im Vergleich zu Pixel-Diffusion und Stable Diffusion.

• Ablationsstudien:

  • Die Verwendung eines festen Rauschpegels im Encoder ist entscheidend für Stabilität und Leistung.
  • Ein reiner VAE-Ansatz mit normaler Prior-Verteilung führt zu schlechteren Ergebnissen, da er höhere Bitraten erfordert, die schwerer zu modellieren sind.
  • Die sigmoid-Gewichtung im Decoder ist essenziell, um hochfrequente Details zu erhalten, ohne die Latent-Modellierung zu überlasten.

5. Bedeutung und Fazit

Das Paper „Unified Latents" bietet einen prinzipiellen Ansatz zur Gestaltung latenter Räume für Diffusionsmodelle.

• Effizienz: Es ermöglicht das Training von Generationsmodellen mit weniger Rechenaufwand bei höherer Qualität.
• Kontrolle: Es macht den Trade-off zwischen Informationsgehalt und Modellierungskomplexität explizit und steuerbar.
• Skalierbarkeit: Die Ergebnisse deuten darauf hin, dass größere Basis-Modelle von informativeren Latents profitieren, was neue Skalierungsgesetze für zukünftige Foundation Models eröffnet.

Obwohl Diffusions-Decodierer aktuell rechenintensiver beim Sampling sind als GAN-basierte Decoder, überwiegt der Vorteil der effizienteren Trainingsphase und der überlegenen generativen Qualität. Das Framework ist zudem potenziell auf andere Datenarten (z. B. Text) übertragbar.