Open Materials Generation with Inference-Time Reinforcement Learning

Ce document introduit OMatG-IRL, un nouveau cadre d'apprentissage par renforcement par gradient de politique qui opère directement sur les champs de vitesse de modèles génératifs en temps continu afin de permettre une prédiction de structures cristallines efficace et alignée sur les propriétés cibles sans nécessiter de calcul explicite du score.

L'essentiel

Imaginez que vous soyez un maître architecte essayant de concevoir les blocs de construction parfaits pour un nouveau type de gratte-ciel. Dans le monde de la science des matériaux, ces « blocs » sont des cristaux. Pendant longtemps, les ordinateurs ont été doués pour apprendre à quoi ces blocs ressemblent en étudiant des millions d'exemples existants. Ils peuvent générer de nouvelles structures cristallines stables qui ressemblent beaucoup à la réalité.

Cependant, il y a un bémol : l'ordinateur est excellent pour copier la forme, mais il n'est pas très doué pour suivre des instructions spécifiques comme : « Rends ce cristal super solide » ou « Fais en sorte qu'il conduise mieux l'électricité ». C'est comme si vous aviez un robot capable de dessiner une maison parfaite, mais si vous lui demandez de « dessiner une maison qui ne prend pas feu », il redessine simplement la même maison parce qu'il ne sait pas comment donner la priorité à cet objectif spécifique.

Ce document présente une nouvelle méthode appelée OMatG-IRL pour corriger cela. Voici comment elle fonctionne, décomposée en concepts simples :

1. Le Problème : Le « Score » vs la « Vélocité »

La plupart des modèles d'IA avancés qui génèrent des formes fonctionnent de l'une des deux manières suivantes :

• La méthode du « Score » : L'IA apprend un « score » (comme un gradient sur une colline) qui lui indique exactement quelle direction prendre pour obtenir une meilleure forme. C'est comme avoir un GPS qui dit : « Tournez à gauche pour vous rapprocher de la destination ».
• La méthode de la « Vélocité » : L'IA apprend une « vélocité » (vitesse et direction) pour passer d'un amas de bruit aléatoire à une forme cristalline. C'est comme une rivière qui coule d'une montagne vers la mer. L'IA connaît la direction du courant, mais elle ne connaît pas nécessairement le « score » ou le gradient mathématique exact de la colline.

Le problème est que les outils les plus puissants pour apprendre à l'IA à suivre des objectifs spécifiques (appelés Apprentissage par Renforcement) nécessitent généralement la méthode du « Score ». Si vous n'avez que la méthode de la « Vélocité », vous ne pouvez pas facilement apprendre à l'IA à optimiser des propriétés comme l'efficacité énergétique.

2. La Solution : Enseigner au fleuve à couler différemment

Les auteurs ont créé un contournement ingénieux. Ils ont réalisé que même si vous n'avez que la « vélocité » (le flux de la rivière), vous pouvez toujours apprendre à l'IA à suivre de nouveaux objectifs en ajoutant une petite dose de hasard (bruit) au flux.

Voyez cela comme ceci :

• Imaginez que l'IA essaie de faire rouler une bille en bas d'une colline pour trouver le point le plus bas (le cristal le plus stable).
• Normalement, la bille roule parfaitement droit sur le chemin conçu par l'IA.
• OMatG-IRL ajoute une « brise » légère et contrôlée qui pousse la bille légèrement hors de sa trajectoire.
• Grâce à cette brise, la bille finit parfois dans un endroit légèrement différent. L'ordinateur vérifie : « Est-ce que ce nouvel endroit avait une énergie plus basse ? Était-ce un meilleur cristal ? »
• Si la réponse est « Oui », l'IA apprend : « D'accord, la prochaine fois, pousse la bille un peu plus dans cette direction. »

Cela permet à l'IA d'apprendre de ses erreurs et de ses succès sans avoir besoin de la carte complexe du « score ». Elle apprend en expérimentant avec le flux lui-même.

3. Le Tour de Magie du « Voyage dans le Temps » (Vélocité d'Annealing)

Les auteurs ont également découvert quelque chose de surprenant sur la vitesse à laquelle l'IA génère ces cristaux. Habituellement, pour obtenir un cristal parfait, l'IA doit effectuer des centaines de petites étapes lentes (comme descendre prudemment un escalier escarpé). Cela prend beaucoup de temps.

Les auteurs ont utilisé leur nouvelle méthode d'apprentissage pour enseigner à l'IA un nouveau programme pour sa vitesse. Au lieu de marcher lentement tout le long, l'IA a appris à :

1. Commencer avec une vitesse spécifique.
2. Accélérer ou ralentir aux moments les plus opportuns.
3. Terminer la tâche en une fraction du temps habituel.

C'est comme apprendre à un coureur qui a l'habitude de faire des joggings de 10 miles à soudainement sprinter le dernier kilomètre parfaitement, ou à prendre un raccourci qui ne fonctionne que s'il court à un rythme précis. Le résultat ? L'IA peut générer des cristaux de haute qualité 10 fois plus vite (ou plus encore) qu'auparavant, avec le même niveau de précision.

4. Pourquoi cela importe pour les cristaux

Dans la tâche spécifique de la Prédiction de Structure Cristalline (CSP) — où vous donnez à l'IA une liste d'ingrédients (comme du Carbone et de l'Oxygène) et lui demandez de construire le meilleur cristal possible — les auteurs ont montré que :

• Ils pouvaient apprendre à l'IA à construire des cristaux avec une énergie plus basse (ce qui signifie qu'ils sont plus stables et susceptibles d'exister dans la nature).
• Ils ont fait cela sans avoir besoin de calculer le « score » complexe que les autres méthodes exigent.
• Ils ont fait cela tout en maintenant une grande variété de cristaux (pour que l'IA ne se contente pas de mémoriser une seule réponse).
• Ils ont rendu le processus beaucoup plus rapide, réduisant le temps nécessaire pour générer un cristal de centaines d'étapes à seulement quelques dizaines.

Résumé

Le papier présente une nouvelle façon d'entraîner l'IA pour concevoir de meilleurs matériaux. C'est comme prendre une rivière qui coule naturellement dans une certaine direction et lui apprendre à changer occasionnellement de cours pour trouver une meilleure destination, le tout sans avoir besoin d'une carte détaillée de l'ensemble du paysage. Cela permet aux scientifiques de concevoir de nouveaux matériaux plus rapidement et avec des propriétés plus spécifiques que jamais.

Résumé technique

Résumé Technique : Génération de Matériaux Ouverte avec Apprentissage par Renforcement au Moment de l'Inférence (OMatG-IRL)

1. Énoncé du Problème

Les modèles génératifs en temps continu sont devenus des outils puissants pour la conception inverse de matériaux, capables de prédire des structures cristallines stables. Cependant, une limitation significative persiste : l'intégration de propriétés cibles explicites (par exemple, des objectifs mécaniques, électroniques ou énergétiques spécifiques) dans le processus génératif reste difficile. Bien que l'apprentissage par renforcement (RL) par gradient de politique offre un mécanisme rigoureux pour aligner les modèles génératifs avec des objectifs en aval, son application aux modèles de flux a été entravée par une contrainte technique.

Les méthodes standards de RL par gradient de politique nécessitent généralement l'accès au score (le gradient de la densité de probabilité logarithmique) pour calculer les ratios de politique et effectuer les mises à jour. De nombreux modèles de flux modernes, particulièrement ceux utilisant des interpolants stochastiques (SI) ou le Flow Matching, apprennent uniquement des champs de vitesse et ne calculent ni ne stockent explicitement le score. Par conséquent, ces modèles ont été inaccessibles aux cadres de RL standards, limitant leur capacité à optimiser des objectifs spécifiques non implicites au-delà de la stabilité inhérente à la distribution d'entraînement.

2. Méthodologie : OMatG-IRL

Les auteurs introduisent OMatG-IRL (Open Materials Generation with Inference-Time Reinforcement Learning), un cadre de RL par gradient de politique conçu pour opérer directement sur les champs de vitesse appris des modèles génératifs en temps continu, éliminant ainsi le besoin de calcul explicite du score.

Mécanisme Central

OMatG-IRL tire parti de l'observation empirique selon laquelle les métriques d'évaluation standard de la prédiction de structure cristalline (CSP) sont robustes aux petites perturbations stochastiques introduites dans la dynamique de l'équation différentielle ordinaire (ODE) sous-jacente. La méthode procède comme suit :

1. Processus Stochastique de Substitution : Pour les modèles qui apprennent uniquement un champ de vitesse $\hat{v}_\theta(t, x_t)$, l'intégration de l'ODE déterministe est augmentée d'un petit programme de bruit $\sigma_{ref}(t)$. Cela crée une équation différentielle stochastique (SDE) de substitution qui préserve la performance de base du modèle pré-entraîné tout en permettant l'exploration nécessaire.
   $$x_{t+\Delta t} = x_t + \hat{v}_{\theta_{ref}}(t, x_t)\Delta t + \sigma_{ref}(t)\sqrt{\Delta t}\xi$$
   Cette substitution définit une politique de référence pour la régularisation de Kullback-Leibler (KL).

2. Exploration au Moment de l'Inférence : Pendant le RL, le modèle explore en utilisant un champ de vitesse renforcé $\hat{v}_\theta(t, x_t)$ et potentiellement un programme de bruit $\sigma(t)$ différent pour améliorer l'exploration.

3. Optimisation de la Politique (GRPO) : Le cadre utilise l'optimisation de politique relative par groupe (GRPO). Pour une composition donnée, plusieurs trajectoires sont générées. Les récompenses terminales (par exemple, l'énergie négative par atome) sont calculées, et les avantages relatifs au groupe sont calculés pour mettre à jour la politique. Cette approche évite le besoin d'une fonction de valeur apprise et stabilise l'optimisation à travers des échelles de récompense hétérogènes.

4. Apprentissage par Recuit de Vitesse : Une application novatrice d'OMatG-IRL implique l'apprentissage d'un programme de recuit de vitesse dépendant du temps $s_\theta(t)$. Au lieu d'utiliser des programmes de recuit faits à la main, le modèle apprend une correction résiduelle au champ de vitesse gelé :
   $$x_{t+\Delta t} = x_t + [1 + s_\theta(t)]\hat{v}_{\theta_{ref}}\Delta t + \sigma(t)\hat{v}_{\theta_{ref}}\sqrt{\Delta t}\xi$$
   Cela permet au modèle d'adapter l'échelle du champ de vitesse pour améliorer l'efficacité de l'échantillonnage.

Applicabilité

Le cadre est conçu pour être flexible :

• Basé sur la Vitesse : Opère sur des modèles apprenant uniquement des champs de vitesse (aucun score requis).
• Basé sur le Score : Peut également être appliqué à des modèles prédisant à la fois la vitesse et le débruiteur (denoiser), mettant à jour conjointement les deux composantes.

3. Contributions Clés

• Première Application du RL à la CSP : Ce travail présente la première application du RL par gradient de politique spécifiquement à la tâche de prédiction de structure cristalline (CSP), où la composition est fixe et la structure est générée.
• RL sans Score pour les Modèles de Flux : OMatG-IRL permet le RL pour les modèles génératifs de flux qui apprennent uniquement des champs de vitesse, surmontant la limitation qui restreignait auparavant le RL aux modèles de diffusion basés sur le score.
• Renforcement Basé sur l'Énergie sans Récompenses de Diversité : Contraênment aux tâches de génération De Novo (DNG) qui nécessitent des récompenses de diversité explicites pour éviter l'effondrement de mode, la tâche CSP maintient naturellement la diversité grâce au conditionnement de la composition. Les auteurs démontrent que les objectifs basés sur l'énergie peuvent être renforcés efficacement sans pénalités de diversité supplémentaires.
• Programmes de Recuit Appris : Le papier introduit une méthode pour apprendre des programmes de recuit de vitesse dépendant du temps via le RL, remplaçant les heuristiques faites à la main.

4. Résultats Expérimentaux

Les auteurs ont évalué OMatG-IRL sur le jeu de données MP-20 (Materials Project) en utilisant le cadre OMatG.

• Renforcement de l'Énergie : Les variantes de l'OMatG-IRL, tant basées sur le score que sur la vitesse, ont réussi à renforcer l'énergie relative par atome, atteignant des réductions d'environ 0,5 eV par atome par rapport au modèle de base pré-entraîné.
• Parité de Performance : L'approche basée sur la vitesse (qui ne nécessite pas de calcul de score) a atteint une performance comparable à l'approche basée sur le score, validant l'efficacité du processus stochastique de substitution.
• Efficacité de l'Échantillonnage :
  • Le cadre de RL a permis une CSP précise avec une réduction drastique des étapes d'intégration.
  • Plus précisément, la variante Velocity-Annealing OMatG-IRL a récupéré la performance d'un modèle de base nécessitant $N_t = 950$ étapes d'intégration en utilisant seulement $N_t = 100$ étapes.
  • Remarquablement, le programme de recuit appris est resté robuste même lorsque les étapes étaient réduites à $N_t = 10$, alors que le modèle de base fait main a échoué rapidement sous une discrétisation temporelle agressive.
• Robustesse : La méthode a maintenu les taux de correspondance (match rates) et réduit l'erreur quadratique moyenne (RMSE) tout en diminuant considérablement le coût de calcul de la génération (d'un ordre de grandeur).

5. Signification et Revendications

Les auteurs affirment qu'OMatG-IRL représente une avancée significative dans la conception inverse de matériaux cristallins en :

1. Démocratisant le RL pour les Modèles de Flux : En supprimant la dépendance à un calcul explicite du score, le cadre étend les bénéfices du RL (optimisation d'objectifs en aval spécifiques) à une classe plus large de modèles génératifs en temps continu, y compris ceux basés sur le Flow Matching et les interpolants stochastiques généraux.
2. Améliorant l'Efficacité : La capacité d'apprendre des programmes de recuit de vitesse optimaux permet une prédiction de structure précise avec beaucoup moins d'étapes d'intégration, répondant directement au goulot d'étranglement computationnel du criblage de matériaux.
3. Optimisation Spécifique à la Tâche : Ce travail démontre que le RL peut efficacement optimiser les objectifs physiques (comme la minimisation de l'énergie) dans la CSP sans compromettre la diversité structurelle inhérente à la tâche, offrant une voie plus directe vers la découverte de matériaux possédant des propriétés ciblées.

Les auteurs notent des limites, notamment le fait que le processus stochastique de substitution n'est pas exactement préservateur des marges (bien que la divergence soit bornée et négligeable pour un faible bruit) et que l'objectif actuel basé sur l'énergie n'optimise pas directement les métriques de correspondance de structure comme le taux de correspondance, bien que ces métriques restent corrélées. Le code est publié dans le cadre de la mise à jour du framework Open Materials Generation (OMatG).