Sintesi Tecnica: Decomposizione Tensoriale Bayesiana con Prior di Modello di Diffusione

Definizione del Problema

La decomposizione tensoriale (TD) a basso rango è uno strumento standard per l'analisi di dati multidimensionali, efficace quando i dati sono puliti e completamente osservati. Tuttavia, le prestazioni degradano significativamente in presenza di una severa mancanza di dati o rumore. Sebbene l'assunzione di basso rango funga da prior strutturale, essa è spesso insufficiente per il recupero di dati reali. I tentativi precedenti di potenziare la TD con prior manuali (ad esempio, sparsità, regolarità/smoothness) non riescono a catturare la ricca statistica delle distribuzioni di dati complessi. Al contrario, sebbene i modelli di diffusione allo stato dell'arte offrano potenti prior di dati appresi, essi non sono facilmente compatibili con gli attuali framework di TD o trattabili per l'inferenza del posteriore, particolarmente quando accoppiati con vincoli di basso rango e selezione automatica del rango.

Metodologia: DiffBCP

Gli autori propongono DiffBCP, un framework di decomposizione CP bayesiano con prior ibrido progettato per integrare un prior di diffusione appreso con vincoli strutturali di basso rango.

1. Formulazione Probabilistica

Il modello definisce una distribuzione di probabilità congiunta sui dati osservati $Y$, il tensore ricostruito $X$ e i fattori CP latenti ($\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)$$

I componenti chiave includono:

• Verosimiglianza (Likelihood): Assume che le voci osservate seguano una distribuzione Normale con precisione $\tau$.
• Vincolo Strutturale: $p(X | A^{(1:N)}, \lambda)$ impone la struttura della decomposizione CP ($X = \text{CP}(\lambda, A^{(1:N)})$) tramite una funzione delta di Dirac.
• Prior dei Fattori Latenti:
  • Prior CUSP: I pesi CP $\lambda$ seguono un prior di Cumulative Shrinkage Process (CUSP). Questo prior gerarchico consente la selezione automatica del rango riducendo i componenti ridondanti verso una piccola costante $\theta_\infty$ all'aumentare del rango.
  • Prior Gaussiani: Le matrici di fattore $A^{(n)}$ seguono prior gaussiani standard per regolarizzare le norme.
• Prior di Diffusione: Al tensore ricostruito $X$ viene assegnato un potente prior di modello di diffusione pre-addestrato $p(X)$. Questo agisce come un regolarizzatore implicito e guidato dai dati, capace di catturare texture e strutture complesse che le sole assunzioni di basso rango non possono catturare.

2. Inferenza: Split Gibbs Sampler (SGS)

L'inferenza diretta del posteriore è intrattabile a causa dell'accoppiamento tra la verosimiglianza, il vincolo di basso rango e l'implicito prior di diffusione. Gli autori utilizzano un Split Gibbs Sampler con una variabile ausiliaria $Z$ per disaccoppiare questi termini:

• Augmentation: Viene introdotta una variabile ausiliaria $Z$ con un termine di accoppiamento $\phi(Z, X; \rho) = \frac{1}{2\rho^2}\|Z - X\|_F^2$.
• Aggiornamenti delle Variabili Latenti: Dato $Z$, i fattori CP ($\lambda, A^{(n)}$) e gli iperparametri ($\tau, \theta, \nu$) ammettono aggiornamenti coniugati in forma chiusa (Gaussiana, Gamma, Beta, ecc.).
• Aggiornamento del Blocco di Diffusione: Dato $X$, la variabile ausiliaria $Z$ viene campionata da una distribuzione definita dal prior di diffusione e dal termine di accoppiamento. Questo viene risolto come un problema di denoising utilizzando la rete di score pre-addestrata $s_\psi$ tramite il processo inverso della SDE (Equazione Differenziale Stocastica), guidato dalla stima di basso rango $X$.
• Accoppiamento Adattivo al Rumore: Viene introdotto uno schema in cui il parametro di accoppiamento $\rho$ viene regolato dinamicamente in base alla precisione del rumore inferita $\tau$ (specificamente $\tau\rho^2 = c$). Ciò riduce la sensibilità agli schemi di annealing manuali.

3. Adattamento del Rango

Il framework include uno schema di adattamento del rango bidirezionale. Il prior CUSP restringe i componenti inattivi, mentre un meccanismo separato aggiunge nuovi componenti se il rango attuale è insufficiente. Ciò garantisce robustezza contro l'errata specificazione iniziale del rango.

Contributi Chiave

1. Framework con Prior Ibrido: Il primo modello di decomposizione tensoriale bayesiana che integra un modello di diffusione pre-addestrato "off-the-shelf" come prior di dati sul tensore ricostruito, combinato con un prior CP strutturale di basso rango.
2. Inferenza Trattabile: Sviluppo di uno Split Gibbs Sampler che consente l'inferenza del posteriore nonostante il complesso accoppiamento tra verosimiglianza, vincoli di basso rango e prior di diffusione. Ciò include aggiornamenti in forma chiusa per i fattori latenti e denoising guidato dal basso rango per il blocco di diffusione.
3. Inferenza Automatica di Rango e Rumore: L'uso del prior CUSP e la stima bayesiana completa del rumore eliminano la necessità di una sintonizzazione manuale di rango e livelli di rumore.
4. Robustezza: Il metodo dimostra robustezza rispetto alla errata specificazione iniziale del rango e ai livelli di rumore variabili attraverso il suo schema di accoppiamento adattivo.

Risultati Sperimentali

Il metodo è stato valutato su compiti di inpainting e denoising di immagini utilizzando i dataset FFHQ e ImageNet, nonché immagini ad alta risoluzione fuori distribuzione (OOD).

• Prestazioni: DiffBCP ha superato costantemente i baseline, inclusi la TD bayesiana tradizionale (BCP, BTR), la TD basata su deep learning (HLRTF, DeepTensor) e i metodi plug-and-play (GLON, PnP-DM).
  • Su FFHQ con il 90% di dati mancanti (maschera Uniforme), DiffBCP ha raggiunto 28,28 dB PSNR, superando il miglior baseline (DeepTensor) di 2,17 dB.
  • Su maschere irregolari, i guadagni sono stati ancora più pronunciati (+2,33 dB rispetto a DeepTensor).
• OOD e Alta Risoluzione: Il metodo ha ricostruito con successo immagini 2048×2048 (es. Marseille, Tokyo) utilizzando un modello di diffusione addestrato su immagini 256×256, mantenendo coerenza e dettaglio dove altri metodi (come PuTT) introducevano artefatti o fallivano.
• Studi di Ablazione:
  • La rimozione del prior di diffusione ha degradato significativamente le prestazioni su maschere strutturate (Stripe, Irregular), confermando il ruolo del prior di diffusione nel recuperare la continuità spaziale.
  • Il prior CUSP si è rivelato critico per il recupero quando il rango iniziale era severamente sotto-specificato (ad esempio, partendo da un rango 10 invece di 200).
• Confronto con i Campionatori di Diffusione: Rispetto a PnP-DM, DiffBCP ha mostrato una maggiore stabilità riguardo allo scheduling del livello di rumore e migliori prestazioni su immagini ad alta risoluzione, sebbene PnP-DM potesse ottenere risultati competitivi su maschere uniformi se i suoi parametri di rumore venissero sintonizzati manualmente ai valori ottimali.

Significato e Rivendicazioni

L'articolo sostiene che DiffBCP riesce a colmare il divario tra il bias induttivo strutturale della decomposizione tensoriale a basso rango e il potere espressivo dei moderni modelli generativi di diffusione. Formulando questa integrazione all'interno di un framework completamente probabilistico, gli autori ottengono:

• Recupero allo stato dell'arte per dati tensoriali rumorosi e incompleti.
• Adattamento automatico a ranghi e livelli di rumore sconosciuti, eliminando la necessità di un'estesa sintonizzazione degli iperparametri.
• Robustezza ai cambiamenti di distribuzione e agli input ad alta risoluzione, sfruttando la struttura di basso rango come un bias induttivo stabilizzante quando il prior di diffusione affronta dati OOD.

Gli autori notano che le prestazioni dipendono dal fatto che il segnale sottostante ammetta una significativa struttura a basso rango; se questa assunzione viene violata, il contributo strutturale diminuisce, sebbene il prior di diffusione rimanga attivo. Si suggerisce come lavoro futuro l'estensione del framework ad altri problemi inversi (ad esempio, compressed sensing) e ad altre decomposizioni tensoriali (Tensor Train, Tensor Ring).