[논문리뷰] MITS: Enhanced Tree Search Reasoning for LLMs via Pointwise Mutual Information

링크: 논문 PDF로 바로 열기

저자: Jiaxi Li, Yucheng Shi, Jin Lu, Ninghao Liu

핵심 연구 목표

대규모 언어 모델(LLM)의 다단계 추론 과정에서 중간 단계의 품질을 효율적이고 신뢰성 있게 평가하고, 계산 비용이 높은 경로 탐색 문제를 해결하고자 합니다. 특히, 기존 방법론이 일반적인 추론 경로를 선호하는 한계를 극복하고 질문별로 특화된(question-specific) 추론 경로 를 식별하는 원칙적인 프레임워크를 개발하는 것을 목표로 합니다.

핵심 방법론

제안하는 Mutual Information Tree Search (MITS) 는 정보 이론적 원리를 활용하여 추론을 안내합니다. 주요 기술은 Pointwise Mutual Information (PMI) 기반의 스코어링 함수로, 중간 추론 단계의 품질을 평가하고 빔 서치(Beam Search) 를 통해 검색 트리를 확장합니다. 또한, 엔트로피 기반 동적 샘플링 전략 으로 불확실한 추론 단계에 계산 리소스를 적응적으로 할당하며, 최종 예측은 PMI 점수 와 예측 합의를 결합한 가중 평균 투표(Weighted Average Voting) 방식을 통해 이루어집니다.

주요 결과

MITS 는 다양한 추론 벤치마크에서 강력한 성능 향상을 보였습니다. 예를 들어, StrategyQA(QWEN2.5-3B) 데이터셋에서 MITS는 68.45% 의 정확도를 달성하여 CoT(47.34%) 대비 21.11% 향상되었고, ARC-Challenge(QWEN2.5-7B) 에서는 92.55% 로 최강 베이스라인인 rStar(87.24%) 를 5.31% 능가했습니다. 또한, MITS 는 64.41초 의 계산 비용으로 RAP(203.42초) 나 rStar(815.67초) 보다 훨씬 빠르면서도 더 높은 정확도를 제공하며 최적의 효율성을 입증했습니다.

AI 실무자를 위한 시사점

PMI 기반의 정보 이론적 접근법 은 LLM의 복잡한 추론 과정에서 비용이 많이 드는 룩어헤드 시뮬레이션 없이 추론 경로의 품질을 효율적으로 평가할 수 있는 강력한 방법을 제시합니다. 동적 샘플링 과 가중 평균 투표 는 LLM 시스템의 불확실성 관리와 예측 신뢰성 향상에 실용적인 도움을 주어, 실제 애플리케이션에서 LLM 기반 시스템의 성능과 견고성을 높일 수 있습니다. 이러한 방법론은 다양한 추론 태스크에 적용 가능한 일반화된 프레임워크 로서 AI 개발 및 응용에 새로운 방향을 제시합니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.