Resumo Técnico: Modelos Fundamentais Atômicos Compactos e Equivariantes a SO(3) via Poda Estrutural

Declaração do Problema

Redes neurais em grafos (GNNs) equivariantes a SO(3) estabeleceram-se como o paradigma dominante para modelos fundamentais atômicos, alcançando alta precisão e eficiência de dados ao codificar explicitamente a simetria rotacional. No entanto, esses modelos dependem de operações de produto tensorial de ordem superior para acoplar características geométricas, resultando em custos computacionais que escalam até $O(L_{max}^6)$, onde $L_{max}$ é a ordem máxima das representações irredutíveis. Isso cria uma compensação significativa entre a precisão do modelo e a eficiência de inferência.

Estratégias existentes de compressão de modelos — especificamente Distilação de Conhecimento (KD) e Quantização — apresentam limitações neste domínio. A KD frequentemente exige o treinamento de um estudante a partir do zero e não reutiliza diretamente os pesos do professor, enquanto a quantização reduz a precisão de bits, mas deixa inalterada a contagem de parâmetros e os caminhos de produto tensorial. Além disso, métodos de poda estrutural padrão desenvolvidos para Grandes Modelos de Linguagem (LLMs) não podem ser aplicados diretamente a modelos equivariantes a SO(3). Remover dimensionalidades de características de forma ingênua quebra a simetria rotacional do caminho de produto tensorial, violando a propriedade fundamental de equivariância necessária para a consistência física. Tentativas anteriores de poda de modelos atômicos limitaram-se à remoção de blocos de granularidade grosseira (por exemplo, camadas de interação inteiras), o que reduz o campo receptivo do modelo e é particularmente prejudicial para arquiteturas rasas.

Metodologia

Os autores propõem uma estrutura de poda estrutural especificamente projetada para modelos fundamentais atômicos equivariantes a SO(3). O método opera em checkpoints pré-treinados para produzir modelos menores e densos que mantêm a equivariância a SO(3) e alta precisão.

1. Granularidade da Poda Equivariante

A inovação central reside na definição da unidade de poda. Em modelos equivariantes a SO(3), as características são decompostas em representações irredutíveis de ordem $l \in [0, L_{max}]$. Um vetor de características para um canal específico $k$ e ordem $l$ consiste em $2l+1$ componentes indexados por $m \in [-l, l]$.

• Restrição: A operação de rotação em SO(3) acopla todos os componentes $m$ dentro de uma ordem específica $l$. Remover componentes individuais (poda elemento a elemento) quebra a equivariância.
• Solução: O método trata o subespaço abrangido pelos $2l+1$ componentes de uma ordem específica $l$ e canal $k$ como uma unidade atômica. A poda é realizada na tupla $(k, l)$, representada por uma máscara binária $z^{(t)}_{k,l}$. Isso garante que, se um bloco de características for retido, todos os seus componentes rotacionais sejam preservados, mantendo estritamente a equivariância a SO(3).

2. Critério de Importância Informado pela Física

Para determinar quais blocos $(k, l)$ reter, os autores propõem uma pontuação de importância baseada na sensibilidade da superfície de energia potencial.

• Objetivo Duplo: MLIPs devem prever tanto a energia potencial ($E$) quanto as forças ($F_i = -\nabla_{r_i} E$).
• Métrica: A pontuação de importância $I^{(t)}_{k,l}$ é derivada de uma expansão de Taylor de primeira ordem da mudança de energia ao remover um bloco de características. Ela agrega o produto da magnitude da característica e o gradiente da energia em relação a essa característica sobre todos os átomos e dimensões de representação:
  $$I^{(t)}_{k,l} = \mathbb{E}_{R,Z \sim \mathcal{D}} \left| \sum_{j,m} \frac{\partial E}{\partial h^{(t)}_{j,klm}} \cdot h^{(t)}_{j,klm} \right|$$
• Vantagem: Este critério é invariante a SO(3) e incorre em sobrecarga negligenciável, pois os gradientes $\frac{\partial E}{\partial h}$ já são computados durante a inferência padrão de forças.

3. Alinhamento Estrutural e Pipeline

A estrutura segue um pipeline de quatro etapas:

1. Calibração: Um pequeno subconjunto de dados de pré-treinamento é passado pelo modelo grande para calcular pontuações de importância.
2. Poda Orientada por Objetivo: Uma arquitetura-alvo é definida. Os blocos $(k, l)$ com maior classificação são retidos com base nas pontuações de importância.
3. Alinhamento Estrutural: A estrutura do modelo é fisicamente alterada para corresponder às dimensões reduzidas. Isso envolve:
   • Reconfiguração do Produto Tensorial: Fatia de pesos radiais e caminhos de Clebsch-Gordan para alinhar com os canais retidos.
   • Fatiamento de Pesos: Remoção física de linhas e colunas dos pesos de transformação linear para criar um tensor de pesos menor e denso.
4. Retreinamento e Ajuste Fino: O modelo podado passa por um retratamento leve em um subconjunto de dados de pré-treinamento para recuperar de erros de aproximação, seguido de ajuste fino em tarefas específicas de domínio.

Contribuições Principais

1. Método de Poda Equivariante: Introdução de um método de poda estrutural para tensores equivariantes a SO(3) de ordem superior nas dimensões de canal e ordem, que preserva estritamente a equivariância a SO(3) ao tratar blocos $(k, l)$ como unidades atômicas.
2. Estrutura de Quatro Etapas: Um pipeline completo (Calibração, Poda, Retreinamento, Ajuste Fino) para comprimir modelos fundamentais.
3. Generalizabilidade: Demonstração de que o método se aplica a múltiplas arquiteturas, incluindo MACE, SevenNet e eSCN.
4. Combinabilidade: Validação de que a poda estrutural pode ser combinada com quantização e distilação de conhecimento para ganhos de eficiência compostos.

Resultados Experimentais

O método foi avaliado principalmente em MACE-MP (materiais inorgânicos) e MACE-OFF (moléculas orgânicas), com testes de generalizabilidade em SevenNet e eSCN.

Precisão vs. Treinamento a Partir do Zero

• Superioridade: Modelos podados consistentemente superam modelos do mesmo tamanho treinados a partir do zero.
  • Em MACE-MP, modelos podados alcançaram menor MAE de Força e melhor desempenho fora da distribuição no conjunto de dados sAlex em comparação com baselines treinadas a partir do zero.
  • Em Matbench Discovery, um modelo podado MACE-MP (usando apenas 25% dos dados de pré-treinamento) superou o pequeno modelo oficial treinado a partir do zero em 7 de 9 métricas.
  • Em MACE-OFF, modelos podados igualaram ou excederam as baselines treinadas a partir do zero em diversos subconjuntos moleculares, mesmo ao usar apenas 50% dos dados de pré-treinamento.
• Desempenho em Domínio Específico: O ajuste fino de modelos podados em oito conjuntos de dados inorgânicos downstream reduziu os erros de energia e força em 70,1% e 34,4%, respectivamente, em comparação com o treinamento de modelos específicos de tarefa a partir do zero.

Ganhos de Eficiência

• Redução de Parâmetros: Modelos comprimidos contêm 1,5× a 4× menos parâmetros.
• Custo de Treinamento: O pipeline completo requer 2,5× a 4× menos computação de pré-treinamento do que treinar um modelo pequeno a partir do zero (por exemplo, 181 horas de GPU versus 750 horas de GPU para uma configuração específica).
• Velocidade de Inferência: Modelos podados alcançam um aceleração de inferência de 2,7× a 3,0× e até 5,7× de redução no uso de pico de memória.

Combinação com Outros Métodos

• Quantização: A combinação de poda com quantização FP32 resultou em uma melhoria de 9,6× no throughput (relativa ao modelo base FP64).
• Distilação de Conhecimento: A combinação de poda com KD baseada em Hessiano refinou ainda mais o desempenho do modelo, demonstrando que os dois métodos são complementares.

Significado e Alegações

O artigo afirma que a poda estrutural oferece um caminho natural e eficaz para preencher a lacuna entre precisão e eficiência em modelos fundamentais atômicos. Ao herdar diretamente pesos de checkpoints grandes e pré-treinados e comprimi-los enquanto preservam simetrias físicas, o método produz modelos compactos que são:

1. Mais precisos do que modelos do mesmo tamanho treinados a partir do zero.
2. Significativamente mais baratos de treinar e implantar.
3. Fisicamente consistentes, mantendo a estrita equivariância a SO(3) necessária para simulações confiáveis de dinâmica molecular (validado via simulações de MD de água líquida mostrando deriva de energia estável e funções de distribuição radial precisas).

Os autores enfatizam que o princípio central — avaliar a importância de representações irredutíveis acopladas — é amplamente aplicável a qualquer arquitetura equivariante que represente características como somas de representações irredutíveis de SO(3). O trabalho sugere que a poda permite que profissionais implantem modelos fundamentais de alta fidelidade em ambientes com recursos limitados, sem sacrificar a transferabilidade aprendida durante o pré-treinamento em grande escala.