[논문 리뷰] AoT : Atom of Thoughts for Markov LLM Test-Time Scaling

Atom of Thoughts for Markov LLM Test-Time Scaling

Chain-of-Thought의 새로운 지평이라고 하여 AoT, Atom of Thoughts가 새롭게 제안되었다. Question에 대해 Decomposition / Contraction 과정을 반복하면서 상태를 업데이트하는 방식이라고 할 수 있을 것인데, 자세한 내용은 논문 내용을 보면서 같이 살펴보도록 하자.

Abstract

Large Language Models (LLMs) achieve superior performance through training-time scaling, and test-time scaling further enhances their bilities by conducting effective reasoning during inference. However, as the scale of reasoning increases, existing test-time scaling methods suffer from accumulated historical information, which not only wastes computational resources but also interferes with effective reasoning. To address this issue, we observe that complex reasoning progress is often achieved by solving a sequence of independent subquestions, each being self-contained and verifiable. These subquestions are essentially atomic questions, relying primarily on their current state rather than accumulated history, similar to the memoryless transitions in a Markov process. Based on this observation, we propose Atom of Thoughts (AOT), where each state transition in the reasoning process consists of decomposing the current question into a dependency-based directed acyclic graph and contracting its subquestions, forming a new atomic question state. This iterative decomposition-contraction process continues until reaching directly solvable atomic questions, naturally realizing Markov transitions between question states. Furthermore, these atomic questions can be seamlessly integrated into existing test-time scaling methods, enabling AOT to serve as a plug-in enhancement for improving reasoning capabilities. Experiments across six benchmarks demonstrate the effectiveness of AOT both as a standalone framework and a plug-in enhancement. Notably, on HotpotQA, when applied to gpt-4omini, AOT achieves an 80.6% F1 score, surpassing o3-mini by 3.4% and DeepSeek-R1 by 10.6%. The code will be available at https://github.com/qixucen/atom.

중요한 내용은 볼드체 처리 된 부분만 보면 될 것이라 생각한다. 실제 우리가 문제를 작은 문제로 나누어 생각하듯, directed acyclic graph(이하 DAG) 를 이용해서 질문을 atomic questions로 변경한 후 각각에 대한 해경릉 통해  최종 goal을 달성할 수 있다는 것이 주요한 contribution이 될 것이다.

이와 같이 Computational Cost가 낮아진다고 한다. MDP를 사용하기 때문이 아닐까 싶다는 생각이 든다.

Introduction

현재의 LLM들은 이제 데이터 규모와 품질을 높일 만큼 높인 상황에서, Test time Scaling(이 뭔지 모른다면 이 링크를 참조하기 바란다)을 통해 성능을 향상시키고 있다. 하지만 이런 Test Time Scaling은 추론 과정에서 계속해서 구조적인 dependency를 유지하여야 하기 때문에, 계속해서 정보가 누적되어 Computational Cost 측면에서 문제가 발생한다. 그래서 이 논문에서 제안하는 AoT 방식을 사용하면, 이 방식에서는  Markov Process에서의 memoryless transition과 비슷한 방법을 사용하여 훨씬 복잡도를 낮출 수 있다는 것이 주요 contribution이라고 생각을 한다. 방법론에 대해 이야기하기 전, 저자가 말하는 AoT의 장점을 보고 넘어가자. 과연 이게 그렇게 타당한 말일지는 후에 살펴보도록 하자.

1. Computational Resource에 대해 scaling을 진행할 때, 계속해서 과거의 정보를 유지할 필요가 없다.
2. 단순히 혼자 사용되는 것이 아닌, Test Time Scaling을 진행할 때 함께 진행될 수 있다.

Overview - Method of AoT

앞에서도 언급하였듯, AoT는 문제를 독립적인 하위 문제들로 분해(decomposition)한 후, 이후 이를 atomic unit으로 다시 contraction하는 방식을 반복하면서 Test-Time Scaling을 진행한다. 하나하나 살펴보도록 하자.

Decomposition(분해)

이 과정에서는 DAG 구조를 사용한다. 주어진 질문을 지금까지는 reasoning 과정을 통해(CoT) 해결하였다면, 이 방식은 전체 질문에 대해서 한 질문을 subquestions들의 그래프로 만들어낸다는 것이다.

DAG of Subquestions

질문을 여러 subquestion으로 분해 후, 각 question 사이를 edge로 이어 하나의 그래프를 만든다고 생각하면 될 것 같다. 이 방식에서, 전체 subquestion들에 대해서 저자는 질문을 Independent Subquestion, Dependent Subquestion으로 나누게 된다. 이는 이름만 봐도 얼추 알 수 있지만 '다른 질문이 먼저 해결되어야 해결될 수 있는가?', 그러므로 위 그래프의 개념을 이용하자면 ' Incoming Edge가 있는가? ' 의 개념이다. 이를 식으로 나타내면 아래와 같다.

이는 사람이 생각하는 과정과 상당히 유사하고, 이는 저자가 얻고자 하는 최종 목표이기도 하다. dependent하다는 것을 좀더 쉽게 설명해 보자면, 위의 식에서 'Q_j가 해결되지 않는다면 Q_i 또한 해결이 시작될 수 없다' 정도로 이해하면 될 것 같다.

수축(Contraction)

수축은, 말 그대로 이 문제를 해결해나가는 과정을 의미한다. 이는 사람이 실제로 생각하고 문제를 해결해 나가는 과정과 상당히 유사하다.

DAG가 모두 완성되게 된다면, 이제 이를 전체 Graph의 Leaf Node부터 해결해나가게 된다. 작은 질문이 해결되면, 거기에서 edge를 뻗어나가 그 다음 node를 계속해서 풀이해 나가는 과정이라고 보면 된다.

def contract(question: str, decompose_result: dict, independent: list, dependent: list):
    instruction = """
        You are a multiple choice question solver specializing in optimizing step-by-step reasoning processes. Your task is to optimize the existing reasoning trajectory into a more efficient, single self-contained question.
        
        For the original question: {question}
        
        Here are step-by-step reasoning process:
        {response}
        
        {sub_questions}
        
        Here are explanations of key concepts:
        1. self-contained: The optimized question must be solvable independently, without relying on any external information
        2. efficient: The optimized question must be simpler than the original, requiring fewer reasoning steps and having a clearer reasoning process (these steps are reduced because some solved sub-problems become known conditions in the optimized question or are excluded as incorrect explorations)
        
        Note: Since this is a multiple choice question, the optimized question must completely retain the options of the original question.
        
        You can freely reason in your response, but please enclose the your optimized question within <question></question> tags
    """
    sub_questions = """
        The following sub-questions and their answers can serve as known conditions:
        {independent}

        The descriptions of the following questions can be used to form the description of the optimized problem:
        {dependent}
        
        """
    answer = decompose_result["answer"]
    for sub_q in independent:
        sub_q.pop("depend", None)
    for sub_q in dependent:
        sub_q.pop("depend", None)
        
    sub_questions = sub_questions.format(independent=independent, dependent=dependent)
    return instruction.format(question=question, answer=answer, response=decompose_result["response"], sub_questions=sub_questions)

실제 코드 구현은 다음과 같이 되어 있다. independent, dependent 질문 2가지로 나누어 진행한다고 보면 된다. 아무래도 LLM 관련 코드다 보니 프롬프팅 형식으로 되어 있기는 하지만, 어떤 느낌인지는 알 수 있을 것이라 생각한다. 이 과정에서는 AoT가 CoT보다 적은 Computational Cost를 가지는 이유가 나오는데, 바로 Markov Property의 Memoryless 방식을 차용하기 때문이다. 그로 인해서 과거의 모든 정보가 부하가 걸리지 않고, 이전의 데이터와 업데이트된 서브그래프만 데이터에 남게 되면서 사용량이 적어진다는 것이다. 결국 최종 결과는 똑같이 나오기 때문에 과거 정보의 손실로 인한 정확도 감소 문제만 해결된다면 기존 test-time scaling 방식을 뒤바꿔 놓을 수 있지 않을까 하는 생각이 든다.

Integration

이건 뭐 전체 구조에 대한 이야기는 아니고, 이 전체 과정이 다른 프레임워크와 통합될 수 있는 방법론을 제시한다고 보면 된다.
쉽게 말하면 'Decomposition / Contract를 진행' - 'Contracted Question 생성' - '이를 기존 Test-Time Scaling에 사용하여 리소스 절약'정도로 생각하면 될 것 같다. 이에 대해서는 뒤의 Experiment에서 조금더 보도록 하자.

Experiments

본 장에서는 AoT 프레임워크의 유효성을 입증하기 위한 실제적인 분석을 수행하였다. AoT의 성능 및 효율성을 평가하기 위해 다양한 벤치마크 실험을 설계하였으며, 다음의 세 가지 연구 질문을 설정하여 실험을 진행하였다.

1. AoT 프레임워크가 실제로 대규모 언어 모델(LLM)의 추론 성능 향상에 기여하는가?
2. AoT는 기존의 대표적인 추론 방식과 비교하여 computational cost 측면에서 보다 효율적인가?
3. AoT가 기존의 Test-Time Scaling 방법론과 통합 시 어떤 효과를 나타내는가?

실험 설계 및 환경

실험 평가를 위해 추론 능력을 다양한 관점에서 검증할 수 있는 여섯 가지 벤치마크 데이터셋을 선정하였다.

• HotpotQA: 다중단계(multi-hop) 추론 능력 평가
• StrategyQA: 상식 기반의 이진(yes/no) 판단 능력 평가
• GSM8K: 수학적 문제해결 및 논리적 추론 능력 평가
• SVAMP: 자연어 기반 수학적 추론 능력 평가
• Date Understanding: 날짜 및 시간 관련 정보의 이해 및 추론 능력 평가
• CommonsenseQA: 상식 기반 추론 능력 평가

실험 결과 및 상세 분석

1. 추론 성능(Performance) 개선

• HotpotQA 벤치마크 실험에서 AoT를 적용한 GPT-4-o-mini 모델은 80.6%의 F1 점수를 기록하였다. 이는 기존의 최첨단 모델(o3-mini 대비 3.4%, DeepSeek-R1 대비 10.6%)보다 유의미한 성능 향상을 나타냈다.
• GSM8K와 같은 복합적인 수학 문제에서도 AoT 기반의 접근법이 단순한 Chain-of-Thought(CoT) 방식에 비해 보다 정확한 결과를 도출하였으며, 이는 AoT가 복잡한 문제 상황에서도 효과적으로 작동함을 입증한다.

2. 계산 비용(Computational Cost) 효율성

• 실험 결과, AoT는 기존의 Chain-of-Thought(CoT) 접근법 대비 평균적인 추론 단계 수가 감소하였다. 이는 AoT의 DAG(Directed Acyclic Graph) 기반의 구조적 문제 분해(decomposition)와 수축(contraction) 전략 덕분으로 분석된다.
• 특히, AoT의 Markov 특성으로 인해 이전 단계의 불필요한 정보를 유지하지 않게 되면서 메모리 사용량이 현저히 감소하였으며, 이에 따라 전체 추론 시간도 단축되는 것으로 확인되었다.

3. 기존 방법과의 통합성(Integration)

• AoT는 기존의 Chain-of-Thought(CoT) 및 Tree-of-Thought(ToT)와 같은 대표적 Test-Time Scaling 기법들과 통합하여 사용할 때 독립적으로 사용하는 경우보다 더욱 뛰어난 성능을 나타냈다.
• 구체적으로 AoT와 CoT가 통합된 방식은 단독의 CoT보다 성능 향상 및 계산 효율성 면에서 우수한 결과를 보였다. 이는 AoT가 제공하는 'contracted question' 생성이 기존의 Test-Time Scaling 기법을 보다 효과적이고 간소화된 형태로 개선시키기 때문이다.

종합적 결론 및 시사점

본 연구에서 수행된 실험을 통해 AoT 프레임워크는 LLM의 성능 향상과 동시에 계산 비용의 효율적 절감을 실현할 수 있음을 실증하였다. 특히 기존 Test-Time Scaling 방법론과의 통합 시 더욱 뛰어난 결과를 제공할 수 있다는 사실은 AoT가 앞으로 LLM 연구 및 응용 분야에서 핵심적인 역할을 담당할 수 있음을 시사하는 것 같다. LLM을 엄청 많이 공부했지는 않아서 이게 그정도의 혁신인지는 아직 잘 모르겠지만, 일단 박수 한 번 치고 갈만한 일이 아닐까 싶다.

하지만 limitation 부분에서는 다음과 같은 한계점을 갖는다. 
첫째, 실험에 사용된 데이터셋의 범위가 특정 추론 유형에 한정되어 있어, 보다 다양한 실세계 문제로의 일반화 가능성에 대해서는 추가적인 연구가 필요하다.
둘째, AoT의 효과성이 특정 LLM 구조 및 파라미터 설정에 따라 달라질 수 있다는 점에서, 다양한 모델 구조와 규모를 포함한 후속 실험을 통해 AoT의 범용성을 더욱 면밀히 검증할 필요가 있다.
이러한 한계를 극복하고 AoT의 실무적 응용 가능성을 확대하기 위해, 다양한 도메인과 모델 환경에서의 추가적인 연구가 요구된다. 특히 QA에 관련된 데이터셋이 아니라, 일반화성능이 잘 나오는지부터 봐야할 것 같다는 생각이 들었다. 너무 이게 모든 부분에서 실험을 해보지 않은 것 가인ㄴ가 하는 생각이 들기도 해서... 이에 대해서 한번 코드를 까봐야 할 것 같다.