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

Resumo Técnico do ConRetroBert

Formulação do Problema
A retrosíntese de passo único baseada em templates visa prever conjuntos de reagentes para um produto-alvo selecionando e aplicando um template de reação explícito (codificado como uma transformação SMARTS). Embora essa abordagem ofereça interpretabilidade e rastreabilidade inerentes — cruciais para o planejamento de síntese —, os métodos baseados em templates foram historicamente considerados menos competitivos do que os modelos sem templates. Os autores argumentam que essa limitação não decorre do uso dos próprios templates, mas da formulação de aprendizado dominante: classificação global sobre uma biblioteca de longo alcance de regras de reação. Essa formulação cria três problemas: templates frequentes dominam a função de perda, o denominador softmax trata templates não observados, mas quimicamente válidos, como rótulos negativos concorrentes (incompletude de rótulos), e a camada de saída escala mal com bibliotecas grandes de templates.

Metodologia: ConRetroBert
O artigo apresenta o ConRetroBert, um framework de codificador duplo que reformula a retrosíntese baseada em templates como um problema de recuperação e classificação densa, em vez de classificação global. A arquitetura emprega dois codificadores Transformer independentes (um para SMILES do produto e outro para SMARTS do template) que mapeiam as entradas para um espaço de incorporação compartilhado. O processo de treinamento ocorre em duas etapas:

1. Etapa 1: Pré-treinamento Contrastivo
   O modelo aprende um espaço de representação conjunta produto-template usando pré-treinamento contrastivo simétrico (perda InfoNCE). Para um minibatch de exemplos pareados, o codificador do produto alinha-se com seu template correspondente, enquanto outros exemplos no batch servem como negativos. Isso estabelece um espaço de incorporação pesquisável onde os templates podem ser recuperados por meio de busca por vizinho mais próximo.

2. Etapa 2: Classificação Listwise Multi-Positiva
   Esta etapa refina o espaço de recuperação em uma política de proposta de templates. Em vez de classificar contra o vocabulário completo, o modelo classifica um conjunto candidato $C(x)$ para cada produto. Este conjunto inclui:

   • Positivos Observados: Todos os templates registrados com o produto nos dados de treinamento.
   • Negativos Difíceis: Templates quimicamente similares recuperados de um índice FAISS.
   • Negativos no Batch: Templates associados a outros produtos no batch.

   O modelo otimiza um objetivo de verossimilhança negativa listwise multi-positiva, incentivando a separação dos positivos observados de negativos difíceis quimicamente similares. Essa formulação corresponde melhor ao problema de decisão no momento da inferência, onde o modelo deve escolher entre um pequeno conjunto de candidatos recuperados, em vez de toda a biblioteca.

Adaptação de Templates Estabilizada por EMA
Um desafio técnico chave abordado é a instabilidade da mineração de negativos difíceis quando o codificador de templates é atualizado. Se o codificador de templates mudar muito rapidamente, o índice de recuperação (usado para minerar negativos difíceis) torna-se desatualizado. O ConRetroBert resolve isso usando um mecanismo de Média Móvel Exponencial (EMA):

• Um codificador de templates ativo é atualizado via gradientes da função de perda de classificação.
• Um codificador sombra EMA de movimento lento (atualizado com $\alpha = 0,999$) é usado para reconstruir o índice de recuperação e minerar negativos difíceis no nível da época.
  Isso permite que as representações de templates se adaptem ao objetivo de classificação sem desestabilizar o processo de construção de candidatos.

Principais Contribuições

1. Formulação Baseada em Recuperação: Os autores reformulam a retrosíntese baseada em templates como recuperação densa produto-template seguida de classificação listwise de conjunto candidato, substituindo a classificação global por um objetivo alinhado com decisões no momento da inferência.
2. Treinamento Estabilizado por EMA: Eles introduzem um codificador de templates treinável estabilizado por uma sombra EMA, permitindo adaptação controlada das representações de templates enquanto mantêm um índice de recuperação estável para mineração de negativos difíceis.
3. Pipeline de Avaliação Rigoroso: O artigo fornece um protocolo de avaliação abrangente, incluindo classificação de conjunto de reagentes, aplicação de template, deduplicação e transferência controlada de vazamento do USPTO-Full para o USPTO-50k.
4. Validação Empírica: Através de análises locais e de cauda longa, os autores demonstram que a classificação de conjunto candidato fortalece substancialmente a retrosíntese baseada em templates, mantendo ao mesmo tempo a rastreabilidade explícita das regras.

Resultados
Os experimentos foram conduzidos no benchmark USPTO-50k (regime local) e via aprendizado por transferência a partir de um checkpoint USPTO-Full controlado quanto a vazamentos.

• Desempenho Local USPTO-50k:

  • A transição da Etapa 1 (recuperação contrastiva) para a Etapa 2 (classificação com codificador de templates congelado) melhorou a precisão de reação top-1 de 50,5% para 61,4%.
  • Adicionar adaptação de templates estabilizada por EMA (Etapa 2 EMA Mid) melhorou ainda mais a precisão top-1 para 62,4%.
  • O modelo alcançou precisões top-3, top-5 e top-10 de 81,6%, 85,3% e 87,8%, respectivamente.
  • O ganho de desempenho é atribuído principalmente à formulação de classificação de conjunto candidato da Etapa 2, que separa melhor os positivos dos negativos difíceis do que o alinhamento contrastivo isolado.

• Aprendizado por Transferência:

  • O ajuste fino a partir de um checkpoint USPTO-Full controlado quanto a vazamentos atingiu 75,4% de precisão top-1 no USPTO-50k, demonstrando a escalabilidade da abordagem com dados de pré-treinamento maiores.

• Cauda Longa e Confiabilidade:

  • A previsão baseada em recuperação mostrou força particular na cauda longa de templates raros, superando a pontuação baseada em classificação em classes raras.
  • A análise revelou que 42,8% das previsões corretas de reagentes top-1 foram "sucessos secundários", gerados por um template diferente do rótulo de verdade fundamental registrado, mas quimicamente equivalente após aplicação e deduplicação. Isso destaca que o modelo aprende a recuperar alternativas quimicamente válidas, e não apenas o rótulo específico registrado.

Significado e Alegações
O artigo argumenta que a retrosíntese baseada em templates não precisa sacrificar o desempenho preditivo pela interpretabilidade. Ao mudar da classificação global para a classificação baseada em recuperação com mineração de negativos difíceis e estabilização por EMA, o ConRetroBert desafia a suposição de que alta precisão exige abandonar templates de reação explícitos. Os autores afirmam que seu método produz previsões que não apenas são competitivas com modelos sem templates e semi-templates de última geração, mas também permanecem inspecionáveis quimicamente, rastreáveis a regras de transformação específicas e compatíveis com restrições de planejamento de síntese a jusante. O trabalho sugere que a "fraqueza" dos métodos baseados em templates é uma questão de formulação solucionável através de técnicas modernas de aprendizado por recuperação, e não uma limitação inerente ao uso de regras explícitas.