PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Large Language Models for Black-Box Optimization

Ye Yuan, Can, Chen, Zipeng Sun, Dinghuai Zhang, Christopher Pal, Xue Liu

arXiv:2601.14446v1cs.CEcs.AI

TL;DR

Offline BBO seeks high-performing designs from limited labeled data, while autoregressive LLMs can struggle with bidirectional dependencies. The paper conditions diffusion LLMs on task and dataset prompts and guides denoising with tree search; dLLM achieves state-of-the-art results in few-shot design-bench settings.

  • Problem

    Offline BBO must find high-performing designs from an offline dataset, often with few labels, while autoregressive LLMs struggle to model bidirectional dependencies in complex designs.

    Source: §1 · pp. 1–2

  • Method

    dLLM conditions diffusion LLM denoising on task descriptions and offline data, then uses masked diffusion tree search with GP expected improvement to balance exploration and exploitation.

    Source: §1 · pp. 2–3

  • Results

    dLLM achieves state-of-the-art results on design-bench in the few-shot setting.

    Source: §6 · p. 10

  • Takeaways & Limitations

    The results indicate the potential of diffusion LLMs as general-purpose optimizers.

    Source: §6 · p. 10

    Extending classical UCT guarantees to this setting would require strong additional assumptions and is beyond the work’s scope.

    Source: §3.2 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper uses diffusion LLMs for offline BBO, where few labeled examples and bidirectional design dependencies challenge existing approaches. It combines in-context denoising with masked diffusion tree search to generate and guide improved candidates.

1 INTRODUCTION

  • Traditional proxy and conditional generative methods are task-specific and data-limited, and may face out-of-distribution generalization issues.The paper presents pre-trained LLMs’ in-context learning as an alternative to these approaches.Source: §1 · p. 1
  • Autoregressive LLMs prompted with task descriptions and offline data struggle with designs whose components depend on both prefixes and suffixes.The paper gives DNA units as an example of such bidirectional dependencies.Source: §1 · pp. 1–2
  • Diffusion LLMs iteratively denoise masked designs conditioned on natural-language task descriptions and offline datasets.The prompt also instructs the model to generate improved designs.Source: §1 · p. 2
  • Masked diffusion tree search treats denoising as step-wise MCTS, using masked designs as nodes and denoising steps as actions.UCT selects nodes, GP expected improvement evaluates completed designs, and rewards are backpropagated to balance exploration and exploitation.Source: §1 · p. 3

2 PRELIMINARIES

The preliminaries define offline BBO and its few-shot setting, then describe diffusion LLMs’ masking and denoising process. They also motivate Gaussian Processes for limited-data optimization and explain how diffusion models use bidirectional context and iterative sampling.

2.1 OFFLINE BLACK-BOX OPTIMIZATION

  • Offline BBO seeks a design maximizing an unknown objective using an offline dataset of designs and scores; the paper focuses on few-shot cases with small N, such as N = 10.A Gaussian Process is described as often more effective than a learned DNN in this setting because data are limited and it provides principled uncertainty estimates.Source: §2.1 · p. 3
  • The method adopts a Gaussian Process with an RBF kernel to model new designs from offline design and label data.The passage introduces its predictive posterior distribution.Source: §2.1 · p. 3

2.2 DIFFUSION LARGE LANGUAGE MODELS

  • Compared with autoregressive models, diffusion LLMs use bidirectional context and iterative sampling that scales at test time.Source: §2.2 · p. 3
  • Diffusion LLMs progressively mask tokens in a forward process, then recover masked tokens as t decreases from 1 to 0.At intermediate t, each token is independently masked with probability t.Source: §2.2 · pp. 3–4
  • A mask predictor takes a partially masked sequence and predicts all masked tokens simultaneously, with training loss computed only at masked positions.The passage specifies that the loss uses a uniformly sampled timestep and the forward masking process.Source: §2.2 · p. 4

3 METHOD

The method combines natural-language in-context denoising with masked diffusion tree search to generate and guide candidate designs. The search evaluates denoised candidates with Gaussian Process expected improvement and returns high-scoring designs.

3.1 IN-CONTEXT DENOISING

  • The diffusion LLM receives a prompt containing the task description, few-shot offline design–label pairs, and an instruction to generate improved candidates.At each denoising step, it completes masked tokens, remasks the least-confident tokens at ratio s/t, and restricts outputs to valid task-specific symbols.Source: §3.1 · p. 4

3 METHOD

  • Algorithm 1 combines in-context denoising and masked diffusion tree search, then returns fully denoised designs with the highest expected improvement scores among explored candidates.The search is initialized with an RBF-kernel Gaussian Process fit to the offline dataset.Source: §3 · p. 4, §3.1 · p. 5, §3.2 · p. 6

3.2 MASKED DIFFUSION TREE SEARCH

  • Masked diffusion tree search treats partially masked designs as nodes and denoising steps as actions, using selection, expansion, evaluation, and backpropagation to balance exploration and exploitation.UCT selection favors promising value estimates while its exploration term favors less-visited nodes; diffusion completions and remasking generate diverse children.Source: §3.2 · p. 5
  • Each child is evaluated through multiple denoising runs, with its reward set to average Gaussian Process expected improvement and propagated to ancestor nodes.The expected-improvement calculation uses the GP predictive mean and standard deviation and the best score in the offline dataset as its reference.Source: §3.2 · pp. 5–6
  • The tree search is presented as a practical UCT-based heuristic; theoretical analysis for diffusion-generated expansions and GP-based expected-improvement rewards is left beyond scope.The paper instead reports ablations and sensitivity studies showing effectiveness and robustness across tasks.Source: §3.2 · p. 6

4 EXPERIMENTS

Across four few-shot design tasks, dLLM ranks first on all four and leads 15 baselines overall; experiments also test its components and sensitivity to search and data settings.

4.1 BENCHMARKS

  • Experiments use 10 uniformly sampled examples for two continuous morphology tasks and two discrete DNA-binding tasks.The tasks are Ant, D’Kitty, TF8, and TF10; each method is evaluated on 128 candidates using maximum normalized ground-truth score, with mean and median ranks also reported.Source: §4.1 · p. 6

4.2 BASELINES

  • Baselines span proxy-based and generative methods, alongside CMA-ES and MCTS-transfer.The comparison includes gradient-guided proxies, conditional generators, autoregressive methods, and diffusion models.Source: §4.2 · pp. 6–7

4.3 IMPLEMENTATION DETAILS

  • The implementation uses LLaDA-8B-Instruct, an RBF-kernel GP, depth 4, branching factor 5, 16 search steps, and 5 fully unmasked designs.A full search episode takes about 30 minutes on Ant and 40 minutes on TF8.Source: §4.3 · p. 7

4.4 RESULTS AND ANALYSIS

  • dLLM achieves the highest average and median rank among 15 baselines and ranks first on all four tasks.The tasks span continuous and discrete design spaces.Source: §4.4 · p. 7, Table 1 · p. 7

4.5 ABLATION STUDIES

  • Removing masked diffusion tree search substantially decreases performance across all tasks, while removing task descriptions also lowers performance across all tasks.Replacing the diffusion LLM with continuous diffusion causes a clear drop, especially on discrete sequence tasks; classifier-free guidance gives similar results but doubles inference cost.Source: §4.5 · p. 8

4.6 HYPERPARAMETER SENSITIVITY

  • Performance improves steadily with greater tree depth and generally with larger branching factors, while gains are ultimately limited by GP epistemic uncertainty.Depth varies over {1, 2, 4, 6, 8}; branching factor varies over {1, 3, 5, 7, 9}.Source: §4.6 · p. 8
  • Performance drops sharply with only 2 offline examples, then improves and stabilizes with more data; MMaDA and LLaDA backbone results are close on Ant and TF8.The tested dataset sizes are {2, 5, 10, 20}; the paper also reports robustness to prompt variations.Source: §4.6 · pp. 8–9

5 RELATED WORK

Prior work uses LLMs both to predict design properties and to generate or optimize designs directly. This work differs from earlier autoregressive approaches by applying diffusion LLMs, while related MCTS methods span black-box optimization, diffusion planning, and molecular optimization.

5 RELATED WORK

  • LLMs have been used to predict design properties and directly generate promising designs, with bidirectional encoders used for regression and autoregressive models explored for prediction.Examples include embedding-based regression in high-dimensional design spaces (Nguyen et al., 2024; Tang et al., 2024) and molecular property prediction (Nguyen & Grover, 2024).Source: §5 · p. 9
  • Direct LLM optimization includes prompting labeled designs for generation and selecting parents for evolutionary crossover and mutation (Zhang et al., 2023; Liu et al., 2024).Program-synthesis methods FunSearch (Veliˇckovi´c et al., 2024) and AlphaEvolve (Novikov et al., 2025) use LLM-guided evolutionary strategies.Source: §5 · p. 9
  • This work follows direct LLM optimization but investigates diffusion LLMs rather than the predominantly used autoregressive models (Zhang et al., 2023).The paper cites diffusion LLMs' bidirectional modeling and iterative sampling as suited to black-box design problems.Source: §5 · p. 9
  • MCTS has guided black-box optimization through space partitioning, transferred search-space refinement, and reward-guided protein design (Wang et al., 2020; Wang et al., 2024; Liu et al., 2025).These approaches focus on concentrating sampling in promising regions, adapting search spaces using related tasks, or discovering proteins that fold to a target backbone.Source: §5 · p. 9
  • Diffusion MCTS methods target planning, image generation, and language modeling, while Tang et al. (2025) apply it to molecular optimization using SMILES.The paper distinguishes its approach from the molecular method because that method lacks in-context learning and is limited to small-molecule representations.Source: §5 · p. 9

6 CONCLUSION

The paper applies diffusion LLMs to offline BBO through in-context denoising and masked diffusion tree search. It reports state-of-the-art few-shot results on design-bench and presents diffusion LLMs as potential general-purpose optimizers.

6 CONCLUSION

  • dLLM achieves state-of-the-art results on design-bench in the few-shot setting.The authors describe these results as demonstrating the potential of diffusion LLMs as general-purpose optimizers.Source: §6 · p. 10
  • In-context denoising conditions diffusion LLMs on task descriptions and offline datasets to refine masked designs.The inputs are used to generate improved designs through iterative denoising.Source: §6 · p. 10
  • Masked diffusion tree search formulates denoising as Monte Carlo Tree Search to balance exploration and exploitation.The conclusion identifies this and in-context denoising as the method's two key modules.Source: §6 · p. 10
  • The authors frame dLLM as opening directions for data-efficient design generation in scientific and engineering domains.Source: §6 · p. 10

ETHICS STATEMENT

The authors describe dLLM as potentially useful for scientific discovery in biomedicine and robotics, while warning that generated designs could be misused. They call for safeguards, oversight, and regulatory frameworks for responsible use.

ETHICS STATEMENT

  • The authors say dLLM could support scientific discovery in biomedicine and robotics through optimization of complex design spaces.Source: ETHICS STATEMENT · p. 11
  • They warn that design generation and refinement could be misapplied to harmful contexts, including biological weapons development.Source: ETHICS STATEMENT · p. 11
  • The authors call for safeguards, oversight, and regulatory frameworks to promote responsible use for beneficial purposes.Source: ETHICS STATEMENT · p. 11

A.1 LLM USAGE

The authors report using LLMs only to polish the paper's writing. They specify that LLM assistance did not extend to method design, implementation, or results analysis.

A.1 LLM USAGE

  • LLMs were used to refine wording and check grammar for clarity and readability.Source: §A.1 · p. 16
  • The authors state that LLMs did not assist with method design or implementation.Source: §A.1 · p. 16
  • The authors state that LLMs were not used for experimental results analysis.Source: §A.1 · p. 16

A.2 IN-CONTEXT PROMPT

The in-context prompts pair task-specific goals and offline examples with instructions for proposing improved DNA sequences or robot morphologies. They specify the relevant design representation and, where shown, output constraints.

In-Context Prompt for Ant

  • The Ant prompt targets a quadruped morphology that runs fast, representing each design as 60 continuous parameters grouped into four legs with 15 parameters each.It asks for a new design with float-valued features rounded to two decimal places.Source: §A.2 · p. 16

In-Context Prompt for D’Kitty

  • The D’Kitty prompt targets locomotion toward a fixed location and represents designs with 56 continuous morphology parameters grouped into four legs with 14 parameters each.Source: §A.2 · p. 17

In-Context Prompt for TF8

  • The TF8 prompt requests a length-8 DNA sequence maximizing binding affinity with SIX6 REF R1, supplies existing sequences and labels, and requires a different sequence.The requested sequence uses A, C, G, and T and must differ from existing sequences in at least one position.Source: §A.2 · p. 17

In-Context Prompt for TF10

  • The TF10 prompt requests a length-10 DNA sequence maximizing binding affinity with Pho4 and provides existing sequences with their binding affinities.Source: §A.2 · p. 18

A.3 MEDIAN RESULTS

At the 50th percentile, dLLM has the best overall mean and median rank across 16 methods, placing first on D’Kitty, TF8, and TF10 and third on Ant. Its normalized scores remain stable on Ant and TF8 under five GPT-5-generated task-description templates.

A.3 MEDIAN RESULTS

  • dLLM achieves the best overall mean and median rank across 16 methods, ranking first on D’Kitty, TF8, and TF10 and third on Ant.Table 3 reports 50th-percentile normalized scores across four tasks.Source: §A.3 · p. 18

A.4 DIFFERENT TASK DESCRIPTIONS

  • Normalized scores remain stable across five GPT-5-generated task-description templates: 0.642–0.654 on Ant and 0.869–0.876 on TF8.Source: Figure 6 · p. 18, §A.4 · p. 19
Loading 2601.14446v1…