Explanation of Dynamic Physical Field Predictions using WassersteinGrad: Application to Autoregressive Weather Forecasting

Resumo Técnico: WassersteinGrad para Previsão de Campos Físicos Dinâmicos

O artigo apresenta o WassersteinGrad, um novo framework de Inteligência Artificial Explicável (XAI) projetado especificamente para modelos de aprendizado profundo aplicados à previsão de fenômenos físicos dinâmicos, como a previsão meteorológica.

1. O Problema: Deslocamento Geométrico de Atribuições

Atualmente, métodos baseados em gradiente (como SmoothGrad) são o padrão para explicar previsões em dados de alta dimensão. Esses métodos aplicam ruído gaussiano à entrada e calculam a média dos gradientes para reduzir o ruído visual. No entanto, os autores identificam uma falha fundamental quando esses métodos são aplicados a campos físicos (como ventos ou pressão):

• Ruído Geométrico vs. Estacionário: Em imagens, o ruído costuma afetar apenas a amplitude da explicação. Em campos físicos, o ruído na entrada causa um deslocamento geométrico das atribuições. Ou seja, a "massa de importância" migra para localizações espacialmente incorretas, mas fisicamente estruturadas.
• Erro de Fase: Esse fenômeno é análogo ao "erro de fase" na meteorologia (prever o evento certo, mas no lugar errado).
• Acúmulo Autoregressivo: Em modelos de previsão que utilizam rollouts autoregressivos (onde a previsão de $t+1$ é usada para prever $t+2$), esse deslocamento geométrico se acumula, tornando as explicações em horizontes de longo prazo (ex: $t+5$) completamente incoerentes.

Os autores demonstram mecanisticamente que isso ocorre devido a dois componentes arquiteturais:

1. CNNs: A instabilidade do operador argmax em camadas de pooling desloca fisicamente o gradiente.
2. Transformers: A natureza não-Lipschitz da atenção e a existência de múltiplos atratores de energia (metaestabilidade) fazem com que diferentes ruídos levem o mecanismo de atenção a diferentes configurações de agrupamento (clustering) de tokens.

2. Metodologia: O Framework WassersteinGrad

Para resolver o problema da média de mapas deslocados, os autores propõem substituir a média aritmética ponto a ponto (que resulta em um borrão/blur das características) por uma média geométrica baseada em transporte ótimo.

O processo consiste em:

1. Mapeamento para Medidas de Probabilidade: Os mapas de gradiente bruto ($G_i$) são convertidos em distribuições de probabilidade espacial ($\mu_i$) através da aplicação de valor absoluto e normalização $\ell_1$. Isso garante que a "massa" de importância seja preservada.
2. Cálculo do Baricentro de Wasserstein: Em vez de tirar a média dos valores nos pixels, o método calcula o Baricentro de Wasserstein Entrópico. Isso encontra uma distribuição que minimiza o custo de transporte para todas as distribuições perturbadas, preservando a estrutura geométrica e a localização das características.
3. Variantes de Saída:
   • WGBary: Fornece o baricentro diretamente como o mapa de atribuição (robusto, mas perde o sinal/direção do gradiente).
   • WGBary $\times$ Grad: Utiliza o baricentro como uma máscara espacial sobre o gradiente original para recuperar a direção (positivo/negativo) e o contraste de amplitude.

3. Principais Contribuições

• Análise Mecanística: Uma investigação teórica de como o ruído de entrada se transforma em deslocamento espacial através de camadas de convolução e mecanismos de atenção.
• Novo Framework de Agregação: A introdução do WassersteinGrad, que utiliza transporte ótimo para obter um consenso geométrico de explicações perturbadas.
• Validação em Dados Reais: A primeira caracterização empírica do deslocamento de gradientes em modelos de previsão meteorológica neural.

4. Resultados e Significância

Os experimentos foram realizados utilizando o dataset TITAN (Météo-France) e um modelo de previsão híbrido (UNetRPP).

• Qualitativo: Enquanto o SmoothGrad e o BaseGrad produzem mapas fragmentados e ruidosos (especialmente em previsões de longo prazo $t+5$), o WassersteinGrad produz atribuições geograficamente coerentes e localizadas.
• Quantitativo: O método superou os baselines em métricas cruciais:
  • Fidelidade (Faithfulness): Maior capacidade de identificar quais regiões realmente impulsionam a previsão.
  • Robustez (Robustness): Redução drástica na sensibilidade à rotação e deslocamento espacial do padrão de atribuição (medido por $LLE_{cos}$).
  • Esparsidade (Sparsity): A variante WGBary $\times$ Grad conseguiu manter explicações concentradas e interpretáveis.

Significância: Este trabalho é crucial para a adoção de IA em ambientes de alto risco (como a meteorologia operacional). Ele demonstra que a transferência direta de métodos de XAI de visão computacional para ciências físicas é falha e propõe uma solução matematicamente rigorosa que respeita a natureza contínua e geométrica dos fenômenos físicos.