Technisches Resümee: Qualitätsbewusste Modulation für Diffusion Transformer

Problemstellung

Moderne Text-zu-Bild-Diffusionsmodelle, insbesondere Diffusion Transformer (DiT) wie Stable Diffusion 3.5 (SD3.5), verlassen sich auf Konditionierungssignale – speziell auf Zeitschritt-Embeddings und gepoolte Prompt-Embeddings –, um den Denoising-Prozess via Adaptive Layer Normalization (AdaLN) zu modulieren. Während diese Signale das aktuelle Rauschniveau und die semantische Intention effektiv kommunizieren, fehlt ihnen eine qualitätsbewusste Information. Infolgedessen kann der Modulationsmechanismus nicht direkt auf samplespezifische Fehlermodi reagieren, wie etwa Fehlerakkumulation, falsche Interpretation des Prompts oder halluzinierte Merkmale, die sich im Verlauf des latenten Bildes entwickeln. Die Basis-Konditionierung ist unabhängig vom aktuellen latenten Zustand ($\partial M_{base} / \partial x_t = 0$), was verhindert, dass das Modell auf Fidelity-Probleme basierend auf der tatsächlichen Qualität des Zwischenbildes reagieren kann.

Methodik

Die Autoren schlagen das Quality Representation Module (QRM) vor, einen leichtgewichtigen Transformer-basierten Adapter, der darauf ausgelegt ist, qualitätsbewusste Signale in den Diffusionsprozess zu injizieren, ohne den vortrainierten Backbone zu modifizieren.

Architektur

QRM ist als Transformer-Decoder implementiert, der additive Modulations-Offsets ($M_{qrm}$) zu den Basis-AdaLN-Parametern ($M_{base}$) vorhersagt.

• Eingaben: Das Modul erhält das aktuelle latente Bild ($x_t$), das gepoolte Prompt-Embedding ($p$), das Zeitschritt-Embedding ($t$) und die Basis-Modulationsparameter ($M_{base}$).
• Verarbeitung:
  1. Das latente Bild $x_t$ wird mittels einer gelernten $1\times1$-Konvolution und adaptivem Pooling in latente Tokens projiziert.
  2. Die Basis-Modulationsparameter werden normalisiert und projiziert, um Query-Tokens ($Z_0$) zu bilden.
  3. Ein Key-Value-Speicher ($KV$) wird aus den latenten Tokens, dem Prompt und den Zeitschritt-Embeddings konstruiert.
  4. Durch Cross-Attention- und Self-Attention-Layer attendieren die Query-Tokens auf den $KV$-Speicher, um eine verfeinerte Repräsentation zu generieren.
  5. Der Output wird zurück auf die native AdaLN-Breite projiziert, um den Offset-Vektor $M_{qrm}$ zu erzeugen.
• Anwendung: Die finale Modulation, die auf die Transformer-Blöcke angewendet wird, ist $M_t = M_{base,t} + M_{qrm,t}$.

Trainingsstrategie (ReFL)

QRM wird unter Verwendung einer Reward Feedback Learning (ReFL) Strategie trainiert:

• Gefrorene Komponenten: Der Diffusions-Backbone (Denoiser, Text-Encoder, VAE), der Sampler und das Reward-Modell bleiben eingefroren. Nur die QRM-Parameter werden aktualisiert.
• Zielsetzung: Das Modul lernt, einen skalaren Reward-Score zu maximieren, der von einem eingefrorenen Reward-Modell (z. B. HPSv2.1 oder CLIPScore) abgeleitet wird, welches auf ein aus einem Zwischenschritt des Denoising-Prozesses dekodiertes Bild angewendet wird.
• Loss-Funktionen: Das Paper nutzt entweder einen Reward-Maximierungs-Loss oder einen Margin-Seeking-Loss, welcher das QRM-modulierte Bild dazu drängt, das Baseline-Modell um eine spezifische Differenz zu übertreffen.
• Zeitschritt-Selektion: Um die Effizienz und den Fokus auf die strukturelle Formung zu verbessern, wird QRM nur während einer vordefinierten Teilmenge von High-Noise-Zeitschritten (z. B. $t \in \{50, \dots, 26\}$) aktiviert, anstatt über den gesamten Diffusions-Zeitplan hinweg.

Inferenz

Während der Inferenz fungiert das trainierte QRM-Modul als Plug-in-Adapter. Es sagt Modulations-Offsets basierend auf dem aktuellen latenten Zustand voraus und addiert diese zu den Basis-Parametern. Dies ermöglicht es dem Modell, die Denoising-Trajektorie hin zu Zuständen höherer Qualität zu steuern, ohne den Backbone neu zu trainieren oder den Sampling-Zeitplan zu ändern.

Kernbeiträge

1. QRM-Architektur: Einführung eines leichtgewichtigen Transformer-Moduls, das aus bestehenden Modell-Inputs (Latent, Prompt, Zeitschritt) qualitätsbewusste Repräsentationen lernt, um die Modulationsparameter zu verfeinern.
2. Nicht-invasive Integration: Ein Mechanismus zur Erweiterung der Diffusion-Transformer-Modulation mit latenten Bildqualitäts-Signalen, ohne den zugrunde liegenden Diffusions-Backbone, den VAE oder den Sampling-Zeitplan zu verändern.
3. Reward-gesteuerte Modulation: Demonstration, dass vorhergesagte Modulations-Updates ($M_{qrm}$) den Denoising-Prozess in jedem Zeitschritt effektiv verfeinern können, was die Prompt-Adhärenz und die Bild-Fidelity verbessert.
4. Umfassende Evaluierung: Umfangreiche Ablationsstudien, welche die Auswirkungen der QRM-Architekturgröße, der aktiven Zeitschrittbereiche und der Trainings-Reward-Metriken auf die Bildqualität analysieren.

Experimentelle Ergebnisse

Die Autoren evaluierten QRM mit Stable Diffusion 3.5 (SD3.5) auf dem Parti-Prompts Benchmark.

• Quantitative Leistung:
  • Im Vergleich zur SD3.5-Baseline erreichte QRM, trainiert mit HPSv2.1 als Reward-Signal, die höchsten Scores unter allen evaluierten Methoden (einschließlich Prompt-Refinement, Fine-Tuning und ReNeg-Ansätzen).
  • Verbesserungen: +0,20 im Aesthetic Score und +0,82 im HPSv2.1 Score.
  • Wenn mit CLIPScore trainiert wurde, verbesserte QRM die menschliche Präferenz-Ausrichtung (HPSv2.1), während es eine vergleichbare perzeptuelle Qualität beibehielt.
• Qualitative Analyse:
  • Visualisierungen der Denoising-Trajektorien zeigen, dass die Anwendung von QRM während der frühen, hochrauschbehafteten Zeitschritte dazu führt, dass erkennbare semantische Strukturen signifikant früher entstehen als in der Baseline.
  • Die Anwendung von QRM nur in mittleren oder späten Phasen lieferte geringere Verbesserungen, was bestätigt, dass die globale Struktur früh im Diffusionsprozess etabliert wird.
• Ablationsbefunde:
  • Kapazität: Größere QRM-Modelle optimierten stärker auf den spezifischen Trainings-Reward (z. B. HPSv2.1), zeigten jedoch leichte Trade-offs in anderen Metriken (z. B. CLIPScore).
  • Zeitschritt-Bereich: Die Modulation während früher Zeitschritte ($50 \to 26$) war konsistent effektiver für die Ausrichtung als die Modulation in späteren Phasen.
  • Reward-Metrik: Die Wahl des Reward-Modells (CLIPScore vs. HPSv2.1) prägte direkt die durch QRM erlernten perzeptuellen Biases.

Bedeutung und Ansprüche

Das Paper behauptet, dass Modulationsebene-Interventionen einen skalierbaren und effektiven Weg bieten, um große Diffusionsmodelle mit menschlichen Präferenz-Signalen auszurichten. Im Gegensatz zu Methoden, die ein teures Fine-Tuning des gesamten Backbones erfordern (z. B. LoRA oder Full Fine-Tuning), erreicht QRM konsistente Qualitätssteigerungen, indem es vollständig durch additive Modulations-Updates operiert, die auf der sich entwickelnden latenten Repräsentation konditioniert sind.

Die Autoren positionieren QRM als einen flexiblen, qualitätsbewussten Adapter, der durch verschiedene Reward-Signale gesteuert und auf verschiedene Diffusion-Transformer-Architekturen angewendet werden kann. Sie räumen ein, dass QRM die Biases der Trainings-Reward-Modelle erbt, und schlagen vor, dass zukünftige Arbeiten Multi-Reward-Training und adaptive Zeitschritt-Selektion explorieren könnten, um Robustheit und Generalisierung weiter zu erhöhen. Die Arbeit zeigt, dass signifikante Gewinne in der Bild-Fidelity und Prompt-Ausrichtung erzielt werden können, ohne die vortrainierten Gewichte von State-of-the-Art-Modellen zu verändern.