🤖 AI

Can I Have Your Order? Monte-Carlo Tree Search for Slot Filling Ordering in Diffusion Language Models

본 논문은 마스킹 확산 모델에서 슬롯 채우기 순서를 최적화하기 위해 몬테카를로 트리 탐색을 활용하는 McDiffuSE 프레임워크를 소개하며, 비순차적 완성 전략을 효과적으로 탐색함으로써 수학적 및 코드 추론 작업에서 생성 품질을 크게 향상시킵니다.

원저자: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

게시일 2026-05-28

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Joshua Ong Jun Leang, Yu Zhao, Mihaela Cătălina Stoian, Wenda Li, Shay B. Cohen, Eleonora Giunchiglia

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

복잡한 케이크를 굽고 있다고 상상해 보세요. 하지만 레시피를 단계별로 따라 하는 대신 (반죽을 섞은 다음 굽고, 그 다음에 아이싱을 바르는 식), 원하는 순서대로 케이크를 채울 수 있는 마법 오븐을 가지고 있다고 가정해 봅시다. 먼저 윗부분에 아이싱을 바르고, 그 다음에 중간 부분을 굽고, 마지막으로 스프링클을 뿌릴 수도 있습니다. 이것이 **마스크된 확산 모델 (Masked Diffusion Models, MDMs)**이 작동하는 방식입니다. 이들은 누락된 정보의 "슬롯"을 어떤 순서로든 채워 넣음으로써 텍스트 (코드나 수학 해답 등) 를 생성할 수 있는 강력한 AI 도구들입니다.

하지만 함정이 하나 있습니다. 반죽이 팬에 들어가기 전에 아이싱을 바르면, 전체가 무너져 내립니다. 마찬가지로 AI 가 잘못된 "슬롯"을 먼저 채워 넣으면, AI 가 그 특정 부분에 대해 매우 확신하더라도 나머지 텍스트는 nonsensical(말도 안 되는) 것이 될 수 있습니다.

이 논문은 이 순서 문제를 해결하기 위해 MCDIFFUSE라는 새로운 방법을 소개합니다. 간단한 비유를 들어 그 작동 원리를 설명해 보겠습니다.

1. 문제: "확신하지만 틀린" 요리사

현재의 AI 모델들은 종종 첫 번째 단계를 매우 확신하는 요리사처럼 행동합니다. 만약 모델이 "이 문장은 함수 정의로 시작해야 한다고 90% 확신한다"고 생각하면, 즉시 그렇게 합니다. 하지만 때로는 정의로 시작하는 것이 나중에 혼란을 초래합니다. 모델이 코드가 작동하려면 먼저 주석이나 특정 구문 마커를 작성해야 한다는 사실을 깨닫지 못했기 때문입니다.

이 논문은 이를 "계획하고 채우기 (plan-and-infill)" 문제라고 부릅니다. AI 는 텍스트의 어떤 부분을 다음에 작성할지 결정해야 합니다. 잘못된 순서를 선택하면 최종 결과가 나빠집니다.

2. 해결책: "만약에" 시뮬레이터 (MCTS)

저자들은 **MCTS(Monte Carlo Tree Search)**라는 기법을 사용하는 MCDIFFUSE를 소개합니다.

MCTS 를 슈퍼 플래너나 시뮬레이터라고 생각하세요. AI 가 텍스트의 특정 부분을 작성하기로 결정하기 전에, 이 플래너는 머릿속에서 수천 가지의 "만약에" 시나리오를 실행합니다.

시나리오 A: 함수 정의를 먼저 작성하면 어떨까? (나머지 코드를 시뮬레이션해 보니... 아뇨, 나중에 충돌이 발생합니다.)
시나리오 B: 구문 마커를 먼저 작성하면 어떨까? (나머지를 시뮬레이션해 보니... 아하, 코드가 완벽하게 흐릅니다.)

이 플래너는 단순히 즉각적인 확신 (AI 가 다음 단어에 대해 얼마나 확신하는지) 만 보지 않습니다. 장기적인 결과를 봅니다. "이 경로를 선택하면 전체 케이크가 잘 나올까?"라고 묻는 것입니다.

3. 결정 방식: "탐색"과 "시뮬레이션"의 균형

이 논문은 이 플래너를 조정하는 방법에 대한 흥미로운 발견을 제시합니다. 보통 사람들은 더 나은 답을 얻으려면 더 많은 시뮬레이션을 실행해야 한다고 생각합니다. 하지만 MCDIFFUSE 는 그것이 가장 중요한 것이 아니라는 것을 발견했습니다.

옛 방식: 같은 경로를 100 번 시뮬레이션하여 확실히 합니다.
MCDIFFUSE 방식: 시뮬레이션 횟수는 줄이지만, AI 가 보통 무시하는 이상하고 unlikely(가능성이 낮은) 경로를 시도할 준비를 합니다.

비유: 숨겨진 보물을 찾고 있다고 상상해 보세요.

높은 시뮬레이션, 낮은 탐색: 지도에 "보물 여기"라고 적혀 있어 같은 자리에서 100 번을 파봅니다. 아무것도 찾지 못했지만, 충분히 열심히 찾았다는 확신은 매우 큽니다.
높은 탐색 (MCDIFFUSE): 눈에 띄는 곳에서 몇 번 파보지만, 지도가 틀릴 수도 있으므로 옆의 이상하고 덤불이 우거진 곳에서도 파기로 결정합니다. 알고 보니 보물은 그 덤불 속에 있었습니다!

이 논문은 플래너에 "큰 탐색 상수 (large exploration constant)"를 부여하여 (용감해지고 unlikely 한 순서들을 시도하도록 지시하는 것) 단순히 시뮬레이션을 더 많이 실행하는 것보다 훨씬 효과적임을 발견했습니다. 이는 AI 가 자신의 과도한 확신에 빠지는 함정에서 벗어나도록 도와줍니다.

4. 결과: 더 나은 코드와 수학

파이썬 코드 작성과 수학 문제 해결과 같은 작업에서 이를 테스트했을 때:

작동합니다: MCDIFFUSE 는 이전 방법들을 크게 능가했습니다. 일부 코딩 테스트에서는 정확도가 거의 20% 향상되었습니다.
효율적입니다: 추가적인 계획을 수행하지만, 표준 AI 모델들보다 더 짧고 간결한 답변을 실제로 생성합니다. 표준 모델들은 종종 장황하게 말하거나 긴 루프에 갇히곤 하지만, MCDIFFUSE 는 직접적인 경로를 찾습니다.
"비밀 소스": AI 는 주로 인간이 읽는 것처럼 정상적인 왼쪽에서 오른쪽 순서를 따르지만, 필요할 때 그 규칙을 전략적으로 깨뜨립니다. 마치 연대기적으로 쓰기를 하되, 이야기의 논리를 맞추기 위해 언제 뒤로 돌아가서 중요한 세부 사항을 삽입해야 하는지 정확히 아는 작가와 같습니다.

요약

MCDIFFUSE는 AI 텍스트 생성을 위한 똑똑한 "교통 통제관"입니다. AI 가 그 순간에 옳다고 느껴지는 대로 맹목적으로 텍스트를 채워 넣게 하는 대신, 다양한 작업 순서를 테스트하기 위해 시뮬레이션 엔진을 사용합니다. AI 가 보통 무시하는 "이상한" 순서들을 시도할 용기를 냄으로써 실수를 피하고 더 깔끔하고 정확한 코드 및 수학 해답을 생성합니다.

기술 요약: MCDIFFUSE – 확산 언어 모델의 슬롯 채우기 순서 최적화를 위한 MCTS

1. 문제 정의

마스크드 확산 모델 (MDM) 은 비순차적, 반복적 탈노이즈를 가능하게 함으로써 추론 작업에 있어 자기회귀 모델 (ARM) 의 유망한 대안을 제공합니다. 그러나 MDM 은 출력 품질이 슬롯 채우기 순서에 민감하게 반응하는 경향으로 인해 종종 ARM 보다 낮은 성능을 보입니다. "계획 및 채우기 (plan-and-infill)"프레임워크 (예: ReFusion) 는 슬롯 내의 국소적 의존성을 해결하지만, 슬롯 생성의 전역적 순서를 결정하기 위해 휴리스틱 계획자 (종종 신뢰도 기반) 에 의존합니다. 이러한 휴리스틱은 종종 장기 의존성을 고려하지 못해 오류 전파와 비일관된 전역 출력을 초래합니다. 핵심적인 과제는 가능한 슬롯 순열의 방대한 조합 공간으로, 전역적 일관성을 극대화하는 순서를 식별하는 것이 포괄적 탐색 없이는 어렵다는 점입니다.

2. 방법론: MCDIFFUSE

저자들은 MCDIFFUSE를 소개합니다. 이는 슬롯 선택을 **몬테카를로 트리 탐색 (MCTS)**을 통해 해결되는 순차적 의사결정 문제로 공식화하는 훈련이 불필요한 프레임워크입니다. 이 접근법은 추가적인 모델 훈련을 요구하지 않으며, 대신 추론 시 생성 순서를 최적화합니다.

2.1 MDP 로서의 문제 공식화

슬롯 계획 과정은 결정적 마르코프 의사결정 과정 (MDP) $\langle S, A, T, R \rangle$ 으로 모델링됩니다:

상태 ( $S$ ): 생성된 슬롯 집합과 현재 부분적으로 채워진 텍스트 시퀀스를 인코딩합니다. 채워지지 않은 슬롯은 마스크됩니다.
행동 ( $A$ ): 다음에 채울 마스크된 슬롯의 인덱스를 선택합니다.
전이 ( $T$ ): 결정적입니다. 행동을 선택하면 모델의 argmax 디코딩을 사용하여 해당 슬롯이 채워집니다.
보상 ( $R$ ): 슬롯 수준 신뢰도로 정의되며, 이는 해당 특정 슬롯에서 생성된 토큰에 대해 모델이 부여하는 평균 확률로 계산됩니다.

2.2 MCTS 적응

MCDIFFUSE 는 확산 생성의 고유한 제약을 처리하기 위해 MCTS 를 특정 메커니즘으로 적응시킵니다:

사전 지향 확장: 모델의 고유한 신뢰도 점수를 정규화하여 PUCT(예측기 - 상한 신뢰도 트리) 선택 기준을 위한 사전 확률 $P(a|s)$ 를 형성합니다. 이는 탐색을 허용하면서도 의미상 타당한 행동으로 탐색을 유도합니다.
하이브리드 가치 추정: 탐욕적 국소 신뢰도의 단시성을 피하기 위해 상태 - 행동 쌍 $V(s, a)$ $V (s, a)$ 의 가치는 다음 두 가지의 가중 조합입니다:
1. 즉시 보상: 현재 슬롯의 신뢰도.
2. 롤아웃 반환: 남은 슬롯을 채워 완료할 때까지의 확률적 시뮬레이션 (롤아웃) 으로, 장기적 궤적 일관성을 추정합니다.
- 공식: $V(s, a) = \lambda \cdot R(s, a) + (1 - \lambda) \cdot E_{\pi_{roll}}[G]$ 이며, 여기서 $\lambda$ 는 즉각적 신뢰도와 장기적 계획 사이의 균형을 맞춥니다.
선택 전략: 알고리즘은 탐색 중 탐색과 활용의 균형을 맞추기 위해 PUCT 를 사용합니다. 루트에서의 최종 행동 선택은 강건한 자식 기준(최대 방문 횟수) 에 의존하여 안정성을 보장합니다.

3. 주요 기여

MCDIFFUSE 프레임워크: MDM 에서 슬롯 순서를 탐색하기 위해 MCTS 를 특별히 맞춤화한 훈련 불필요 방법으로, 자기회귀 모델과의 성능 격차를 효과적으로 해소합니다.
전략적 비순차적 생성: 실증 분석에 따르면, MCDIFFUSE 는 주로 순차적(왼쪽에서 오른쪽) 순서를 따르지만, 복잡한 추론 작업에서 성능을 극대화하기 위해서는 비순차적 생성으로의 전략적 이탈이 필수적입니다.
탐색 대 시뮬레이션 깊이: 논문은 시뮬레이션 횟수 ( $N_{sim}$ ) 를 늘리는 것이 일관되게 성능을 향상시키지 않는다고 보여줍니다. 대신, 모델의 편향된 신뢰도 사전 지식을 극복하고 효과적인 비탐욕적 슬롯 순서를 발견하기 위해서는 큰 탐색 상수( $c_{uct}$ )가 필요합니다.

4. 실험 결과

MCDIFFUSE 는 자기회귀 베이스라인 (Qwen2.5, Qwen3) 과 확산 베이스라인 (LLaDA, Dream, ReFusion) 에 대해 여섯 가지 추론 벤치마크 (GSM8K, MATH500, ARC, GPQA-Diamond, MBPP, HumanEval) 에서 평가되었습니다.

성능 향상: MCDIFFUSE 는 자기회귀 베이스라인 대비 평균 **3.2%**의 개선을 달성했으며, 베이스라인 계획 및 채우기 방법 (ReFusion) 대비 **8.0%**의 개선을 보였습니다.
구체적 벤치마크: 주목할 만한 개선은 코드 생성인 MBPP 에서 19.5%, MATH500 에서 **4.9%**였습니다.
효율성: MCDIFFUSE 는 훨씬 더 컴팩트한 추론 시퀀스를 생성했습니다. 예를 들어, MATH500 에서 Qwen2.5 가 436.0 개의 토큰을 생성한 반면, 평균 길이 152.2 개의 토큰으로 더 높은 정확도를 달성했습니다.
코드 생성 대 추론: 이 방법은 다중 선택 추론에 비해 구조적 의존성이 있는 코드 생성 작업에서 상당한 개선을 보였으며, 이는 슬롯 순서가 의미적 유효성에 결정적으로 영향을 미치는 영역에서 MCTS 가 특히 효과적임을 시사합니다.

5. 중요성 및 주장

이 논문은 전략적 슬롯 계획이 마스크드 확산 모델의 성능에 있어 결정적인 요소임을 확립합니다. 저자들은 다음과 같이 주장합니다:

MCTS 의 효과성: MCTS 기반 계획은 모델 재훈련 없이 MDM 의 생성 품질을 향상시키기 위한 실현 가능하고 효과적인 접근법입니다.
비순차성의 필수성: 확산 모델이 엄격하게 순차적 궤적을 따라야 한다는 가정과 달리, 왼쪽에서 오른쪽 순서에서 이탈할 수 있는 능력은 복잡한 제약을 해결하고 정확도를 극대화하는 데 필수적입니다.
깊이보다 탐색 폭: 슬롯 계획의 주요 과제는 탐색 깊이 (시뮬레이션 횟수) 가 아니라 탐색 폭입니다. 모델의 국소적 신뢰도 편향을 극복하고 전역적으로 일관된 경로를 방문하도록 보장하기 위해서는 높은 탐색 상수가 필요합니다.

이 연구는 MCDIFFUSE 를 전역적 일관성이 가장 중요한 구조화된 다단계 추론이 필요한 도메인에서 확산 모델과 자기회귀 모델 간의 격차를 해소하는 한 걸음으로 위치시킵니다.