Toward Better Geometric Representations for Molecule Generative Models

Resumo Técnico: LENSEs – Aprimoramento Latente para Codificações Estruturais Não Suaves

Declaração do Problema
A geração condicionada a representações geométricas (GeoRCG) emergiu como um paradigma eficiente para a geração de moléculas 3D, desacoplando a modelagem de representações moleculares da geração de estruturas 3D. Neste arcabouço, um codificador molecular pré-treinado (por exemplo, UniMol, Frad) define um espaço latente, um Modelo de Difusão de Representação (RDM) amostra um vetor ao nível da molécula a partir deste espaço, e um gerador reconstrói a estrutura 3D condicionada a este vetor.

No entanto, a eficácia do GeoRCG é fundamentalmente limitada pela qualidade do espaço latente fornecido pelos codificadores pré-treinados. Os autores identificam dois gargalos críticos:

1. Variedades de Representação Não Suaves: Codificadores pré-treinados, otimizados principalmente para previsão de propriedades via remoção de ruído de coordenadas, produzem representações com constantes de Lipschitz elevadas. Pequenas perturbações na conformação molecular levam a deslocamentos desproporcionalmente grandes no espaço latente. Esta não suavidade torna a distribuição difícil de modelar para o RDM, frequentemente necessitando de mecanismos compensatórios como injeção de ruído ajustada manualmente.
2. Semântica Multinível Subexplorada: Codificadores pré-treinados exibem características hierárquicas onde camadas rasas capturam motivos estruturais locais (por exemplo, grupos funcionais) e camadas mais profundas codificam semânticas globais. Pipelines existentes de GeoRCG tipicamente utilizam apenas a representação da camada final, descartando a rica informação semântica preservada em camadas anteriores. Consequentemente, o gerador deve reaprender essas semânticas do zero, aumentando a dificuldade de otimização.

Metodologia: LENSEs
Para abordar essas questões, os autores propõem LENSEs (Aprimoramento Latente para Codificações Estruturais Não Suaves), um arcabouço que refina ativamente o espaço latente de condicionamento durante o treinamento generativo. Em vez de depender passivamente de um codificador congelado, LENSEs introduz três mecanismos complementares:

1. Cabeça de Representação com Pooling de Camadas Aprendível:
   Em vez de utilizar uma única saída determinística do codificador, LENSEs agrega características de múltiplas camadas. Emprega um mecanismo de pooling de camadas aprendível para ponderar representações de diferentes profundidades do codificador ($H^{(l)}$). Essas características agrupadas são mapeadas para os parâmetros de uma posterior Gaussiana diagonal ($\mu, \sigma$) via cabeças de projeção. O código latente $z$ é amostrado usando o truque de reparametrização. Esta cabeça é treinada conjuntamente com o gerador para extrair representações que são simultaneamente semanticamente ricas e favoráveis à geração, regularizadas por uma perda de divergência KL para garantir um espaço latente compacto.

2. Perda Perceptiva Molecular ($L_{PERC}$):
   Inspirada em perdas perceptivas em visão computacional, esta perda alinha a representação semântica de alto nível da molécula removida de ruído (gerada) com a da molécula limpa de verdade fundamental dentro do espaço de características do codificador congelado. Ao passar tanto a molécula gerada quanto a molécula limpa pelo codificador e minimizar a distância entre suas representações agrupadas ao nível do grafo, o gerador é incentivado a preservar padrões estruturais significativos (por exemplo, grupos funcionais) em vez de apenas minimizar erros de coordenadas de baixo nível. Uma operação de stop-gradient é aplicada à representação da verdade fundamental.

3. Perda de Alinhamento de Representação (REPA) ($L_{REPA}$):
   Motivada pela observação de que os estados ocultos do gerador se alinham naturalmente com as características do codificador, esta perda alinha explicitamente as características ao nível dos nós do gerador com as do codificador pré-treinado. Um projetor aprendível mapeia as características do gerador para o espaço de características do codificador, e a similaridade de cosseno é maximizada entre os nós correspondentes. Isso reduz a lacuna semântica entre pré-treinamento e geração, acelerando o treinamento e melhorando a aquisição de conhecimento quimicamente significativo.

O objetivo total de treinamento combina a perda de geração base ($L_{gen}$) com a perda KL, perda perceptiva e perda REPA.

Resultados Chave
Os autores avaliam LENSEs nos conjuntos de dados GEOM-DRUG e QM9, comparando-o com baselines de última geração, incluindo GeoRCG, CanonFlow e SemlaFlow.

• Qualidade de Geração: No desafiador conjunto de dados GEOM-DRUG, LENSEs alcança desempenho de última geração. Alcança 97,28% de validade e 98,51% de estabilidade molecular, superando o melhor GeoRCG anterior (95,3% de validade, 97,20% de estabilidade) e o CanonFlow (95,9% de validade). Também alcança as métricas de energia e tensão mais baixas, indicando plausibilidade física superior.
• Suavidade de Representação: A análise da constante de Lipschitz revela que o espaço latente refinado é significativamente mais suave que o espaço do codificador bruto. No GEOM-DRUG, a constante de Lipschitz empírica é reduzida em 4,6×, confirmando que a variedade de representação é mais adequada para modelagem por difusão.
• Riqueza de Representação: A rank efetiva da representação aumenta de 65,32 (codificador bruto) para 88,42 (LENSEs), indicando uma redução no colapso dimensional e um espaço latente mais rico e estruturado.
• Utilidade em Tarefas Descendentes: Quando utilizado para tarefas de regressão descendente em propriedades do QM9 com um codificador e cabeça de representação congelados, LENSEs supera o codificador bruto em todas as propriedades, demonstrando que o próprio processo de treinamento generativo atua como um paradigma de pré-treinamento eficaz para refinar representações moleculares.
• Eficiência de Amostragem: LENSEs demonstra eficiência de amostragem aprimorada, alcançando alta validade e estabilidade com menos etapas de geração em comparação com a baseline.

Significado e Alegações
O artigo afirma que a geração de alta qualidade de moléculas 3D depende não apenas da arquitetura do gerador, mas criticamente da qualidade da representação de condicionamento. LENSEs demonstra que o espaço latente de codificadores pré-treinados pode ser sistematicamente melhorado para ser mais suave, mais informativo semanticamente e melhor alinhado com objetivos generativos.

Os autores postulam que seu trabalho estabelece uma relação bidirecional entre qualidade de geração e qualidade de representação. Especificamente, sugerem que o treinamento generativo com objetivos de alinhamento pode servir como um paradigma potencial de pré-treinamento para codificadores moleculares, refinando representações para serem mais preditivas e úteis para tarefas descendentes. O trabalho destaca que representações moleculares úteis não são inerentemente favoráveis à geração e requerem refinamento explícito para desbloquear seu potencial pleno na modelagem generativa.