Résumé technique : Distillation de modèles de fondation pour les EDP dépendantes du temps

1. Définition du problème

Les équations aux dérivées partielles (EDP) sont fondamentales pour la modélisation des processus physiques spatio-temporels tels que la dynamique des fluides. Bien que les réseaux de neurones soient apparus comme des modèles de substitution (surrogates) efficaces pour les EDP, ils sont confrontés à un compromis critique :

• Modèles de fondation (FMs) : Les modèles à grande échelle préentraînés sur divers ensembles de données physiques (ex: Poseidon, Walrus) font preuve d'une forte capacité de généralisation aux tâches en aval avec peu de données. Cependant, ils sont coûteux en termes de calcul, s'appuyant souvent sur de larges architectures de type Transformer qui manquent de biais inductifs spécifiques (ex: l'équivariance) et nécessitent un temps d'inférence et une mémoire importants. Cela limite leur utilité dans les applications en temps réel comme le contrôle et la planification.
• Substituts compacts : Les modèles plus petits (ex: opérateurs de Fourier neuronaux) sont efficaces mais nécessitent généralement de grandes quantités de données d'entraînement spécifiques à la tâche pour atteindre une précision élevée, ce qui est souvent indisponible dans les systèmes expérimentaux ou partiellement observés.

Le défi central consiste à transférer les capacités de généralisation des grands modèles de fondation préentraînés vers des modèles substituts compacts et efficaces, sans nécessiter l'accès à la distribution complète des conditions initiales du domaine cible, laquelle est souvent inconnue ou coûteuse à échantillonner.

2. Méthodologie : Teacher Rollout Extension (TREX)

Les auteurs proposent Teacher Rollout Extension (TREX), un cadre de distillation de connaissances conçu pour combler l'écart entre les grands modèles de fondation et les substituts efficaces en aval.

2.1 Concept central

TREX répond à la limitation de la distillation standard, qui nécessite souvent d'échantillonner de nouvelles conditions initiales à partir d'une distribution $p_0$ connue. Dans de nombreux scénarios en aval (ex: données expérimentales), $p_0$ est inconnue. Au lieu de cela, TREX augmente les trajectoires limitées disponibles grâce à la génération de longues trajectoires synthétiques par des déroulements de l'enseignant (teacher rollouts).

2.2 La procédure TREX

1. Ajustement fin (Fine-tuning) : Un modèle de fondation préentraîné (Enseignant, $T$) est d'abord ajusté sur un petit ensemble de trajectoires disponibles en aval.
2. Génération de déroulements (Rollout Generation) : En partant des états initiaux des quelques trajectoires de vérité terrain disponibles, l'enseignant ajusté est déroulé de manière autorégressive pour générer de longues trajectches synthétiques.
   • Mesure d'occupation : Ce processus échantillonne à partir de la distribution de déroulement induite par l'enseignant (états rencontrés lors de la prédiction autorégressive) plutôt que de se limiter aux courtes trajectoires observées.
   • Injection de bruit : Pour élargir la couverture de l'espace d'états et apprendre au substitut comment le système évolue à partir d'états perturbés, un bruit gaussien périodique est injecté dans l'état de l'enseignant pendant les déroulements. L'enseignant prédit ensuite l'état suivant basé sur cet état perturbé.
3. Entraînement du substitut (Student) : Un modèle substitut compact ($S$) est entraîné sur un mélange de :
   • Transitions de vérité terrain originales.
   • Transitions générées par l'enseignant (segments sans bruit et avec bruit).
   • La perte du substitut combine la supervision de vérité terrain ($L_{GT}$) et la perte de distillation de l'enseignant ($L_{TREX}$).

2.3 Biais inductifs architecturaux

Contrairement à l'enseignant, qui est généralement un Transformer large et flexible, l'architecture du substitut est choisie pour incorporer des biais inductifs spécifiques à la tâche. Les auteurs utilisent un Opérateur de Fourier Neural Tensorisé (TFNO), qui impose l'équivariance aux translations et supprime le bruit à haute fréquence, des propriétés qui ne sont pas explicitement imposées à l'enseignant.

3. Contributions clés

1. Première étude systématique : Le papier présente la première investigation systématique de la distillation de modèles de fondation de PDE préentraînés généraux vers des modèles substituts compacts.
2. Cadre TREX : Introduction d'une nouvelle méthode de distillation qui génère des données de supervision via des déroulements d'enseignants bruyants, éliminant le besoin d'un échantillonneur explicite de conditions initiales.
3. Intégration de biais inductifs : Démonstration que la distillation d'un modèle de fondation général vers une petite architecture dotée de biais inductifs appropriés (ex: l'équivariance) peut respecter les propriétés physiques tout en maintenant la précision prédictive de l'enseignant.

4. Résultats expérimentaux

Les auteurs ont évalué TREX en utilisant Poseidon et Walrus comme enseignants à travers diverses épreuves de dynamique des fluides (équations de Navier-Stokes et d'Euler compressibles).

• Précision : Les substituts entraînés par TREX ont égalé ou surpassé la précision des grands modèles de fondation enseignants dans des régimes de faibles données (4 à 32 trajectoires). Dans certains cas, le substitut a surpassé l'enseignant lorsqu'il était entraîné uniquement sur les données générées par l'enseignant (sans vérité terrain).
• Efficacité :
  • Paramètres : Les modèles substituts ont réduit le nombre de paramètres de plusieurs ordres de grandeur (ex: de ~628M de paramètres pour Poseidon-L à ~174K pour TFNO, soit une réduction d'environ ~3600x).
  • Vitesse d'inférence : Le temps d'inférence a été réduit de plus d'un ordre de grandeur (ex: de ~270ms à ~22ms par étape).
  • Mémoire : La consommation de mémoire GPU pour une passe avant a été réduite d'un facteur d'environ 3,2.
• Stabilité et Équivariance :
  • Les substituts ont présenté des déroulements à long horizon plus stables avec une croissance d'erreur plus lente par rapport aux enseignants.
  • Le substitut TFNO a démontré des erreurs d'équivariance par translation significativement plus faibles que l'enseignant non équivariant, maintenant la cohérence lors des décalages spatiaux durant les déroulements autorégressifs.
• Ablations :
  • La méthode est robuste à la longueur de déroulement et aux hyperparamètres de bruit.
  • L'injection de bruit lors des déroulements s'est avérée bénéfique, particulièrement avec des amplitudes plus élevées.
  • La distillation s'est révélée efficace même lorsque l'architecture du substitut (UNet, FNO standard) ne pouvait pas surpasser l'enseignant en utilisant uniquement la vérité terrain.

5. Signification et affirmations

Le papier affirme que TREX réconcilie efficacement le compromis entre la forte capacité de généralisation des grands modèles de fondation et les exigences d'efficacité des modélisations substituts pratiques. En étendant les données limitées en aval via les déroulements de l'enseignant, TREX permet aux modèles compacts d'apprendre la dynamique à long horizon et les symétries physiques qui nécessiteraient autrement des ensembles de données massifs ou des architectures complexes.

Les auteurs positionnent ce travail comme une étape vers la rendre déployable pour des applications de contrôle et d'ingénierie en temps réel (ex: contrôle de fabrication, systèmes énergétiques) où le coût computationnel et la vitesse d'inférence sont des contraintes critiques. Ils notent que bien que la méthode dépende de la qualité de l'enseignant ajusté, elle transfère avec succès les connaissances même lorsque l'enseignant manque de biais inductifs spécifiques, permettant au substitut d'imposer des contraintes physiques telles que l'équivariance.