Likelihood Matching for Diffusion Models

Technische Samenvatting: Likelihood Matching voor Diffusiemodellen

Probleemstelling
Huidige mainstream diffusiemodellen, met name Score-Based Generative Models (SGM's), vertrouwen op Score Matching (SM) als hun primaire trainingsdoelstelling. Hoewel dit empirisch effectief is, minimaliseert SM de discrepantie tussen een geparametriseerde scorefunctie en de ware scorefunctie bij verschillende ruisniveaus. Theoretisch gezien minimaliseert deze aanpak slechts een bovengrens op de negatieve log-likelihood van de oorspronkelijke datadistributie, in plaats van de likelihood zelf. Deze indirecte verbinding kan leiden tot een verlies aan statistische efficiëntie vergeleken met Maximum Likelihood Estimation (MLE), wat bekend staat om het produceren van consistente en efficiënte schatters. Bovendien richt standaard SM zich uitsluitend op informatie van de eerste-orde momenten (de score/gradiënt), waardoor hogere-orde structurele informatie die aanwezig is in de transitiedichtheden mogelijk wordt verwaarloosd.

Methodologie
De auteurs stellen Likelihood Matching (LM) voor, een raamwerk dat direct mikt op de likelihood van de datadistributie door gebruik te maken van de equivalentie tussen de likelihood van de doeldata en de pad-likelihood van het reverse diffusieproces.

1. Theoretische Fundering:
   De methode is gebaseerd op de eigenschap dat de gezamenlijke likelihood van het reverse diffusiepad equivalent is aan de likelihood van de oorspronkelijke dataverdeling (op constante factoren na die gerelateerd zijn aan het forward proces). De auteurs stellen vast dat het maximaliseren van de exacte pad-likelihood van het reverse proces equivalent is aan het maximaliseren van de data log-likelihood.

2. Quasi-Maximum Likelihood Estimation (QMLE):
   Omdat de exacte reverse transitiedichtheden onhandelbaar zijn, benaderen de auteurs deze met een Quasi-Maximum Likelihood-aanpak. Zij modelleren elke reverse transitiedichtheid als een Gaussische distributie waarbij de conditionele gemiddelde en covariantie worden gematcht aan het ware reverse proces.

   • Kerninzicht: Het conditionele gemiddelde en de covariantie van het reverse proces hangen niet alleen af van de scorefunctie ($\nabla \log q_t$), maar ook van de Hessian-functie ($\nabla^2 \log q_t$).
   • Parametrisatie: De auteurs parametriseren zowel de score als de Hessian met behulp van neurale netwerken (specifiek U-Nets). Om de computationele complexiteit te beheersen, wordt de Hessian vaak benaderd met een diagonaal-plus-low-rank structuur.

3. De Likelihood Matching Doelstelling:
   De voorgestelde doelstelling, Likelihood Matching (LM), wordt afgeleid door het te minimaliseren van de negatieve quasi-log-likelihood van de reverse transities.

   • De doelstelling verenigt score matching en covariantie matching.
   • Het incorporeert een impliciet likelihood-wegingsmechanisme via de gematchte covariantie-matrix $(I + \sigma^2 H_t)^{-1/2}$, wat de weging van de score matching-term automatisch aanpast op basis van lokale kromming.
   • De training maakt gebruik van een stochastische sampler die willekeurig tijdstappen langs het diffusietraject selecteert om de loss te berekenen, wat zorgt voor een efficiënt gebruik van de evolutionaire informatie van de forward SDE.

4. Stochastische Sampler:
   Een nieuwe samplingsprocedure wordt geïntroduceerd die zowel de geleerde score- als de Hessian-informatie gebruikt. De reverse stap wordt gedefinieerd als:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   waarbij het gemiddelde $\hat{\mu}$ en de covariantie $\hat{\Sigma}$ functies zijn van respectievelijk de geschatte score en de Hessian.

Belangrijkste Bijdragen

• Nieuwe Trainingsdoelstelling: Het artikel introduceert een quasi-likelihood-gebaseerde trainingsdoelstelling (LM) die de reverse pad log-likelihood benadert, waardoor score matching en covariantie matching effectief worden gecombineerd met impliciete likelihood-weging.
• Hessian-Verrijkte Sampling: Een stochastische sampler is afgeleid die zowel de score- als de Hessian-informatie benut, wat een nauwkeuriger benadering biedt van de reverse transitiedichtheid dan standaard DDPM-type samplers.
• Theoretische Garanties:
  • Convergentie: Er worden non-asymptotische convergentiegaranties gegeven in totale variatieafstand. De foutmarge schaalt als $O(d^3 \log^{4.5} T / T)$ voor de reverse stap-fout, met extra termen voor score- en Hessian-schattingsfouten.
  • Consistentie: Het artikel bewijst de consistentie van de voorgestelde quasi-MLE onder een oracle parametrische setting, waarbij wordt aangetoond dat de schatter convergeert naar de ware dataparameters naarmate de steekproefomvang en de diffusiestappen toenemen.
• Efficiënte Implementatie: Het gebruik van low-rank Hessian-benaderingen en de Sherman-Morrison-Woodbury-formule maakt een efficiënte berekening van de likelihood en de samplingstappen mogelijk.

Experimentele Resultaten
De auteurs valideren de methode op zowel synthetische als real-world datasets:

• Synthetische Data: Op Gaussian en t-distributie mengsels presteerde LM consequent beter dan standaard Score Matching (SM), met name bij het vastleggen van goed gescheiden modi. Op een sterk gekrulde Swiss-roll distributie (die een gekromde manifold representeert), demonstreerde LM superieure prestaties, wat het voordeel van de Hessian-geïnformeerde covariantie valideert.
• Parameterestimatie: In een parametrische Gaussian mixture setting bereikte LM een lagere Mean Absolute Error (MAE) en standaardfout in parameterestimatie vergeleken met SM, wat de claims over theoretische consistentie ondersteunt.
• Beelddatasets (MNIST, CIFAR-10, CelebA):
  • LM behaalde een lagere Fréchet Inception Distance (FID) en betere Negative Log-Likelihood (NLL) scores vergeleken met SM-baselines.
  • De prestaties verbeterden met de rang van de Hessian-benadering ($r$), met een piek rond $r=20-30$.
  • Speed-Quality Trade-off: LM behaalde een lagere FID voor vergelijkbare wall-clock tijd en vereiste minder reverse sampling stappen om een hoge visuele kwaliteit te bereiken. Bijvoorbeeld, op MNIST produceerde LM herkenbare cijfers na slechts 20 reverse iteraties, terwijl SM-samples wazig bleven.

Betekenis en Claims
Het artikel claimt dat Likelihood Matching een statistisch meer rigoureus alternatief biedt voor Score Matching door direct de datalikelihood te targeten via een quasi-likelihood raamwerk. Door tweede-orde informatie (de Hessian) te incorporeren, legt de methode rijkere transitiedichtheden vast en verbetert het de sampling-efficiëntie, met name in lage-stap regimes en op complexe, gekromde datamanifolds. De auteurs benadrukken dat hoewel LM hogere computationele kosten per stap met zich meebrengt door de Hessian-modellering, de reductie in het aantal vereiste sampling stappen en de verbetering in generatiekwaliteit (FID en NLL) de effectiviteit ervan valideren voor het trainen van hoogwaardige generatieve modellen. Het werk overbrugt de kloof tussen de statistische efficiëntie van MLE en het praktische succes van diffusiemodellen.