Technical Summary: Diffusion Large Language Models for Black-Box Optimization

Problem Statement

Offline Black-Box Optimization (BBO) seeks to identify optimal designs ($x^*$) that maximize an unknown objective function $f(\cdot)$ using only a limited offline dataset of design-label pairs. This scenario is prevalent in domains such as DNA sequence design and robotics, where querying properties is expensive and labeled data is scarce (few-shot regime).

Traditional approaches typically fall into two categories:

1. Proxy-based methods: Training a surrogate model (e.g., Gaussian Processes or Deep Neural Networks) to approximate $f(\cdot)$ and using its gradients to guide optimization.
2. Generative methods: Training conditional generative models to directly produce high-scoring designs.

However, these methods often suffer from out-of-distribution generalization issues in data-limited settings. Furthermore, recent adaptations of autoregressive Large Language Models (LLMs) to BBO, which frame tasks as natural language prompts, face a critical limitation: they generate sequences left-to-right. This unidirectional approach struggles to capture bidirectional dependencies inherent in many design tasks (e.g., a DNA unit's function depends on both its prefix and suffix).

Methodology: dLLM

The authors propose dLLM, a framework that leverages Diffusion Large Language Models for offline BBO. The method combines two core modules: In-Context Denoising and Masked Diffusion Tree Search.

1. In-Context Denoising

Instead of fine-tuning, the method utilizes the in-context learning capabilities of pre-trained diffusion LLMs.

• Prompt Construction: The input to the diffusion LLM consists of a structured natural language prompt containing:
  • Task Description: Defines the design format (e.g., DNA alphabet), the objective, and constraints.
  • Offline Dataset: A few-shot set of design-label pairs (e.g., DNA sequences and binding affinities).
  • Instruction: A directive to generate improved candidates.
• Denoising Process: The model starts with a fully masked design ($x_{t=1}$). It iteratively predicts masked tokens to produce a complete candidate ($\hat{x}_0$). The least confident tokens are re-masked to form a new partially masked state ($x_s$) for the next step, effectively refining the design through iterative denoising.

2. Masked Diffusion Tree Search (MDTS)

To guide the generation toward high-performing regions and balance exploration with exploitation, the denoising process is cast as a step-wise Monte Carlo Tree Search (MCTS).

• Tree Structure: Each node represents a partially masked design ($x_t$).
• Selection: The algorithm traverses the tree from the root (fully masked) using the Upper Confidence Bound (UCT) score. This balances exploiting high-value nodes and exploring less-visited branches. The UCT score incorporates the node's value estimate and the model's likelihood of the unmasking step.
• Expansion: Selected nodes are expanded by applying the diffusion LLM to generate diverse completions, which are then re-masked to create child nodes.
• Evaluation: Child nodes are unrolled into fully denoised designs. Their rewards are estimated using Expected Improvement (EI) derived from a Gaussian Process (GP) trained on the offline dataset. The GP provides principled uncertainty estimates crucial for few-shot settings.
• Backpropagation: The reward signal is propagated up the tree to update the value estimates and visit counts of ancestor nodes, informing future selection decisions.

Key Contributions

1. First Application of Diffusion LLMs to BBO: The paper introduces the first framework utilizing diffusion LLMs for offline black-box optimization, capitalizing on their bidirectional modeling and iterative refinement capabilities to handle complex design dependencies.
2. In-Context Denoising Module: A novel conditioning mechanism that allows diffusion LLMs to refine masked designs based on natural language task descriptions and few-shot offline datasets without parameter updates.
3. Masked Diffusion Tree Search: A search algorithm that treats diffusion denoising as a sequential decision-making process, dynamically balancing exploration and exploitation via MCTS and GP-based Expected Improvement.

Experimental Results

The authors evaluated dLLM on the Design-Bench benchmark, covering both continuous (Ant Morphology, D'Kitty Morphology) and discrete (TF Bind 8, TF Bind 10) tasks in a few-shot setting (10 labeled examples).

• Performance: dLLM achieved state-of-the-art results across all four tasks, ranking first in both mean and median ranks among 15 baselines. It demonstrated consistent gains in both continuous and discrete design spaces.
• Comparison:
  • Outperformed pure gradient-based proxy methods (e.g., Grad-mean, Grad-EI), indicating that generative models coupled with search are superior for exploration.
  • Surpassed other generative baselines, including autoregressive LLMs (OPRO) and other diffusion models (DDOM, GTG). The authors attribute this to the bidirectional modeling of diffusion LLMs and the broad pre-trained knowledge they leverage.
  • Outperformed the strong autoregressive baseline UniSO-T, suggesting that bidirectional context is critical for tasks with complex dependencies.
• Ablation Studies:
  • Removing the tree search (w/o MDTS) caused a substantial performance drop, highlighting the necessity of guided exploration.
  • Replacing the diffusion LLM with a vanilla diffusion model trained only on the few-shot data led to significant degradation, confirming the importance of pre-trained knowledge.
  • Removing the explicit task description in the prompt reduced performance, validating the role of semantic framing.

Significance and Claims

The paper claims that dLLM demonstrates the potential of diffusion LLMs as general-purpose optimizers for data-efficient design generation. By addressing the limitations of autoregressive models in capturing bidirectional dependencies and integrating them with a robust search strategy, the method offers a versatile alternative for scientific and engineering domains where labeled data is scarce. The authors position this work as opening new directions for leveraging pre-trained foundation models in offline black-box optimization.