Multi-Scale Wavelet Transformers for Operator Learning of Dynamical Systems

Resumo Técnico: Transformers de Onda Multiescala para Aprendizado de Operadores de Sistemas Dinâmicos

Declaração do Problema
Surrogados baseados em dados para sistemas dinâmicos, particularmente operadores neurais, oferecem vantagens significativas de velocidade sobre solvers numéricos para equações diferenciais parciais (EDPs). No entanto, esses modelos frequentemente sofrem de viés espectral (ou suavização excessiva), onde redes neurais aprendem componentes de baixa frequência mais facilmente do que componentes de alta frequência. No contexto do aprendizado de operadores, isso se manifesta como uma subestimação sistemática do conteúdo de alto número de onda. Essa limitação é crítica para sistemas multiescala e caóticos, onde erros em pequena escala podem degradar rapidamente previsões em grande escala, levando a instabilidade de longo horizonte e soluções não físicas em aplicações como previsão do tempo. Estratégias de mitigação existentes, como arquiteturas multiescala ou representações conscientes de frequência, frequentemente falham em preservar características de alta frequência durante o subamostragem agressiva ou dependem de mecanismos de atenção computacionalmente caros aplicados diretamente a grades físicas de alta resolução.

Metodologia
Os autores propõem o Transformador de Onda Multiescala (MSWT), um framework de aprendizado de operadores projetado para aprender a dinâmica do sistema dentro de um domínio de onda tokenizado. A arquitetura aborda o viés espectral separando explicitamente o conteúdo de baixa e alta frequência através de escalas, mantendo ao mesmo tempo a eficiência computacional. O MSWT consiste em três componentes principais:

1. Tokenizador de Patch: O campo físico é mapeado para um espaço de tokens via particionamento em patches não sobrepostos. Isso reduz a resolução efetiva para operações de atenção, suavizando o landscape de otimização e melhorando a eficiência para EDPs de alta resolução.
2. Estrutura de Onda Multiescala: O modelo emprega uma topologia em U com $L$ escalas de subamostragem e superamostragem.
   • Subamostragem Preservadora de Ondas: Diferente de pooling padrão ou convoluções com stride que descartam informações de alta frequência, o codificador utiliza a Transformada Discreta de Onda (DWT) (especificamente ondas de Haar) para decompor características em quatro sub-bandas: Baixa-Baixa (LL), Baixa-Alta (LH), Alta-Baixa (HL) e Alta-Alta (HH). Todas as sub-bandas são retidas e passadas para a próxima escala, garantindo que detalhes de alta frequência não sejam perdidos.
   • Superamostragem Preservadora de Ondas: O decodificador utiliza a DWT Inversa (iDWT) para reconstruir características de maior resolução a partir das sub-bandas, fundindo-as com conexões de salto do codificador.
3. Bloco de Atenção de Onda (WAttn): Operando dentro do espaço de ondas, este bloco modela dependências entre diferentes sub-bandas de frequência. Ele aplica uma compressão de dimensão, realiza DWT para extrair coeficientes de onda e, em seguida, utiliza uma combinação de convolução e autoatenção (especificamente atenção de janela local) para misturar características entre sub-bandas. Isso permite que o modelo amplifique componentes de alta frequência que tipicamente possuem baixa energia, mas são cruciais para a estrutura de pequena escala. Finalmente, uma iDWT mapeia os coeficientes de volta para o espaço de tokens.

O modelo é treinado usando uma perda relativa padrão $\ell_2$ (tipo MSE) em pares de previsão de um passo, sem exigir funções objetivo especializadas para penalizar erros de alta frequência, isolando a contribuição do viés indutivo arquitetural.

Principais Contribuições

1. Mitigação do Viés Espectral via Atenção de Onda: O artigo introduz uma arquitetura transformer que modela conteúdo de ondas multibanda através de escalas. Ao operar no domínio de ondas, o modelo preserva e mistura explicitamente sub-bandas de alta frequência, resultando em fidelidade espectral aprimorada em comparação com operadores baseados em Fourier ou convolucionais padrão.
2. Eficiência Aprimorada e Estabilidade de Longo Horizonte: Ao combinar tokenização com sub/superamostragem baseada em ondas, o método reduz significativamente o custo computacional dos mecanismos de atenção (operando em resoluções espaciais reduzidas) enquanto mantém a estabilidade. Isso permite execuções precisas de longo horizonte de sistemas dinâmicos caóticos onde outros modelos tipicamente divergem ou acumulam erro.
3. Avaliação de Longo Horizonte em Dados Climáticos Reais: Os autores validam o método no conjunto de dados de reanálise climática ERA5. Os resultados demonstram viés climatológico reduzido em previsões de execução livre, indicando consistência física de longo prazo aprimorada em comparação com baselines de última geração.

Resultados Experimentais
O MSWT foi avaliado em três benchmarks: Fluxo Caótico de Kolmogorov (CKF), Equação de Água Rasa (SWE) e reanálise climática ERA5.

• Fluxo Caótico de Kolmogorov (CKF): O MSWT alcançou reduções substanciais de erro sobre as baselines (FNO, U-Net, WNO, SAOT, HFS). No passo 64, o MSWT reduziu o erro L2 Relativo em mais de 60% comparado ao segundo melhor modelo estável (FNO). Também mostrou desempenho superior em métricas espectrais (EMLR e EMAE), com melhorias de 76,5% e 80,6% respectivamente, e manteve fidelidade espectral correspondendo estreitamente à verdade fundamental, enquanto as baselines exibiram superestimação de energia ou soluções não físicas.
• Equação de Água Rasa (SWE): O MSWT demonstrou menor acumulação de erro em execuções de longo horizonte (passos 41 e 81), melhorando a segunda melhor baseline estável (FNO) em aproximadamente 30% em L2 Relativo, EMLR e EMAE. Evitou com sucesso artefatos de fronteira observados em outros métodos.
• Reanálise Climática ERA5: Em execuções autoregressivas em escala decenal, o MSWT consistentemente superou o Emulador Climático Desacoplado Leve (LUCIE). Alcançou menor Erro Quadrático Médio (RMSE) em todas as variáveis (temperatura, umidade, ventos, pressão, precipitação) e reduziu notavelmente o viés climatológico, particularmente na pressão de superfície perto de regiões polares e viés de vento meridional sobre o Polo Sul.

Significado
O artigo afirma que o framework MSWT isola efetivamente o papel da representação e do viés indutivo na mitigação do viés espectral. Ao aprender dinâmicas em um espaço de ondas tokenizado, o método alcança um equilíbrio entre eficiência computacional e preservação de estruturas de pequena escala. Os resultados sugerem que amostragem baseada em ondas explícita e mecanismos de atenção são superiores para estabilidade de longo prazo em sistemas caóticos e multiescala, oferecendo uma solução robusta para tarefas de previsão do mundo real onde manter consistência física ao longo de horizontes longos é crítico. Os autores enfatizam que essas melhorias são alcançadas sem modificar objetivos de treinamento padrão, destacando o design arquitetural como o principal impulsionador do desempenho.