Training Diffusion Language Models for Black-Box Optimization

1. Problem Definition: Offline Black-Box Optimization (BBO)

The paper addresses Offline Black-Box Optimization, a critical challenge in fields like robotics, DNA synthesis, and materials science.

• Goal: Discover optimal designs ($x^*$) that maximize an unknown objective function $f(x)$ (e.g., binding affinity, robot speed).
• Constraint: Direct interaction with the objective function (experimentation) is prohibitively expensive or time-consuming. Therefore, optimization must rely solely on a static, offline dataset of design-label pairs $\mathcal{D} = \{(x_i, y_i)\}$.
• Data Scarcity: Real-world applications often suffer from small-data regimes (e.g., $N=500$ labeled samples), making standard iterative online optimization impossible.
• Limitations of Existing Methods:
  • Surrogate Models: Traditional approaches (GPs, DNNs) struggle with out-of-distribution (OOD) generalization.
  • Autoregressive (AR) LLMs: Recent methods use AR LLMs (like GPT) to generate designs via natural language prompts. However, AR models are inherently unidirectional (left-to-right), failing to capture the strong bidirectional dependencies common in design problems (e.g., a DNA base is constrained by both its prefix and suffix).

2. Methodology: DiBO Framework

The authors propose DiBO, a framework that adapts Diffusion Language Models (Diffusion LLMs) for offline BBO. The core innovation is leveraging the bidirectional modeling capabilities of diffusion models while bridging the domain gap between natural text and heterogeneous design data.

A. Domain Adaptation via Unified Corpus

Diffusion LLMs are pre-trained on natural text, but BBO data involves heterogeneous signals: natural language prompts, discrete/continuous designs, and scalar labels.

• Unified Prompt-Response Corpus: The authors construct a dataset where inputs are unified sequences containing:
  1. Task Description: Natural language instructions.
  2. Offline Examples: Few-shot design-label pairs.
  3. Delimiters: Special tokens (|design-start|, |design-end|, |label-start|, |label-end|) are introduced to explicitly mark semantic boundaries between text, designs, and labels.
• Joint Masked Prediction: During the Domain Adaptation (DA) stage, the model is trained to reconstruct masked tokens in both the prompt and the response simultaneously. This forces the model to learn the semantic roles of the delimiters and the relationship between instructions and design constraints.

B. Two-Stage Post-Training Framework

To align the diffusion model specifically with high-performing designs, a two-stage fine-tuning process is employed:

1. Supervised Fine-Tuning (SFT):

   • Objective: Masked-response prediction. The model learns to generate a response design $o$ given a prompt $q$, where the response's label is strictly higher than the best label in the prompt.
   • Goal: Instills an inductive bias toward high-label regions of the design space and initializes the model to follow the prompt context.

2. Reinforcement Learning (RL):

   • Reward Definition: The reward $r(q, o)$ is defined as the improvement in label value: $r = y(o) - y(q)$. Unlike SFT, this stage allows for negative rewards (if the generated design is worse), providing fine-grained feedback.
   • Efficient Optimization: Instead of iterative denoising for every RL step (which is computationally expensive), the authors use a one-step log-probability approximation. They unmask the response tokens in a single step from a fully masked state to estimate the probability $p_\theta(o|q)$.
   • Loss Function: The model is optimized to maximize the expected reward, integrating fine-grained feedback directly into the diffusion generation process.

3. Key Contributions

1. Adaptation of Diffusion LLMs to BBO: First work to apply diffusion LLMs to offline BBO, exploiting their bidirectional modeling to capture complex structural dependencies that AR models miss.
2. Semantic Delimiter Strategy: Introduction of explicit delimiter tokens to bridge the domain gap between natural language pre-training and the heterogeneous signals of BBO (prompts, designs, labels).
3. Two-Stage Post-Training Pipeline: A novel framework combining SFT (for coarse alignment) and RL (for fine-grained reward maximization) to effectively navigate the optimization landscape in small-data settings.
4. State-of-the-Art Performance: Demonstrated superior results on the Design-Bench suite, particularly in small-data regimes.

4. Experimental Results

The method was evaluated on four Design-Bench tasks: two discrete (TF Bind 8, TF Bind 10) and two continuous (Ant Morphology, D'Kitty Morphology).

• Performance: DiBO achieved State-of-the-Art (SOTA) results, ranking 1st in both mean and median rank across all tasks.
  • TF Bind 8: Outperformed the best baseline by a significant margin (>0.108 normalized score).
  • Ant Morphology: Achieved a normalized score of 0.944, surpassing the next best method (ExPT at 0.929).
• Comparison:
  • vs. Forward Methods (Surrogates): DiBO is more robust to distributional shifts and does not rely on potentially inaccurate surrogate gradients.
  • vs. Inverse Methods (VAE, GAN, AR-LLM): DiBO outperforms AR-LLMs (like OPRO) by capturing bidirectional dependencies and outperforms standard diffusion models by effectively integrating textual instructions.
• Ablation Studies:
  • Delimiters: Removing delimiter tokens caused significant performance drops, especially in the Domain Adaptation stage (e.g., a 0.182 drop on Ant Morphology).
  • Training Stages: Removing any stage (DA, SFT, or RL) led to performance degradation. Notably, RL-only failed to produce valid outputs on continuous tasks, highlighting the necessity of the SFT initialization.
  • Context Construction: Selecting few-shot examples based on design similarity (rather than random selection) was critical for performance.

5. Significance and Impact

• Scientific Discovery: The framework offers a cost-effective alternative to physical experimentation in molecular design and robotics, accelerating the discovery of high-performance materials and sequences.
• Generalization: By treating BBO as a language modeling problem with bidirectional context, DiBO provides a general-purpose paradigm for high-dimensional optimization where data is scarce.
• Safety & Ethics: The authors acknowledge the dual-use nature of generative models (e.g., potential for harmful biochemical designs) and emphasize that the tool is intended to augment, not replace, expert judgment, requiring responsible oversight.

In summary, DiBO successfully bridges the gap between natural language processing and scientific optimization, demonstrating that Diffusion LLMs, when properly adapted with semantic delimiters and a reward-aligned training pipeline, are superior to autoregressive models for complex, bidirectional design problems.