🤖 AI

Diffusion Large Language Models for Black-Box Optimization

이 논문은 인컨텍스트 디노이징 모듈과 마스크드 디퓨전 트리 서치를 갖춘 디퓨전 거대 언어 모델을 활용하여 설계를 반복적으로 개선함으로써 퓨샷 시나리오에서 최첨단 성능을 달성하는 오프라인 블랙박스 최적화를 위한 새로운 접근 방식인 dLLM을 소개한다.

원저자: Ye Yuan, Can, Chen, Zipeng Sun, Dinghuai Zhang, Christopher Pal, Xue Liu

게시일 2026-01-22

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Ye Yuan, Can, Chen, Zipeng Sun, Dinghuai Zhang, Christopher Pal, Xue Liu

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신은 완벽한 새 레시피를 발명하려는 마스터 셰프라고 상상해 보세요. 하지만 당신에게는 큰 문제가 있습니다. 오직 10개의 오래된 레시피와 그 평점이 적힌 아주 작은 수첩 하나뿐이라는 점입니다. 새로운 아이디어를 테스트할 수 있는 주방도 없습니다. 테스트 비용이 너무 비싸거나 위험하기 때문입니다. 당신은 그 작은 수첩을 들여다보는 것만으로 가장 좋은 새로운 레시피를 찾아내야 합니다.

이것이 바로 **오프라인 블랙박스 최적화(Offline Black-Box Optimization)**의 과제입니다. 여기서 "블랙박스"란 어떤 설계(예: DNA 서열이나 로봇의 형태)가 좋은지 나쁜지를 결정하는 미지의 규칙을 의미합니다. "오프라인"이라는 부분은 실시간으로 새로운 아이디어를 테스트할 수 없으며, 오직 과거의 데이터만을 사용할 수 있다는 뜻입니다.

이 논문의 새로운 방법인 dLLM이 이 문제를 어떻게 해결하는지, 쉬운 비유를 통해 설명해 드리겠습니다.

기존 방식의 문제점

이 문제를 해결하려는 이전의 시도들은 **자기회귀 모델(Autoregressive Models)**을 사용했습니다. 이것은 마치 글을 쓸 때 왼쪽에서 오른쪽으로, 한 단어씩 써 내려가는 작가를 생각하면 됩니다.

결함: 문장을 쓰고 있다면, 문장의 끝부분이 처음에 어떻게 써야 했을지를 결정하는 경우가 많습니다. 하지만 왼쪽에서 오른쪽으로만 쓰는 작가는 시작 부분을 쓰는 동안 끝을 볼 수 없습니다. 그들은 "전체적인 그림"의 연결 고리를 놓치게 됩니다. 복잡한 설계(DNA와 같은)에서는 모든 부분이 서로에게 의존하기 때문에, 엄격하게 왼쪽에서 오른쪽으로만 쓰는 것은 엉망이거나 수준 미달의 결과를 초래하기 쉽습니다.

새로운 솔루션: "디퓨전(Diffusion)" 셰프

저자들은 **디퓨전 대규모 언어 모델(Diffusion Large Language Models, dLLMs)**을 소개합니다. 한 단어씩 써 내려가는 대신, **빈 재료들(masking된 [M]으로 표현됨)**이 담긴 그릇을 앞에 두고, 단계별로 재료를 조금씩 드러내며 완성해 나가는 셰프를 상상해 보세요.

이 접근 방식에는 두 가지 초능력이 있습니다.

양방향 시야 (Bidirectional Vision): 셰프는 전체 그릇을 한꺼번에 볼 수 있습니다. 레시피의 "끝"이 "시작"에 어떤 영향을 미치는지 볼 수 있으므로, 국소적인 수정이 아닌 전역적인 수정을 할 수 있습니다.
반복적 정교화 (Iterative Refinement): 그들은 최종 요리를 즉석에서 바로 추측하지 않습니다. 거친 스케치에서 시작하여, 빈칸을 천천히 채워 나가며 매 단계마다 점점 더 나은 결과물을 만들어냅니다.

작동 원리: 두 단계의 춤

논문은 이 셰프를 더욱 뛰어나게 만들기 위해 두 가지 영리한 기술을 결합했습니다.

1. 인컨텍스트 디노이징 (In-Context Denoising, "스마트 프롬프트")

셰프가 요리를 시작하기 전에, 당신은 특별한 지시 카드를 줍니다. 이 카드에는 다음 내용이 포함됩니다:

목표: "이 특정 단백질에 가장 잘 달라붙는 DNA 서열을 만드세요."
수첩: 당신의 오프라인 데이터셋에 있는 10개의 예시.
명령: "더 나은 새로운 서열을 제안해 주세요."

모델은 이 카드를 읽고 자신의 방대한 사전 학습 지식(수백만 권의 요리책을 읽은 셰프와 같은 지식)을 사용하여 마스킹된 재료들을 "디노이징(denoising)"하기 시작합니다. 이는 빈 [M]들을 실제 글자(A, C, G, T)로 바꾸며, 후보 설계를 서서히 드러내는 과정입니다.

2. 마스크 디퓨전 트리 서치 (Masked Diffusion Tree Search, "분기형 탐험가")

단순히 하나의 레시피를 추측하는 것만으로는 충분하지 않습니다. 만약 셰프가 잘못된 경로에 빠진다면 어떻게 될까요? 저자들은 트리 서치(Tree Search) 메커니즘을 추가했습니다. 셰프가 단순히 하나의 요리를 만드는 것이 아니라, 여러 갈래의 가능성을 펼쳐놓는다고 상상해 보세요.

선택 (Selection): 셰프는 현재 진행 중인 모든 부분적인 레시피들을 살펴보고, 다음에 작업할 가장 유망한 것을 고릅니다 (새로운 것을 시도하는 것과 기존의 것을 고수하는 것 사이의 균형을 맞춘 점수를 사용합니다).
확장 (Expansion): 그 유망한 부분 레시피를 가져와서 다음 몇 개의 빈칸을 채울 수 있는 여러 가지 다른 방법들을 생성합니다. 이제 하나의 경로가 아니라, 하나의 나무(tree) 형태의 많은 가능성이 생깁니다.
평가 (Evaluation): 각 새로운 가지(branch)에 대해, 그들은 "예측기"(가우시안 프로세스, 즉 당신의 10개 예시를 바탕으로 훈련된 스마트 계산기)를 사용하여 만약 요리를 완성한다면 최종 결과물이 얼마나 좋을지 예측합니다.
역전파 (Backpropagation): 만약 어떤 가지가 맛있어 보인다면, 셰프는 그 경로를 기억하고 더 깊이 탐구합니다. 만약 어떤 가지가 맛이 없다면, 그 가지를 쳐내고(prune) 시간을 낭비하지 않습니다.

이 과정은 몬테카를로 트리 서치(Monte Carlo Tree Search)(AI가 바둑을 두는 데 사용하는 전략)와 비슷하지만, 체스를 두는 대신 AI가 완벽한 설계를 찾기 위해 "빈칸 채우기" 게임을 하는 것입니다.

결과

논문은 네 가지 서로 다른 과제를 통해 이 방법을 테스트했습니다:

Ant Morphology (개미 형태학): 빠르게 기어가는 로봇 개미 설계하기.
D'Kitty Morphology (D'Kitty 형태학): 빠르게 기어가는 로봇 고양이 설계하기.
TF Bind 8 & 10: 특정 단백질에 잘 달라붙는 짧은 DNA 서열 설계하기.

모든 테스트에서 dLLM 방식은 경사 하강법(gradients), 생성 모델(generative models), 또는 표준 언어 모델(standard language models)을 사용하는 기존의 모든 방법을 능가했습니다. 특히 단 10개의 예시만 가지고 학습할 때도 절대적인 최적의 설계(상위 1%)를 찾아내는 데 탁받은 성능을 보였습니다.

요 요약

기존의 방식들이 왼쪽에서 오른쪽으로만 쓸 수 있고 쉽게 막히는 작가라면, 새로운 dLLM 방식은 다음과 같은 마스터 셰프와 같습니다:

당신의 제한된 수첩과 지시 사항을 읽습니다.
빈 캔버스에서 시작합니다.
전체적인 그림을 보면서 설계를 서서히 드러냅니다.
완벽한 해결책을 놓치지 않도록 트리 서치 전략을 사용하여 동시에 여러 가지 "만약에" 시나리오를 탐색합니다.

이를 통해 데이터가 부족하고 규칙이 숨겨져 있는 상황에서도 최선의 설계를 찾아낼 수 있습니다.

기술 요약: 블랙박스 최적화를 위한 확산 대규모 언어 모델 (Diffusion Large Language Models)

문제 정의

오프라인 블랙박스 최적화(Black-Box Optimization, BBO)는 제한된 오프라인 데이터셋(설계-레이블 쌍)만을 사용하여 미지의 목적 함수 $f(\cdot)$ 를 최대화하는 최적의 설계( $x^*$ )를 식별하는 것을 목표로 합니다. 이러한 시나리오는 DNA 서열 설계 및 로보틱스와 같이 속성을 쿼리하는 비용이 많이 들고 레이블이 있는 데이터가 부족한(few-shot regime) 도메인에서 흔히 발생합니다.

전통적인 접근 방식은 일반적으로 다음 두 가지 범주로 나뉩니다:

대리 모델 기반 방법(Proxy-based methods): $f(\cdot)$ 를 근사하기 위해 대리 모델(예: 가우시안 프로세스 또는 심층 신경망)을 학습시키고, 그 기울기(gradient)를 사용하여 최적화를 유도합니다.
생성 모델 기반 방법(Generative methods): 고득점 설계를 직접 생성하도록 조건부 생성 모델을 학습시킵니다.

그러나 이러한 방법들은 데이터가 제한적인 환경에서 분포 외 일반화(out-of-distribution generalization) 문제를 겪는 경우가 많습니다. 또한, 최근 BBO를 위해 **자기회귀(autoregressive) 대규모 언어 모델(LLM)**을 자연어 프롬프트 형태로 적응시킨 방식들은 결정적인 한계에 직면해 있습니다. 이들은 시퀀스를 왼쪽에서 오른쪽으로 생성하는 단방향 방식을 취하는데, 이러한 방식은 많은 설계 작업에 내재된 양방향 의존성(예: DNA 단위의 기능은 접두사와 접미사 모두에 의존함)을 포착하는 데 어려움을 겪습니다.

방법론: dLLM

저자들은 오프라인 BBO를 위해 **확산 대규모 언어 모델(Diffusion Large Language Models)**을 활용하는 프레임워크인 dLLM을 제안합니다. 이 방법은 **인컨텍스트 디노이징(In-Context Denoising)**과 **마스크드 확산 트리 탐색(Masked Diffusion Tree Search)**이라는 두 가지 핵심 모듈을 결합합니다.

1. 인컨텍스트 디노이징 (In-Context Denoising)

이 방법은 파라미터 미세 조정(fine-tuning) 대신, 사전 학습된 확산 LLM의 인컨텍스트 학습 능력을 활용합니다.

프롬프트 구성: 확산 LLM의 입력은 다음을 포함하는 구조화된 자연어 프롬프트로 구성됩니다:
- 작업 설명(Task Description): 설계 형식(예: DNA 알파벳), 목적 및 제약 조건을 정의합니다.
- 오프라인 데이터셋(Offline Dataset): 설계-레이블 쌍(예: DNA 서열 및 결합 친화도)의 few-shot 세트입니다.
- 지시문(Instruction): 개선된 후보를 생성하라는 지시입니다.
디노이징 과정: 모델은 완전히 마스킹된 설계( $x_{t=1}$ )에서 시작합니다. 모델은 마스킹된 토큰을 반복적으로 예측하여 완전한 후보( $\hat{x}_0$ )를 생성합니다. 신뢰도가 가장 낮은 토큰들은 다시 마스킹되어 다음 단계를 위한 새로운 부분 마스킹 상태( $x_s$ )를 형성하며, 이를 통해 반복적인 디노이징을 거쳐 설계를 정교화합니다.

2. 마스크드 확산 트리 탐색 (Masked Diffusion Tree Search, MDTS)

고성능 영역으로 생성을 유도하고 탐색(exploration)과 이용(exploitation) 사이의 균형을 맞추기 위해, 디노이징 과정을 단계별 **몬테카를로 트리 탐색(Monte Carlo Tree Search, MCTS)**으로 구성합니다.

트리 구조: 각 노드는 부분적으로 마스킹된 설계( $x_t$ )를 나타냅니다.
선택(Selection): 알고리즘은 상한 신뢰 구간(Upper Confidence Bound, UCT) 점수를 사용하여 루트(완전 마스킹 상태)로부터 트리를 탐색합니다. 이는 높은 가치를 가진 노드를 이용하는 것과 방문 횟수가 적은 브랜치를 탐색하는 것 사이의 균형을 맞춥니다. UCT 점수는 노드의 가치 추정치와 언마스킹(unmasking) 단계에 대한 모델의 가능도(likelihood)를 통합합니다.
확장(Expansion): 선택된 노드는 확산 LLM을 적용하여 다양한 완성을 생성함으로써 확장되며, 생성된 결과물은 다시 마스킹되어 자식 노드를 생성합니다.
평가(Evaluation): 자식 노드들은 완전히 디노이징된 설계로 전개됩니다. 이들의 보상은 오프라인 데이터셋에 대해 학습된 가우시안 프로세스(GP)로부터 도출된 **기대 개선량(Expected Improvement, EI)**을 사용하여 추정됩니다. GP는 few-shot 설정에서 중요한 원칙적인 불확실성 추정치를 제공합니다.
역전파(Backpropagation): 보상 신호는 트리 상단으로 전파되어 조상 노드의 가치 추정치와 방문 횟수를 업데이트함으로써 향러의 선택 결정에 정보를 제공합니다.

주요 기여

BBO에 확산 LLM을 최초로 적용: 본 논문은 확산 LLM의 양방향 모델링 및 반복적 정교화 능력을 활용하여 복잡한 설계 의존성을 처리할 수 있는, 오프라인 BBO를 위한 최초의 프레임워크를 소개합니다.
인컨텍스트 디노이징 모듈: 파라미터 업데이트 없이 자연어 작업 설명과 few-shot 오프라인 데이터셋을 기반으로 마스킹된 설계를 정교화할 수 있는 새로운 컨디셔닝 메커니즘을 제공합니다.
마스크드 확산 트리 탐색: 확산 디노이징을 순차적 의사결정 과정으로 취급하며, MCTS와 GP 기반의 기대 개선량을 통해 탐색과 이용을 동적으로 조절하는 탐색 알고리즘을 제시합니다.

실험 결과

저자들은 Design-Bench 벤치마크를 통해 연속형(Ant Morphology, D'Kitty Morphology) 및 이산형(TF Bind 8, TF Bind 10) 작업을 포함한 few-shot 설정(10개의 레이블된 예시)에서 dLLM을 평가했습니다.

성능: dLLM은 네 가지 작업 모두에서 최첨단(SOTA) 결과를 달eric했으며, 15개의 베이스라인 중 평균 및 중앙값 순위에서 1위를 차지했습니다. 이는 연속형 및 이산형 설계 공간 모두에서 일관된 이득을 보여주었습니다.
비교:
- 순수 그래디언트 기반 프록시 방법(예: Grad-mean, Grad-EI)보다 우수한 성능을 보였으며, 이는 생성 모델이 탐색을 위해 결합되었을 때 더 우월함을 나타냅니다.
- 자기회귀 LLM(OPRO) 및 다른 확산 모델(DDOM, GTG)을 포함한 다른 생성형 베이스라인을 능가했습니다. 저자들은 이를 확산 LLM의 양방향 모델링과 활용되는 광범위한 사전 학습 지식 덕분으로 돌렸습니다.
- 강력한 자기회귀 베이스라인인 UniSO-T를 능가하였으며, 이는 양방향 컨텍스트가 복잡한 의존성을 가진 작업에 매우 중요하다는 것을 시사합니다.
절제 연구(Ablation Studies):
- 트리 탐색을 제거했을 때(w/o MDTS) 성능이 크게 하락하여, 유도된 탐색의 필요성을 입증했습니다.
- 확산 LLM을 few-shot 데이터로만 학습된 일반적인 확산 모델로 교체했을 때 성능이 현저히 저하되었으며, 이는 사전 학습된 지식의 중요성을 확인시켜 줍니다.
- 프롬프트에서 명시적인 작업 설명을 제거하면 성능이 감소하여, 의미론적 프레이밍(semantic framing)의 역할을 검증했습니다.

의의 및 주장

본 논문은 dLLM이 데이터 효율적인 설계 생성을 위한 범용 최적화 도구로서 확산 LLM의 잠재력을 보여준다고 주장합니다. 자기회귀 모델이 양방향 의존성을 포착하는 데 겪는 한계를 해결하고 이를 강력한 탐색 전략과 결합함으로써, dLLM은 레이블된 데이터가 부족한 과학 및 공학 분야에서 다재다능한 대안을 제공합니다. 저자들은 이 연구가 사전 학습된 파운데이션 모델을 오프라인 블랙박스 최적화에 활용하는 새로운 방향을 열었다고 평가합니다.