Diffusion Large Language Models for Black-Box Optimization

TL;DR

This work introduces diffusion large language models (dLLMs) for offline black-box optimization, addressing bidirectional dependencies and limited labeled data by combining in-context denoising with a masked diffusion tree search (MDTS). The method prompts a diffusion LLM with task descriptions and offline datasets to iteratively denoise masked designs, while MDTS guides exploration via a step-wise Monte Carlo Tree Search using GP-based Expected Improvement as the reward. Empirical results on Design-Bench show state-of-the-art few-shot performance across continuous and discrete design tasks, with ablations confirming the necessity of MDTS and careful prompting. The approach demonstrates the practical potential of diffusion LLMs as general-purpose, data-efficient optimizers for complex scientific and engineering design spaces.

Abstract

Offline black-box optimization (BBO) aims to find optimal designs based solely on an offline dataset of designs and their labels. Such scenarios frequently arise in domains like DNA sequence design and robotics, where only a few labeled data points are available. Traditional methods typically rely on task-specific proxy or generative models, overlooking the in-context learning capabilities of pre-trained large language models (LLMs). Recent efforts have adapted autoregressive LLMs to BBO by framing task descriptions and offline datasets as natural language prompts, enabling direct design generation. However, these designs often contain bidirectional dependencies, which left-to-right models struggle to capture. In this paper, we explore diffusion LLMs for BBO, leveraging their bidirectional modeling and iterative refinement capabilities. This motivates our in-context denoising module: we condition the diffusion LLM on the task description and the offline dataset, both formatted in natural language, and prompt it to denoise masked designs into improved candidates. To guide the generation toward high-performing designs, we introduce masked diffusion tree search, which casts the denoising process as a step-wise Monte Carlo Tree Search that dynamically balances exploration and exploitation. Each node represents a partially masked design, each denoising step is an action, and candidates are evaluated via expected improvement under a Gaussian Process trained on the offline dataset. Our method, dLLM, achieves state-of-the-art results in few-shot settings on design-bench.

Figures (6)

• Figure 1: In-context denoising.
• Figure 2: Masked diffusion tree search.
• Figure 3: Sensitivity to different tree depths.
• Figure 4: Sensitivity to different branching factors.
• Figure 5: Sensitivity to different offline dataset sizes.