Riepilogo Tecnico: KRONOS – Diffusione Latente Autoregressiva per la Generazione di Molecole 3D

Definizione del Problema

La generazione di molecole tridimensionali (3D) è un compito critico nella scoperta di farmaci e nella progettazione molecolare. Gli attuali approcci allo stato dell'arte si affidano principalmente a modelli di diffusione e flow matching. Sebbene questi metodi raggiungano un'elevata qualità di generazione, richiedono tipicamente che la dimensione molecolare (numero di atomi) sia specificata a priori, limitando la loro flessibilità per la generazione a lunghezza variabile.

Al contrario, i modelli autoregressivi (AR) supportano naturalmente la generazione a lunghezza variabile e il condizionamento su strutture molecolari parziali (ad esempio, l'elaborazione di scaffold noti), il che è estremamente prezioso nella chimica medicinale. Tuttavia, gli esistenti modelli AR affrontano due sfide principali:

1. Gap di Prestazioni: Storicamente, i modelli AR sono rimasti indietro rispetto ai modelli di diffusione in termini di qualità di generazione, in particolare per quanto riguarda la fedeltà geometrica e la validità chimica.
2. Bilanciamento degli Obiettivi: Raggiungere prestazioni elevate sia nella generazione incondizionata (design de novo) che in quella condizionata da frammenti (completamento molecolare) all'interno di un'unica architettura rimane difficile. Molti modelli AR faticano a generalizzare tra questi paradigmi senza modifiche architettoniche o compromessi sulle prestazioni.

Inoltre, i generatori molecolari latenti esistenti si sono quasi esclusivamente affidati alla diffusione, lasciando in gran parte inesplorato il potenziale della generazione autoregressiva nello spazio latente.

Metodologia: KRONOS

Gli autori introducono KRONOS, un framework di diffusione autoregressiva latente che genera molecole 3D direttamente nello spazio latente di un Autoencoder Unificato (UAE) pre-addestrato. Il metodo combina tre componenti principali:

1. Generazione Autoregressiva Latente

Invece di operare nello spazio dei dati grezzi (predicendo simultaneamente tipi di atomi/legami discreti e coordinate continue), KRONOS opera nello spazio latente continuo di un UAE pre-addestrato (Luo et al., 2025).

• Rappresentazione: Una molecola 3D viene codificata in una sequenza di token latenti $Z = \{z_1, \dots, z_n\}$, dove ogni token codifica congiuntamente l'identità atomica, l'ambiente di legame e la geometria 3D.
• Fattorizzazione: Il modello apprende una distribuzione autoregressiva da sinistra a destra:
  $$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$$
• Lunghezza Variabile: Poiché lo spazio latente non possiede un naturale token di fine sequenza, KRONOS impiega un classificatore di stop dedicato (un MLP leggero) addestrato insieme al generatore per predire dinamicamente la terminazione della sequenza, eliminando la necessità di un numero di atomi predefinito.

2. Predizione dei Token Basata sulla Diffusione

Per modellare la distribuzione condizionata dei token latenti continui $p_\theta(z_i | z_{<i})$, KRONOS utilizza una formulazione di Loss di Diffusione (Li et al., 2024).

• Un transformer causale codifica i token precedenti in una rappresentazione di contesto $c_i$.
• Una rete di denoising per diffusione $\epsilon_\theta$ predice il rumore aggiunto al corrente token latente $z_i$, condizionato da $c_i$ e dal timestep di diffusione.
• Questo approccio sostituisce la predizione di token discreti con un obiettivo di denoising, consentendo la modellazione flessibile di variabili continue senza discretizzazione.

3. Strategia di Addestramento Misto (Generazione Condizionata da Frammenti)

Per supportare sia la generazione incondizionata che quella condizionata da frammenti all'interno di un singolo modello, gli autori introducono una strategia di addestramento ispirata al paradigma Fill-in-the-Middle (FIM).

• Decomposizione in Frammenti: Le molecole vengono decomposte in frammenti chimicamente significativi utilizzando uno schema basato su BRICS.
• Costruzione della Sequenza: Durante l'addestramento, ogni molecola viene trasformata in una:
  • Sequenza di ordine predefinito (per la generazione incondizionata).
  • Sequenza condizionata da frammenti dove un sottoinsieme casuale di frammenti funge da "seed" (contesto), e i restanti frammenti sono il target di generazione.
• Meccanismo: Il modello è addestrato a generare autoregressivamente i frammenti target condizionati dai frammenti seed. Ciò viene ottenuto assegnando un unico ordine deterministico ai frammenti e agli atomi al loro interno, piuttosto che marginalizzare su tutte le permutazioni (come nei modelli AR-Agnostici rispetto all'ordine). Questo permette al modello di apprendere capacità di completamento molecolare senza modifiche architettoniche.

Contributi Chiave

1. Framework KRONOS: Un nuovo framework di diffusione autoregressiva latente che modella congiuntamente la topologia del grafo molecolare e la geometria in uno spazio latente unificato.
2. Strategia di Addestramento Unificata: Un approccio di addestramento misto che consente la generazione incondizionata e condizionata da frammenti all'interno di un unico modello autoregressivo da sinistra a destra, ispirato al paradigma FIM.
3. Generazione a Lunghezza Variabile: Un framework unificato che supporta naturalmente la generazione a lunghezza variabile tramite un classificatore di stop, evitando la necessità di una specifica del numero di atomi a priori.
4. Prestazioni: Dimostrazione che la generazione autoregressiva latente può raggiungere prestazioni d'eccellenza tra i metodi autoregressivi, rimanendo competitiva con i modelli di diffusione, senza sacrificare la qualità incondizionata quando viene abilitato il condizionamento da frammenti.

Risultati Sperimentali

Gli autori hanno valutato KRONOS sui dataset QM9 e GEOM-Drugs.

Generazione Incondizionata

• QM9: KRONOS ha ottenuto la più alta Validità 3D (97,3%) e la prestazione complessiva più forte tra i metodi autoregressivi. È rimasto competitivo con i principali modelli di diffusione (es. JODO, EDM) su metriche quali Stabilità Atomica 3D, Lunghezza di Legame e MMD di Angolo di Legame.
• GEOM-Drugs: KRONOS ha ottenuto la prestazione complessiva più forte tra tutti i baseline valutati (sia autoregressivi che di diffusione).
• Plausibilità Fisica: Utilizzando l'ottimizzazione della geometria semi-empirica (GFN2-xTB), le conformazioni generate da KRONOS hanno mostrato bassi valori di Root Mean Square Distance (RMSD) e differenze di energia ($\Delta E$) rispetto alle strutture ottimizzate, indicando che si trovano vicino a regioni fisicamente plausibili dello spazio conformazionale.

Generazione Condizionata da Frammenti

• Su un dataset di circa 15k coppie molecola-frammento, KRONOS ha superato sostanzialmente il baseline autoregressivo QUETZAL sia nelle metriche Mean@3 che Best@3 per validità e stabilità.
• Crucialmente, abilitare la generazione condizionata da frammenti ha avuto un impatto trascurabile sulla generazione incondizionata, dimostrando che entrambi i paradigmi possono essere supportati all'interno di un unico modello.

Significato e Rivendicazioni

Il paper sostiene che KRONOS fornisca una alternativa promettente agli attuali approcci autoregressivi e basati sulla diffusione nello spazio dei dati. La sua importanza risiede nel:

• Colmare il Gap: Riduce sensibilmente il divario di prestazioni tra modelli autoregressivi e di diffusione, raggiungendo risultati d'eccellenza tra i metodi AR pur eguagliando i baseline di diffusione.
• Flessibilità: Unifica i punti di forza dei modelli autoregressivi (lunghezza variabile, condizionamento parziale) con l'alta qualità di generazione dei modelli di diffusione, il tutto all'interno di una singola formulazione nello spazio latente.
• Utilità Pratica: Supportando la generazione condizionata da frammenti senza modifiche architettoniche, offre uno strumento pratico per compiti di chimica medicinale dove le molecole vengono spesso progettate elaborando scaffold noti.

Gli autori segnalano alcune limitazioni, tra cui la dipendenza dall'espressività dello spazio latente pre-addestrato e l'uso di uno schema di decomposizione BRICS fisso, suggerendo che il lavoro futuro potrebbe esplorare l'apprendimento congiunto delle astrazioni dei frammenti con il modello generativo o l'estensione del framework al design guidato dalle proprietà.