Adaptive Guidance for Retrieval-Augmented Masked Diffusion Models

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 대규모 언어 모델 (LLM) 은 정적 파라미터 지식에 의존하여 환각 (Hallucination) 이나 구식 정보를 생성하는 문제가 있습니다. 이를 해결하기 위해 외부 지식을 활용하는 **검색 증강 생성 (RAG)**이 널리 사용되지만, 검색된 문맥이 노이즈가 있거나 모델의 내부 지식과 상충될 때 성능이 저하되는 지식 충돌 (Knowledge Conflict) 문제가 발생합니다.
• 현재의 한계:
  • 기존 RAG 연구는 주로 자기회귀 (Autoregressive, AR) 모델에 집중되어 있습니다.
  • 최근 주목받는 **마스크 확산 모델 (Masked Diffusion Models, MDM)**은 이터레이션 기반의 비 causal(양방향) 생성 방식을 가지며, AR 모델보다 컨텍스트 활용 능력이 뛰어나다는 잠재력이 있지만, RAG 환경에서의 지식 충돌 해결은 거의 연구되지 않았습니다.
  • 기존 확산 모델에 적용되는 **클래스프리 가이드 (Classifier-Free Guidance, CFG)**는 고정된 가이드 스케일 ($\lambda$) 을 사용합니다. 이는 검색된 문맥이 신뢰할 수 없을 때 (노이즈) 과도하게 가이드하여 환각을 유발하거나, 신뢰할 수 있는 문맥일 때 가이드가 부족하여 오류를 수정하지 못하는 딜레마를 초래합니다.

2. 제안 방법: ARAM (Methodology)

저자들은 적응형 검색 증강 마스크 확산 (Adaptive Retrieval-Augmented Masked Diffusion, ARAM) 프레임워크를 제안합니다. 이는 학습이 필요 없는 (Training-free) 적응형 가이드 메커니즘으로, **신호 - 대 - 잡음비 (Signal-to-Noise Ratio, SNR)**를 기반으로 각 디노이징 (Denoising) 단계와 토큰 위치에서 가이드 스케일을 동적으로 조절합니다.

핵심 아이디어 및 수학적 유도

1. 정보 이론적 관점: 검색된 문맥의 영향력을 '사전 분포 (Prior)'에서 '조건부 분포 (Posterior)'로의 분포 변화로 해석합니다.
2. 신호 (Signal) 와 잡음 (Noise) 의 정의:
   • 신호: 검색된 문맥이 모델의 믿음을 얼마나 효과적으로 변화시키는지 (정보 획득량). $DKL(p_{cond} || p_{prior}) + DKL(p_{prior} || p_{cond})$로 계산됩니다.
   • 잡음: 검색된 문맥이 가져오는 불확실성. 기존 연구에서 제안된 분산 (Variance) 대신, **조건부 엔트로피 ($H(p_{cond})$)**를 잡음의 대용량 (Proxy) 으로 사용합니다. 이는 문맥이 명확할 때 엔트로피가 낮아지고, 모호하거나 상충될 때 높아지는 특성을 이용합니다.
3. 적응형 가이드 스케일 ($\lambda_t$) 계산:
   • 최적의 가이드 스케일 $\lambda^*$를 SNR ($\frac{Signal}{Noise}$) 로 근사화합니다.
   • 실제 적용 시 안정성을 위해 다음과 같이 조정합니다:
     $$\lambda_t = \lambda_{max} \cdot \tanh\left(\beta \cdot \frac{Signal}{Noise + \epsilon}\right)$$
   • 동작 원리:
     • 신뢰할 수 있는 문맥 (Gold Doc): 신호가 강하고 잡음 (엔트로피) 이 낮아 $\lambda_t$가 커짐 $\rightarrow$ 검색된 정보를 강력하게 반영.
     • 노이즈/상충 문맥: 신호가 약하거나 잡음 (엔트로피) 이 높아 $\lambda_t$가 작아짐 $\rightarrow$ 모델의 내부 지식을 유지하며 환각 방지.

3. 주요 기여 (Key Contributions)

1. ARAM 프레임워크 제안: RAG 환경의 MDM 을 위한 최초의 적응형 CFG 프레임워크로, SNR 기반의 동적 가이드 스케일 조절을 통해 지식 충돌을 해결합니다.
2. 정보 이론적 분석: 확산 기반 생성에서 검색 문맥과 파라미터 사전 지식 간의 상호작용을 정보 이론적 관점 (상호 정보량, 엔트로피) 으로 분석하고, 신호와 불확실성에 기반한 가이드 규칙을 유도했습니다.
3. 광범위한 실험 검증: 다양한 지식 집약적 QA 벤치마크에서 ARAM 이 기존 RAG 베이스라인 및 다른 디코딩 전략 (CAD, COIECD, SPREAD 등) 보다 우수한 성능을 보임을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋: NQ, TriviaQA, MARCO QA, HotpotQA, T-REx 등 5 가지 지식 집약적 QA 데이터셋.
• 모델: LLaDA, Dream (MDM) 및 LLaMA, Qwen (AR 모델) 과 비교.
• 성능 향상:
  • 정확도 (EM) 및 F1 점수: ARAM 은 모든 벤치마크에서 기존 고정 가이드 RAG 및 다른 적응형 방법들보다 높은 EM 과 F1 점수를 기록했습니다. (예: LLaDA 기반에서 EM 19.3 $\rightarrow$ 30.3, Dream 기반에서 EM 35.9 $\rightarrow$ 37.7)
  • 지식 충돌 해결: 검색된 문맥이 정답을 포함하는 경우 (Positive) 에는 가이드를 강화하여 정확도를 높이고, 문맥이 오답을 유도하거나 관련 없는 경우 (Negative) 에는 가이드를 억제하여 환각을 줄였습니다.
  • 비교 분석: AR 모델용 적응형 방법 (CAD 등) 을 MDM 에 적용하면 성능이 저하되는 반면, ARAM 은 MDM 의 이터레이션 특성에 맞춰 토큰별/단계별 가이드를 조절하여 최적의 성능을 냈습니다.

5. 의의 및 결론 (Significance)

• MDM 의 RAG 적용 가능성 증대: 확산 기반 언어 모델이 RAG 환경에서도 외부 지식을 효과적으로 활용할 수 있음을 보여주었습니다.
• 동적 신뢰도 조절: 고정된 가이드 스케일의 한계를 극복하고, 검색된 정보의 신뢰도에 따라 모델이 스스로 '믿을지 말지' 결정하는 메커니즘을 제시했습니다.
• 실용성: 학습이 필요 없는 (Training-free) 방식이므로 기존 MDM 모델에 즉시 적용 가능하지만, 매 단계에서 전체 어휘에 대한 계산을 수행해야 하므로 추론 시 약간의 지연 (Latency) 이 발생한다는 점은 한계로 남습니다.

결론적으로, ARAM 은 검색된 문맥의 품질과 불확실성을 실시간으로 평가하여 생성 과정을 최적화함으로써, RAG 기반의 확산 언어 모델이 환각을 줄이고 사실 기반 생성의 정확도를 높이는 데 중요한 기여를 합니다.