Sintesi Tecnica: Modellazione Generativa di Strutture Latenti Discrete tramite Gradienti di Politica Dinamici

Definizione del Problema

Il paper affronta la sfida di inferire stati latenti meccanicistici non osservati ($S^*$) da osservazioni indirette ($X$) in domini scientifici dove lo spazio latente è combinatoriamente grande. A differenza dell'apprendimento non supervisionato generale (ad esempio, clustering o topic modeling) che apprende rappresentazioni latenti artificiali, questo lavoro si concentra sulla ricostruzione di stati meccanicistici reali come percorsi di reazione chimica, reti di trasporto o isoforme geniche.

Il problema centrale è formulato come un compito di apprendimento (Problema 2.1) per trovare i parametri del modello $\theta$ per una distribuzione generativa $Pr(S|\theta)$ che massimizzi la verosimiglianza dei dati osservati $X_{1:N}$, dove $Pr(X|\theta) = \sum_{S} Pr(X|S)Pr(S|\theta)$. Un compito di inferenza secondario (Problema 2.2) consiste nel stimare lo specifico stato latente $\hat{S}_i$ per ogni osservazione $X_i$.

Gli approcci classici come l'Aspetto-Massimizzazione (EM) faticano in questo contesto perché l'E-step richiede il calcolo delle aspettative su spazi di stato esponenzialmente grandi, spesso intrattabili. Sebbene l'Apprendimento per Rinforzo (RL) sia naturalmente adatto alla generazione sequenziale di stati strutturati, i gradienti di politica standard ottimizzano le ricompense attese sotto funzioni di ricompensa stazionarie, non la verosimiglianza marginale di osservazioni indirette.

Metodologia: GReinSS

Gli autori introducono GReinSS (Generative Reinforcement Learning of Structured States), un framework che adatta i gradienti di politica per ottimizzare la verosimiglianza marginale delle osservazioni indirette.

Ricalibratura Dinamica della Ricompensa

L'innovazione centrale è una funzione di ricompensa dinamicamente ricalibrata. In un RL standard, il gradiente della politica è $\mathbb{E}_{\tau}[r(\tau) \nabla_\theta \log Pr(\tau|\theta)]$. GReinSS definisce la ricompensa $r(\tau)$ per una traiettoria $\tau$ (che produce uno stato latente $S(\tau)$) come:
$$r(\tau) = \sum_{i=1}^{N} \frac{Pr(X_i | \tau)}{Pr(X_i | \theta)}$$
dove $Pr(X_i | \theta) = \mathbb{E}_{\tau}[Pr(X_i | \tau)]$.

Teorema 3.1 stabilisce che questa specifica funzione di ricompensa rende il gradiente della politica un estimatore non distorto del gradiente della log-verosimiglianza dei dati $\nabla_\theta \log Pr(X_{1:N} | \theta)$.

• Meccanismo: Il denominatore $Pr(X_i | \theta)$ ricalibra il contributo di ogni osservazione. Ciò impedisce alla politica di collassare su una singola traiettoria ad alta ricompensa (che massimizzerebbe la verosimiglianza di un'osservazione ma fallirebbe per le altre) e incoraggia invece la politica a imparare una distribuzione $Pr(S|\theta)$ che bilanci la verosimiglianza tra tutte le osservazioni.
• Ciclo di Addestramento: L'algoritmo alterna il campionamento di traiettorie, il calcolo delle ricompense dinamiche basate sull'attuale $\theta$ e l'aggiornamento di $\theta$ tramite la salita del gradiente.

Apprendimento Off-Policy

Per affrontare la scarsità di traiettorie che producono una verosimiglianza non nulla per specifiche osservazioni, il paper impiega il campionamento off-policy. Il Teorema 3.3 identifica la distribuzione di proposta off-policy ottimale che minimizza la varianza come $q(\tau | X_{1:N}, \theta) = \frac{1}{N} \sum_{i=1}^N Pr(\tau | X_i, \theta)$. Sebbene il campionamento esatto da questa distribuzione sia spesso intrattabile, gli autori utilizzano euristiche (ad esempio, campionamento distorto basato sulle caratteristiche delle osservazioni) combinate con correzioni di importance sampling per mantenere stime del gradiente non distorte.

Mini-batching

Il framework supporta il mini-batching (Corollario 3.2), dove le ricompense sono calcolate su un sottoinsieme di osservazioni $B \subset [N]$, permettendo al metodo di scalare su grandi dataset.

Contributi Chiave

1. Formulazione: Il paper formula il problema dell'apprendimento di distribuzioni su stati latenti discreti da osservazioni indirette come un problema di ottimizzazione della politica, colmando il divario tra modellazione probabilistica e apprendimento per rinforzo.
2. Connessione Teorica: Fornisce una prova teorica (Teorema 3.1) che i gradienti di politica con ricompense ricalibrate dinamicamente sono equivalenti alla massimizzazione della verosimiglianza marginale delle osservazioni indirette, un risultato non raggiunto dalle formulazioni standard di RL o di Inferenza Variazionale (VI).
3. Generalizzazione: A differenza delle applicazioni precedenti di tecniche simili (ad esempio, nella filogenesi del cancro o nella variazione del numero di copie), questo lavoro generalizza la procedura a arbitrari strutture discrete (grafi, insiemi, sequenze) ed stabilisce connessioni con metodi esistenti come la Generalized Expectation Maximization (GEM) e i GFlowNets in condizioni specifiche.

Risultati Sperimentali

Gli autori valutano GReinSS su dataset simulati e reali, confrontandolo con la ricerca locale, i gradienti di politica standard, i GFlowNets e le implementazioni basate su GEM utilizzando VAE, modelli autoregressivi e diffusione discreta.

1. Inferenza di Grafi di Processo Simulati

• Task: Ricostruire grafi diretti da endpoint di cammini casuali (random walk).
• Risultati: GReinSS ha costantemente ottenuto i punteggi F1 mediani più elevati. Notevolmente, con bassa informazione per osservazione ($k=10$ random walks), GReinSS ha raggiunto un F1 di 0,891, mentre tutti i baseline (inclusi i gradienti di politica ingenui e i GFlowNets) hanno ottenuto punteggi inferiori a 0,55. Il baseline del gradiente di politica ingenuo spesso collassava predicendo grafi vuoti, evidenziando la necessità della ricalibrazione dinamica della ricompensa.

2. Inferenza di Sottoinsiemi Simulati

• Task: Ricostruire sottoinsiemi da misurazioni di elementi rumorose.
• Risultati: GReinSS ha superato tutti i baseline al variare dei livelli di rumore ($\sigma$) e delle dimensioni dell'universo ($|U|$).
  • Per universi piccoli ($|U|=10$), la maggior parte dei metodi ha performato bene.
  • Per universi grandi ($|U|=1000$), i baseline (VAE, diffusione, autoregressione) sono collassati (F1 mediano < 0,4), mentre GReinSS ha mantenuto un F1 mediano di 0,938.
  • I risultati dimostrano che GReinSS scala efficacemente a grandi spazi combinatori dove altri modelli generativi falliscono.

3. Applicazione nel Mondo Reale: Ricostruzione di Isoforme di RNA

• Task: Inferire le strutture e le proporzioni delle isoforme di RNA da dati di sequenziamento short-read (dataset GTEx), validati contro la verità di base (ground truth) di sequenziamento long-read.
• Confronto: Confrontato con RSEM, l'algoritmo standard basato su Expectation Maximization utilizzato nel progetto GTEx.
• Risultati: GReinSS ha superato significativamente RSEM.
  • Per il gene MBD2, RSEM ha predetto un'isoforma maggiore con proporzione del 18% (rispetto al 96% della ground truth) e isoforme minori errate, risultando in un errore di predizione di 0,5374. GReinSS ha predetto il 92% per l'isoforma maggiore con un errore di 0,0067.
  • Su 14.390 geni, GرeinSS ha superato RSEM con un margine $\ge 0,05$ nel 46,6% dei casi, mentre RSEM ha superato GReinSS in solo il 9,4% dei casi. La differenza mediana di errore è stata di $-0,0405$ a favore di GReinSS.

Significato e Rivendicazioni

Il paper sostiene che GReinSS sia un "approccio fondato e praticamente efficace" per la modellazione generativa in spazi combinatori. La sua importanza primaria risiede nel:

• Superare l'Intrattabilità: Fornisce una soluzione per l'inferenza dello stato latente in spazi dove l'EM esatto è intrattabile a causa della complessità esponenziale.
• Superiorità rispetto al RL Standard: Dimostra che i gradienti di politica standard falliscono in questo contesto senza la ricalibrazione dinamica della ricompensa, e che i GFlowNets (che bilanciano le probabilità delle traiettorie) sono insufficienti quando la struttura della ricompensa non si allinea perfettamente con la verosimiglianza delle osservazioni.
• Impatto nel Mondo Reale: L'applicazione di successo allo splicing dell'RNA, superando un metodo standard ampiamente utilizzato (RSEM) su dati biologici reali, valida l'utilità del metodo oltre i benchmark sintetici.

Gli autori osservano che, sebbene il metodo attualmente si basi su euristiche per il campionamento off-policy e assuma una $Pr(X|S)$ nota, il framework offre una base per estensioni future, inclusi i framework Q-learning, actor-critic e l'apprendimento del modello di osservazione stesso.