An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

Technische Zusammenfassung: Eine unbedingte Darstellung des bedingten Scores bei linearen inversen Problemen in unendlichdimensionalen Räumen

Problemstellung
Score-basierte Diffusionsmodelle (SDMs) haben sich zu einem leistungsfähigen Werkzeug für das Sampling aus Posterior-Verteilungen bei bayesschen inversen Problemen entwickelt und bieten flexible Priori-Verteilungen, die herkömmliche, manuell entworfene Regularisierer übertreffen. Allerdings stehen bestehende Methoden für das Posterior-Sampling vor einem kritischen Zielkonflikt zwischen Recheneffizienz und Skalierbarkeit.

1. Korrekturbasierte Methoden (z. B. das Hinzufügen von Likelihood-Gradienten zum Score oder das Erzwingen von Datenkonsistenz) erfordern wiederholte Auswertungen des Vorwärtsoperators $A$ während des Sampling-Prozesses. Dies stellt eine unüberwindbare Rechenbarriere für großskalige, hochdimensionale Probleme dar.
2. Methoden des direkten bedingten Trainings lernen eine Score-Funktion, die auf Messungen konditioniert ist ($s(x, t; \mu_y)$). Zwar vermeidet dies Vorwärtsauswertungen während des Samplings, doch erhöht dies drastisch die Eingabedimensionalität des neuronalen Netzwerks (von $x$ auf $(x, y)$), was umfangreiche Trainingsdaten und komplexe Architekturen erfordert, um den bedingten Score genau zu approximieren, insbesondere in unendlichdimensionalen Settings.

Die Arbeit adressiert die Herausforderung, effizientes und skalierbares Posterior-Sampling bei linearen inversen Problemen in unendlichdimensionalen Räumen durchzuführen, ohne dabei die hohen Online-Kosten für Vorwärtsmodell-Auswertungen oder die hohe Sample-Komplexität des direkten bedingten Trainings in Kauf nehmen zu müssen.

Methodik: Unbedingte Darstellung des bedingten Scores (UCoS)
Die Autoren schlagen UCoS vor, ein Framework, das die Rechenlast der Konditionierung von der Online-Sampling-Phase in eine Offline-Trainingsphase verlagert. Der Kerngedanke besteht darin, dass für lineare inverse Probleme mit Gaußschem Rauschen die bedingte Score-Funktion exakt aus einer unbedingten Score-Funktion eines modifizierten, aufgabenabhängigen Diffusionsprozesses mittels affiner Transformationen abgeleitet werden kann.

• Theoretische Grundlage:

  • Die Methode operiert in einem separablen Hilbertraum $\mathcal{H}$, was Invarianz gegenüber der Diskretisierung sicherstellt.
  • Anstatt ein Netzwerk zu trainieren, das $s(x, t; \mu_y)$ approximiert, definieren die Autoren einen neuen Diffusionsprozess $\tilde{X}_t$, dessen Verteilung eine Mischung aus der Priori-Verteilung und einem vom Vorwärtsoperator $A$ abhängigen Gaußschen Term ist.
  • Sie beweisen (Satz 3.7), dass der wahre bedingte Score $s(x, t; \mu_y)$ mit dem unbedingten Score $\tilde{s}$ dieses neuen Prozesses durch die folgende Identität verknüpft ist:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    wobei $r$ eine transformierte Version des unbedingten Scores ist, $\lambda(t)$ ein zeitabhängiger Skalar und $\xi_t(x, y)$ eine affine Transformation, die die Messung $y$ und den adjungierten Operator $A^*$ beinhaltet.
  • Entscheidend ist, dass der Term $A^*\Gamma^{-1}y$ (wobei $\Gamma$ die Rausch-Kovarianz ist) einmalig vor Beginn des Samplings berechnet werden kann. Der nachfolgende Sampling-Prozess erfordert lediglich die Auswertung des trainierten Netzwerks $r_\theta$ und einfache Vektoradditionen, wodurch die Notwendigkeit weiterer Anwendungen von $A$ oder $A^*$ während der Schritte der inversen Diffusion entfällt.

• Implementierung:

  • Offline-Phase: Ein neuronales Netzwerk (speziell ein Fourier-Neural-Operator, FNO) wird trainiert, um die Funktion $r$ zu approximieren. Das Training nutzt ein Denoising-Score-Matching-Ziel auf dem modifizierten Prozess $\tilde{X}_t$. Diese Phase erfordert Zugriff auf $A$ und $A^*$, jedoch keine spezifischen Messdaten $y$.
  • Online-Phase: Posterior-Proben werden durch Lösen der SDE in umgekehrter Zeit unter Verwendung des transformierten Scores generiert. Der Prozess ist „matrixfrei" und erfordert nur den vorkalkulierten Verschiebungsvektor $A^*\Gamma^{-1}y$ sowie die Netzwerkauswertung.

Hauptbeiträge

1. Exakte Identität in unendlichen Dimensionen: Die Arbeit etabliert eine rigorose theoretische Identität (Satz 3.7), die den bedingten Score über affine Transformationen mit einem aufgabenabhängigen unbedingten Score in unendlichdimensionalen Hilberträumen verbindet. Dies verallgemeinert das Prinzip über spezifische endlichdimensionale Diskretisierungen hinaus.
2. Konvergenzanalyse: Die Autoren liefern eine quantitative Fehlerabschätzung (Satz 4.1) für den generativen Prozess, die die Fehlerbeiträge aus Trainingsverlust, numerischer Diskretisierung und Initialisierung explizit charakterisiert. Dies rechtfertigt die Konvergenz der Methode zum wahren Posterior-Maß.
3. Skalierbarkeit und Effizienz: Durch die Entkopplung der Konditionierung von der Sampling-Schleife erreicht UCoS die Recheneffizienz unbedingter Methoden (keine Vorwärtsauswertungen während des Samplings) bei gleichzeitiger Beibehaltung der Genauigkeit bedingter Methoden. Es vermeidet den Fluch der Dimensionalität, der mit dem Training von Netzwerken auf dem gemeinsamen Raum $(x, y)$ verbunden ist.

Experimentelle Ergebnisse
Die Autoren validieren UCoS an drei inversen Problemen: Bildinpainting, sparse-view Computertomographie (CT) und Bildentwackelung.

• Inpainting (Gaußscher Prior): In einem Setting, bei dem das Posterior Gaußsch ist und Scores in geschlossener Form berechnet werden können, schneidet UCoS vergleichbar mit der exakten bedingten Methode und anderen Baselines ab, was die theoretische Identität bestätigt.
• CT-Bildgebung und Entschärfung (Gelernte Priors): Unter Verwendung der Datensätze LIDC-IDRI bzw. CelebA wird UCoS mit SDE ALD, Diffusion Posterior Sampling (DPS), projektionsbasierten Methoden (Proj) und direktem bedingtem Training verglichen.
  • Genauigkeit: UCoS erzeugt Proben mit einer Verzerrung und Standardabweichung, die mit unbedingten Methoden vergleichbar oder besser sind und die der bedingten Methode bei Verwendung von Architekturen mit geringer Komplexität (z. B. 32 Knoten pro Schicht) deutlich überlegen sind.
  • Effizienz: UCoS erreicht die Sampling-Geschwindigkeit der bedingten Methode (beide vermeiden Vorwärtsauswertungen), übertrifft jedoch korrektionbasierte Methoden (DPS, Proj, SDE ALD), die wiederholte Vorwärtspässe erfordern, deutlich.
  • Robustheit: Die Leistung der bedingten Methode verschlechtert sich bei weniger Netzwerkparametern merklich aufgrund der erhöhten Komplexität der Approximation der gemeinsamen Verteilung, während UCoS auch mit einfacheren Architekturen hohe Qualität beibehält.

Bedeutung und Behauptungen
Die Arbeit behauptet, dass UCoS eine „Best-of-Both-Worlds"-Lösung bietet: Sie eliminiert den Rechenengpass der Vorwärtsmodell-Auswertungen während des Samplings (eine Einschränkung korrektionbasierter Methoden), vermeidet gleichzeitig jedoch die hohe Sample-Komplexität und die architektonischen Anforderungen des direkten bedingten Trainings.

Die Autoren betonen, dass ihr Ansatz diskretisierungsinvariant ist, was ihn für unendlichdimensionale inverse Probleme geeignet macht, bei denen Leistungsgarantien in hochdimensionalen endlichdimensionalen Approximationen oft versagen. Sie räumen ein, dass die Methode aufgabenabhängig ist und ein Neutraining erfordert, wenn sich der Vorwärtsoperator $A$ ändert, und dass die exakte Korrespondenz auf der Linearität des Vorwärtsmodells und der Gaußschheit der Likelihood-Funktion beruht (was den Schritt des „Quadrats vervollständigens" ermöglicht). Die Arbeit wird als Proof of Concept für eine neue Klasse von diffusionsbasierten Samplern präsentiert, die die Komplexität der Konditionierung in die Trainingsphase verlagern, ohne in der theoretischen Formulierung Approximationsfehler zu erzeugen.