Tensor Train Diffusion: Leveraging Low-Rank Structures for High-Dimensional Score-Based Sampling

Technische Samenvatting: Tensor Train Diffusion

Probleemstelling
Het artikel behandelt de uitdaging van het bemonsteren van complexe, hoogdimensionale waarschijnlijkheidsdensiteiten $p_{\text{target}} = \rho_{\text{target}}/Z$, waarbij de ongenormaliseerde densiteit $\rho_{\text{target}}$ puntgewijs evalueerbaar is, maar de normaliserende constante $Z$ onhandelbaar is en geen monsters beschikbaar zijn van $p_{\text{target}}$. Dit is een centraal probleem in computationele wetenschappen, waaronder moleculaire dynamica en Bayesiaanse statistiek.

Hoewel diffusiemodellen een krachtig kader bieden voor dergelijke taken door een ruisend stochastisch differentiaalvergelijking (SDE) te reverseren, vereisen ze doorgaans het leren van een scorefunctie van monsters via denoising score matching. In deze specifieke setting moet de score direct worden afgeleid van de ongenormaliseerde densiteit zonder toegang tot monsters. De tijd-omgekeerde dynamica wordt beheerst door een Hamilton-Jacobi-Bellman (HJB) partiële differentiaalvergelijking (PDE). Bestaande methoden voor het oplossen van deze hoogdimensionale PDE, zoals Physics-Informed Neural Networks (PINNs) of trajectgebaseerde technieken, lijden aan lange trainingstijden, hoge gevoeligheid voor hyperparameters en de noodzaak van kostbare evaluaties van $\rho_{\text{target}}$ en de hogere afgeleiden daarvan.

Methodologie
De auteurs stellen Tensor Train Diffusion (TTD) voor, een nieuwe solver voor de onderliggende HJB-vergelijking die gebruikmaakt van Functional Tensor Train (FTT) representaties.

1. Wiskundige Formulering:

   • Het probleem wordt geformuleerd als het vinden van de optimale controle (drift) voor een bemonsterings-SDE die een voorwaarts ruisproces reverseert.
   • De log-densiteit van het voorwaartse proces, $V = -\log \vec{p}_Y$, voldoet aan een specifieke HJB PDE.
   • Door gebruik te maken van de verbinding tussen HJB-vergelijkingen en Backward Stochastic Differential Equations (BSDEs), herformuleren de auteurs de oplossing van de PDE als een reeks regressieproblemen. Dit maakt het mogelijk om de oplossing te berekenen via een iteratief schema dat achteruit in de tijd werkt.

2. Functional Tensor Train (FTT) Approximatie:

   • In plaats van neurale netwerken wordt de waardefunctie $V$ en zijn gradiënt benaderd met behulp van FTT's. Dit formaat maakt gebruik van latente laag-rang structuren in hoogdimensionale functies, door ze te representeren als een product van laag-rang tensor cores.
   • De methode maakt gebruik van Extended Tensor Trains (xTT), die het TT-formaat combineren met specifieke univariate basisfuncties (bijv. Legendre-polynomen, B-splines, Fourier-modi).
   • De optimalisatie van de FTT-coëfficiënten wordt uitgevoerd met behulp van Alternating Least Squares (ALS), een regressie-gebaseerde aanpak die de instabiliteiten van stochastic gradient descent (SGD) vermijdt.

3. Algoritmische Kenmerken:

   • Iteratieve Verfijning: Het algoritme maakt gebruik van een outer loop waarbij het controlebeleid iteratief wordt bijgewerkt. Beginnemend vanuit een initiële schatting (bijv. Langevin-dynamica), verfijnt de methode de bemonsterings-trajecten om de trainingsdistributie beter af te stemmen op de doel-distributie, waardoor effectief nieuwe modi worden ontdekt.
   • Adaptiviteit: Het framework bevat adaptieve strategieën voor:
     • Rangselectie: Het dynamisch aanpassen van de TT-rang op basis van de singulariteitswaarde-decay om een evenwicht te vinden tussen nauwkeurigheid en computationele kosten.
     • Basisgraden: Het aanpassen van de polynoomgraad van de basisfuncties op basis van de gladheid van de oplossing in verschillende tijdstappen.
     • Regularisatie: Het automatisch afstemmen van de regularisatie-magnitude om het leren van afgeleiden te stabiliseren.
   • Domeinbehandeling: De methode gaat om met bewegende benaderingsdomeinen en breidt de geleerde policy uit naar regio's buiten het trainingsdomein met behulp van affine lineaire extensies om stabiliteit tijdens het bemonsteren te garanderen.

Belangrijkste Bijdragen

• Nieuwe Solver: Introductie van TTD, een solver voor HJB-vergelijkingen die diffusie-gebaseerd bemonsteren, HJB-theorie en BSDE's integreert om een gunstige afruil te bereiken tussen computationele kosten en nauwkeurigheid.
• Robuustheid en Flexibiliteit: Demonstratie van de flexibiliteit van de methode door diverse basisfuncties (Legendre, B-splines, Fourier) te integreren en de robuustheid te verbeteren ten opzichte van eerdere tensor-train-gebaseerde PDE-solvers, met name bij het afhandelen van trajecten buiten het primaire trainingsdomein.
• Prestaties: Aantonen dat TTD kan bemonsteren van complexe, multimodale distributies van variërende dimensies, waarbij het bestaande diffusie-gebaseerde samplers overtreft in zowel snelheid als nauwkeurigheid.

Experimentele Resultaten
De auteurs evalueren TTD op verschillende uitdagende benchmarks:

• Hoogdimensionale Multiwell Problemen: In dimensies $d=10$ en $d=50$ met meerdere modi, bereikt TTD een hoge getrouwheid in bemonstering. De methode vermijdt mode collapse, een veelvoorkomend probleem bij diffusie-gebaseerde bemonstering. De runtime schaalt lineair met het aantal tijdstappen en is aanzienlijk sneller dan neurale netwerk-gebaseerde alternatieven (bijv. DIS).
• Anisotrope Gaussische Targets: Experimenten bevestigen het vermogen van het algoritme om adaptief rangen te verlagen naarmate de doelverdeling convergeert naar een standaard normale verdeling, wat de rang-adaptiviteit mechanisme valideert.
• $\phi^4$ Scalar Field Theory: De methode bemonster succesvol de gediscretiseerde Ginzburg-Landau model, een systeem met exponentieel veel metastabiele modi en sterke naburige correlaties, waar standaard MCMC-methoden doorgaans falen vanwege mode collapse.
• Vergelijking: TTD presteert consistent beter dan state-of-the-art baselines (inclusclusief PIS, DDS, DIS en diverse MCMC-varianten) wat betreft de Effective Sample Size (ESS), log-variantie divergentie en de nauwkeurigheid van geschatte normaliserende constanten, waarbij vaak betere resultaten worden behaald in minder tijd.

Betekenis en Claims
Het artikel claimt dat TTD een snelle, nauwkeurige en stabiele alternatief biedt voor neurale netwerktechnieken voor het oplossen van hoogdimensionale HJB-vergelijkingen in bemonsteringscontexten. Door gebruik te maken van laag-rang tensorstructuren, omzeilt de methode de "vloek van dimensionaliteit" die gepaard gaat met traditionele mesh-gebaseerde methoden en de hoge trainingskosten en hyperparameter-gevoeligheid van SGD-gebaseerde neurale solvers. De auteurs positioneren TTD als een robuuste fundering voor het oplossen van HJB-vergelijkingen, die minder evaluaties van de doel-densiteit vereist en een verbeterde efficiëntie biedt voor hoogdimensionale bemonsteringsproblemen. Ze erkennen beperkingen met betrekking tot de garantie voor het identificeren van alle modi in zeer complexe distributies en de expressiviteitslimieten van FTT's voor bepaalde doel-densiteiten, maar benadrukken de superieure prestaties van de methode in de geteste regimes.