Résumé technique : Modèles de langage à diffusion pour l'optimisation en boîte noire

Énoncé du problème

L'optimisation en boîte noire (BBO) hors ligne cherche à identifier des conceptions optimales ($x^*$) qui maximisent une fonction objectif inconnue $f(\cdot)$ en utilisant uniquement un ensemble limité de données de paires conception-étiquette. Ce scénario est prévalent dans des domaines tels que la conception de séquences d'ADN et la robotique, où l'interrogation des propriétés est coûteuse et les données étiquetées sont rares (régime de few-shot).

Les approches traditionnelles se répartissent généralement en deux catégories :

1. Méthodes basées sur des substituts (proxies) : Entraîner un modèle de substitution (par exemple, des processus gausiens ou des réseaux de neurones profonds) pour approximer $f(\cdot)$ et utiliser ses gradients pour guider l'optimisation.
2. Méthodes génératives : Entraîner des modèles génératifs conditionnels pour produire directement des conceptions à haut score.

Cependant, ces méthodes souffrent souvent de problèmes de généralisation hors distribution dans des contextes de données limitées. De plus, les adaptations récentes des modèles de langage de grande taille (LLM) autorégressifs à la BBO, qui cadrent les tâches comme des prompts en langage naturel, font face à une limitation critique : elles génèrent des séquences de gauche à droite. Cette approche unidirectionnelle peine à capturer les dépendances bidirectionnelles inhérentes à de nombreuses tâches de conception (par exemple, la fonction d'une unité d'ADN dépend à la fois de son préfixe et de son suffixe).

Méthodologie : dLLM

Les auteurs proposent dLLM, un cadre qui exploite les Modèles de Langage à Diffusion (Diffusion LLMs) pour la BBO hors ligne. La méthode combine deux modules de base : le Débruitage en Contexte (In-Context Denoising) et la Recherche par Arbre de Diffusion Masquée (Masked Diffusion Tree Search).

1. Débruitage en Contexte

Au lieu de procéder à un ajustement fin (fine-tuning), la méthode utilise les capacités d'apprentissage en contexte des LLM de diffusion pré-entraînés.

• Construction du Prompt : L'entrée du LLM de diffusion consiste en un prompt structuré en langage naturel contenant :
  • Description de la tâche : Définit le format de conception (ex: alphabet de l'ADN), l'objectif et les contraintes.
  • Jeu de données hors ligne : Un ensemble de quelques exemples (few-shot) de paires conception-étiquette (ex: séquences d'ADN et affinités de liaison).
  • Instruction : Une directive pour générer des candidats améliorés.
• Processus de débruitage : Le modèle commence par une conception entièrement masquée ($x_{t=1}$). Il prédit de manière itérative les jetons (tokens) masqués pour produire un candidat complet ($\hat{x}_0$). Les jetons les moins confiants sont re-masqués pour former un nouvel état partiellement masqué ($x_s$) pour l'étape suivante, affinant ainsi la conception par un débruitage itératif.

2. Recherche par Arbre de Diffusion Masquée (MDTS)

Pour guider la génération vers des régions de haute performance et équilibrer l'exploration et l'exploitation, le processus de débruitage est traité comme une recherche arborescente de Monte Carlo (MCTS) étape par étape.

• Structure de l'arbre : Chaque nœud représente une conception partiellement masquée ($x_t$).
• Sélection : L'algorithme parcourt l'arbre depuis la racine (entièrement masquée) en utilisant le score Upper Confidence Bound (UCT). Cela équilibre l'exploitation des nœuds de haute valeur et l'exploration des branches moins visitées. Le score UCT intègre la estimation de la valeur du nœud et la probabilité du modèle pour l'étape de démasquage.
• Expansion : Les nœuds sélectionnés sont étendus en appliquant le LLM de diffusion pour générer des complétions diverses, qui sont ensuite re-masquées pour créer des nœuds enfants.
• Évaluation : Les nœuds enfants sont déroulés en conceptions entièrement débruites. Leurs récompenses sont estimées à l'aide de l'Amélioration Attendue (Expected Improvement - EI) dérivée d'un processus gaussien (GP) entraîné sur le jeu de données hors ligne. Le GP fournit des estimations d'incertitude rigoureuses, cruciales dans les contextes de few-shot.
• Rétropropagation : Le signal de récompense est propagé vers le haut de l'arbre pour mettre à jour les estimations de valeur et les comptes de visites des nœuds ancêtres, informant ainsi les futures décisions de sélection.

Contributions Clés

1. Première application des LLM de diffusion à la BBO : L'article présente le premier cadre utilisant les LLM de diffusion pour l'optimisation en boîte noire hors ligne, capitalisant sur leurs capacités de modélisation bidirectionnelle et de raffinement itératif pour gérer des dépendances de conception complexes.
2. Module de Débruitage en Contexte : Un mécanisme de conditionnement novateur qui permet aux LLM de diffusion d'affiner les conceptions masquées sur la base de descriptions de tâches en langage naturel et de jeux de données hors ligne en few-shot, sans mise à jour des paramètres.
3. Recherche par Arbre de Diffusion Masquée : Un algorithme de recherche qui traite le débruitage de diffusion comme un processus de prise de décision séquentiel, équilibrant dynamiquement l'exploration et l'exploitation via MCTS et l'amélioration attendue basée sur les GP.

Résultats Expérimentaux

Les auteurs ont évalué dLLM sur le benchmark Design-Bench, couvrant des tâches continues (Ant Morphology, D'Kitty Morphology) et discrètes (TF Bind 8, TF Bind 10) dans un contexte de few-shot (10 exemples étiquetés).

• Performance : dLLM a obtenu des résultats de pointe (state-of-the-art) sur les quatre tâches, se classant premier en rang moyen et médian parmi 15 baselines. Il a démontré des gains constants dans les espaces de conception continus et discrets.
• Comparaison :
  • A surpassé les méthodes de substitution purement basées sur le gradient (ex: Grad-mean, Grad-EI), indiquant que les modèles génératifs couplés à la recherche sont supérieurs pour l'exploration.
  • A surpassé d'autres baselines génératives, y compris les LLM autorégressifs (OPRO) et d'autres modèles de diffusion (DDOM, GTG). Les auteurs attribuent cela à la modélisation bidirectionnelle des LLM de diffusion et au vaste savoir pré-entraîné qu'ils exploitent.
  • A surpassé la forte baseline autorégressive UniSO-T, suggérant que le contexte bidirectionnel est critique pour les tâches avec des dépendances complexes.
• Études d'ablation :
  • La suppression de la recherche par arbre (sans MDTS) a provoqué une chute substantielle de la performance, soulignant la nécessité d'une exploration guidée.
  • Le remplacement du LLM de diffusion par un modèle de diffusion classique entraîné uniquement sur les données few-shot a entraîné une dégradation significative, confirmant l'importance du savoir pré-entraîné.
  • La suppression de la description explicite de la tâche dans le prompt a réduit la performance, validant le rôle du cadrage sémantique.

Signification et Revendications

L'article affirme que dLLM démontre le potentiel des LLM de diffusion en tant qu'optimisateurs polyvalents pour la génération de conceptions efficaces en données. En abordant les limites des modèles autorégressifs pour capturer les dépendances bidirectionnelles et en les intégrant à une stratégie de recherche robuste, la méthode offre une alternative versatile pour les domaines scientifiques et d'ingénierie où les données étiquetées sont rares. Les auteurs positionnent ce travail comme ouvrant de nouvelles directions pour l'utilisation des modèles de fondation pré-entraînés dans l'optimisation en boîte noire hors ligne.