Modelos típicos de aprendizagem profunda para visão computacional, como Redes Reurais Convolucionais (CNNs) e transformadores de visão (ViT), processam sinais assumindo espaços planares (planos). Por exemplo, as imagens digitais são representadas como uma grade de pixels em um plano. No entanto, este tipo de dados representa apenas uma fração dos dados que encontramos em aplicações científicas.

Variáveis ​​amostradas da atmosfera terrestre, como temperatura e umidade, são naturalmente representadas na esfera. Alguns tipos de dados cosmológicos e fotos panorâmicas também são sinais esféricos e são melhor tratados como tal.

O uso de métodos projetados para imagens planares para processar sinais esféricos é problemático por alguns motivos. Primeiro, existe um problema de amostragem, ou seja, não há como definir grades uniformes na esfera, que são necessárias para CNNs e ViTs planares, sem forte distorção.

Ao projetar a esfera em um plano, a área representada pelo círculo vermelho fica fortemente distorcida perto dos pólos. Este problema de amostragem prejudica a precisão dos CNNs e ViTs convencionais em entradas esféricas.

Em segundo lugar, os sinais e padrões locais na esfera são muitas vezes complicados pelas rotações, pelo que os modelos precisam de uma forma de resolver isso. Gostaríamos de equivalência às rotações 3D, o que garante que os recursos aprendidos sigam as rotações da entrada. Isso leva a uma melhor utilização dos parâmetros do modelo e permite o treinamento com menos dados.

A equivalência com rotações 3D também é útil na maioria dos ambientes onde as entradas não têm uma orientação preferencial, como formas e moléculas 3D.

Corridas de drones com câmeras panorâmicas. Aqui, as curvas acentuadas resultam em grandes rotações 3D da imagem esférica. Gostaríamos que nossos modelos fossem robustos a tais rotações. Fonte:  https://www.youtube.com/watch?v=_J7qXbbXY80

 

 

Na atmosfera, é comum ver padrões semelhantes aparecendo em diferentes posições e orientações. Gostaríamos que nossos modelos compartilhassem parâmetros para reconhecer esses padrões.

Com os desafios acima em mente, em “ Scaling Spherical CNNs ”, apresentado no ICML 2023 , apresentamos uma biblioteca de código aberto em JAX para aprendizado profundo em superfícies esféricas. Demonstramos como as aplicações desta biblioteca correspondem ou superam o desempenho de última geração em previsão do tempo e benchmarks de previsão de propriedades moleculares, tarefas que normalmente são abordadas com transformadores e redes neurais gráficas.

Antecedentes das CNNs esféricas

CNNs esféricas resolvem os problemas de amostragem e de robustez à rotação , aproveitando a convolução esférica e as operações de correlação cruzada, que normalmente são calculadas por meio de transformadas generalizadas de Fourier . Para superfícies planas, entretanto, a convolução com filtros pequenos é mais rápida, porque pode ser realizada em grades regulares sem usar transformadas de Fourier.

O maior custo computacional para entradas esféricas restringiu até agora a aplicação de CNNs esféricas a pequenos modelos e conjuntos de dados e conjuntos de dados de baixa resolução.

Nossas contribuições

Implementamos as convoluções esféricas de CNNs esféricas ponderadas por spin em JAX com foco na velocidade e habilitamos o treinamento distribuído em um grande número de TPUs usando paralelismo de dados .

Também introduzimos uma nova camada de ativação de colapso de fase e normalização de lote espectral, e um novo bloco residual que melhora a precisão e a eficiência, o que permite treinar modelos mais precisos até 100x maiores do que antes. Aplicamos esses novos modelos na regressão de propriedades moleculares e na previsão do tempo.

 

Dimensionamos CNNs esféricos em até duas ordens de grandeza em termos de tamanhos de recursos e capacidade do modelo, em comparação com a literatura:  Cohen’18 ,  Esteves’18 ,  Esteves’20 e  Cobb’21 . VGG-19  está incluído como referência convencional da CNN. Nosso maior modelo para previsão do tempo tem 256 x 256 x 78 entradas e saídas e executa 96 camadas convolucionais durante o treinamento com uma resolução interna mais baixa de 128 x 128 x 256.

Regressão de propriedades moleculares

A previsão de propriedades de moléculas tem aplicações na descoberta de medicamentos , onde o objetivo é rastrear rapidamente inúmeras moléculas em busca daquelas com propriedades desejáveis. Modelos semelhantes também podem ser relevantes na concepção de medicamentos direcionados à interação entre proteínas.

Os métodos atuais em química quântica computacional ou experimental são caros, o que motiva o uso do aprendizado de máquina.

As moléculas podem ser representadas por um conjunto de átomos e suas posições no espaço 3D; as rotações da molécula alteram as posições, mas não as propriedades moleculares. Isso motiva a aplicação de CNNs esféricas devido à sua equivariância de rotação.

No entanto, as moléculas não são definidas como sinais na esfera, então o primeiro passo é mapeá-las para um conjunto de funções esféricas. Fazemos isso aproveitando as interações baseadas na física entre os átomos da molécula.

Cada átomo é representado por um conjunto de sinais esféricos acumulando interações físicas com outros átomos de cada tipo (mostrados nos três painéis à direita). Por exemplo, o átomo de oxigênio (O; painel superior) possui um canal para oxigênio (indicado pela esfera rotulada “O” à esquerda) e hidrogênio (“H”, à direita).  As  forças de Coulomb acumuladas  no átomo de oxigênio em relação aos dois átomos de hidrogênio são indicadas pelas regiões sombreadas em vermelho na parte inferior da esfera denominadas “H”. Como o átomo de oxigênio não contribui com forças para si mesmo, a esfera “O” é uniforme. Incluímos canais extras para as  forças de Van der Waals  .

CNNs esféricas são aplicadas às características de cada átomo e os resultados são posteriormente combinados para produzir as previsões de propriedades. Isto resulta num desempenho de última geração na maioria das propriedades, conforme normalmente avaliado no benchmark QM9 :

Previsão do tempo

Previsões climáticas precisas servem como ferramentas valiosas para fornecer avisos oportunos sobre eventos climáticos extremos, permitindo uma gestão eficaz dos recursos hídricos e orientando o planeamento informado das infraestruturas. Num mundo cada vez mais ameaçado por catástrofes climáticas , há uma urgência em fornecer previsões muito mais rápidas e precisas num horizonte temporal mais longo do que os modelos de circulação geral.

Os modelos de previsão também serão importantes para prever a segurança e a eficácia dos esforços destinados a combater as alterações climáticas, tais como as intervenções climáticas . O estado da arte atual utiliza modelos numéricos dispendiosos baseados em dinâmica de fluidos e termodinâmica, que tendem a mudar após alguns dias.

Tendo em conta estes desafios, é urgente que os investigadores da aprendizagem automática resolvam os problemas de previsão climática, uma vez que as técnicas baseadas em dados têm o potencial de reduzir o custo computacional e de melhorar a precisão a longo prazo.

CNNs esféricas são adequadas para esta tarefa, uma vez que os dados atmosféricos são apresentados nativamente na esfera. Eles também podem lidar com eficiência com padrões repetidos em diferentes posições e orientações que são comuns em tais dados.

Aplicamos nossos modelos a vários benchmarks de previsão do tempo e superamos ou igualamos modelos meteorológicos neurais baseados em CNNs convencionais (especificamente, 1 , 2 e 3 ). Abaixo mostramos os resultados em um ambiente de teste onde o modelo toma uma série de variáveis ​​atmosféricas como entrada e prevê seus valores com seis horas de antecedência.

O modelo é então aplicado iterativamente em suas próprias previsões para produzir previsões mais longas. Durante o treinamento, o modelo prevê até três dias de antecedência e é avaliado em até cinco dias. Keisler propôs uma rede neural gráfica para esta tarefa, mas mostramos que CNNs esféricos podem corresponder à precisão do GNN na mesma configuração.

Previsão meteorológica iterativa com até cinco dias (120h) de antecedência com CNNs esféricas. As animações mostram a previsão específica de umidade a uma determinada pressão e seu erro.

Previsões de velocidade e temperatura do vento com CNNs esféricas.

Recursos adicionais

Nossa biblioteca JAX para CNNs esféricas eficientes já está disponível . Mostramos aplicações para regressão de propriedades moleculares e previsão do tempo, e acreditamos que a biblioteca será útil em outras aplicações científicas, bem como em visão computacional e visão 3D.

A previsão do tempo é uma área ativa de pesquisa no Google com o objetivo de construir modelos mais precisos e robustos – como Graphcast , um modelo recente de previsão de médio alcance baseado em ML – e construir ferramentas que permitam avanços adicionais em toda a comunidade de pesquisa, como o recentemente lançado WeatherBench 2 .