Training Diffusion Language Models for Black-Box Optimization

Motivation

• Known: 기존 연구는 offline BBO를 자연어 프롬프트로 포맷하여 autoregressive LLM으로 후보 design을 생성하거나, surrogate 모델을 학습하여 라벨이 부족한 상황에서 최적화를 수행해왔다.
• Gap: autoregressive LLM의 좌에서 우로의 unidirectional 생성 방식은 DNA 서열 등 design 문제에 내재된 강한 bidirectional dependency를 포착하지 못하며, diffusion LLM을 적용하려 해도 natural text pre-training과 BBO의 이질적 신호(prompt, design, label) 사이에 domain gap이 존재한다.
• Why: 라벨이 부족한 로보틱스, DNA, 재료과학 등 실세계 문제에서 offline 데이터만으로 우수한 design을 효율적으로 발굴할 수 있는 방법은 실험 비용을 크게 절감할 수 있어 실질적 파급력이 크다.
• Approach: diffusion LLM의 bidirectional modeling 능력을 활용하기 위해 unified prompt–response corpus와 delimiter token을 도입해 domain gap을 해소하고, SFT와 RL로 구성된 2단계 post-training 프레임워크로 high-label design 생성에 정렬시킨다.

Achievement

Figure 1. Overview of the DiBO framework. (a) Unified Prompt–Response Corpus: Heterogeneous BBO signals (natural-languag

1. State-of-the-art 성능: Design-Bench의 small-data(N=500) 설정에서 기존 방법들 대비 우수한 성능을 달성했다.
2. 효율적인 학습: discrete task 기준 1.5 H100 GPU 시간만으로 학습이 가능해 매우 높은 학습 효율성을 보였다.
3. Domain gap 해소: delimiter token 기반 unified prompt–response corpus를 통해 diffusion LLM을 BBO의 이질적 신호에 성공적으로 적응시켰다.

How

Figure 1. Overview of the DiBO framework. (a) Unified Prompt–Response Corpus: Heterogeneous BBO signals (natural-languag

• design과 label을 각각 |design-start|/|design-end|, |label-start|/|label-end| delimiter token으로 감싸 unified prompt–response corpus를 구축
• diffusion LLM이 prompt와 response 양쪽에서 masked token을 joint하게 예측하도록 domain adaptation 수행
• 1단계(SFT): masked-response prediction을 통해 high-label 영역으로의 inductive bias를 주입
• 2단계(RL): prompt 대비 response의 label 개선치를 reward로 정의하고, 효율적인 one-step log-probability approximation으로 RL loss를 계산해 세밀한 reward feedback을 통합

Originality

• autoregressive LLM 대신 diffusion LLM의 bidirectional modeling을 BBO에 최초로 도입한 시도
• 이질적 BBO 신호(prompt, design, label)를 delimiter token으로 명시적으로 구분하는 unified corpus 설계
• 기존 diffusion LLM 연구가 few-shot에서 frozen-model prompting에 그친 것과 달리, domain adaptation과 SFT+RL 기반 post-training을 결합해 모델 자체를 최적화 문제에 맞게 적응시킴
• RL 단계에서 label improvement 기반 reward와 one-step log-probability approximation을 활용한 효율적 RL loss 설계

Limitation & Further Study

• 본문 발췌만으로는 다양한 Design-Bench task 전반(continuous/discrete)에 대한 세부 성능 및 baseline 비교의 폭이 충분히 드러나지 않아 일반화 가능성 검증이 제한적일 수 있음
• one-step log-probability approximation이 실제 policy gradient 대비 근사 오차를 얼마나 유발하는지에 대한 이론적/실증적 분석이 부족해 보임
• reward가 label improvement로만 정의되어 있어 다목적(multi-objective) 최적화나 제약 조건이 있는 design 문제로의 확장성은 추가 연구가 필요함
• 향후 연구로 더 큰 diffusion LLM 백본이나 다양한 도메인(로보틱스, 재료과학)으로의 확장, RL reward 설계의 정교화가 기대됨

Evaluation

총평: Diffusion LLM의 bidirectional modeling을 offline BBO에 접목한 참신하고 실용적인 접근으로, delimiter 기반 domain adaptation과 SFT+RL 결합 설계가 효율성과 성능을 동시에 입증한 인상적인 연구이다.