Flow-Based Generative Modeling for Optimizing Sampling Policies in Compressed Sensing Applications

Sintesi Tecnica: Modellazione Generativa Basata su Flow per l'Ottimizzazione delle Politiche di Campionamento nel Compressed Sensing

Enunciato del Problema
Le moderne applicazioni nell'elaborazione dei segnali e nell'imaging medico richiedono spesso l'acquisizione di segnali ad alta dimensionalità sotto rigorosi vincoli di risorse. La teoria classica del campionamento stabilisce che una ricostruzione accurata richiede misurazioni proporzionali alla dimensione ambiente del segnale, un requisito che è frequentemente impraticabile. Il Compressed Sensing (CS) affronta questo problema consentendo il recupero di segnali sparsi da meno misurazioni, a condizione che l'operatore di misurazione soddisfi specifiche condizioni (ad esempio, l'incoerenza). Tuttavia, la qualità della ricostruzione dipende fortemente dalla scelta dello schema di campionamento. Trovare un pattern di sottocampionamento ottimale è un problema combinatorio NP-hard perché la maschera di campionamento è discreta e non differenziabile, impedendo un'ottimizzazione end-to-end ingenua. Inoltre, lo spazio di ricerca cresce esponenzialmente con il numero di posizioni dei sensori, rendendo impraticabile l'esplorazione esaustiva. Gli approcci generativi esistenti, come il Deep Probabilistic Subsampling (DPS) o l'Active Diffusion Subsampling (ADS), spesso producono maschere fisse o soffrono di un'alta latenza computazionale, limitando la loro adattabilità e l'utilità clinica in tempo reale.

Metodologia
Gli autori propongono un framework generativo basato su flow, consapevole del compito (task-aware), che riformula il paradigma convenzionale di addestramento del Flow Matching (FM). Inveza di modellare direttamente il manifold dei dati, il modello opera sul manifold dello spazio di campionamento per apprendere maschere di sottocampionamento ottimali.

• Flow Matching sul Manifold di Campionamento: Il framework utilizza un Ordinary Differential Equation (ODE) neurale per trasportare punti da una distribuzione di rumore Gaussiano standard ($p_0$) a una distribuzione target di operatori di campionamento ($p_1$). Il modello apprende un campo vettoriale dipendente dal tempo $v_\phi(x, t)$ che definisce una traiettoria generativa.
• Ottimizzazione Condizionata al Compito: A differenza del FM standard che si basa su un accoppiamento lineare tra rumore e dati, questo framework accoppia il flusso generativo con un risolutore di task CS a valle ($f_\theta$). L'obiettivo di addestramento minimizza una funzione di perdita composita che combina la perdita di regressione FM standard ($L_{FM}$) e una perdita specifica del task ($L_{task}$), pesate da parametri di incertezza apprendibili ($\sigma_1, \sigma_2$). Ciò consente al modello di ottimizzare congiuntamente la politica di campionamento e le prestazioni del task.
• Gating Sigmoideo Differenziabile: Per affrontare la non-differenziabilità delle maschere binarie discrete, l'output continuo del modello di flow ($x_t$) viene passato attraverso un gate sigmoideo differenziabile personalizzato. Questo gate impiega una funzione di pendenza dipendente dal tempo $s(t)$ e una soglia percentile $\tau$ per produrre maschere morbide ($g(x_t) \in (0, 1)$). Man mano che il passo temporale $t$ aumenta, la transizione diventa più netta, regolarizzando efficacemente l'entropia del processo di generazione della maschera.
• Modalità di Inferenza: Il framework supporta due modalità operative:
  1. Da Random a Ottimizzato: Vengono effettuate misurazioni casuali iniziali, vengono estratte le caratteristiche e il modello di flow genera una maschera ottimizzata per le misurazioni successive. Questo è adatto per scenari in cui le misurazioni sono economiche o i segnali sono statici (ad esempio, MNIST, CelebA).
  2. Sensing Adattivo: I dati dello spazio k a bassa frequenza (LF) vengono acquisiti per primi. Un estrattore di caratteristiche condiziona il modello di flow per determinare il pattern di campionamento ottimale per le restanti bande ad alta frequenza. Questa modalità è progettata per misurazioni costose, specificamente l'accelerazione MRI.

Contributi Chiave

• Nuovo Paradigma: Il documento introduce una prova di concetto in cui il Flow Matching è addestrato all'interno del manifold degli operatori di campionamento piuttosto che nello spazio del segnale, progettato specificamente per task di CS a valle.
• Politica Continua e Adattiva: Il framework fornisce una parametrizzazione continua e tempo-continua dello schema di campionamento, offrendo un'evoluzione fluida da strutture grossolane a raffinamenti specifici del task, a differenza dei metodi di generazione di maschere statiche.
• Ottimizzazione Congiunta: Il metodo consente un'ottimizzazione end-to-end congiunta della strategia di campionamento e del task di ricostruzione/classificazione, sfruttando il prior del task per guidare il processo generativo.
• Framework Unificato: Offre un approccio flessibile applicabile a vari problemi inversi, inclusa la classificazione di immagini, la ricostruzione di immagini e l'accelerazione MRI, senza richiedere il riaddestramento per diversi operatori di misurazione.

Risultati Sperimentali
Il framework è stato valutato su MNIST, CelebA e il dataset fastMRI knee:

• Classificazione MNIST: L'approccio basato su flow ha migliorato l'accuratezza della classificazione di oltre il 20% rispetto ai classificatori addestrati su maschere casuali. Sebbene sia rimasto leggermente al di sotto del modello A-DPS in specifici rapporti di campionamento fissi, ha dimostrato una maggiore versatilità attraverso una gamma di rapporti (dall'1% all'8%) senza necessità di addestramento specifico per il rapporto.
• Ricostruzione Immagini (CelebA & MNIST): Il metodo ha raggiunto prestazioni state-of-the-art nella ricostruzione delle immagini. Sul dataset CelebA, ha raggiunto un Peak Signal-to-Noise Ratio (PSNR) di 25,17 dB con un tasso di sottocampionamento del 5%. Ha superato il metodo state-of-the-art ADS, in particolare ai tassi di campionamento aggressivi (bassi).
• Accelerazione MRI: Negli esperimenti fastMRI, il framework ha raggiunto un Structural Similarity Index (SSIM) di 0,7163 con accelerazione 8×, superando tutti i baseline (LOUPE, DPS, ADS, A-DPS) a questo alto fattore di accelerazione. Con accelerazione 4×, ha performato in modo comparabile agli altri metodi.
• Efficienza Computazionale: Il tempo di inferenza per generare una maschera di campionamento è stato di circa 127–130 ms su una NVIDIA GeForce RTX 4070. Questo è significativamente più veloce del tempo di acquisizione per ogni linea di k-space in MRI (500–2000 ms) e drasticamente più veloce del tempo di acquisizione per linea riportato dal metodo ADS (3040 ms), suggerendo la fattibilità per un impiego clinico in tempo reale.

Significatività e Rivendicazioni
Gli autori affermano che questo lavoro stabilisce la fattibilità fondamentale dell'uso di modelli generativi basati su flow condizionati dal task per ottimizzare le politiche di campionamento nel compressed sensing. Lo studio evidenzia che:

1. Il Condizionamento al Task è Efficace: Integrare un prior di un task a valle direttamente nel processo di addestramento del flow matching migliora significativamente le prestazioni nei problemi inversi.
2. Vincoli del Manifold di Campionamento: Lo spazio di campionamento possiede vincoli strutturali intrinseci (binarietà, sparsità) che possono essere sfruttati per migliorare l'efficienza dell'addestramento rispetto alla modellazione del manifold dell'immagine naturale.
3. Viabilità Clinica: Il basso overhead computazionale e la natura adattiva del framework lo rendono un candidato promettente per applicazioni di imaging medico nel mondo reale, in particolare per ridurre i tempi di scansione MRI e migliorare la qualità dei dati diagnostici.
4. Generalizzabilità: L'approccio offre una strategia unificata e flessibile per progettare schemi di sensing guidati dai dati e dai task, applicabili a una vasta gamma di problemi inversi oltre ai compiti specifici testati.

Il documento si pone come uno studio di prova di concetto, riconoscendo che, sebbene abbia raggiunto risultati state-of-the-art in compiti specifici di ricostruzione, il contributo primario è la dimostrazione di un nuovo e flessibile framework per apprendere politiche di campionamento che possono essere adattate a vari problemi inversi.