Who Should Lead Decoding Now? Tracking Reliable Trajectories for Ensembling Masked Diffusion Language Models

기술 요약: 누가 지금 디코딩을 주도해야 하는가? 마스크드 확산 언어 모델(MDLM) 앙상블을 위한 신뢰할 수 있는 궤적 추적

문제 정의
마스크드 확산 언어 모델(MDLM)은 순차적인 자기회귀 방식이 아닌, 반복적인 재마스킹(remasking)을 통해 시퀀스를 병렬적으로 노이즈 제거(denoising)하는 독특한 패러다임으로 부상했습니다. 다양한 강점과 디코딩 역학을 가진 모델들이 등장함에 따라, 이질적인 MDLM들로부터 지식을 효과적으로 조율하거나 융합하는 방법이 중요한 과제로 떠오르고 있습니다. 차기 토큰 확률을 평균하거나 가장 확신이 높은 모델으로 라우팅하는 기존의 자기회귀 모델용 앙상블 전략은 MDLM에 직접 적용하기 어렵습니다. MDLM에서 시퀀스는 유연하고 비선형적인 순서로 생성되므로, 서로 다른 모델들은 각 단계에서 서로 다른 부분적으로 노이즈가 제거된 시퀀스 상에서 작동할 수 있습니다. 이러한 '공통된 다음 토큰' 위치의 부재는 표준적인 앙상블 정의를 적용하기 어렵게 만들며, 이로 인해 MDLM 앙상블 연구는 여전히 미개척 영역으로 남아 있습니다.

방법론: TIE (궤적 기반 반복 앙상블)
이를 해결하기 위해 저자들은 MDLM에 특화된 지식 융합 프레임워크인 TIE를 제안합니다. TIE는 디코딩 역학 분석에서 도출된 두 가지 핵심 관찰 결과에 기반하여 작동합니다:

1. 정답 궤적의 안정성: 정답을 맞힌 샘플은 토큰이 여전히 마스킹된 상태일지라도, 정답 관련 토큰(answer-relevant tokens)에 대해 오답 샘플보다 더 안정적이고 확신 있는 역학(낮은 엔트로피, 높은 top-1 확률, 적은 토큰 변화)을 보입니다.
2. 주입을 통한 교정 가능성: 오답을 생성하는 모델들도 만약 더 신뢰할 수 있는 궤적을 따르는 모델로부터 유망한 부분적 노이즈 제거 중간 상태를 전달받는다면, 정답으로 유도될 수 있습니다.

이러한 통찰을 바탕으로, TIE는 생성 과정 중에 다음과 같은 반복적인 3단계 사이클을 구현합니다:

1. 궤적 생성 (Trajectory Generation): 각 구성 모델($M$)은 고정된 단계 수(앙상블 간격, $n$) 동안 독립적으로 디코딩하여 부분적인 디코딩 궤적을 생성합니다.
2. 궤적 평가 (Trajectory Assessment): 프레임워크는 정답 토큰 위치(예: "Answer:"와 같은 접미사 뒤에 나타나는 토큰)에 대한 확신 역학을 점수화하여 각 궤적의 신뢰도를 평가합니다. 점수 산출에는 네 가지 지표가 사용됩니다:
   • 음의 토큰 변화 횟수 (Negative Token Change Count): 마스킹된 정답 위치에 대한 top-1 토큰 변화의 이력을 추적합니다. 낮은 횟수는 높은 안정성을 의미합니다.
   • Top-1 확률 (Top-1 Probability): 가장 가능성이 높은 토큰의 확률입니다.
   • 음의 엔트로피 (Negative Entropy): 분포의 확실성을 측정합니다.
   • 확률 마진 (Probability Margin): top-1과 top-2 확률 간의 격차입니다.
   • 교차 모델 점수 산출 (Cross-Model Scoring): 로짓(logit) 기반 지표의 경우, 프레임워크는 각 궤적을 모든 구성 모델에 통과시켜 점수를 평균합니다. 이는 모델 간의 캘리브레이션 차이를 완화하여, 선택된 궤적이 단순히 해당 모델에 의해 선호되는 것이 아니라 앙상블 전체에서 일관되게 확신을 갖도록 보장합니다.
3. 궤적 전달 (Trajectory Relay): 가장 높은 총합 점수를 받은 궤적이 선택됩니다. 이 "승리한" 궤적은 모든 구성 모델로 전달되어, 그들의 현재 부분 상태를 대체합니다. 모델들은 이 공유된, 더 신뢰할 수 있는 중간 상태로부터 독립적인 디코딩을 재개합니다.

이 사이클은 생성이 종료될 때까지 반복됩니다. 마지막으로, 시스템은 가장 안정적인 정답 토큰 역학(가장 낮은 정규화된 토큰 변화 횟수)을 가진 궤적을 기준으로 $M$개의 후보 중 최종 응답을 선택합니다.

주요 기여

• 새로운 디코딩 분석: 본 논문은 정답을 맞힌 MDLM 생성에서 정답 관련 토큰이 안정적인 확신 궤적을 따른다는 점을 식별하여, 전체 응답이 생성되기 전에 신뢰할 수 있는 경로를 식별할 수 있는 신호를 제공합니다.
• MDLM 전용 앙상블 프레임워크: TIE는 부분 궤적을 반복적으로 전달함으로써 지식을 융합하는 메커니즘을 도입하며, 이를 통해 모델들이 생성 과정 중에 서로의 편차를 교정할 수 있도록 합니다.
• 교차 모델 점수 산출 메커니즘: 저자들은 과도하게 확신하는 모델에 의한 편향된 선택을 방지하기 위해, 궤적을 모든 구성 모델에 걸쳐 평가하는 것의 필요성을 입증했습니다.

실험 결과
저자들은 네 가지 도메인에 걸친 8개의 벤치마크(일반 추론: MMLU, ARC-Challenge, Winogrande; 수학적 추론: GSM8K, MATH500; 코딩: HumanEval, MBPP; 계획: Countdown)에서 TIE를 평가했습니다.

• 성능 향상: TIE는 개별 MDLM(LLaDA-1.5, Dream-7B, DreamCoder-7B, DiffuCoder-7B) 및 사후 생성 앙상블 베이스라인보다 일관되적으로 우수한 성능을 보였습니다. 예를 들어, GSM8K에서 TIE를 통한 LLaDA와 Dream의 앙상블은 정확도를 78.39%(Dream)에서 83.47%로 향상시켰습니다.
• 점수 지표: 점수 함수 중 토큰 변화 횟수(특히 정규화된 버전)가 가장 견고한 것으로 나타났으며, 8개 벤치마크 중 4개에서 가장 좋은 결과를 달성했습니다.
• 모델 시너지: 프레임워크는 구성 모델들의 역량이 비슷하고 강력할 때 가장 효과적입니다. 성능 차이가 클 경우(예: >15%), 약한 모델이 노이즈 섞인 신호를 유입시켜 앙상블의 효과가 떨어집니다.
• 동적 리더십: 분석 결과, 가장 높은 점수의 궤적을 생성하는 모델은 디코딩 과정 내내 빈번하게 변경되었습니다(태스크에 따라 모델 변화율 12~44%). 이는 서로 다른 모델이 생성의 각기 다른 단계에서 보완적인 강점을 기여함을 확인시켜 줍니다.
• 호환성: TIE는 임계값 설정(thresholding) 및 top-$k$ 언마스킹과 같은 표준 MDLM 가속 기술과 결합했을 때도 효과적입니다.

의의 및 주장
본 논문은 TIE가 추가 학습 없이도 MDLM 앙상블이라는 미개척 문제에 대한 실용적인 접근법을 제공한다고 주장합니다. 정답 관련 토큰에 대한 확신 역학을 활용함으로써, TIE는 이질적인 모델들이 협력적으로 응답을 정교화할 수 있게 하며, 약한 모델이 하위 수준의 궤적에서 회복할 수 있도록 합니다. 저자들은 이 연구가 확산 기반 언어 모델의 보다 효과적인 오케스트레이션을 향한 의미 있는 단계라고 위치시키며, 궤적 안정성에 의해 유도되는 협력적 추론이 다양한 작업에서 추론 능력을 크게 향상시킬 수 있음을 시사합니다. 또한, 모델 간의 성능 차이가 클 때 효과가 감소한다는 점과 현재 연구가 두 모델의 앙상블에 국한되어 있다는 점 등 한계점도 겸허히 언급하였습니다.