Domain decomposition of large neural network surrogate models

Resumen Técnico: Descomposición de Dominio para Modelos Sustitutos de Redes Neuronales

1. Planteamiento del Problema

En la optimización de diseño basada en mecánica computacional, la evaluación eficiente de cantidades de interés (QoIs) espaciotemporales (como campos de tensión y deformación) es un desafío crítico. Tradicionalmente, estos campos se calculan mediante solvers del Método de Elementos Finitos (FEM), lo cual es computacionalmente costoso, especialmente cuando se introducen dimensiones paramétricas adicionales para la optimización estocástica.

Aunque las Redes Neuronales (NN) han surgido como modelos sustitutos (surrogates) eficientes para aproximar estas mapeos, presentan limitaciones significativas:

• Aproximación Global: Las NN estándar intentan aproximar todo el dominio globalmente. Para capturar no linealidades locales con alta precisión, requieren una sobre-parametrización masiva, lo que conduce a tiempos de entrenamiento excesivos.
• Curse of Dimensionality: En espacios paramétricos de alta dimensión, los métodos tradicionales (como la expansión del caos polinomial) requieren un número prohibitivo de simulaciones FEM.
• Ineficiencia: Entrenar una única red neuronal grande para problemas complejos es lento y a menudo ineficiente en términos de recursos computacionales.

El objetivo del artículo es superar estas barreras mediante la Descomposición de Dominio (DDM), dividiendo el espacio de entrada en subdominios locales más simples, entrenados en paralelo, para reducir la complejidad y los tiempos de cómputo sin sacrificar la precisión.

2. Metodología Propuesta

Los autores proponen un enfoque de descomposición de dominio no superpuesto que divide el dominio espacial $G$ en $M$ subdominios $G_i$. En lugar de una red global, se entrena un conjunto de redes neuronales locales ($\hat{u}_{\theta_i}$), una por subdominio.

Para garantizar que la solución global sea continua y físicamente consistente, se imponen restricciones de interfaz entre los subdominios adyacentes. Se requieren dos tipos de continuidad:

1. Continuidad $C^0$: La función de aproximación debe ser continua en la interfaz.
2. Continuidad $C^1$: El gradiente normal de la función también debe ser continuo en la interfaz.

Para resolver el problema de optimización con restricciones, se proponen y comparan dos algoritmos basados en multiplicadores:

• Algoritmo de Multiplicadores de Lagrange (LMA):

  • Reformula el problema como una optimización no restringida mediante la función Lagrangiana.
  • Utiliza un enfoque de ascenso dual (dual ascent).
  • Dado que las NN son altamente no convexas, las restricciones de interfaz se linealizan iterativamente alrededor de puntos de referencia usando expansiones de Taylor.
  • Utiliza el optimizador L-BFGS (quasi-Newton) para actualizar los parámetros de la red dentro de las restricciones linealizadas.
  • Desventaja: La linealización es computacionalmente costosa y puede fallar en encontrar pasos factibles en problemas altamente no lineales, requiriendo múltiples intentos.

• Algoritmo de Multiplicadores de Lagrange Aumentados (ALMA):

  • Introduce un término de penalización (regularización) adicional en la función de pérdida, además de los multiplicadores de Lagrange.
  • La función objetivo incluye un término cuadrático de penalización ($\rho \|Q_i(\theta_i)\|^2$) que castiga el incumplimiento de las restricciones.
  • Permite actualizar los multiplicadores y los parámetros de la red directamente mediante métodos basados en gradientes sin necesidad de linealizar las restricciones de la NN en cada paso.
  • Utiliza también L-BFGS para la actualización primal.
  • Ventaja: Es más robusto, converge más rápido y es menos sensible a la no linealidad local que el LMA.

3. Contribuciones Clave

1. Enfoque Generalizado: A diferencia de trabajos previos limitados a Redes Neuronales Informadas por Física (PINNs), este método es basado en datos, aplicable a cualquier mapeo entrada-salida generado por FEM.
2. Nuevas Restricciones de Interfaz: Se implementan restricciones para la continuidad tanto del valor de la función como de sus gradientes normales ($C^1$), utilizando multiplicadores de Lagrange en lugar de solo términos de penalización por MSE (Error Cuadrático Medio).
3. Comparativa Rigurosa: Se evalúan exhaustivamente dos estrategias de optimización (LMA vs. ALMA) en términos de convergencia, precisión y escalabilidad.
4. Análisis de Escalabilidad: Se demuestra que la descomposición de dominio permite el entrenamiento paralelo de redes, reduciendo drásticamente el tiempo total de entrenamiento para problemas grandes.

4. Resultados Experimentales

Los métodos se validaron mediante dos casos de estudio:

• Caso 1: Compresión de un cilindro 2D (Lineal Elástico)

  • Configuración: Se dividió el dominio en subdominios verticales y multidireccionales.
  • Rendimiento de Tiempo:
    • ALMA demostró una escalabilidad superior. El tiempo de entrenamiento fue de 5.8 horas para 2 dominios y creció casi linealmente.
    • LMA mostró un crecimiento exponencial del tiempo de cómputo al aumentar los subdominios (hasta 140 horas), debido a la dificultad de encontrar pasos de linealización factibles.
  • Precisión:
    • Ambos métodos mejoraron significativamente la continuidad en las interfaces en comparación con NNs sin restricciones.
    • ALMA alcanzó el error relativo máximo mínimo (aprox. 4%) con una descomposición de 3 subdominios verticales.
    • LMA ofreció una precisión ligeramente superior en las interfaces pero a un costo computacional prohibitivo.
  • Robustez: Se probó la sensibilidad ante la falta de datos (huecos en la interfaz) y la reducción de puntos de muestreo. ALMA mantuvo la estabilidad, aunque la falta de datos en la interfaz introduce incertidumbre estocástica en la solución de la interfaz.

• Caso 2: Propagación de incertidumbre en un cilindro 3D

  • Objetivo: Optimización de parámetros de material (módulo bulk y shear) con incertidumbre.
  • Resultado: El modelo sustituto basado en ALMA con dos subdominios espaciales logró aproximar con alta precisión los campos de desplazamiento y sus estadísticas (media y desviación estándar).
  • Errores: El error relativo medio fue bajo (<2.5% para la media, <10% para la desviación estándar). La continuidad en la interfaz se mantuvo efectiva, validando el uso del método para problemas de optimización estocástica de alta dimensión.

5. Significado y Conclusiones

El estudio concluye que la Descomposición de Dominio con Multiplicadores de Lagrange Aumentados (ALMA) es la estrategia preferida para entrenar modelos sustitutos de redes neuronales a gran escala.

• Eficiencia: ALMA supera a LMA en velocidad de convergencia y escalabilidad, haciendo viable el entrenamiento de redes para dominios complejos y no lineales mediante computación paralela.
• Precisión vs. Costo: Aunque LMA puede ofrecer una precisión teórica ligeramente mayor en las interfaces, su alto costo computacional y su inestabilidad en problemas no lineales lo hacen menos práctico. ALMA ofrece el mejor equilibrio entre tiempo de entrenamiento y precisión.
• Aplicabilidad: El método permite abordar problemas de optimización de diseño que antes eran intratables debido a la complejidad computacional, permitiendo el uso de modelos sustitutos precisos en lugar de simulaciones FEM directas en cada iteración de optimización.

En resumen, el trabajo presenta un marco robusto para descomponer problemas de aprendizaje profundo complejos en subproblemas manejables, resolviendo los desafíos de continuidad mediante técnicas avanzadas de optimización con restricciones, lo que abre nuevas vías para la ingeniería computacional eficiente.