Resumo Técnico: Cache de Características Preditivo para Aceleração sem Treinamento de Geração de Geometria Molecular

Declaração do Problema

Modelos generativos profundos, particularmente modelos de flow matching e difusão, tornaram-se o estado da arte para a geração de geometrias moleculares de alta fidelidade (coordenadas 3D, tipos de átomos e ordens de ligação). No entanto, sua aplicação prática na descoberta de fármacos é dificultada pelos custos computacionais significativos durante a inferência. Gerar uma única molécula pode exigir centenas de avaliações de redes neurais (passos) para integrar a equação diferencial ordinária (ODE) subjacente. Ao escalar para as centenas de milhares ou milhões de amostras necessárias para triagem combinatória ou triagem virtual, essa latência de inferência torna-se o principal gargalo.

Métodos de aceleração existentes, como reparametrização de trajetória, destilação progressiva e geração em espaço latente, geralmente exigem treinamento adicional, incorrendo em custos de dados, computação e tempo. Além disso, simplesmente reduzir o número de passos de inferência (redução uniforme de passos) frequentemente leva a uma degradação imediata na qualidade das amostras, como estruturas moleculares inválidas ou propriedades físicas ruins.

Metodologia

Os autores propõem o Cache de Características Preditivo (Predictive Feature Caching), uma estratégia livre de treinamento que acelera a geração de geometria molecular ao explorar a suavidade temporal das trajetórias dos estados ocultos durante o processo de amostragem.

Mecanismo Central

O método opera diretamente no backbone SE(3)-equivariante de modelos de flow matching pré-treinados. Ele reconhece que, conforme o solver percorre o tempo ($t$), as ativações intermediárias (características) da rede neural evoluem suavemente. Em vez de recomputar a passagem direta (forward pass) completa da rede em cada passo, o método:

1. Armazena em Cache: Avalia o backbone completo em passos de tempo de "checkpoint" específicos.
2. Prediz: Para passos intermediários dentro de um intervalo de cache $D$, substitui a avaliação cara do backbone por uma função de predição barata $\hat{F}$ baseada em características previamente computadas.
3. Preserva a Equivariância: Os esquemas de predição são combinações lineares de características armazenadas em cache e diferenças finitas, que comutam com a ação do grupo SE(3), garantindo que as moléculas geradas permaneçam invariantes à rotação e translação.

Estratégias de Cache Específicas

O artigo investiga dois esquemas preditivos adaptados da literatura de geração de imagens, aplicados especificamente ao último bloco da rede (um "previsão do último bloco" para evitar a recomputação de todo o prefixo):

• TaylorSeer (TS): Utiliza expansões de série de Taylor locais para prever características futuras. As derivadas são aproximadas via diferenças finitas de saídas anteriores armazenadas em cache.
• Adams–Bashforth (AB): Emprega uma recursão linear de múltiplos passos usando os últimos $j$ outputs armazenados em cache para prever o output atual.

A implementação garante que o passo final de inferência seja sempre computado via backbone completo para garantir a convergência, enquanto os passos intermediários utilizam as predições em cache.

Principais Contribuições

1. Aceleração sem Treinamento: O artigo introduz o primeiro esquema livre de treinamento para acelerar a geração de geometria molecular com pouca ou nenhuma perda de qualidade. É uma substituição direta ("drop-in") para modelos pré-treinados, não exigindo retreinamento ou ajuste fino (fine-tuning).
2. Adaptação SE(3)-Equivariante: Os autores transferiram e analisaram com sucesso estratégias de cache de características para operar diretamente em geometrias moleculares e representações SE(3)-equivariantes, provando que as operações de cache preservam as simetrias geométricas necessárias.
3. Complementaridade: O trabalho demonstra que o cache preditivo é ortogonal a outras acelerações baseadas em treinamento e otimizações de nível de sistema (ex: compilação de grafos, TensorFloat-32), permitindo acelerações compostas.
4. Capacidade de Refinamento: Além da aceleração, os autores mostram que o cache pode atuar como um mecanismo de refinamento, melhorando a qualidade da amostra em relação ao modelo base quando operando com o mesmo orçamento de inferência efetivo.

Resultados Experimentais

O método foi avaliado usando geração molecular incondicional e baseada em estrutura (condicionada por bolsa de proteína) usando modelos como SemlaFlow, Tabasco, FLOWR e FLOWR.root em conjuntos de dados incluindo GEOM, SPINDR e CrossDocked2020.

• Troca entre Velocidade vs. Qualidade: O cache alcança uma redução de 2× no tempo de inferência de tempo real (wall-clock) mantendo a qualidade da amostra equivalente. Com uma degradação mínima, acelerações de até 3× são alcançáveis.
• Superioridade sobre a Redução de Passos: Em orçamentos de passos mais baixos (ex: 1/3 ou 1/4 dos passos originais), o cache preditivo supera significativamente a redução uniforme de passos. Por exemplo, o modelo Tabasco falhou em produzir moléculas válidas em 14 passos sem o cache, enquanto as variantes com cache mantiveram alta validade e plausibilidade física.
• Melhoria de Qualidade: Em algumas configurações, modelos com cache com um orçamento efetivo de 100 passos (alcançado via 200 passos totais com cache) superaram o modelo base padrão de 100 passos em métricas como Vina Score e validade de Posebusters.
• Otimizações Compostas: Quando combinado com compilação de grafos e kernels TF32, a abordagem gerou um aceleração total de até 7×, reduzindo o tempo para gerar 10.000 moléculas de >14 minutos para ~2 minutos.

Significância e Alegações

O artigo afirma que o cache de características preditivo define uma nova fronte de Pareto de velocidade-qualidade para a geração molecular. Ao alavancar a suavidade das trajetórias de características, o método permite escalar modelos generativos para os milhões de amostras necessárias para a descoberta de fármacos prática sem os custos proibitivos de retreinamento ou as perdas de qualidade associadas à redução de passos ingênua.

Os autores posicionam este trabalho como um passo sistemático para melhorar a eficiência do tempo de inferência, destacando que o cache não é apenas uma heurística, mas um mecanismo robusto que preserva a validade física e as restrições geométricas das moléculas geradas. O código é disponibilizado para facilitar a pesquisa adicional e a adoção pela comunidade de design molecular.