Resumo Técnico: SinAE – Um Autoencoder de Flow-Matching de Arquitetura Única para Sistemas Atômicos Cross-Domain

1. Definição do Problema

A geração de estruturas atômicas 3D é crítica para a descoberta de fármacos, design de materiais e engenharia de proteínas. No entanto, os pipelines generativos atuais estão fragmentados em três domínios distintos: pequenas moléculas, cristais e proteínas. Cada domínio depende de arquiteturas especializadas adaptadas aos seus respectivos priors estruturais:

• Moléculas: Tipicamente utilizam redes de passagem de mensagens equivariantes para lidar com a conectividade de ligações e regras de valência.
• Cristais: Requerem o tratamento de condições de contorno periódicas e invariância de base de rede (lattice-basis invariance).
• Proteínas: Frequentemente dependem de parametrizações ao nível de resíduos (ex: frames SE(3)) ou mecanismos de atenção triangular.

Essa fragmentação exacerba a escassez de dados, uma vez que os conjuntos de dados por domínio (ex: QM9 com 134k moléculas, MP-20 com 45k cristais) são pequenos para os padrões modernos. Embora o treinamento cross-domain possa mitigar isso, as abordagens existentes enfrentam dois obstáculos primários:

1. Limitações de Geração Direta: Gerar diretamente no espaço de coordenadas 3D força um único modelo a aprender simultaneamente composições válidas, geometria precisa e priors estruturais heterogêneos (periodicidade, backbones de resíduos, valência), o que geralmente exige tokenização complexa e específica de domínio ou geração em estágios.
2. Limitações de Autoencoders Latentes: Os autoencoders latentes cross-domain existentes sofrem de problemas de fidelidade de reconstrução (decoders com perda de informação introduzem erros geométricos que os priors não conseguem corrigir) e complexidade de treinamento (exigindo combinações de perda e ajuste de pesos específicos para cada domínio). Nenhum método anterior alcançou um autoencoder simultaneamente sem perdas (lossless), arquiteturalmente geral e com objetivo unificado para todos os três domínios.

2. Metodologia: SinAE

Os autores introduzem o SinAE, um autoencoder de flow-matching de arquitetura única projetado para moléculas, cristais e proteínas.

Arquitetura

• Backbone Unificado: O SinAE utiliza um Transformer vanilla tanto para o encoder quanto para o decoder. Ele evita explicitamente operadores equivariantes, passagem de mensagem em grafos ou módulos específicos de domínio (ex: camadas conscientes de rede ou atenção triangular).
• Featurização de Entrada:
  • Coordenadas: Coordenadas cartesianas brutas são utilizadas sem descritores invariantes (ex: distâncias par a par) ou construção de frames canônicos.
  • Periodicidade: Para cristais, os três vetores da rede são anexados como "tokens de átomos virtuais" à sequência, permitindo que o modelo processe a matriz da rede de forma idêntica às coordenadas atômicas.
  • Equivariância: Em vez de equivariância arquitetural, o modelo baseia-se em aumento por rotação e translação SO(3) aleatória durante o treinamento para garantir robustez de orientação.
• Tokenização: Átomos são representados pela soma de embeddings de elementos aprendidos e projeções lineares de suas coordenadas. Para proteínas, os resíduos são tratados como tokens únicos em suas posições $C_\alpha$.

Objetivos de Treinamento

O modelo é treinado em duas etapas:

1. Autoencoder de Flow-Matching:
   • Encoder: Um Transformer estocástico mapeia a estrutura de entrada $G=(X, A)$ para um código latente por token $z$.
   • Decoder: Um Transformer não-equivariante reconstrói a estrutura a partir de $z$ usando flow matching iterativo. Diferente dos decoders VAE de passagem única, o decoder refina a estrutura ao longo de múltiplos passos de ODE.
   • Função de Perda: Um objetivo unificado combina flow matching condicional contínuo para coordenadas ($L_{coord}$) e flow matching discreto para tipos de átomos ($L_{atom}$), além de um regularizador KL leve ($L_{KL}$). Não são necessários pesos de perda ou perdas auxiliares específicos de domínio.
2. Prior de Flow Matching Latente:
   • Um Diffusion Transformer (DiT) vanilla aprende a distribuição $p(z)$ sobre as saídas do encoder.
   • Como o termo KL empurra a posterior em direção a uma Gaussiana padrão, o prior opera em um espaço latente suave e aproximadamente Gaussiano, sem restrições equivariantes.
   • Durante a geração, um rótulo de domínio é usado como condição (via classifier-free guidance) para direcionar o DiT para moléculas, cristais ou proteínas.

3. Principais Contribuições

1. Tokenização de Arquitetura Única: O SinAE é a primeira arquitetura a lidar com moléculas, cristais e backbones de proteínas usando um único Transformer vanilla sem operadores equivariantes ou específicos de domínio.
2. Reconstrução Próxima de Sem Perdas (Near-Lossless): Ao transferir o fardo da reconstrução para um decoder de flow-matching iterativo, o SinAE atinge um RMSD sub-mili-Ångström entre os domínios:
   • Moléculas: 0.0002 Å no QM9, 0.0006 Å no GEOM-Drugs.
   • Cristais: 0.0017 Å no MP-20.
   • Proteínas: ~0.013 Å de CA-RMSD de backbone no CASP15.
   • Isso representa uma melhoria de 10 a 100 vezes na fidelidade de reconstrução em relação aos baselines latentes anteriores, com taxas de correspondência de estrutura de 100% em moléculas e cristais.
3. Transferência Cross-Domain Positiva: O treinamento conjunto em QM9 e MP-20 melhora estritamente o desempenho de reconstrução e geração em ambos os domínios comparado ao treinamento de domínio único, fornecendo evidências de que um latente atômico compartilhado captura regularidades geométricas transferíveis (ex: distribuições de comprimento de ligação, padrões de coordenação).

4. Resultados Experimentais

Os autores avaliaram o SinAE em benchmarks padrão (QM9, GEOM-Drugs, MP-20, AFDB-FS/CASP15, QMOF).

• Fidelidade de Reconstrução: O SinAE supera todos os baselines (incluindo GeoLDM, CDVAE, ADiT, UAE-3D, ProteinAE) em RMSD e taxas de correspondência de estrutura: o decoder iterativo resolve o posicionamento atômico precocemente e refina a geometria de ligação/estéricas posteriormente, alcançando alta precisão sem relaxação post-hoc.
• Geração Molecular: No QM9 e GEOM-Drugs, o prior DiT latente alcança alta validade (escores PB Valid de 0.97 e 0.93, respectivamente), superando ou igualando geradores equivariantes especializados (ex: EQGAT-diff, SemlaFlow).
• Geração de Cristais: No MP-20, o modelo SinAE conjunto alcança os maiores escores S.U.N. (Estável, Único, Novo) e M.S.U.N. entre os métodos comparados, incluindo aqueles com arquiteturas invariantes de rede. Demonstra superior estabilidade termodinâmica (estruturas metaestáveis) apesar de usar um Transformer genérico.
• Geração de Proteínas: O SinAE alcança alta designabilidade (0.96 scRMSD $\le$ 2.0 Å) e diversidade simultaneamente, superando ProteinAE e LatentDiff. Em tarefas de co-design sequência-estrutura, atinge escores competitivos de scTM e scRMSD em comprimentos de 100 e 200.

5. Significância e Alegações

O artigo afirma que o SinAE demonstra que a geração atômica cross-domain pode ser unificada sob uma única arquitetura simples sem sacrificar a fidelidade.

• Simplicidade vs. Desempenho: Os autores argumentam que operadores complexos específicos de domínio não são estritamente necessários para uma geração de alta qualidade se o decoder for suficientemente poderoso (flow matching iterativo) e o espaço latente for adequadamente regularizado.
• Eficiência de Dados: Os resultados fornecem evidência direta de que latentes atômicos compartilhados permitem que diferentes populações químicas se beneficiem mutuamente de seus dados, melhorando o desempenho em regimes de poucos dados.
• Desacoplamento de Geometria e Priors: O trabalho sugere que a complexidade geométrica pode ser tratada pelo refinamento iterativo do decoder, permitindo que o encoder e o prior operem em um espaço latente simplificado e agnóstico ao domínio.

Limitações Reconhecidas:

• A avaliação de proteínas é atualmente limitada à geração de nível de backbone; a modelagem de todos os átomos ainda não foi abordada.
• O modelo atual não suporta geração condicionada por propriedade, composição ou comprimento em um único checkpoint através de todos os três domínios simultaneamente.
• Cadeias de proteínas mais longas (>256 resíduos) requerem retreinamento com orçamentos de comprimento maiores.