Training Diffusion Language Models for Black-Box Optimization

이 논문은 오프라인 블랙박스 최적화 (Offline Black-Box Optimization, BBO) 문제를 해결하기 위해 확산 언어 모델 (Diffusion Language Models, Diffusion LLMs) 을 적용한 새로운 프레임워크인 DiBO (Diffusion for Black-box Optimization) 를 제안합니다.

주요 내용은 다음과 같습니다.

1. 문제 정의 (Problem)

• 배경: 로봇 설계, DNA 합성, 신소재 발견 등 다양한 과학 분야에서 성능 지표를 극대화하는 새로운 설계를 찾는 것이 목표입니다.
• 제약: 실제 실험은 시간과 비용이 많이 들기 때문에, 기존에 수집된 레이블이 있는 설계 데이터셋 (오프라인 데이터) 만을 사용하여 최적의 설계를 찾아야 하는 오프라인 BBO가 필수적입니다.
• 기존 방법의 한계:
  • 기존 방법은 주로 자기회귀 (Autoregressive, AR) LLM 을 사용하거나, 서브레이트 모델 (Surrogate Model) 을 학습합니다.
  • AR 모델의 문제: 왼쪽에서 오른쪽으로만 생성하는 특성상, DNA 시퀀스나 복잡한 설계 문제에서 필수적인 양방향 의존성 (Bidirectional Dependencies) 을 포착하는 데 한계가 있습니다.
  • 도메인 간극 (Domain Gap): 기존 확산 모델은 자연어 텍스트로만 사전 학습되었으나, BBO 작업은 자연어 프롬프트, 설계 데이터 (Design), 그리고 성능 레이블 (Label) 이라는 이질적인 신호를 포함하고 있어 직접 적용하기 어렵습니다.

2. 방법론 (Methodology)

저자들은 확산 LLM 의 양방향 모델링 능력을 활용하고 도메인 간극을 해결하기 위해 다음과 같은 세 가지 핵심 기법을 제안합니다.

A. 도메인 적응 (Domain Adaptation, DA)

• 통합 프롬프트 - 응답 코퍼스 (Unified Prompt-Response Corpus): 자연어 프롬프트, 설계, 레이블을 하나의 시퀀스로 통합합니다.
• 구분자 토큰 (Delimiter Tokens): 이질적인 신호를 명확히 구분하기 위해 |design-start|, |design-end|, |label-start|, |label-end|와 같은 구분자 토큰을 어휘집 (Tokenizer) 에 추가합니다. 이를 통해 모델이 텍스트, 설계, 레이블의 의미적 역할을 명확히 인식하도록 합니다.
• 공동 손실 함수 (Joint Loss): 프롬프트와 응답 (설계 + 레이블) 의 모든 마스킹된 토큰을 동시에 복원하도록 학습하여 도메인 적응을 수행합니다.

B. 2 단계 후학습 (Two-Stage Post-Training)

모델이 고품질 (High-Label) 설계를 생성하도록 정렬하기 위해 두 단계의 후학습을 거칩니다.

1. 지도 미세 조정 (Supervised Fine-Tuning, SFT):
   • 통합 데이터셋을 사용하여 프롬프트에 기반한 응답 (설계) 만을 마스킹 예측하도록 학습합니다.
   • 이는 모델이 프롬프트 컨텍스트를 기반으로 고품질 설계 영역으로 편향 (Inductive Bias) 되도록 초기화하는 역할을 합니다.
2. 강화 학습 (Reinforcement Learning, RL):
   • 보상 (Reward): 프롬프트 내 기존 설계의 최고 레이블과 생성된 응답 설계의 레이블 차이 ($y(o) - y(q)$) 를 보상으로 정의합니다.
   • 효율성: 확산 모델의 반복적 디노이즈 (Denoising) 과정은 계산 비용이 크므로, 단일 단계 로그 확률 근사 (One-step log-probability approximation) 기법을 사용하여 RL 손실을 효율적으로 계산합니다.
   • 이 단계를 통해 모델은 단순한 모방을 넘어, 구체적인 보상 신호를 통해 설계 공간을 탐색하고 최적화합니다.

3. 주요 기여 (Key Contributions)

1. 확산 LLM 의 BBO 적용: 자기회귀 모델의 한계를 극복하고, 설계 문제의 양방향 의존성을 포착하기 위해 확산 LLM 을 오프라인 BBO 에 최초로 적용했습니다.
2. 도메인 적응 프레임워크: 이질적인 신호 (프롬프트, 설계, 레이블) 를 처리하기 위한 구분자 토큰과 통합 코퍼스 구축 방법을 제안했습니다.
3. 효율적인 후학습 파이프라인: SFT 와 RL 을 결합한 2 단계 파이프라인을 통해 모델이 고품질 설계를 생성하도록 효과적으로 정렬했습니다.
4. 성능 입증: Design-Bench 벤치마크의 소규모 데이터 (Small-data) 설정에서 기존 최첨단 (SOTA) 방법들을 능가하는 결과를 달성했습니다.

4. 실험 결과 (Results)

• 벤치마크: Design-Bench 의 4 가지 태스크 (Ant Morphology, D'Kitty Morphology, TF Bind 8, TF Bind 10) 에서 평가되었습니다.
• 성능:
  • 제안된 DiBO는 4 개 태스크 중 3 개 (Ant, TF Bind 8, TF Bind 10) 에서 가장 높은 성능을 기록했습니다.
  • 특히 TF Bind 8에서는 기존 최강 베이스라인보다 정규화된 점수에서 0.108 이상 향상된 결과를 보여주었습니다.
  • 전체 평균 순위 (Mean Rank) 와 중앙값 순위 (Median Rank) 에서 모든 비교 대상 방법 중 1 위를 차지했습니다.
• 비교:
  • 전향적 방법 (Forward Methods, 서브레이트 모델 기반): DiBO 는 서브레이트 모델의 오차와 분포 이동 (Distribution Shift) 에 덜 민감하여 더 강건한 성능을 보였습니다.
  • 역방향 방법 (Inverse Methods, VAE, GAN, AR-LLM 등): 양방향 컨텍스트 모델링과 텍스트 지시사항의 통합 능력으로 인해 AR 기반 LLM 이나 기존 확산 모델보다 우수한 성능을 발휘했습니다.

5. 의의 및 결론 (Significance)

• 과학적 발견 가속화: 고비용의 물리적 실험을 줄이고, 제한된 오프라인 데이터만으로도 고품질의 설계 (DNA, 로봇 형태, 소재 등) 를 생성할 수 있는 강력한 도구를 제공합니다.
• 범용 최적화 패러다임: 확산 언어 모델이 고차원 설계 공간에서의 오프라인 최적화를 위한 강력하고 일반적인 패러다임이 될 수 있음을 입증했습니다.
• 데이터 효율성: 레이블이 부족한 소규모 데이터 환경에서도 뛰어난 성능을 발휘하여, 실제 과학 및 공학 응용 분야에서의 실용성을 높였습니다.

요약하자면, 이 논문은 확산 언어 모델의 양방향성과 구체적인 도메인 적응 기법 (구분자 토큰, SFT+RL) 을 결합하여, 제한된 데이터로 복잡한 설계 문제를 최적화하는 새로운 표준을 제시한 연구입니다.