Resumen Técnico: Destilación de Modelos de Fundación para EDP dependientes del tiempo

1. Definición del Problema

Las ecuaciones diferenciales parciales (EDP) son fundamentales para modelar procesos físicos espaciotemporales, como la dinámica de fluidos. Si bien las redes neuronales han surgido como modelos sustitutos (surrogates) eficientes para las EDP, enfrentan un compromiso crítico:

• Modelos de Fundación (FMs): Los modelos a gran escala preentrenados en diversos conjuntos de datos físicos (por ejemplo, Poseidon, Walrus) demuestran una fuerte generalización en tareas de seguimiento con datos limitados. Sin embargo, son computacionalmente costosos, ya que a menudo dependen de arquitecturas Transformer de gran tamaño que carecen de sesgos inductivos específicos (por ejemplo, la equivariancia) y requieren un tiempo de inferencia y memoria significativos. Esto limita su utilidad en aplicaciones en tiempo real, como el control y la planificación.
• Sustitutos Compactos: Los modelos más pequeños (por ejemplo, los Operadores Neuronales de Fourier) son eficientes, pero típicamente requieren grandes cantidades de datos de entrenamiento específicos de la tarea para lograr una alta precisión, los cuales a menudo no están disponibles en sistemas experimentales o parcialmente observados.

El desafío central es transferir las capacidades de generalización de los grandes modelos de fundación de EDP preentrenados hacia modelos sustitutos compactos y eficientes sin requerir acceso a la distribución completa de condiciones iniciales del dominio objetivo, la cual suele ser desconocida o costosa de muestrear.

2. Metodología: Extensión de Rollout del Profesor (TREX)

Los autores proponen Teacher Rollout Extension (TREX), un marco de destilación de conocimiento diseñado para cerrar la brecha entre los grandes modelos de fundación y los sustitutos eficientes para tareas específicas.

2.1 Concepto Central

TREX aborda la limitación de la destilación estándar, que a menudo requiere el muestreo de nuevas condiciones iniciales a partir de una distribución conocida $p_0$. En muchos escenarios de aplicación (por ejemplo, datos experimentales), $p_0$ es desconocida. En su lugar, TREX aumenta las trayectorias limitadas disponibles mediante la generación de trayectorias sintéticas largas a través de rollouts del profesor.

2.2 El Procedimiento TREX

1. Ajuste Fino (Fine-tuning): Un modelo de fundación preentrenado (Profesor, $T$) se ajusta primero utilizando un pequeño conjunto de trayectorias disponibles de la tarea específica.
2. Generación de Rollouts: Partiendo de los estados iniciales de las pocas trayectorias reales disponibles, el profesor ajustado se ejecuta de forma autorregresiva para generar trayectorias sintéticas largas.
   • Medida de Ocupación: Este proceso muestrea de la distribución de rollout inducida por el profesor (estados encontrados durante la predicción autorregresiva) en lugar de solo las trayectorias cortas observadas.
   • Inyección de Ruido: Para ampliar la cobertura del espacio de estados y enseñar al estudiante cómo evoluciona el sistema desde estados perturbados, se inyecta ruido gaussiano periódico en el estado del profesor durante los rollouts. El profesor predice entonces el siguiente estado basándose en este estado perturbado.
3. Entrenamiento del Estudiante: Un modelo compacto (Estudiante, $S$) se entrena sobre una mezcla de:
   • Transiciones de la verdad de campo (ground-truth) originales.
   • Transiciones generadas por el profesor (segmentos tanto libres de ruido como con ruido).
   • La pérdida del estudiante combina la supervisión de la verdad de campo ($L_{GT}$) y la pérdida de destilación del profesor ($L_{TREX}$).

2.3 Sesgos Inductivos Arquitectónicos

A diferencia del profesor, que es típicamente un Transformer grande y flexible, la arquitectura del estudiante se elige para incorporar sesgos inductivos específicos de la tarea. Los autores utilizan un Operador Neuronal de Fourier Tensorizado (TFNO), el cual impone la equivariancia a las traslaciones y suprime el ruido de alta frecuencia, propiedades que no son forzadas explícitamente por el profesor.

3. Contribuciones Clave

1. Primer Estudio Sistemático: El artículo presenta la primera investigación sistemática sobre la destilación de modelos de fundación de EDP preentrenados generales en modelos sustitutos compactos para tareas específicas.
2. Marco TREX: Introducción de un novedoso método de destilación que genera datos de supervisión mediante rollouts ruidosos del profesor, eliminando la necesidad de un muestreador explícito de condiciones iniciales.
3. Integración de Sesgo Inductivo: Demostración de que destilar un modelo de fundación general en una arquitectura pequeña con sesgos inductivos adecuados (por ejemplo, la equivariancia) puede respetar las propiedades físicas manteniendo la precisión predictiva del profesor.

4. Resultados Experimentales

Los autores evaluaron TREX utilizando Poseidon y Walrus como profesores a través de varios benchmarks de dinámica de fluidos (ecuaciones de Navier-Stokes y Euler compresibles).

• Precisión: Los estudiantes entrenados con TREX igualaron o superaron la precisión de los modelos de fundación (profesores) de gran tamaño en regímenes de pocos datos (de 4 a 32 trayectorias). En algunos casos, el estudiante superó al profesor incluso cuando fue entrenado únicamente con datos generados por el profesor (sin verdad de campo).
• Eficiencia:
  • Parámetros: Los modelos estudiante redujeron el recuento de parámetros en varios órdenes de magnitud (por ejemplo, de ~628M de parámetros en Poseidon-L a ~174K en TFNO, una reducción de ~3600x).
  • Velocidad de Inferencia: El tiempo de inferencia se redujo en más de un orden de magnitud (por ejemplo, de ~270ms a ~22ms por paso).
  • Memoria: El consumo de memoria GPU para un pase hacia adelante (forward pass) se redujo en un factor de ~3.2.
• Estabilidad y Equivarianza:
  • Los estudiantes exhibieron rollouts de largo horizonte más estables con un crecimiento de error más lento en comparación con los profesores.
  • El estudiante TFNO demostró errores de equivariancia de traslación significativamente menores en comparación con el profesor no equivariante, manteniendo la consistencia bajo desplazajes espaciales durante los rollouts autorregresivos.
• Ablaciones:
  • El método es robusto a la longitud del rollout y a los hiperparámetros de ruido.
  • Se encontró que la inyección de ruido durante los rollouts es beneficiosa, particularmente con amplitudes más altas.
  • Se demostró que la destilación es efectiva incluso cuando la arquitectura del estudiante (UNet, FNO estándar) no podía superar al profesor utilizando únicamente la verdad de campo.

5. Significado y Reivindicaciones

El artículo afirma que TREX reconcilia eficazmente el compromiso entre la fuerte generalización de los grandes modelos de fundación y los requisitos de eficiencia de los modelos sustitutos prácticos. Al extender los datos limitados de la tarea mediante los rollouts del profesor, TREX permite que los modelos compactos aprendan la dinámica de largo horizonte y las simetrías físicas que de otro modo requerirían conjuntos de datos masivos o arquitecturas complejas.

Los autores posicionan este trabajo como un paso hacia la implementación de los sustitutos basados en modelos de fundación en aplicaciones científicas y de ingeniería en tiempo real (por ejemplo, control de fabricación, sistemas de energía) donde el costo computacional y la velocidad de inferencia son restricciones críticas. Señalan que, si bien el método depende de la calidad del profesor ajustado, logra transferir el conocimiento incluso cuando el profesor carece de sesgos inductivos específicos, permitiendo que el estudiante imponga restricciones físicas como la equivariancia.