Training Diffusion Language Models for Black-Box Optimization

Fuente: arXiv
Saved in:
Bibliographic Details
Main Authors: Sun, Zipeng, Chen, Can, Yuan, Ye, Wu, Haolun, Gu, Jiayao, Pal, Christopher, Liu, Xue
Format: Preprint
Published: 2026
Subjects:
Online Access:
Tags: Add Tag
No Tags, Be the first to tag this record!
_version_ 1866917548909395968
author Sun, Zipeng
Chen, Can
Yuan, Ye
Wu, Haolun
Gu, Jiayao
Pal, Christopher
Liu, Xue
author_facet Sun, Zipeng
Chen, Can
Yuan, Ye
Wu, Haolun
Gu, Jiayao
Pal, Christopher
Liu, Xue
contents We study offline black-box optimization (BBO), aiming to discover improved designs from an offline dataset of designs and labels, a problem common in robotics and DNA with limited labeled samples. While recent work applies autoregressive LLMs to BBO by formatting tasks as natural-language prompts, their left-to-right design generation struggles to capture the strong bidirectional dependencies inherent in design problems. To address this, we propose adapting diffusion LLMs to offline BBO to leverage their bidirectional modeling capabilities. However, a domain gap exists between the natural text pre-training of diffusion LLMs and the heterogeneous signals in BBO (prompts, designs, and labels). To bridge this gap, we construct a unified prompt--response corpus and introduce delimiter tokens to explicitly mark field boundaries for domain adaptation. We further propose a two-stage post-training framework to align the diffusion LLM generation with high-label designs. The first stage performs supervised fine-tuning on the unified dataset via masked-response prediction, and the second stage adopts reinforcement learning with rewards defined by label improvements. Our method achieves state-of-the-art results on Design-Bench under small-data settings with highly efficient training, requiring only $1.5$ H100 GPU hours for discrete tasks. Code for our work is available here: https://github.com/zpointS/DiBO.
format Preprint
id arxiv_https___arxiv_org_abs_2603_17919
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Training Diffusion Language Models for Black-Box Optimization
Sun, Zipeng
Chen, Can
Yuan, Ye
Wu, Haolun
Gu, Jiayao
Pal, Christopher
Liu, Xue
Computational Engineering, Finance, and Science
We study offline black-box optimization (BBO), aiming to discover improved designs from an offline dataset of designs and labels, a problem common in robotics and DNA with limited labeled samples. While recent work applies autoregressive LLMs to BBO by formatting tasks as natural-language prompts, their left-to-right design generation struggles to capture the strong bidirectional dependencies inherent in design problems. To address this, we propose adapting diffusion LLMs to offline BBO to leverage their bidirectional modeling capabilities. However, a domain gap exists between the natural text pre-training of diffusion LLMs and the heterogeneous signals in BBO (prompts, designs, and labels). To bridge this gap, we construct a unified prompt--response corpus and introduce delimiter tokens to explicitly mark field boundaries for domain adaptation. We further propose a two-stage post-training framework to align the diffusion LLM generation with high-label designs. The first stage performs supervised fine-tuning on the unified dataset via masked-response prediction, and the second stage adopts reinforcement learning with rewards defined by label improvements. Our method achieves state-of-the-art results on Design-Bench under small-data settings with highly efficient training, requiring only $1.5$ H100 GPU hours for discrete tasks. Code for our work is available here: https://github.com/zpointS/DiBO.
title Training Diffusion Language Models for Black-Box Optimization
topic Computational Engineering, Finance, and Science
url https://arxiv.org/abs/2603.17919