[논문 리뷰] Training Diffusion Language Models for Black-Box Optimization

Zipeng Sun, Can Chen외 5명|arXiv (Cornell University)|2026. 03. 18.

Machine Learning in Materials Science인용 수 0

한 줄 요약

이 논문은 오프라인 블랙박스 최적화를 위한 확산 LLM 프레임워크 DiBO를 소개합니다. 이 프레임워크는 구분 토큰(delimiter tokens)으로 도메인 적응을 수행하고 두 단계 후학습(SFT then RL)을 통해 이질적인 프롬프트, 디자인, 라벨로부터 고라벨 디자인을 생성하며, 소데이터 설정에서 Design-Bench에서 최첨단 성과를 달성합니다.

ABSTRACT

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, DNA, and materials science 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 small-data settings.

연구 동기 및 목표

• 자연어 사전 학습과 오프라인 디자인 데이터 간의 격차를 줄여 BBO를 위한 양방향 모델링을 가능하게 한다.
• 확산 LLM을 활용하여 디자인 공간의 양방향 의존성을 포착한다.
• 도메인 적응 및 후학습 파이프라인을 개발하여 확산 생성이 고라벨 디자인과 정렬되도록 한다.
• Discrete 및 Continuous 작업에서 소데이터 설정의 Design-Bench에서 강력한 성능을 보여준다.

제안 방법

• 디자인과 라벨을 명시적으로 구분 토큰으로 표시하는统一 프롬프트–응답 코퍼스를 구성한다.
• 프롬프트와 응답에서 마스킹된 토큰을 함께 예측하는 방식으로 도메인 적응을 수행한다(공동 DA 손실).
• 统一 코퍼스에서의 지도 학습 미세조정(마스킹된 응답 예측)으로 1단계 후학습을 수행하고, 고라벨 개선으로 정의된 보상을 사용하는 강화학습으로 이어진다.
• 효율성을 위해 한 단계 로그-확률 RL 근사치를 사용하고 보상 r(q,o)=y(o)−y(q)와 보상의 표준편차로 정규화한다.
• 128개의 후보를 각 태스크에 대해 사용하고 제약: Design-Bench 태스크(TF8, TF10, Ant Morphology, D’Kitty Morphology)에서 안정성과 강건성을 분석한다.

Figure 1 : Overview of the DiBO framework. (a) Unified Prompt–Response Corpus: Heterogeneous BBO signals (natural-language prompts, offline designs and their associated labels) are unified using explicit delimiter tokens. (b) Domain Adaptation (DA): The diffusion LLM is domain-adapted via joint mask

실험 결과

연구 질문

• RQ1 오프라인 BBO에서 확산 LLM이 자기회귀 LLM보다 양방향 의존성을 더 잘 포착할 수 있는가?
• RQ2 Delimiter를 강화한 프롬프트를 통한 도메인 적응이 이질적 오프라인 데이터 학습을 향상시키는가?
• RQ3 세 단계 훈련 파이프라인(DA, SFT, RL)이 소데이터 체제에서 고라벨 디자인에 대해 더 나은 정합성을 내는가?
• RQ4 DiBO가 이산 및 연속 디자인 작업에서 다양한 베이스라인과 비교하여 어떤 성과를 보이는가?
• RQ5 프롬프트 유사성, 구분 토큰, 훈련 단계가 성능에 미치는 영향은 무엇인가?

주요 결과

Method	Ant Morphology	D’Kitty Morphology	TF Bind 8	TF Bind 10	Mean Score ↑	Rank Mean ↓	Rank Median ↓
D (ours) DiBO (full)	0.944±0.016	0.923±0.002	0.965±0.038	0.755±0.012	0.897±0.017	2.5	1.0
OPRO	0.517±0.039	0.856±0.046	0.758±0.017	0.500±0.013	0.657±0.028	13.5	14.5
GTG	0.603±0.039	0.917±0.023	0.762±0.016	0.730±0.026	0.753±0.026	7.25	9.0
DDOM	0.590±0.026	0.929±0.037	0.739±0.016	0.497±0.002	0.689±0.020	10.75	11.5
MIN	0.570±0.003	0.886±0.017	0.764±0.008	0.517±0.030	0.684±0.015	11.75	12.5
ExPT	0.929±0.049	0.950±0.041	0.810±0.044	0.703±0.022	0.848±0.039	3.0	3.0
OPRO (alternate)	0.517±0.039	0.856±0.046	0.758±0.017	0.500±0.013	0.657±0.028	13.5	14.5
BONET	0.632±0.042	0.920±0.040	0.776±0.007	0.492±0.043	0.705±0.033	9.5	7.5
CMA-ES	0.592±0.010	0.711±0.045	0.784±0.029	0.658±0.031	0.686±0.029	9.75	8.5
UniSO-T	0.636±0.045	0.939±0.007	0.836±0.027	0.522±0.017	0.733±0.024	6.0	5.0
Grad-mean	0.644±0.039	0.907±0.016	0.666±0.011	0.695±0.027	0.728±0.023	9.5	8.5
Grad-EI	0.626±0.002	0.901±0.045	0.673±0.012	0.689±0.013	0.722±0.018	10.5	10.5

• DiBO는 소데이터 설정에서 여러 작업에 걸쳐 최첨단 성능을 달성한다.
• 구분 토큰은 plain-text 경계에 비해 모든 학습 단계에서 성능을 크게 향상시킨다.
• 세 단계 파이프라인(DA + SFT + RL)은 두 단계 변형보다 우수한 정합성을 보이며, RL이 미세 보상 최적화를 제공한다.
• 유사도 조건부 맥락 구성(디자인 유사성으로 선택된 프롬프트 예시)이 무작위 맥락에 비해 성능을 크게 높인다.
• DiBO는 정방향 대리자 기반 방법 및 다수의 확산 계열 베이스라인을 꾸준히 능가하며, 특히 Ant Morphology 및 TF Bind 태스크에서 두드러진 성과를 보인다.
• 하이퍼파라미터(RL 학습률, 프롬프트 템플릿 변화)에도 로버스트한 편이다.

Figure 2 : Hyperparameter sensitivity on TF Bind 8 and Ant Morphology at the RL stage: (a) the number of few-shot examples in the prompt (context length), (b) size of the offline dataset, (c) learning rate used in RL stage, and (d) variations of prompt templates. Results are reported as relative per