Resumen Técnico: Samudra 2: Escalamiento de Emuladores Oceánicos a través de Resoluciones

Planteamiento del Problema

Los Modelos de Circulación General del Océano (OGCM, por sus siglas en inglés) son fundamentales para la ciencia climática al simular el transporte de calor, sal, carbono y cantidad de movimiento. Sin embargo, su costo computacional limita severamente la capacidad de ejecutar grandes conjuntos (ensembles) o explorar diversos escenarios de forzamiento, particularmente cuando se requiere una alta resolución espacial para resolver remolinos de mesoescala y estructuras frontales nítidas. Aunque los emuladores neuronales ofrecen potenciales aceleraciones, los emuladores oceánicos existentes han tenido dificultades para combinar una resolución espacial fina con estabilidad dinámica a largo plazo.

Los autores identifican tres modos de falla específicos en su predecesor, Samudra, que estaba limitado a una resolución de $1^\circ$:

1. Colapso de la Varianza: La acumulación de errores autorregresivos provoca que las predicciones regresen hacia la media climatológica, suprimiendo la variabilidad temporal, especialmente en el océano profundo.
2. Impronta de Artefactos: Patrones espaciales de los campos de velocidad se "filtran" en los campos de temperatura y salinidad, creando bandas no físicas y ruido de alta frecuencia que se amplifica en despliegues (rollouts) de varios años.
3. Limitación de Resolución: La incapacidad de operar en resoluciones más finas que $1^\circ$ impide la emulación de remolinos de mesoescala y estructuras de corrientes de contorno occidental que son críticas para la variabilidad climática.

Metodología

Samudra 2 aborda estos desafíos mediante dos modificaciones arquitectónicas y de entrenamiento complementarias aplicadas a una arquitectura U-Net basada en bloques modificados de estilo ConvNeXt.

1. Escalamiento Arquitectónico

Para soportar resoluciones más altas ($1/2^\circ$ y $1/4^\circ$), los autores ensancharon la arquitectura U-Net en lugar de profundizarla.

• Ancho de Canales: Los anchos de canal a través de las etapas de codificador-decodificador se incrementaron de $[200, 250, 300, 400]$ a $[280, 380, 480, 520]$.
• Factor de Expansión: El factor de expansión interna dentro de los bloques ConvNeXt modificados se redujo de 4 a 2. Este diseño desplaza la capacidad de los parámetros desde el cuello de botella interno del bloque hacia las dimensiones de características entre etapas, equilibrando la capacidad de representación con el costo computacional.

2. Función de Pérdida Dinámica

Para abordar el colapso de la varianza y la impronta de artefactos, los autores reemplazaron la función de pérdida estándar de Error Cuadrático Medio (MSE) con una pérdida dinámica ponderada por canales.

• Mecanismo: El MSE estándar está dominado por los campos superficiales de alta varianza, lo que causa que el optimizador descuide las señales de evolución lenta del océano profundo. La nueva función de pérdida repondera cada canal de forma inversamente proporcional a su error de predicción actual.
• Implementación: Los pesos se actualizan en línea utilizando un Promedio Móvil Exponencial (EMA) del MSE inverso por canal. Esto amplifica adaptativamente la señal de gradiente para los canales con errores pequeños (típicamente los campos del océano profundo) mientras evita que los canales de alta varianza de la superficie dominen la pérdida.
• Limitación (Clamping): Los pesos se limitan a una razón (establecida en $L=20$) para evitar la sobreamplificación de canales con una señal aprendible insignificante.

Configuración de Entrenamiento y Evaluación

• Datos: Los modelos se entrenan con la salida del modelo oceánico GFDL OM4p25 ($1/4^\circ$ nativo), re-dimensionada (regridded) a $1^\circ$, $1/2^\circ$ y $1/4^\circ$.
• Variables: El emulador predice 77 canales (temperatura potencial, salinidad, velocidad zonal/meridional en 19 niveles de profundidad, y altura de la superficie del mar) impulsados por el forzamiento atmosférico.
• Despliegues (Rollouts): El entrenamiento utiliza despliegues autorregresivos cortos ($K=4$ pasos, 40 días). La evaluación involucra despliegues autorregresivos de largo horizonte de aproximadamente 8 años (~580 pasos) sin retroalimentación de la verdad de campo (ground-truth), evaluando métricas como la varianza temporal, distribuciones espectrales e índices climáticos (ej. Niño 3.4).

Principales Contribuciones

1. Modelo Samudra 2: Un emulador oceánico de IA mejorado que integra una U-Net ConvNeXt más ancha con una pérdida ponderada por varianza dinámica. Mitiga con éxito el colapso de la varianza y los artefactos de impronta en la emulación autorregresiva de largo horizonte.
2. Escalamiento Multi-Resolución: La primera demostración de la emulación oceánica multianual a resoluciones de $1/2^\circ$ y $1/4^\circ$. Los modelos recuperan estructuras de mesoescala (remolinos, meandros, filamentos) y corrientes de contorno occidental nítidas que no se resuelven a $1^\circ$.
3. Marco de Diagnóstico: Una evaluación exhaustiva que muestra que, si bien la precisión a corto plazo es necesaria, los diagnósticos climáticos de largo horizonte (varianza, distribución espectral) son esenciales para validar los emuladores oceánicos destinados a la proyección climática.

Resultados

Desempeño a Resolución de $1^\circ$

• Océano Superior: El $R^2$ para la temperatura del océano superior de media global mejoró de 0.56 (Samudra) a 0.87 (Samudra 2).
• Océano Profundo: El error de temperatura en el océano profundo se redujo aproximadamente siete veces. Aunque los valores de $R^2$ siguieron siendo negativos (indicando un desempeño peor que la línea base de la media temporal), la magnitud del error se redujo significativamente (de $-111.36$a$-16.14$).
• Impronta: El análisis visual y estadístico confirmó la supresión de los artefactos de impronta de velocidad en los campos de temperatura y salinidad del océano profundo.

Desempeño Multi-Resolución ($1/2^\circ$ y $1/4^\circ$)

• Fidelidad Espacial: Las resoluciones más altas resolvieron progresivamente escalas espaciales más finas. El modelo de $1/4^\circ$ capturó remolinos de mesoescala y puntos calientes de corrientes de contorno occidental (Golfo, Kuroshio, Agulhas) con estructuras de varianza que coinciden estrechamente con la referencia OM4.
• Fidelidad Espectral: El modelo de $1/4^\circ$ extendió el rango de número de onda resuelto por un factor de cuatro en comparación con el modelo de $1^\circ$, accediendo a la dinámica de mesoescala ausente en mallas más gruesas.
• Seguimiento Temporal: Todas las resoluciones rastrearon bien la variabilidad del océano superior ($R^2 = 0.84–0.92$). El desempeño en profundidades intermedias mejoró con la resolución (ej. el $R^2$ a $1/4^\circ$ se acercó a 0.01), aunque el seguimiento del océano profundo ($2000–7000$ m) sigue siendo un desafío abierto con $R^2$ negativos en todas las resoluciones.

Estudios de Ablación

Los experimentos de ablación a $1^\circ$ revelaron una fuerte interacción sinérgica entre las dos modificaciones:

• Arquitectura Más Ancha por sí sola: No mejoró el desempeño del océano superior ($R^2 = 0.50$) y degradó la salinidad de profundidad intermedia.
• Pérdida Dinámica por sí sola: No mejoró el desempeño del océano superior ($R^2 = 0.51$) pero fue el principal motor para reducir el error en el océano profundo.
• Combinado (Samudra 2): Logró un $R^2 = 0.87$ en el océano superior, demostrando que la arquitectura más ancha proporciona la capacidad necesaria para que la pérdida dinámica funcione efectivamente.

Significado y Reivindicaciones

El artículo afirma que Samudra 2 permite realizar conjuntos más grandes para proyecciones del nivel del mar, estudios de absorción de calor oceánico y análisis de variabilidad climática al ejecutarse en una sola GPU. Al resolver características de mesoescala a $1/4^\circ$, el emulador ofrece una alternativa tratable para estudiar procesos dependientes de la dinámica de remolinos.

Los autores mantienen una postura modesta respecto a la fidelidad del océano profundo, reconociendo que las predicciones por debajo de $\sim700$ m siguen siendo un desafío debido a las señales de entrenamiento extremadamente débiles y la dominancia de la varianza de fondo. Posicionan a Samudra 2 como un paso significativo que aborda modos de falla específicos de largo horizonte (colapso de varianza, impronta), mientras señalan que las simulaciones basadas en la física siguen siendo la referencia más confiable para la variabilidad del océano profundo. El trabajo se enmarca como una contribución a una clase más amplia de problemas de aprendizaje automático autorregresivos de largo horizonte, abordando cuestiones como el sesgo espectral y el desequilibrio de pérdida entre variables.