🤖 AI

E2Former-V2: On-the-Fly Equivariant Attention with Linear Activation Memory

O E2Former-V2 é uma arquitetura de rede neural de grafos equivariante e escalável que supera os gargalos computacionais de modelos convencionais ao introduzir a Esparsificação Equivariante Alinhada aos Eixos e um mecanismo de Atenção Equivariante On-the-Fly, alcançando uma melhoria de 20 vezes em TFLOPS enquanto mantém o desempenho preditivo em sistemas atomísticos 3D.

Autores originais: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Publicado 2026-06-08

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Lin Huang, Chengxiang Huang, Ziang Wang, Yiyue Du, Chu Wang, Haocheng Lu, Yunyang Li, Xiaoli Liu, Arthur Jiang, Jia Zhang

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando prever como uma multidão massiva de pessoas (átomos) se moverá e interagirá em uma sala. No mundo da química e da física, isso é chamado de "modelagem atomística". Por muito tempo, os cientistas usaram um tipo especial de IA chamada Rede Neural de Grafo Equivariante (EGNN). Essas redes são ótimas porque entendem que, se você rotacionar o quarto inteiro, a física dentro dele não deve mudar.

No entanto, as versões atuais dessas redes têm um problema importante: elas são desajeitadas e lentas quando a multidão fica grande.

Aqui está uma explicação simples da solução do artigo, o E2Former-V2, usando analogias do cotidiano.

O Problema: O Engarrafamento "Centrado nas Arestas"

Imagine uma festa onde todos querem conversar com seus vizinhos.

Método Antigo (EGNNs Tradicionais): O sistema cria uma lista gigante e física de cada conversa acontecendo na festa. Se houver 1.000 pessoas, e cada uma fala com 50 vizinhos, o sistema tem que escrever 50.000 notas de conversa separadas em um quadro branco enorme (memória) antes de poder fazer qualquer conta.
O Gargalo: À medida que a festa cresce (mais átomos), este quadro branco fica tão grande que o computador fica sem espaço. Ele passa todo o seu tempo apenas movendo essas notas de lugar em vez de realmente pensar. É por isso que os modelos antigos travam ou ficam extremamente lentos ao simular moléculas grandes, como proteínas.

A Solução: E2Former-V2

Os autores construíram um novo sistema que para de escrever cada nota de conversa individualmente. Em vez disso, eles usam dois truques inteligentes para tornar o processo extremamente rápido.

Truque 1: A "Bússola Mágica" (Esparsificação Alinhada ao Eixo Equivariante)

No sistema antigo, calcular como os átomos interagem é como tentar resolver um quebra-cabeça onde cada peça pode ser rotacionada em qualquer direção. É bagunçado e exige muita matemática.

A Analogia: Imagine que você tem um quebra-cabeça 3D complexo. Em vez de tentar resolvê-lo enquanto ele gira, você usa uma "Bússola Mágica" para rotacionar o quebra-cabeça inteiro de modo que uma peça específica esteja sempre apontando para cima.
O Resultado: Uma vez que o quebra-cabeça está alinhado, a maioria das peças torna-se vazia ou zero. Você não precisa mais fazer contas nas partes vazias. O artigo chama isso de EAAS. Isso transforma um cálculo denso e bagunçado em um jogo simples de "reindexação" (apenas mover as peças), o que é muito mais rápido.

Truque 2: O "Chef de Streaming" (Atenção Equivariante On-the-Fly)

Este é o maior avanço.

O Jeito Antigo (O Buffet): O sistema antigo prepara todo o buffet (todas as 50.000 notas de conversa) e o coloca sobre a mesa antes que alguém possa comer. Se a mesa for pequena demais, a comida transborda e a festa para.
O Novo Jeito (O Chef de Streaming): O novo sistema, E2Former-V2, age como um chef que cozinha um prato de cada vez, bem na frente do cliente, e o serve imediatamente.
- Ele não escreve as notas.
- Ele não armazena a lista inteira.
- Ele calcula a interação entre duas pessoas, adiciona o resultado à pontuação final e, em seguida, esquece imediatamente os detalhes para abrir espaço para o próximo par.
O Resultado: Isso é chamado de processamento "On-the-Fly" (em tempo real). Como ele nunca constrói o quadro branco gigante, ele usa quase nada de memória. Ele consegue tudo dentro do cache interno super-rápido do computador (SRAM), tornando-o 20 vezes mais rápido que os métodos antigos.

O Que Eles Realmente Alcançaram

Os autores testaram este novo sistema em duas coisas principais:

Moléculas Pequenas (SPICE): Eles testaram em moléculas semelhantes a medicamentos. O novo sistema foi tão preciso quanto os melhores modelos existentes, mas muito mais rápido.
Moléculas Gigantes (OMol25): Eles testaram em estruturas químicas massivas.
- O Limite: Os modelos antigos travavam (ficavam sem memória) ao tentar simular sistemas com mais de 10.000 átomos.
- A Vitória: O E2Former-V2 simulou com sucesso sistemas com 100.000 átomos sem travar.

Testes do Mundo Real Mencionados

Os autores não apenas rodaram números; eles simularam a física real:

Água Líquida: Eles simularam uma caixa de moléculas de água. O novo modelo previu como as moléculas de água se organizam (a "função de distribuição radial") quase perfeitamente, combinando com experimentos do mundo real melhor do que modelos anteriores.
Proteínas: Eles simularam uma proteína grande (cerca de 30.000 átomos) flutuando em água. O modelo manteve a proteína estável e ajudou a dobrá-la em uma forma que coincidia com dados experimentais reais, provando que pode lidar com estruturas biológicas complexas.

Resumo

E2Former-V2 é como atualizar de um sistema que escreve cada conversa em um quadro branco gigante e lento para um sistema onde um chef inteligente cozinha e serve conversas uma por uma, instantaneamente. Isso permite que cientistas simulem sistemas moleculares massivos (como proteínas gigantes) em computadores padrão sem ficar sem memória, mantendo a física precisa.

Resumo Técnico: E2Former-V2: Atenção Equivariante "On-the-Fly" com Memória de Ativação Linear

1. Definição do Problema

Redes Neurais de Grafos Equivariantes (EGNNs) tornaram-se o padrão para modelagem de sistemas atomísticos 3D devido à sua capacidade de preservar simetrias rotacionais e translacionais. No entanto, as arquiteturas convencionais (ex: MACE, EquiformerV2, eSCN) sofrem de gargalos críticos de eficiência enraizados em seu design centrado em arestas (edge-centric). Esses modelos constroem explicitamente características geométricas ou realizam produtos tensoriais densos em cada aresta, resultando em:

Complexidade de Memória: Complexidade de memória de ativação $O(|E|)$ , onde $|E| \approx kN$ (sendo $k$ o número médio de vizinhos). Para sistemas grandes, isso resulta em pegadas de memória massivas que excedem a capacidade da GPU.
Ineficiência de Hardware: A materialização explícita de tensores de aresta força o movimento de dados entre a Memória de Alta Largura de Banda (HBM) e a SRAM on-chip, criando um "muro de memória" que limita o throughput.
Limites de Escalabilidade: À medida que o tamanho do sistema ( $N$ ) aumenta, a latência escala mal, tornando simulações de dinâmica molecular (MD) de larga escala computacionalmente proibitivas em hardware comum.

Embora avanços recentes em Transformers (ex: FlashAttention) tenham resolvido gargalos semelhantes para dados de sequências via execução em streaming e baseada em blocos (tiles), esses princípios permaneceram inexplorados em arquiteturas equivariantes SO(3).

2. Metodologia

O E2Former-V2 aborda esses problemas integrando esparsificação algébrica com um modelo de execução centrado em nós (node-centric) consciente ao hardware. A arquitetura consiste em duas inovações principais:

2.1. Esparsificação Equivariante Alinhada ao Eixo (EAAS)

As contrações tensoriais SO(3) padrão envolvem somas densas sobre coeficientes de Clebsch–Gordan, resultando em complexidade $O(L^6)$ e padrões de acesso à memória irregulares. O E2Former-V2 introduz a EAAS para transformar essas operações densas em operações esparsas baseadas em permutação:

Mudança de Base SO(3) $\to$ SO(2): O método alinha o sistema de coordenadas globais com o vetor local $\vec{r}$ (o eixo que conecta dois átomos). Neste referencial alinhado, os harmônicos esféricos sólidos tornam-se maximamente esparsos (Lema 4.1), onde apenas o componente $m=0$ é não-nulo.
Operador de Reindexação: Ao comutar rotações através do produto tensorial, a contração densa colapsa em uma regra de reindexação determinística (Eq. 13). Este operador mapeia características de entrada para características de saída via mapas lineares esparsos por blocos, eliminando a necessidade de soma densa sobre índices magnéticos.
Resultado: Isso reduz a complexidade da etapa de convolução e alcança um aumento de velocidade de $\sim 1.5\times$ sobre implementações padrão.

2.2. Atenção Equivariante "On-the-Fly"

Para eliminar o gargalo de memória $O(|E|)$ , os autores propõem um mecanismo de atenção totalmente fundido e centrado em nós, implementado via um kernel Triton customizado:

Fatoração Centrada em Nós: Em vez de materializar características de aresta, o mecanismo de atenção desacopla a computação. Vetores de Query e Key são construídos usando projeções escalares SO(3)-invariantes (Eq. 6), enquanto a informação geométrica equivariante é propagada apenas pelo caminho do valor (value path).
Redução em Streaming: O kernel processa vizinhos em modo streaming. Ele itera sobre vizinhos $j \in \mathcal{N}(i)$ , computando pontuações de atenção e acumulando atualizações de valor on-the-fly dentro da SRAM.
Eliminação da Materialização de Arestas: Tensores intermediários de forma $N \times K \times H \times D$ nunca são escritos na HBM. O kernel mantém estatísticas contínuas (máximo, normalização e acumuladores de valor) on-chip, reduzindo a complexidade de memória de ativação de $O(|E|)$ para $O(|V|)$ .
Softmax Online: Para garantir estabilidade numérica durante o streaming, uma formulação de softmax online é usada para atualizar os acumuladores incrementalmente.

3. Principais Contribuições

Esparsificação Equivariante Aliente ao Eixo (EAAS): Uma técnica algébrica inovadora que unifica o recoupling Wigner-6j com o alinhamento de eixo SO(2), convertendo contrações tensoriais SO(3) densas em operações de reindexação esparsas.
Atenção Equivariante "On-the-Fly": Um kernel Triton totalmente fundido que evita a materialização de arestas, maximizando o reúso de SRAM e alcançando até 20 $\times$ mais TFLOPS em comparação com implementações padrão.
Desempenho Escalável: Capacidade demonstrada de lidar com sistemas de até 100.000 átomos em uma única GPU, um regime onde Transformers equivariantes anteriores falham devido a erros de Out-of-Memory (OOM).
Fidelidade Física: Validação em dinâmica molecular mostrando a reprodução precisa de estruturas de água líquida e o refinamento estável de estruturas de proteínas preditas pelo AlphaFold ( $\sim$ 30k átomos).

4. Resultados Experimentais

Os autores avaliaram o E2Former-V2 nos conjuntos de dados SPICE (química medicinal) e OMol25 (escala de modelo de fundação), bem como em simulações de MD.

Precisão Preditiva:
- No SPICE, o E2Former-V2 alcançou os menores Erros Médios Absolutos (MAE) de Energia e Força em quase todos os subconjuntos, superando o MACE-Large em 48% no MAE de Energia em Dímeros.
- No OMol25, a variante conservadora alcançou um MAE de Energia agregado de 1.27 meV/átomo, igualando o eSEN-small e superando significativamente o UMA-S (3.62 meV/átomo).
- No OC20 S2EF-2M, alcançou precisão comparável a baselines fortes como EquiformerV2 e eSCN.
Eficiência e Escalabilidade:
- Throughput: O kernel fundido alcançou um aumento médio de 20 $\times$ em TFLOPS sobre implementações nativas de PyTorch.
- Memória: O uso de pico de memória foi reduzido em 16.9 $\times$ a 19.5 $\times$ em comparação com os baselines.
- Escalabilidade: Enquanto modelos anteriores (ex: EquiformerV2, MACE-Large) encontravam falhas de OOM além de 10k–50k átomos, o E2Former-V2 escalou de forma confiável para 100k átomos.
- Velocidade de Inferência: Em 100k átomos, o E2Former-V2 (Conservative) atingiu 0.12 passos/s, quase 3 $\times$ mais rápido que o UMA-S. No cenário Direct, foi uma ordem de magnitude mais rápido que Allegro e EquiformerV2 em escalas menores.
Dinâmica Molecular:
- Simulações de água volumétrica periódica (216 moléculas) mostraram que o E2Former-V2 coincide de perto com as funções de distribuição radial (RDF) experimentais, superando o MACE-OFF.
- Uma simulação de 2 ns de uma proteína solvatada ( $\sim$ 30k átomos) manteve altos scores TM em relação à estrutura nativa, demonstrando dinâmica estável.

5. Significância e Alegações

O artigo afirma que o E2Former-V2 demonstra a viabilidade prática de grandes transformers equivariantes em GPUs comuns. Ao mudar o paradigma de execução de centrado em arestas para centrado em nós e aproveitar o streaming consciente ao hardware, o método resolve o gargalo fundamental de memória que limitou a escalabilidade de modelos equivariantes.

Os autores enfatizam que isso não é meramente uma melhoria algorítmica, mas uma realização em nível de sistema que faz a ponte entre a fatoração teórica (reduzindo a complexidade de $O(|E|)$ para $O(|V|)$ ) e a eficiência de hardware. O trabalho sugere que modelos de fundação molecular de larga escala e simulações de MD de longa duração, anteriormente proibitivos computacionalmente, agora são acessíveis usando infraestrutura de GPU padrão. O código é liberado para facilitar a adoção pela comunidade.