"Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search" 논문 리뷰

이 논문 "Wider or Deeper? Scaling LLM Inference-Time Compute with Adaptive Branching Tree Search"는 대규모 언어 모델(LLM)의 추론 시간 계산(inference-time computation)을 확장하여 복잡한 작업에 대한 성능을 향상시키는 새로운 프레임워크인 Adaptive Branching Monte Carlo Tree Search (AB-MCTS)를 제안합니다.

서론 및 배경

최근 연구에 따르면 추론 시 계산량을 늘리는 것이 LLM의 복잡한 작업 성능을 크게 향상시킬 수 있음이 입증되었습니다. 이러한 추론 시간 스케일링 접근 방식은 크게 세 가지로 분류됩니다: (1) 후학습 미세 조정(post-training fine-tuning), (2) 보상 기반 사고 연쇄(chain-of-thought, CoT) 생성, (3) 다중 답변 생성. 이 논문은 세 번째 범주인 다중 답변 생성에 중점을 둡니다.

반복 샘플링(Repeated Sampling)은 이 범주에서 가장 성공적인 접근 방식으로, LLM이 0이 아닌 온도(non-zero temperature)에서 여러 후보 출력을 독립적으로 생성하고 가장 유망한 것을 선택하는 방식입니다. 이는 LLM 생성의 다양하고 방대한 출력 공간을 활용하여 높은 품질의 답변을 얻을 확률을 높이는 데 효과적입니다. 코딩 대회 및 ARC-AGI와 같은 벤치마크에서 그 효과가 입증되었습니다.

그러나 반복 샘플링은 탐색(exploration)에만 초점을 맞추고 활용(exploitation)을 위한 명시적인 메커니즘이 부족합니다. 코딩 작업에서 테스트를 실행하여 생성된 코드의 정확성에 대한 외부 피드백을 얻을 수 있는 것처럼, 특정 실제 시나리오에서는 후보 솔루션에 대한 외부 피드백을 얻을 수 있습니다. 이러한 환경에서는 유망한 솔루션을 선택하고 사용 가능한 피드백을 기반으로 개선하는 것이 자연스러운데, 반복 샘플링만으로는 이를 효과적으로 수행할 수 없습니다.

기존의 몬테카를로 트리 탐색(MCTS) 기반 접근 방식은 탐색과 활용을 모두 통합하지만, 고정된 분기 요소(fixed branching factor)를 사용합니다. LLM의 다양하고 방대한 출력 공간을 활용하는 것이 효과적인 추론 시간 스케일링에 중요하기 때문에, 고정된 너비는 스케일링을 저해할 수 있습니다.

AB-MCTS의 핵심 아이디어 및 방법론

이러한 한계를 극복하기 위해 제안된 AB-MCTS는 반복 샘플링을 다중 턴(multi-turn) 탐색 및 활용으로 일반화하는 새로운 추론 시간 프레임워크입니다.

1. 동적인 "더 넓게" 또는 "더 깊게" 결정: AB-MCTS는 탐색 트리의 각 노드에서 외부 피드백 신호에 기반하여 새로운 후보 응답을 확장하여 "더 넓게"(탐색) 갈 것인지, 아니면 기존 응답을 재검토하여 "더 깊게"(활용) 갈 것인지를 동적으로 결정합니다. 이는 순수하게 넓게 가는 반복 샘플링이나 순수하게 깊게 가는 순차 개선(Sequential Refinement)과는 차별화됩니다.
2. 무한한 분기(Unbounded Branching) 허용: 기존 MCTS와 달리 AB-MCTS는 너비를 고정된 하이퍼파라미터로 설정하지 않습니다. 이는 LLM이 0이 아닌 온도에서 동일한 프롬프트로부터 다른 출력을 생성할 수 있기 때문에 분기 요소가 이론적으로 무한하다는 점을 반영한 것입니다. 이를 위해 GEN 노드(GEN node)라는 개념을 도입하여 이미 한 번 확장된 노드도 다시 확장하고 추가로 분기할 수 있도록 합니다.
3. 베이즈 사후 업데이트(Bayesian Posterior Updates) 기반 의사 결정: AB-MCTS는 베이즈 사후 업데이트를 통해 의사 결정 과정을 공식화하여, 각 확장이 탐색과 활용의 균형을 원칙적인 방식으로 맞출 수 있도록 합니다.
4. 노드 선택을 위한 톰슨 샘플링(Thompson Sampling): 표준 MCTS의 UCT 점수(UCT score)는 GEN 노드에는 할당할 수 없으므로, AB-MCTS는 노드 선택에 톰슨 샘플링을 사용합니다. 톰슨 샘플링은 모델링된 확률 분포에서 점수를 샘플링한 다음 가장 높은 샘플링 점수를 가진 액션을 선택하는 방식으로 진행됩니다.
   • AB-MCTS-M (Mixed Model): 각 노드에서 별도의 노드별 혼합 모델을 사용하여 GEN 노드의 점수를 추정합니다. 혼합 모델에서는 각 노드에 "그룹" 라벨을 할당하고, 동일한 라벨을 공유하는 노드는 모델 파라미터를 공유합니다. 이는 LLM의 응답 다양성에서 발생하는 답변 품질의 변동성을 포착합니다. GEN 노드는 새로운 그룹으로 처리되며, 관찰된 점수가 없으므로 그 사후 분포는 더 불확실하여 탐색을 장려합니다.
   • AB-MCTS-A (Node Aggregation): AB-MCTS-M과 달리 다른 액션 간에 공유되는 모델 파라미터가 없습니다. 모든 자식 노드를 단일 CONT 노드(CONT node) 아래에 통합하여 기존 자식 노드로부터의 개선을 나타냅니다. 확장된 노드의 점수는 GEN 노드에 백업되어 이전 확률 분포를 업데이트하는 데 사용됩니다. 가우시안(Gaussian) 또는 베타(Beta) 분포를 사용하여 점수를 모델링할 수 있습니다.
5. AB-MCTS는 이 노드 선택을 위해 두 가지 전략을 제안합니다:

실험 및 결과

AB-MCTS는 복잡한 코딩 및 기계 학습 엔지니어링 벤치마크, 그리고 ARC-AGI에서 GPT-4o 및 DeepSeek-V3와 같은 최신 모델을 사용하여 평가되었습니다.

1. 벤치마크: LiveCodeBench, CodeContest, ARC-AGI, MLE-Bench. 이 벤치마크들은 LLM이 파이썬 코드를 생성하여 작업을 해결하고, 외부 피드백(예: 테스트 케이스 결과)이 탐색을 안내하는 데 사용됩니다.
2. 모델 및 예산: GPT-4o와 DeepSeek-V3를 사용했으며, 최대 API 호출 예산은 2^7 = 128회로 설정되었습니다.
3. 기준선: 반복 샘플링(Best-of-n), 순차 개선(Sequential Refinement), 표준 MCTS와 비교되었습니다.

주요 결과:

• 전반적인 우수성: AB-MCTS는 다양한 벤치마크와 LLM에 걸쳐 일관되게 우수한 성능을 보이며, 확립된 기준선들을 능가하는 최고 평균 순위를 달성했습니다. 이는 AB-MCTS가 각 문제의 다양한 요구 사항에 따라 탐색 전략을 동적으로 조정하여 탐색과 활용의 균형을 정확하게 맞추는 독특한 능력에서 비롯됩니다.
• LiveCodeBench 및 CodeContest: AB-MCTS 알고리즘은 일반적으로 기준선보다 우수한 성능을 보였습니다. 특히 LiveCodeBench에서는 적은 예산으로도 기준선보다 앞서기 시작했으며, CodeContest에서는 더 큰 예산에서 우수한 성능을 보였습니다.
• ARC-AGI: 반복 샘플링이 강한 기준선으로 나타났음에도 불구하고, AB-MCTS 프레임워크는 이에 필적하는 성능을 달성했습니다. 이는 AB-MCTS가 필요할 때 탐색을 동적으로 확장함으로써 효과적으로 탐색할 수 있음을 시사합니다. 예산을 2^9 = 512까지 늘린 추가 실험에서는 반복 샘플링의 개선율이 정체되기 시작하는 반면, AB-MCTS의 성능은 지속적으로 크게 향상되어 더 큰 계산 규모에서 탐색을 유망한 분기로 더 효과적으로 유도함을 보여주었습니다.
• MLE-Bench: AB-MCTS-M은 다양한 ML 작업에서 가장 강력한 평균 순위를 달성하며 견고한 성능을 입증했습니다. 이는 AB-MCTS-M이 다양한 문제 구조에 맞춰 탐색 전략을 효과적으로 조정하는 내재된 강점을 강조합니다.
• 탐색 행동 분석 (너비 vs. 깊이): AB-MCTS 방법은 표준 MCTS보다 더 넓은 트리를 생성하는 경향을 보였습니다. 이는 AB-MCTS가 표준 MCTS와 달리 기존 노드에서 "더 넓게" 탐색할지(GEN 노드 선택) 동적으로 결정할 수 있기 때문입니다. 이는 탐색과 활용의 강점을 결합하는 적응적인 특성을 보여줍니다.

다중 LLM AB-MCTS (Multi-LLM AB-MCTS)

이 논문은 AB-MCTS를 여러 LLM을 활용하는 시나리오로 확장하는 방법을 제시합니다. LLM A가 더 다양한 초기 답변을 생성하고 LLM B가 개선에 탁월한 경우처럼, 두 모델을 모두 사용하여 답변 트리를 구축하면 성능 향상을 기대할 수 있습니다.

1. 생성기 선택 알고리즘:
   • 단일 GEN 노드: 노드 선택 후 추가적인 생성기 선택 단계가 있어, 전체 답변 트리 내에서 해당 생성기가 생성한 노드들의 점수를 기반으로 사후 분포를 계산하고 톰슨 샘플링을 통해 하나의 생성기를 선택합니다.
   • 다중 GEN 노드: 트리의 모든 노드에 사용 가능한 각 생성기별로 여러 GEN 노드를 자식으로 연결하여, 탐색 트리의 지역적 맥락을 더 잘 포착하고 솔루션 프로세스의 각 특정 단계에서 가장 적합한 생성기를 적응적으로 선택할 수 있도록 합니다.
2. ARC-AGI-2 실험: ARC-AGI-2는 더 높은 수준의 인지 능력을 평가하도록 설계된 벤치마크로, 최전선 LLM에게도 해결하기 어려운 문제입니다 (성공률 5% 미만).
   • 결과: Multi-LLM AB-MCTS는 Gemini-2.5-Pro와 DeepSeek-R1-0528을 통합하여 성능을 더욱 향상시켰고, 궁극적으로 30% 이상의 문제에 대한 올바른 솔루션을 찾아냈습니다. DeepSeek-R1-0528이 개별 성능은 낮았지만, Multi-LLM 프레임워크에 통합되면서 해결된 문제의 수가 증가했습니다.
   • 분석: Multi-LLM AB-MCTS 프레임워크는 문제 특성에 따라 다른 LLM을 효과적으로 할당하는 능력을 입증했습니다. 높은 보상을 빠르게 달성한 시도(쉬운 문제)에는 더 능숙한 모델이 할당되는 경향이 있었고, 높은 보상을 얻기 어려웠던 시도(어려운 문제)에는 모델들이 더 균형 있게 활용되었습니다. 또한, 어떤 단일 LLM으로도 해결할 수 없었던 문제들이 여러 모델의 협업을 통해 해결되는 사례가 관찰되었습니다. 이는 단순히 최상의 모델을 문제에 맞추는 것을 넘어선 시너지 효과를 시사합니다.

결론 및 한계

AB-MCTS는 외부 피드백에 기반한 베이즈 의사 결정을 통해 '더 넓게' 또는 '더 깊게'를 동적으로 결정하여, LLM의 복잡한 작업 성능을 향상시키는 새로운 추론 시간 프레임워크입니다. 실험 결과는 AB-MCTS가 반복 샘플링 및 표준 MCTS를 능가하며, 효과적인 추론 시간 스케일링을 위해 무한한 분기 문제를 적응적으로 처리하는 중요성을 입증했습니다.

한계점: 이 접근 방식은 신뢰할 수 있는 점수 평가자(score evaluator)의 존재를 전제로 하지만, 이러한 평가자 자체를 개발하는 것이 작업에 따라 어려울 수 있습니다. 향후 연구에서는 API 호출 수 이외의 더 세분화된 실제 비용 요소를 통합하는 탐색 전략을 탐색하여 AB-MCTS의 실용성을 더욱 높일 수 있습니다. 또한, Pass@k와 Pass@2 간의 성능 격차를 줄이기 위해 더 정교한 최종 답변 선택 알고리즘을 개발하는 것이 필요합니다 (예: 더 정확한 보상 모델 구축 또는 LLM을 판사로 활용).