Resumo Técnico: KRONOS – Difusão Latente Autorregressiva para Geração de Moléculas 3D

Declaração do Problema

A geração de moléculas tridimensionais (3D) é uma tarefa crítica na descoberta de fármacos e no design molecular. As abordagens de estado da arte atuais baseiam-se principalmente em modelos de difusão e correspondência de fluxo (flow matching). Embora esses métodos alcancem alta qualidade de geração, eles tipicamente exigem que o tamanho molecular (número de átomos) seja especificado a priori, o que limita sua flexibilidade para geração de comprimento variável.

Inversamente, modelos autorregressivos (AR) suportam naturalmente a geração de comprimento variável e o condicionamento em estruturas moleculares parciais (por exemplo, elaboração de scaffolds conhecidos), o que é altamente valioso na química medicinal. No entanto, os modelos AR existentes enfrentam dois desafios principais:

1. Lacuna de Desempenho: Historicamente, os modelos AR ficaram atrás dos modelos de difusão em termos de qualidade de geração, particularmente no que diz respeito à fidelidade geométrica e validade química.
2. Equilíbrio de Objetivos: Alcançar um desempenho forte tanto na geração incondicional (design de novo) quanto na geração condicionada por fragmentos (completude molecular) dentro de uma única arquitetura permanece difícil. Muitos modelos AR têm dificuldade em generalizar entre esses paradigmas sem modificações arquiteturais ou compensações de desempenho.

Adicionalmente, os geradores moleculares latentes existentes têm dependido quase exclusivamente de difusão, deixando o potencial da geração autorregressiva no espaço latente amplamente inexplorado.

Metodologia: KRONOS

Os autores introduzem o KRONOS, um framework de difusão autorregressiva latente que gera moléculas 3D diretamente no espaço latente de um Autoencoder Unificado (UAE) pré-treinado. O método combina três componentes principais:

1. Geração Autorregressiva Latente

Em vez de operar no espaço de dados brutos (prevendo tipos de átomos/ligações discretos e coordenadas contínuas simultaneamente), o KRONOS opera no espaço latente contínuo de um UAE pré-treinado (Luo et al., 2025).

• Representação: Uma molécula 3D é codificada em uma sequência de tokens latentes $Z = \{z_1, \dots, z_n\}$, onde cada token codifica conjuntamente a identidade atômica, o ambiente de ligação e a geometria 3D.
• Fatoração: O modelo aprende uma distribuição autorregressiva da esquerda para a direita:
  $$p_\theta(z_{1:n}) = \prod_{i=1}^n p_\theta(z_i | z_{<i})$$
• Comprimento Variável: Como o espaço latente carece de um token natural de fim de sequência, o KRONOS emprega um classificador de parada dedicado (um MLP leve) treinado junto com o gerador para prever a terminação da sequência dinamicamente, eliminando a necessidade de uma contagem de átomos predefinida.

2. Predição de Tokens Baseada em Difusão

Para modelar a distribuição condicional de tokens latentes contínuos $p_\theta(z_i | z_{<i})$, o KRONOS utiliza uma formulação de Perda de Difusão (Li et al., 2024).

• Um transformer causal codifica os tokens precedentes em uma representação de contexto $c_i$.
• Uma rede de denoising de difusão $\epsilon_\theta$ prevê o ruído adicionado ao token latente atual $z_i$, condicionado a $c_i$ e ao passo de tempo da difusão.
• Esta abordagem substitui a predição de tokens discretos por um objetivo de denoising, permitindo a modelagem flexível de variáveis contínuas sem discretização.

3. Estratégia de Treinamento Misto (Geração Condicionada por Fragmentos)

Para suportar tanto a geração incondicional quanto a condicionada por fragmentos em um único modelo, os autores introduzem uma estratégia de treinamento inspirada no paradigma Fill-in-the-Middle (FIM).

• Decomposição de Fragmentos: As moléculas são decompostas em fragmentos quimicamente significativos usando um esquema baseado em BRICS.
• Construção de Sequência: Durante o treinamento, cada molécula é transformada em uma:
  • Sequência de ordem padrão (para geração incondicional).
  • Sequência condicionada por fragmentos, onde um subconjunto aleatório de fragmentos atua como uma "semente" (contexto), e os fragmentos restantes são o alvo de geração.
• Mecanismo: O modelo é treinado para gerar autorregressivamente os fragmentos alvo condicionados aos fragmentos semente. Isso é alcançado atribuindo uma ordenação determinística única aos fragmentos e átomos dentro deles, em vez de marginalizar sobre todas as permutações (como em modelos AR de Ordem-Agnóstica). Isso permite que o modelo aprenda capacidades de completude molecular sem mudanças arquiteturais.

Principais Contribuições

1. Framework KRONOS: Um novo framework de difusão autorregressiva latente que modela conjuntamente a topologia do grafo molecular e a geometria em um espaço latente unificado.
2. Estratégia de Treinamento Unificada: Uma abordagem de treinamento mista que permite a geração incondicional e condicionada por fragmentos dentro de um único modelo autorregressivo da esquerda para a direita, inspirado pelo paradigma FIM.
3. Geração de Comprimento Variável: Um framework unificado que suporta naturalmente a geração de comprimento variável via um classificador de parada, evitando a necessidade de especificação de tamanho a priori.
4. Desempenho: Demonstração de que a geração autorregressiva latente pode alcançar um desempenho de liderança entre os métodos autorregressivos, permanecendo competitivo com modelos de difusão, sem sacrificar a qualidade incondicional ao habilitar o condicionamento por fragmentos.

Resultados Experimentais

Os autores avaliaram o KRONOS nos conjuntos de dados QM9 e GEOM-Drugs.

Geração Incondicional

• QM9: O KRONOS alcançou a maior Validade 3D (97,3%) e o desempenho geral mais forte entre os métodos autorregressivos. Permaneceu competitivo com modelos de difusão líderes (ex: JODO, EDM) em métricas incluindo Estabilidade Atômica 3D, Comprimento de Ligação e MMD de Ângulo de Ligação.
• GEOM-Drugs: O KRONOS alcançou o desempenho geral mais forte entre todos os baselines avaliados (tanto autorregressivos quanto de difusão).
• Plausibilidade Física: Usando otimização de geometria semi-empírica (GFN2-xTB), as conformações geradas pelo KRONOS mostraram baixo RMSD (Root Mean Square Distance) e diferenças de energia ($\Delta E$) em relação às estruturas otimizadas, indicando que elas residem próximas de regiões fisicamente plausíveis do espaço conformacional.

Geração Condicionada por Fragmentos

• Em um conjunto de dados de ~15k pares molécula-fragmento, o KRONOS superou substancialmente o baseline autorregressivo QUETZAL em ambas as métricas Mean@3 e Best@3 para validade e estabilidade.
• Crucialmente, habilitar a geração condicionada por fragmentos teve um impacto negligenciável na geração incondicional, demonstrando que ambos os paradigmas podem ser suportados dentro de uma única arquitetura.

Significância e Alegações

O artigo afirma que o KRONOS fornece uma alternativa promissora às abordagens existentes de difusão e autorregressão de espaço de dados. Sua significância reside em:

• Redução da Lacuna: Ele reduz substancialmente a lacuna de desempenho entre modelos autorregressivos e de difusão, alcançando resultados de liderança entre métodos AR enquanto iguala os baselines de difusão.
• Flexibilidade: Unifica os pontos fortes dos modelos autorregressivos (comprimento variável, condicionamento parcial) com a alta qualidade de geração dos modelos de difusão, tudo dentro de uma formulação de espaço latente único.
• Utilidade Prática: Ao suportar a geração condicionada por fragmentos sem modificações arquiteturais, oferece uma ferramenta prática para tarefas de química medicinal onde moléculas são frequentemente desenhadas através da elaboração de scaffolds conhecidos.

Os autores observam limitações, incluindo a dependência da expressividade do espaço latente pré-treinado e o uso de um esquema de decomposição BRICS fixo, sugerindo que trabalhos futuros poderiam explorar a aprendizagem de abstrações de fragmentos conjuntamente com o modelo generativo ou estender o framework para o design guiado por propriedades.