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

Sintesi Tecnica: Una Rappresentazione Incondizionata del Punteggio Condizionato in Problemi di Inversione Lineare a Dimensionalità Infinita

Enunciato del Problema
I modelli di diffusione basati sul punteggio (SDM) sono diventati uno strumento potente per il campionamento da distribuzioni a posteriori in problemi di inversione bayesiani, offrendo prior flessibili che superano i regolarizzatori tradizionali costruiti a mano. Tuttavia, i metodi esistenti per il campionamento a posteriori affrontano un compromesso critico tra efficienza computazionale e scalabilità.

1. Metodi basati sulla correzione (ad esempio, aggiungere gradienti di verosimiglianza al punteggio o imporre la coerenza con i dati) richiedono valutazioni ripetute dell'operatore diretto $A$ durante il processo di campionamento. Ciò crea un collo di bottiglia computazionale proibitivo per problemi su larga scala e ad alta dimensionalità.
2. Metodi di addestramento condizionato diretto apprendono una funzione di punteggio condizionata alle misurazioni ($s(x, t; \mu_y)$). Sebbene ciò eviti valutazioni dirette durante il campionamento, aumenta drasticamente la dimensionalità di input della rete neurale (da $x$ a $(x, y)$), richiedendo dati di addestramento estesi e architetture complesse per approssimare accuratamente il punteggio condizionato, in particolare in contesti a dimensionalità infinita.

Il lavoro affronta la sfida di eseguire un campionamento a posteriori efficiente e scalabile in problemi di inversione lineare a dimensionalità infinita senza incorrere nell'alto costo online delle valutazioni del modello diretto o nell'alta complessità di campionamento dell'addestramento condizionato diretto.

Metodologia: Rappresentazione Incondizionata del Punteggio Condizionato (UCoS)
Gli autori propongono UCoS, un framework che sposta il carico computazionale della condizione dalla fase online di campionamento a una fase offline di addestramento. L'idea fondamentale è che, per problemi di inversione lineare con rumore gaussiano, la funzione di punteggio condizionata può essere derivata esattamente da una funzione di punteggio incondizionata di un processo di diffusione modificato, dipendente dal compito, tramite trasformazioni affini.

• Fondamento Teorico:

  • Il metodo opera in uno spazio di Hilbert separabile $\mathcal{H}$, garantendo l'invarianza rispetto alla discretizzazione.
  • Invece di addestrare una rete per approssimare $s(x, t; \mu_y)$, gli autori definiscono un nuovo processo di diffusione $\tilde{X}_t$ la cui distribuzione è una miscela del prior e di un termine gaussiano dipendente dall'operatore diretto $A$.
  • Dimostrano (Teorema 3.7) che il vero punteggio condizionato $s(x, t; \mu_y)$ è legato al punteggio incondizionato $\tilde{s}$ di questo nuovo processo tramite l'identità:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    dove $r$ è una versione trasformata del punteggio incondizionato, $\lambda(t)$ è uno scalare dipendente dal tempo e $\xi_t(x, y)$ è una trasformazione affine che coinvolge la misurazione $y$ e l'operatore aggiunto $A^*$.
  • Crucialmente, il termine $A^*\Gamma^{-1}y$ (dove $\Gamma$ è la covarianza del rumore) può essere calcolato una sola volta prima dell'inizio del campionamento. Il successivo processo di campionamento richiede solo la valutazione della rete addestrata $r_\theta$ e semplici addizioni vettoriali, eliminando la necessità di ulteriori applicazioni di $A$ o $A^*$ durante i passaggi di diffusione inversa.

• Implementazione:

  • Fase Offline: Una rete neurale (nello specifico un Operatore Neurale di Fourier, FNO) viene addestrata per approssimare la funzione $r$. L'addestramento utilizza un obiettivo di matching del punteggio di denoising sul processo modificato $\tilde{X}_t$. Questa fase richiede l'accesso a $A$ e $A^*$ ma non a dati di misurazione specifici $y$.
  • Fase Online: I campioni a posteriori sono generati risolvendo l'SDE a tempo inverso utilizzando il punteggio trasformato. Il processo è "senza matrice", richiedendo solo lo spostamento precalcolato $A^*\Gamma^{-1}y$ e la valutazione della rete.

Contributi Chiave

1. Identità Esatta in Dimensionalità Infinita: Il lavoro stabilisce un'identità teorica rigorosa (Teorema 3.7) che collega il punteggio condizionato a un punteggio incondizionato dipendente dal compito tramite trasformazioni affini in spazi di Hilbert a dimensionalità infinita. Questo generalizza il principio oltre specifiche discretizzazioni a dimensionalità finita.
2. Analisi della Convergenza: Gli autori forniscono un limite di errore quantitativo (Teorema 4.1) per il processo generativo, caratterizzando esplicitamente i contributi di errore derivanti dalla perdita di addestramento, dalla discretizzazione numerica e dall'inizializzazione. Ciò giustifica la convergenza del metodo alla vera misura a posteriori.
3. Scalabilità ed Efficienza: Disaccoppiando la condizione dal ciclo di campionamento, UCoS raggiunge l'efficienza computazionale dei metodi incondizionati (nessuna valutazione diretta durante il campionamento) mantenendo al contempo la precisione dei metodi condizionati. Evita la maledizione della dimensionalità associata all'addestramento di reti sullo spazio congiunto $(x, y)$.

Risultati Sperimentali
Gli autori validano UCoS su tre problemi di inversione: inpainting di immagini, tomografia computerizzata (CT) a viste sparse e deblurring di immagini.

• Inpainting (Prior Gaussiano): In un contesto in cui la posterior è gaussiana e i punteggi possono essere calmati in forma chiusa, UCoS performa in modo comparabile al metodo condizionato esatto e ad altre baseline, confermando l'identità teorica.
• Imaging CT e Deblurring (Prior Appresi): Utilizzando rispettivamente i dataset LIDC-IDRI e CelebA, UCoS viene confrontato con SDE ALD, Diffusion Posterior Sampling (DPS), metodi basati su proiezione (Proj) e addestramento Condizionato diretto.
  • Precisione: UCoS produce campioni con bias e deviazione standard comparabili o migliori rispetto ai metodi incondizionati e significativamente migliori rispetto al metodo Condizionato quando si utilizzano architetture a bassa complessità (ad esempio, 32 nodi per livello).
  • Efficienza: UCoS corrisponde alla velocità di campionamento del metodo Condizionato (entrambi evitano valutazioni dirette) ma supera significativamente i metodi basati sulla correzione (DPS, Proj, SDE ALD) che richiedono passaggi diretti ripetuti.
  • Robustezza: Le prestazioni del metodo Condizionato degradano notevolmente con un minor numero di parametri di rete a causa dell'aumentata complessità nell'approssimare la distribuzione congiunta, whereas UCoS mantiene alta qualità con architetture più semplici.

Significato e Affermazioni
Il lavoro afferma che UCoS offre una soluzione "il meglio di entrambi i mondi": elimina il collo di bottiglia computazionale delle valutazioni del modello diretto durante il campionamento (un limite dei metodi basati sulla correzione) evitando al contempo l'alta complessità di campionamento e le richieste architetturali dell'addestramento condizionato diretto.

Gli autori sottolineano che il loro approccio è invariante rispetto alla discretizzazione, rendendolo adatto a problemi di inversione a dimensionalità infinita dove le garanzie di prestazioni spesso falliscono in approssimazioni a dimensionalità finita ad alta dimensionalità. Riconoscono che il metodo è dipendente dal compito, richiedendo un ri-addestramento se l'operatore diretto $A$ cambia, e che la corrispondenza esatta dipende dalla linearità del modello diretto e dalla gaussianità della verosimiglianza (permettendo il passaggio del "completamento del quadrato"). Il lavoro è presentato come una prova di concetto per una nuova classe di campionatori basati sulla diffusione che spostano la complessità della condizione alla fase di addestramento senza errore di approssimazione nella formulazione teorica.