ConRetroBert: EMA Stabilized Dual Encoders for Template-Based Single-Step Retrosynthesis

Riepilogo Tecnico di ConRetroBert

Formulazione del Problema
La retrosintesi basata su template a singolo passo mira a prevedere gli insiemi di reagenti per un prodotto target selezionando e applicando un template di reazione esplicito (codificato come trasformazione SMARTS). Sebbene questo approccio offra intrinseca interpretabilità e tracciabilità – cruciali per la pianificazione della sintesi – i metodi basati su template sono stati storicamente considerati meno competitivi rispetto ai modelli senza template. Gli autori sostengono che tale limitazione non derivi dall'uso dei template stessi, ma dalla formulazione di apprendimento dominante: la classificazione globale su una libreria a coda lunga di regole di reazione. Tale formulazione crea tre problemi: i template frequenti dominano la funzione di perdita, il denominatore softmax tratta i template non osservati ma chimicamente validi come etichette negative concorrenti (incompletezza delle etichette) e lo strato di output scala male con librerie di template di grandi dimensioni.

Metodologia: ConRetroBert
Il documento presenta ConRetroBert, un framework a doppio encoder che riformula la retrosintesi basata su template come un problema di recupero e classificazione denso, anziché come classificazione globale. L'architettura impiega due encoder Transformer indipendenti (uno per gli SMILES del prodotto e uno per i SMARTS del template) che mappano gli input in uno spazio di embedding condiviso. Il processo di addestramento avviene in due fasi:

1. Fase 1: Pre-addestramento Contrastivo
   Il modello apprende uno spazio di rappresentazione congiunto prodotto-template utilizzando un pre-addestramento contrastivo simmetrico (perdita InfoNCE). Per un minibatch di esempi accoppiati, l'encoder del prodotto si allinea con il suo template corrispondente, mentre gli altri esempi nel batch fungono da negativi. Ciò stabilisce uno spazio di embedding ricercabile in cui i template possono essere recuperati tramite ricerca del vicino più prossimo.

2. Fase 2: Classificazione Listwise Multi-Positiva
   Questa fase affina lo spazio di recupero in una politica di proposta di template. Invece di classificare contro l'intero vocabolario, il modello classifica un insieme candidato $C(x)$ per ogni prodotto. Tale insieme include:

   • Positivi Osservati: Tutti i template registrati con il prodotto nei dati di addestramento.
   • Negativi Difficili: Template chimicamente simili recuperati da un indice FAISS.
   • Negativi nel Batch: Template associati ad altri prodotti nel batch.

   Il modello ottimizza un obiettivo di verosimiglianza negativa logaritmica listwise multi-positiva, incoraggiando la separazione dei positivi osservati dai negativi difficili chimicamente simili. Questa formulazione corrisponde meglio al problema decisionale al momento dell'inferenza, dove il modello deve scegliere tra un piccolo insieme di candidati recuperati piuttosto che l'intera libreria.

Adattamento del Template Stabilizzato da EMA
Una sfida tecnica chiave affrontata è l'instabilità dell'estrazione di negativi difficili quando l'encoder del template viene aggiornato. Se l'encoder del template cambia troppo rapidamente, l'indice di recupero (utilizzato per estrarre negativi difficili) diventa obsoleto. ConRetroBert risolve questo problema utilizzando un meccanismo di Media Mobile Esponenziale (EMA):

• Un encoder del template attivo viene aggiornato tramite gradienti dalla funzione di perdita di classificazione.
• Un encoder "ombra" EMA a movimento lento (aggiornato con $\alpha = 0.999$) viene utilizzato per ricostruire l'indice di recupero ed estrarre negativi difficili a livello di epoca.
  Ciò permette alle rappresentazioni dei template di adattarsi all'obiettivo di classificazione senza destabilizzare il processo di costruzione dei candidati.

Contributi Chiave

1. Formulazione Basata sul Recupero: Gli autori riformulano la retrosintesi basata su template come recupero denso prodotto-template seguito da classificazione listwise dell'insieme candidato, sostituendo la classificazione globale con un obiettivo allineato alle decisioni al momento dell'inferenza.
2. Addestramento Stabilizzato da EMA: Introducono un encoder del template addestrabile stabilizzato da un'ombra EMA, consentendo un adattamento controllato delle rappresentazioni dei template mantenendo al contempo un indice di recupero stabile per l'estrazione di negativi difficili.
3. Pipeline di Valutazione Rigorosa: Il documento fornisce un protocollo di valutazione completo che include la classificazione dell'insieme di reagenti, l'applicazione del template, la deduplicazione e il trasferimento controllato da perdite da USPTO-Full a USPTO-50k.
4. Validazione Empirica: Attraverso analisi locali e a coda lunga, gli autori dimostrano che la classificazione dell'insieme candidato rafforza sostanzialmente la retrosintesi basata su template, mantenendo al contempo la tracciabilità esplicita delle regole.

Risultati
Gli esperimenti sono stati condotti sul benchmark USPTO-50k (regime locale) e tramite apprendimento per trasferimento da un checkpoint USPTO-Full controllato per perdite.

• Prestazioni Locali USPTO-50k:

  • Il passaggio dalla Fase 1 (recupero contrastivo) alla Fase 2 (classificazione con encoder del template congelato) ha migliorato l'accuratezza di reazione top-1 dal 50,5% al 61,4%.
  • L'aggiunta dell'adattamento del template stabilizzato da EMA (Fase 2 EMA Mid) ha ulteriormente migliorato l'accuratezza top-1 al 62,4%.
  • Il modello ha raggiunto accuratezze top-3, top-5 e top-10 rispettivamente del 81,6%, 85,3% e 87,8%.
  • Il guadagno di prestazioni è attribuito principalmente alla formulazione di classificazione dell'insieme candidato della Fase 2, che separa meglio i positivi dai negativi difficili rispetto al solo allineamento contrastivo.

• Apprendimento per Trasferimento:

  • Il fine-tuning da un checkpoint USPTO-Full controllato per perdite ha raggiunto un'accuratezza top-1 del 75,4% su USPTO-50k, dimostrando la scalabilità dell'approccio con dati di pre-addestramento più ampi.

• Coda Lunga e Affidabilità:

  • La previsione basata sul recupero ha mostrato una forza particolare nella coda lunga dei template rari, superando la valutazione basata sulla classificazione per le classi rare.
  • L'analisi ha rivelato che il 42,8% delle previsioni corrette di reagenti top-1 sono stati "successi secondari", generati da un template diverso dall'etichetta ground-truth registrata ma chimicamente equivalente dopo l'applicazione e la deduplicazione. Ciò evidenzia che il modello impara a recuperare alternative chimicamente valide, non solo l'etichetta specifica registrata.

Significato e Affermazioni
Il documento sostiene che la retrosintesi basata su template non deve sacrificare le prestazioni predittive per l'interpretabilità. Spostandosi dalla classificazione globale alla classificazione basata sul recupero con estrazione di negativi difficili e stabilizzazione EMA, ConRetroBert sfida l'assunto che un'alta accuratezza richieda l'abbandono dei template di reazione espliciti. Gli autori affermano che il loro metodo produce previsioni non solo competitive con i modelli senza template e semi-template all'avanguardia, ma anche chimicamente ispezionabili, tracciabili a specifiche regole di trasformazione e compatibili con i vincoli di pianificazione della sintesi a valle. Il lavoro suggerisce che la "debolezza" dei metodi basati su template è un problema di formulazione risolvibile attraverso tecniche moderne di apprendimento per recupero, piuttosto che un limite intrinseco dell'uso di regole esplicite.