Training Diffusion Language Models for Black-Box Optimization

Titel: Training Diffusion Language Models voor Black-Box Optimalisatie

Auteurs: Zipeng Sun, Can (Sam) Chen, Ye Yuan, Haolun Wu, Jiayao Gu, Christopher Pal, Xue Liu.

1. Het Probleem: Offline Black-Box Optimalisatie (BBO)

Het paper richt zich op offline Black-Box Optimalisatie (BBO), een veelvoorkomend probleem in wetenschappelijke domeinen zoals robotica, DNA-synthese en materiaalkunde.

• Doel: Het vinden van een nieuw ontwerp $x^*$ dat een onbekende objectieve functie $f(x)$ maximaliseert (bijv. het vinden van een DNA-sequentie met de hoogste bindingsaffiniteit).
• Beperking: In de praktijk zijn experimenten duur en tijdrovend, waardoor er geen interactie mogelijk is met de objectieve functie tijdens het optimalisatieproces. Men is beperkt tot een offline dataset van gelabelde ontwerpen en hun scores.
• Uitdaging: De dataset is vaak klein ("small-data settings", bijv. $N=500$). Traditionele methoden (zoals surrogate-modellen) lijden onder onzekerheid bij data buiten de verdeling (out-of-distribution).
• Huidige LLM-aanpak: Recent werk gebruikt autoregressieve Large Language Models (LLMs) die taken als natuurlijke taal prompts formatteren. Deze modellen genereren echter tekst van links naar rechts (unidirectioneel), wat problematisch is voor ontwerpproblemen met sterke bidirectionele afhankelijkheden (waarbij een eenheid in een ontwerp wordt beïnvloed door zowel zijn voorafgaande als volgende elementen, zoals in DNA-sequenties).

2. Methodologie: DiBO Framework

De auteurs stellen DiBO voor, een framework dat Diffusion Language Models (Diffusion LLMs) aanpast voor offline BBO. Diffusion LLMs zijn bij uitstek geschikt vanwege hun inherente vermogen om bidirectionele context te modelleren.

Het framework bestaat uit drie hoofdfasen om de kloof tussen natuurlijke taal en heterogene BBO-data te overbruggen:

A. Domeinadaptatie (Domain Adaptation - DA)

Er bestaat een domeinkloof tussen de pre-training van LLMs (natuurlijke tekst) en de BBO-data (prompts, ontwerpen en labels).

• Unificatie: De auteurs bouwen een uniek corpus van prompt-response paren.
• Delimiters: Om de semantische rollen duidelijk te maken, worden speciale delimiter tokens toegevoegd aan de tokenizer:
  • |design-start| / |design-end| om ontwerpen te omhullen.
  • |label-start| / |label-end| om labels te omhullen.
• Verliesfunctie: Een gezamenlijke loss-functie wordt gebruikt om gemaskerde tokens in zowel de prompt als het antwoord te reconstrueren. Dit helpt het model om de structuur van de heterogene data te leren.

B. Twee-staps Post-Training

Om het generatievermogen van het model af te stemmen op ontwerpen met hoge scores, wordt een twee-staps post-training toegepast:

1. Supervised Fine-Tuning (SFT):

   • Het model wordt getraind op het unificatie corpus via masked-response prediction.
   • Het doel is om een inductieve bias te creëren richting gebieden in de ontwerpruimte met hoge labels. Het model leert om, gegeven een prompt met voorbeelden, een verbeterd ontwerp te genereren.

2. Versterkingsleren (Reinforcement Learning - RL):

   • Er wordt een RL-dataset geconstrueerd waarbij de beloning ($r$) wordt gedefinieerd als de verbetering in label: $r(q, o) = y(o) - y(q)$.
   • Efficiëntie: In plaats van iteratief denoising (wat duur is), gebruiken de auteurs een één-staps log-probabiliteit benadering om de RL-loss te berekenen.
   • Dit stadium integreert fijne-granulariteits feedback (rewards) om het model nog nauwkeuriger te sturen naar optimale ontwerpen.

3. Belangrijkste Bijdragen

1. Adaptatie van Diffusion LLMs: Het eerste werk dat diffusion LLMs toepast op offline BBO, gebruikmakend van hun bidirectionele modellering voor complexe ontwerpafhankelijkheden.
2. Unificatie en Delimiters: De introductie van een uniek prompt-response corpus met expliciete delimiter tokens om de domeinkloof tussen tekst en design-data te overbruggen.
3. Twee-staps Post-Training: Een innovatief framework dat Supervised Fine-Tuning combineert met Reinforcement Learning om generaties effectief te aligneren met hoge scores.
4. State-of-the-Art Resultaten: Het bereiken van nieuwe state-of-the-art resultaten op de Design-Bench benchmarks in kleine-data settings.

4. Resultaten

De methode werd geëvalueerd op vier taken uit de Design-Bench suite (twee discrete DNA-taken en twee continue robot-morfologie taken) met een beperkt aantal gelabelde samples ($N=500$).

• Prestaties: DiBO behaalde de hoogste gemiddelde rang en mediaan rang over alle vergelijkingen.
• Vergelijking: Het presteerde consistent beter dan de sterkste bestaande methoden (zoals autoregressieve LLMs zoals OPRO, VAE's, GAN's en andere diffusion-modellen) op drie van de vier benchmarks, met name een aanzienlijke verbetering (+0.108) op de TF Bind 8 taak.
• Ablatie Studies:
  • Het verwijderen van de delimiter tokens leidde tot significante prestatiedalingen, wat aantoont dat deze cruciaal zijn voor het begrijpen van de datastructuur.
  • Het verwijderen van een van de drie trainingsfasen (DA, SFT of RL) resulteerde in een duidelijke prestatiedaling, wat bevestigt dat alle fasen complementair en noodzakelijk zijn.
  • Het gebruik van gelijksoortige context (selectie van few-shot voorbeelden op basis van ontwerp-similariteit) bleek essentieel voor de prestaties.

5. Betekenis en Impact

Dit paper biedt een robuust en algemeen kader voor offline black-box optimalisatie. Door de unieke voordelen van diffusion modellen (bidirectionele context) te combineren met geavanceerde LLM-trainingstechnieken (SFT en RL), overwint de methode de beperkingen van traditionele autoregressieve modellen.

• Wetenschappelijke Impact: Het biedt een krachtig alternatief voor dure fysieke experimenten in domeinen zoals moleculair ontwerp en robotica.
• Technische Innovatie: Het demonstreert dat diffusion LLMs, vaak gezien als een alternatief voor autoregressieve modellen, superieur kunnen zijn in domeinen waar globale afhankelijkheden tussen elementen cruciaal zijn.
• Toekomst: Het framework opent de deur voor het toepassen van generatieve AI in data-schaarse wetenschappelijke optimalisatieproblemen.