Technische Samenvatting: Bayesiaanse Tensor-decompositie met een Diffusiemodel-prior

Probleemstelling

Low-rank tensor-decompositie (TD) is een standaardinstrument voor multidimensionale data-analyse, effectief wanneer data schoon en volledig geobserveerd is. De prestaties verslechteren echter aanzienlijk onder zware ontbrekendeheid of ruis. Hoewel de low-rank assumptie dient als een structurele prior, is deze vaak onvoldoende voor de reconstructie van realistische data. Eerdere pogingen om TD aan te vullen met handmatige priors (bijv. sparsity, gladheid) falen in het vastleggen van de rijke statistieken van complexe datadistributies. Daarentegen bieden state-of-the-art diffusiemodellen weliswaar krachtige, geleerde datapriors, maar zijn ze niet gemakkelijk compatibel met huidige TD-frameworks of tractabele posteriore inferentie, vooral wanneer ze gekoppeld worden aan low-rank restricties en automatische rangselectie.

Methodologie: DiffBCP

De auteurs stellen DiffBCP voor, een hybride-prior Bayesiaans CP-decompositieframework dat is ontworpen om een geleerde diffusie-prior te integreren met structurele low-rank restricties.

1. Probabilistische Formulering

Het model definieert een gezamenlijke waarschijnlijkheidsverdeling over geobserveerde data $Y$, de gereconstrueerde tensor $X$, en latente CP-factoren ($\lambda, A^{(1:N)}$):
$$p(Y, X, \lambda, A^{(1:N)}, \tau) \propto p(Y | X, \tau) \cdot p(X | A^{(1:N)}, \lambda) \cdot p(A^{(1:N)}) \cdot p(\lambda) \cdot p(X) \cdot p(\tau)$$

Kerncomponenten zijn onder meer:

• Likelihood: Veronderstelt dat geobserveerde entries een Normale verdeling volgen met precisie $\tau$.
• Structurele restrictie: $p(X | A^{(1:N)}, \lambda)$ dwingt de CP-decompositiestructuur af ($X = \text{CP}(\lambda, A^{(1:N)})$) via een Dirac-deltafunctie.
• Latente factor priors:
  • CUSP Prior: De CP-gewichten $\lambda$ volgen een Cumulative Shrinkage Process (CUSP) prior. Deze hiërarchische prior maakt automatische rangselectie mogelijk door redundante componenten naar een kleine constante $\theta_\infty$ te krimpen naarmate de rang toeneemt.
  • Gaussiaanse Priors: Factormatrices $A^{(n)}$ volgen standaard Gaussiaanse priors om normen te regulariseren.
• Diffusie Prior: De gereconstrueerde tensor $X$ krijgt een krachtige, vooraf getrainde diffusiemodel-prior $p(X)$. Dit fungeert als een impliciete, data-gedreven regularisator die in staat is complexe texturen en structuren te vangen die een low-rank assumptie alleen niet kan vatten.

2. Inferentie: Split Gibbs Sampler (SGS)

Directe posteriore inferentie is onhandelbaar vanwege de koppeling tussen de likelihood, de low-rank restrictie en de impliciete diffusie-prior. De auteurs gebruiken een Split Gibbs Sampler met een hulpvariabele $Z$ om deze termen te ontkoppelen:

• Augmentatie: Een hulpvariabele $Z$ wordt geïntroduceerd met een koppelingsterm $\phi(Z, X; \rho) = \frac{1}{2\rho^2}\|Z - X\|_F^2$.
• Updates van latente variabelen: Gegeven $Z$, kennen de CP-factoren ($\lambda, A^{(n)}$) en hyperparameters ($\tau, \theta, \nu$) gesloten vorm conjuncte updates toe (Gaussiaans, Gamma, Beta, etc.).
• Diffusie Blok Update: Gegeven $X$, wordt de hulpvariabele $Z$ gesampled uit een distributie gedefinieerd door de diffusie-prior en de koppelingsterm. Dit wordt opgelost als een denoising-probleem met behulp van het vooraf getrainde score-netwerk $s_\psi$ via het Stochastic Differential Equation (SDE) reverse proces, gestuurd door de low-rank schatting $X$.
• Ruis-adaptieve Koppeling: Er wordt een schema geïntroduceerd waarbij de koppelingsparameter $\rho$ dynamisch wordt aangepast op basis van de geïnferrede ruisprecisie $\tau$ (specifiek $\tau\rho^2 = c$). Dit vermindert de gevoeligheid voor handmatige annealing-schema's.

3. Rang Adaptatie

Het framework bevat een bidirectioneel rang-adaptatieschema. De CUSP-prior krimpt inactieve componenten, terwijl een apart mechanisme nieuwe componenten toevoegt als de huidige rang onvoldoende is. Dit zorgt voor robuustheid tegen initiële rang-misspecificatie.

Belangrijkste Bijdragen

1. Hybride-Prior Framework: Het eerste Bayesiaanse tensor-decompositie model dat een kant-en-klare vooraf getrainde diffusie-prior integreert als een data-prior op de gereconstrueerde tensor, gecombineerd met een structurele low-rank CP-prior.
2. Tractable Inferentie: Ontwikkeling van een Split Gibbs Sampler die posteriore inferentie mogelijk maakt ondanks de complexe koppeling tussen likelihood, low-rank restricties en diffusie-priors. Dit omvat gesloten vorm updates voor latente factoren en low-rank-gestuurde denoising voor de diffusie-blok.
3. Automatische Rang en Ruis Inferentie: Het gebruik van de CUSP-prior en volledige Bayesiaanse ruis-estimatie elimineert de noodzaak voor handmatige afstemming van rang en ruisniveaus.
4. Robuustheid: De methode demonstreert robuustheid tegen initiële rang-misspecificatie en variërende ruisniveaus door middel van zijn adaptieve koppelingsschema.

Experimentele Resultaten

De methode werd geëvalueerd op beeld inpainting en denoising taken met FFHQ en ImageNet datasets, evenals hoog-resolutie out-of-distribution (OOD) afbeeldingen.

• Prestaties: DiffBCP presteerde consistent beter dan baselines, inclusief traditionele Bayesiaanse TD (BCP, BTR), deep learning-gebaseerde TD (HLRTF, DeepTensor) en plug-and-play methoden (GLON, PnP-DM).
  • Op FFHQ met 90% ontbrekende data (Uniform mask), bereikte DiffBCP 28.28 dB PSNR, wat een verbetering is van 2.17 dB ten opzichte van de beste baseline (DeepTensor).
  • Op irregulaire maskers waren de winsten nog prominenter (+2.33 dB ten opzichte van DeepTensor).
• OOD en Hoog-Resolutie: De methode reconstrueerde succesvol 2048×2048 afbeeldingen (bijv. Marseille, Tokyo) met behulp van een diffusiemodel getraind op 256×256 afbeeldingen, waarbij coherentie en detail behouden bleven waar andere methoden (zoals PuTT) artefacten introduceerden of faalden.
• Ablatie Studies:
  • Het verwijderen van de diffusie-prior verslechterde de prestaties aanzienlijk op gestructureerde maskers (Stripe, Irregular), wat de rol van de diffusie-prior bij het herstellen van ruimtelijke continuïteit bevestigt.
  • De CUSP-prior bleek cruciaal voor herstel wanneer de initiële rang ernstig ondergespecificeerd was (bijv. startend bij rang 10 versus 200).
• Vergelijking met Diffusie Samplers: Vergeleken met PnP-DM vertoonde DiffBCP superieure stabiliteit met betrekking tot ruisniveau-scheduling en betere prestaties op hoog-resolutie afbeeldingen, hoewel PnP-DM competitieve resultaten kon behalen op uniforme maskers als de ruisparameters handmatig naar optimale waarden werden afgestemd.

Betekenis en Claims

Het artikel claimt dat DiffBCP erin slaagt de kloof te overbruggen tussen de structurele inductieve bias van low-rank tensor-decompositie en de expressieve kracht van moderne generatieve diffusiemodellen. Door deze integratie binnen een volledig probabilistisch framework te formuleren, bereiken de auteurs:

• State-of-the-art herstel voor ruizige en incomplete tensor-data.
• Automatische adaptatie aan onbekende rangen en ruisniveaus, wat uitgebreide hyperparameter-afstemming overbodig maakt.
• Robuustheid tegen distributieverschuivingen en hoog-resolutie inputs, waarbij de low-rank structuur wordt gebruikt als een stabiliserende inductieve bias wanneer de diffusie-prior met OOD-data te maken krijgt.

De auteurs merken op dat de prestaties afhangen van het feit of het onderliggende signaal een betekenisvolle low-rank structuur toelaat; indien deze assumptie wordt geschonden, neemt de structurele bijdrage af, hoewel de diffusie-prior actief blijft. Toekomstig werk wordt gesuggereerd voor het uitbreiden van het framework naar andere inverse problemen (bijv. compressed sensing) en andere tensor-decompties (Tensor Train, Tensor Ring).