💬 NLP

Improving Diffusion Language Model Decoding through Joint Search in Generation Order and Token Space

이 논문은 생성 순서와 토큰 공간을 공동으로 탐색함으로써 수학적 추론 및 코딩 벤치마크에서 기존 베이스라인들을 능가하는 디퓨전 언어 모델을 위한 새로운 디코딩 방법인 Order-Token Search를 소개한다.

원저자: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

게시일 2026-02-06

📖 4 분 읽기☕ 가벼운 읽기

원저자: Yangyi Shen, Tianjian Feng, Jiaqi Han, Wen Wang, Tianlang Chen, Chunhua Shen, Jure Leskovec, Stefano Ermon

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

당신이 복잡한 퍼즐, 예를 들어 수학 문제나 코딩 챌린지를 풀려고 노력하고 있다고 상상해 보세요. 하지만 당신에게는 마법 같은 조수(확산 언어 모델, Diffusion Language Model)가 있습니다. 이 조수는 인간이 글자를 한 글자씩 왼쪽에서 오른쪽으로 타이핑하는 것처럼 작동하는 것이 아니라, "MASK"라는 자리 표시자로 가득 찬 빈 페이지에서 시작하여 모든 곳을 한꺼번에 채우며 어떤 단어가 어디에 들어갈지 추측합니다.

문제는 이 조수가 다소 무질서하다는 점입니다. 조수는 자신이 원하는 어떤 순서로든 빈칸을 채울 수 있습니다. 때로는 첫 번째 빈칸에 맞는 단어를 맞히기도 하지만, 두 번째 빈칸에 잘못된 단어를 추측해서 막혀버리기도 합니다. 또 어떤 때는 단어들은 제대로 맞혔지만 순서가 틀려서 막다른 길에 다다르기도 합니다.

기존 방식: 추측과 가지치기 (Guessing and Pruning)

이전에는 사람들은 이를 해결하기 위해 크게 두 가지 방법을 시도했지만, 둘 다 결함이 있었습니다.

"확신형" 접근 방식 (The "Confident" Approach): 조수가 자신이 가장 확신하는 빈칸부터 먼저 채우는 방식입니다. 이는 마치 가장 단단해 보이는 길로만 걷는 등산객과 같습니다.
- 장점: 보통 정답을 빠르게 찾아냅니다.
- 단점: 만약 그 "단단한 길"이 절벽(오답)으로 이어진다면, 등산객은 그곳에 갇히게 됩니다. 더 나은 경로가 있을지도 모른다는 생각으로 다른 길을 탐색하지 않습니다.
"무작위형" 접근 방식 (The "Random" Approach): 조수가 완전히 무작위로 빈칸을 선택하여 채우는 방식입니다.
- 장점: 매우 다양한 경로를 탐색하므로, 결국에는 올바른 해결책을 우연히 발견할 가능성이 매우 높습니다.
- 단점: 너무 산만해서 첫 시도에 최선의 경로를 선택하는 경우는 드뭅니다. 이는 마치 원을 그리며 헤매는 등산객과 같습니다. 보물을 찾을 수는 있겠지만, 그 과정에서 수많은 빈 구덩이를 파헤치게 될 것입니다.

새로운 해결책: 순서-토큰 탐색 (Order-Token Search)

이 논문의 저자들은 **순서-토큰 탐색(Order-Token Search)**이라는 새로운 방법을 소개했습니다. 이것을 협력하는 탐험가 팀이라고 생각해 보세요.

단 한 명의 등산객이나 혼란스러운 군중을 보내는 대신, 이 방법은 소규모 팀(하나의 "빔", beam)을 보냅니다. 이들은 다음과 같이 작동합니다.

갈라지는 경로 (탐색, The Search): 정기적인 간격으로 팀은 흩어집니다. 각 탐험가는 서로 다른 전략을 시도합니다.
- 탐험가 A는 첫 번째 빠진 단어를 채우기로 결정합니다.
- 탐험가 B는 마지막 빠진 단어를 채우기로 결정합니다.
- 탐험가 C는 중간 지점에 다른 단어를 넣어봅니다.
- 비유: 이들은 다음에 어디에 쓸 것인지(순서)와 무엇을 쓸 것인지(토큰)를 동시에 탐색합니다.
성적표 (우도 추정치, The Likelihood Estimator): 이것이 마법 같은 부분입니다. 팀에는 특별한 심판(우도 추정치)이 있는데, 이 심판은 단순히 최종 답변만 보는 것이 아닙니다. 대신, 심판은 탐험가들이 거친 모든 단계를 살펴봅니다.
- 탐험가가 논리적인 움직임을 보였는가?
- 이 부분적인 문장이 이전에 작성된 내용과 맥락이 맞는가?
- 비유: 계주 경기를 관찰하는 코치를 상상해 보세요. 만약 주자가 초반에 넘어졌다면, 코치는 주자가 경주를 마칠 때까지 기다리지 않고 즉시 중단시킵니다. 왜냐하면 그 단계 자체가 잘못되었기 때문입니다.
막다른 길 차단 (가지치기, Pruning): 심판은 모든 탐험가의 진행 상황을 점수 매깁니다. 만약 어떤 탐험가가 성공 가능성이 낮아 보이는 경로를 걷고 있다면, 팀은 그 경로를 잘라내고 가장 좋은 궤도에 있는 탐험가들에게 자원을 집중합니다.

이것이 왜 중요한가

논문 저자들은 이를 어려운 수학 문제(GSM8K 및 MATH500 데이터셋)와 코딩 작업(HumanEval)에 테스트했습니다.

결과: "탐험가 팀"(순서-토큰 탐색)은 기존의 "확신형" 또는 "무작위형" 방법보다 일관되게 더 자주 정답을 찾아냈습니다.
비교: 이 방법은 비용이 많이 드는 몇 달간의 재학습(예: diffu-GRPO)이 필요한 방법들과 대등하거나 오히려 더 뛰어난 성능을 보였습니다. 이는 AI의 뇌 자체를 재학습할 필요 없이, 테스트 중에 생각하는 방식을 바꾸는 것만으로도 훨씬 더 똑똑한 AI를 얻을 수 있음을 의미합니다.

스도쿠에 대한 특별한 참고 사항

이 논문은 또한 이 방법을 스도쿠 퍼즐에도 적용해 보았습니다. 흥미롭게도, 스도쿠에서는 효과가 좋지 않았습니다. 저자들은 스도쿠가 엄격한 전역 규칙(예: "줄에 숫자가 중복되지 않음")을 요구하는데, AI의 내부 "성적표"가 이를 이해하지 못한다고 설명합니다. 이는 마치 절벽이 표시되지 않은 지도를 가진 등산객 팀을 주는 것과 같습니다. 아무리 잘 탐색하더라도 지도가 잘못되었다면 절벽 아래로 떨어지는 것을 피할 수 없습니다. 이는 어떤 작업의 경우, 탐색 방법뿐만 아니라 AI 자체를 다르게 훈련시켜야 함을 시사합니다.

요약하자면

이 논문은 AI가 여러 가지 다른 방식의 쓰기(순서)와 여러 가지 다른 단어(토큰)를 동시에 탐색하게 하고, 스마트한 점수 시스템을 사용하여 나쁜 아이디어를 조기에 차단함으로써, 재학습 없이도 확산 언어 모델로부터 훨씬 더 나은 결과를 얻을 수 있음을 보여줍니다. 이는 혼란스러운 추측 게임을 진실을 향한 구조적이고 효율적인 탐색으로 바꿉니다.

기술 요약: 생성 순서와 토큰 공간의 결합 탐색을 통한 확산 언어 모델 디코딩 개선

문제 정의

확산 언어 모델(Diffusion Language Models, DLMs), 특히 마스크드 확산 모델(Masked Diffusion Models, MDMs)은 자기회귀(Autoregressive, AR) 모델과 구별되는 뚜렷한 장점을 제공한다: 바로 **순서 불가지론적 생성(order-agnostic generation)**이다. 토큰을 엄격하게 왼쪽에서 오른쪽으로 생성하는 AR 모델과 달리, MDM은 여러 위치를 병렬로 디노이징할 수 있어 적응형 생성 순서를 허용한다. 그러나 현재의 디코딩 방식은 이러한 유연성을 충분히 활용하지 못하고 있다.

기존의 디코딩 전략은 일반적으로 휴리스틱한 리마스킹(remasking) 규칙에 의해 제어되는 **단일 궤적(single trajectory)**에 고착된다:

저신뢰도 리마스킹(Low-Confidence Remasking): 신뢰도가 높은 토큰을 조기에 결정적으로 고정한다. 이는 종종 강력한 단일 샘플 정확도(pass@1)를 제공하지만, 탐색 범위를 좁은 궤적 세트로 제한하여 멀티 샘플 커버리지(pass@k)의 성장을 더디게 하고 잘못된 지역 최적점(local optima)에 빠질 위험을 초래한다.
무작위 리마스킹(Random Remasking): 궤적 다양성을 극대화하기 위해 리마스킹 패턴을 균등하게 샘플링한다. 이는 더 넓은 솔루션 공간을 탐색함으로써 pass@k를 개선하지만, 경쟁하는 후보들 사이에서 올바른 궤적을 선택할 원칙적인 메커니즘이 부족하여 약한 pass@1을 보인다.

이는 탐사(exploration)–활용(exploitation)의 트레이드오프를 생성한다: 현재의 방법들은 지나치게 활용에 치중하여(유효한 솔루션을 놓침) 혹은 맹목적으로 탐사하여(최적의 솔루션을 선택하는 데 실패함) 양극단에 위치한다. 본 논문은 MDM이 다양한 생성 순서와 토큰 값을 탐색할 수 있는 잠재력이 있음에도 불구하고, 표준 디코딩이 이 결합된 공간에서 명시적인 탐색을 수행하지 않기 때문에 그 잠재력이 실현되지 못하고 있다고 주장한다.

방법론: 순서-토큰 탐색 (Order-Token Search)

이를 해결하기 위해 저자들은 생성 순서와 토큰 값의 결합 공간에서 빔 서치(beam search)를 수행하는 디코딩 알고리즘인 순서-토큰 탐색을 제안한다. 이 알고리즘은 후보 궤적을 확장하는 단계와 유망하지 않은 궤적을 가지치기(pruning)하는 단계를 교대로 수행한다.

1. 탐색 과정 (확장)

순서-토큰 탐색은 $K$ 개의 부분적으로 디노이징된 후보들을 유지한다. 고정된 결정 시퀀스를 따르는 표준 샘플링과 달리, 이 알고리즘은 사용자가 지정한 간격에 따라 두 가지 불확실성 소스에 대해 명시적으로 분기한다:

생성 순서: 다음에 어떤 위치(토큰)를 언마스킹(unmask)할 것인가.
토큰 값: 해당 위치에 어떤 토큰을 할당할 것인가.

계산 복잡도를 관리하기 위해, 탐색은 **블록 확산(block diffusion)**을 활용한다. 모든 디노이징 단계마다 분기하는 대신(이는 과도한 오버헤드를 발생시킴), 알고리즘은 연속적인 토큰 블록의 경계에서 후보를 확장한다. 각 블록 내에서는 독립적인 디노이징이 일어난다. 이는 결정 지점을 전체 시퀀스 길이에서 블록의 수로 줄여줌으로써, 추론 비용을 관리 가능한 수준으로 유지하면서도 순서와 토큰 할당 모두에서 상당한 다양성을 확보할 수 있게 한다.

2. 가지치기 과정 (우도 추정)

핵-혁신은 부분 궤적을 점수화하고 가지치기하는 데 사용되는 **우도 추정기(likelihood estimator)**에 있다. 표준 MDM 목적 함수(노이즈가 섞인 상태로부터 깨끗한 전체 시퀀스를 재구성할 확률)는 많은 마스크가 남아 있는 중간 상태에서 불안정하다. 이는 모델이 훈련 중에 최적화되지 않은 어려운 인필링(infilling) 문제를 해결하도록 강요하기 때문이다.

저자들은 **증분적 액션 레벨 점수 함수(incremental action-level scoring function)**를 도입한다:

전체 상태를 점수화하는 대신, 알고리즘은 더 오염된 상태 $x_t$ 에서 덜 오염된 상태 $x_s$ 로 전이하기 위해 취해진 특정 디노이징 액션을 점수화한다.
점수 $s(x_t; x_s)$ 는 모델의 전체 시퀀스 예측인 $x_0$ (컨텍스트 역할을 함)를 조건으로 하여, $x_s$ 에서 새롭게 드러난 토큰들을 모델이 얼마나 자신 있게 재예측할 수 있는지를 측정한다.
이 접근 방식은 궤적의 과거 결정들(순서 및 토큰 선택)이 모델에 의해 뒷받침되는지를 평가하며, 이는 가지치기를 위한 더 낮은 분산과 더 변별력 있는 신호를 제공한다.

후보의 누적 점수는 탐색 간격 동안의 이러한 증분 점수들의 합이다. 알고-리즘은 각 확장 단계 이후 상위 $K$ 개의 후보를 유지하며, 이를 통해 디코딩을 계속하기 전에 잘못되었거나 차선인 경로를 효과적으로 필터링한다.

주요 기여

트레이드오프 식별: 본 논문은 기존 MDM 디코딩 방법들이 탐사-활용 스펙트럼의 반대편 끝에 위치함을 실증적으로 보여준다. 저신뢰도 리마스킹은 과잉 활용(높은 pass@1, 낮은 pass@k 성장)을 하는 반식이고, 무작위 리마스킹은 맹목적 탐사(높은 pass@k, 낮은 pass@1)를 한다.
순서-토큰 탐색 알고리즘: 생성 순서와 토큰 값의 결합 공간을 탐색하는 새로운 디코딩 알고리즘을 제안한다. 이 알고리즘은 안정적인 증분 우도 추정기를 사용하여 부분 궤적을 가지치기함으로써, 궤적의 다양성을 향상된 단일 샘플 정확도로 전환한다.
실증적 검증: 이 방법은 수학적 추론 및 코딩 벤치마크에서 베이스라인을 일관되게 능가하거나 대등한 성능을 보여준다. 이는 diffu-GRPO와 같은 비용이 많이 드는 사후 학습(post-training) 방법들을 통해 얻는 이득과 일치하거나 이를 상회한다.

실험 결과

저자들은 LLaDA-8B-Instruct와 LLaDA-1.5(RL 사후 학습된 변체)를 사용하여 GSM8K, MATH500, Countdown, HumanEval에서 순서-토킨 탐색(OTS)을 평가했다.

성능 향상: OTS는 벤치마크 전반에서 일관되게 가장 높은 평균 정확도를 달성했다.
- LLaDA-8B에서 OTS는 저신뢰도 리마스킹 베이스라인 대비 GSM8K에서 +3.1%, MATH500에서 +3.8%, Countdown에서 +7.9%, HumanEval에서 **+6.8%**의 pass@1 향상을 보였다.
- LLaDA-1.5에서도 유사한 상대적 개선을 보이며, 최상의 베이스라인(Low-conf + MV)인 37.2%에 비해 **40.4%**의 평균 정확도에 도달했다.
사후 학습과의 비교: OTS는 강화 학습 사후 학습(diffu-GRTP)을 거친 모델들의 성능과 일치하거나 이를 능가했다. 이는 추론 시간 탐색이 모델 재학습만큼 효과적일 수 있음을 시사한다.
절제 연구(Ablation Studies):
- 우도 추정기: 특정 증분 점수 함수(식 2)가 결정적이었다. 모든 드러난 블록의 결합 우도나 AR 스타일의 우도를 사용한 변체들은 현저히 낮은 성능을 보였으며, 이는 맞춤형 추정기가 MDM에서의 안정적인 가지치기에 필수적임을 확인시켜 준다.
- 탐색 공간: OTS는 AR 스타일의 우도를 사용하는 "순서 탐색(Order Search, 순서만 탐색)" 및 "토큰 탐색(Token Search, 토큰만 탐색)"보다 우수한 성능을 보였다. 이는 결합 탐색의 필요성과 AR 휴리스틱이 MDM과 맞지 않음을 강조한다.
효율성: OTS는 다수결 투표(majority voting)보다 테스트 시간 연산량(NFE)에 따라 더 효과적으로 확장된다. 동일한 연산 예산에서 OTS는 구조화된 탐색을 활용함으로써 (예: Countdown에서 29.3% vs 19.9%) 더 높은 정확도를 달 Achieved 했다.

의의 및 주장

본 논문은 결합 순서-토큰 탐색이 DLM 디코딩을 발전시키기 위한 핵심 요소라고 주장한다. 생성 순서와 토큰 값을 명시적으로 탐색하고 안정적인 우도 추정기를 기반으로 가지치기를 수행함으로써, 이 방법은 확산 모델의 내재적인 다중 궤적 잠재력을 신뢰할 수 있는 단일 샘플 성능으로 전환한다.

저자들은 이 연구를 현재 MDM 디코딩의 "탐사-활용" 병목 현상에 대한 해결책으로 제시한다. 그들의 접근 방식은 추가적인 모델 훈련 없이도, 탐욕적 방법이 빠지는 지역 최적점에서 벗어나고 무작위 샘플링이 놓치는 정답을 선택할 수 있게 함으로써 DLM의 능력을 높일 수 있다고 단언한다. 이 연구는 특히 병렬적이고 순서 불가지론적인 확산 모델의 특성에 맞춘 추론 시간 탐색 전략이 추론 및 코딩 작업에서 체계적인 개선을 가져올 수 있음을 입증한다.