Training Diffusion Language Models for Black-Box Optimization

1. Problématique : L'Optimisation Boîte Noire Hors Ligne (Offline BBO)

L'article s'intéresse au problème de l'optimisation boîte noire hors ligne (Offline Black-Box Optimization - BBO). Ce défi consiste à découvrir de nouveaux designs optimaux (par exemple, des séquences d'ADN, des morphologies de robots ou des matériaux) à partir d'un jeu de données statique et limité contenant des paires « design-étiquette » (performance), sans pouvoir interroger la fonction objectif réelle (qui est souvent coûteuse ou impossible à évaluer directement).

• Contexte : Ce problème est crucial dans des domaines comme la robotique, la synthèse d'ADN et la science des matériaux, où les données étiquetées sont rares (par exemple, $N=500$ échantillons).
• Limites des approches existantes :
  • Les méthodes traditionnelles utilisent des modèles de substitution (surrogates) qui souffrent d'incertitude épistémique et de problèmes de distribution hors domaine (out-of-distribution).
  • Les approches récentes basées sur les LLM autoregressifs (génération de gauche à droite) peinent à capturer les dépendances bidirectionnelles inhérentes aux problèmes de conception (ex: un nucléotide dans une séquence d'ADN dépend à la fois de son préfixe et de son suffixe).

2. Méthodologie : DiBO (Diffusion Language Models for BBO)

Les auteurs proposent DiBO, un cadre qui adapte les Modèles de Langage à Diffusion (Diffusion LLMs) pour l'optimisation boîte noire, exploitant leur capacité de modélisation bidirectionnelle. La méthodologie se décompose en trois étapes clés :

A. Adaptation de Domaine via des Tokens Délimiteurs

Il existe un écart de domaine entre le pré-entraînement des LLMs sur du texte naturel et les signaux hétérogènes du BBO (prompts textuels, designs discrets/continus, et étiquettes numériques).

• Solution : Construction d'un corpus unifié « Prompt-Réponse ».
• Innovation : Introduction de tokens délimiteurs explicites (|design-start|, |design-end|, |label-start|, |label-end|) pour marquer les frontières sémantiques entre les instructions, les designs et les scores.
• Apprentissage : Une perte conjointe de reconstruction de masques est utilisée pour adapter le modèle à prédire simultanément les tokens masqués dans le prompt et la réponse, permettant au modèle de comprendre la structure du problème BBO.

B. Cadre de Post-Entraînement en Deux Étages

Pour aligner la génération du modèle sur les designs à haute performance, un pipeline en deux étapes est proposé :

1. Affinement Supervisé (SFT) :

   • Le modèle est entraîné sur le corpus unifié en masquant uniquement les tokens de la réponse (le design généré).
   • Cela introduit un biais inductif vers les régions de l'espace de conception à haute étiquette, servant d'initialisation stable.

2. Apprentissage par Renforcement (RL) :

   • Un ensemble de données RL est construit où la récompense est définie par l'amélioration de l'étiquette entre le prompt et la réponse ($r = y_{réponse} - y_{prompt}$).
   • Efficacité : Au lieu d'un débruitage itératif coûteux, les auteurs utilisent une approximation en une étape de la log-probabilité pour calculer la perte RL.
   • Cette étape intègre un feedback de récompense finement granulaire pour affiner la génération au-delà de la simple imitation.

3. Contributions Clés

1. Adaptation des Diffusion LLMs au BBO : C'est la première application des modèles de diffusion (bidirectionnels) à l'optimisation boîte noire hors ligne, surpassant les limitations des modèles autoregressifs unidirectionnels.
2. Corpus Unifié et Délimiteurs : Proposition d'une méthode pour combler l'écart de domaine entre le texte naturel et les données de conception via des tokens sémantiques explicites.
3. Pipeline de Post-Entraînement Hybride : Développement d'une stratégie combinant SFT (pour l'alignement structurel) et RL (pour l'optimisation de la récompense), permettant une exploration efficace de l'espace de conception.
4. Performance État-de-l'Art : Démonstration de résultats supérieurs sur des tâches complexes avec peu de données.

4. Résultats Expérimentaux

Les expériences ont été menées sur la suite de benchmarks Design-Bench, incluant des tâches discrètes (séquences d'ADN TF Bind 8 et 10) et continues (morphologie de robots Ant et D'Kitty).

• Performance Globale : DiBO obtient les meilleurs résultats globaux, se classant premier en moyenne et en médiane des rangs sur tous les benchmarks comparés.
• Comparaison :
  • Sur TF Bind 8, DiBO surpasse la meilleure méthode de référence (baseline) d'un écart significatif (> 0.108 en score normalisé).
  • Il surpasse systématiquement les méthodes « forward » (basées sur des gradients et des modèles de substitution) et les méthodes « inverse » (VAE, GAN, LLM autoregressifs comme OPRO, et autres modèles de diffusion).
• Analyse d'ablation :
  • L'utilisation des tokens délimiteurs est cruciale : leur suppression entraîne une baisse de performance notable, surtout lors de l'adaptation de domaine.
  • Les trois étapes (DA, SFT, RL) sont complémentaires ; retirer l'une d'elles (notamment le SFT intermédiaire avant le RL) dégrade fortement les performances, prouvant que le RL seul ne suffit pas pour apprendre les composantes hétérogènes complexes.
  • La construction du contexte basée sur la similarité des exemples few-shot est essentielle pour la cohérence sémantique.

5. Signification et Impact

Ce travail établit les Diffusion LLMs comme un paradigme robuste et généraliste pour l'optimisation hors ligne dans des espaces de conception à haute dimension.

• Avantages : La capacité de modélisation bidirectionnelle permet de mieux capturer les dépendances structurelles complexes (comme dans l'ADN) que les modèles autoregressifs.
• Efficacité des données : La méthode fonctionne efficacement dans des régimes de données très limitées (small-data), ce qui est critique pour les applications scientifiques réelles où l'expérimentation est coûteuse.
• Perspective : DiBO offre une alternative puissante aux méthodes d'optimisation traditionnelles, réduisant la dépendance aux essais physiques et accélérant la découverte dans des domaines comme la biologie synthétique et la robotique.

En résumé, DiBO démontre que l'adaptation des modèles de diffusion, couplée à une ingénierie de prompt soignée et un post-entraînement par renforcement, permet de résoudre des problèmes d'optimisation complexes là où les approches précédentes échouaient.