AI바라기의 인공지능

LLM : 논문리뷰 : Satori: Reinforcement Learning with Chain-of-Action-Thought Enhances LLM Reasoning via Autoregressive Search 본문

논문리뷰

LLM : 논문리뷰 : Satori: Reinforcement Learning with Chain-of-Action-Thought Enhances LLM Reasoning via Autoregressive Search

AI바라기 2025. 2. 6. 16:34

Overall Summary

Satori는 Chain-of-Action-Thought (COAT) reasoning과 two-stage training paradigm (format tuning + reinforcement learning with RAE)을 통해 single LLM의 autoregressive search 능력을 향상시키는 새로운 접근 방식을 제시한다. Mathematical reasoning benchmarks에서 state-of-the-art performance를 달성하고, out-of-domain tasks에서도 strong generalization을 보이며, iterative self-improvement를 통해 지속적인 성능 향상을 이룬다. 이 연구는 LLM의 reasoning 능력 향상에 대한 새로운 방향을 제시하며, future work를 통해 더 발전된 autoregressive search capabilities를 가진 LLM 개발에 기여할 수 있다.

 

 

쉬운 설명:
Satori는 "스스로 생각하고, 반성하고, 새로운 방법을 탐색하는" 능력을 가진 LLM을 만들기 위한 연구입니다. 기존에는 정답을 맞추기 위해 여러 번 시도하고 외부의 도움을 받는 방식이었다면, Satori는 특별한 "행동 지침" (<|continue|>, <|reflect|>, <|explore|>)을 통해 LLM 스스로 더 나은 답을 찾아가는 과정을 학습합니다. 마치 사람이 문제를 풀 때 "계속 진행", "잠깐, 검토", "다른 방법 시도" 와 같은 생각을 하는 것과 유사합니다. Satori는 적은 양의 힌트 데이터와 많은 양의 self-improvement를 통해 효율적으로 학습하며, 수학 문제뿐만 아니라 다양한 분야의 문제에도 잘 적응하는 능력을 보여줍니다.

 

 

 

Satori: Reinforcement Learning with Chain-of-Action-Thought Enhances LLM Reasoning via Autoregressive Search 논문 학습 노트

Purpose of the Paper

  • Existing Limitation: 기존 Large Language Models (LLMs)의 reasoning 능력 향상은 주로 test-time computation 증가 (extensive sampling)에 의존하며, 이는 external LLM verifier를 사용하는 two-player system을 필요로 함.
  • New Research Problem: Can we internalize the searching capabilities to fundamentally enhance the reasoning abilities of a single LLM? 즉, external guidance 없이 single LLM이 autoregressive search를 통해 복잡한 문제를 해결할 수 있도록 하는 새로운 연구 문제를 제시함.
  • Approach: Post-training LLMs for autoregressive searching (extended reasoning process with self-reflection and self-exploration)에 초점을 맞춤.

Key Contributions

  1. Chain-of-Action-Thought (COAT) reasoning:
    • Novelty: Meta-action tokens (<|continue|>, <|reflect|>, <|explore|>)을 도입하여 LLM이 external intervention 없이 reasoning process를 제어하도록 함.
    • Method: Conventional post-training (large-scale supervised fine-tuning)과 달리, COAT는 LLM이 문제 해결 중 다양한 meta-actions을 취할 수 있도록 함.
  2. Two-stage training paradigm:
    • Novelty: Small-scale format tuning (FT)과 large-scale self-improvement (RL with "Restart and Explore" (RAE) techniques)를 결합한 새로운 training paradigm 제시.
    • Method:
      1. Format Tuning (FT): Imitation learning on small-scale demonstration trajectories를 통해 COAT reasoning format을 internalize.
      2. Self-improvement (RL): RAE techniques를 활용하여 reinforcement learning으로 large-scale self-improvement 진행.
  3. Satori:
    • 7B LLM trained on open-source models and data.
    • Minimal supervision (only 10K samples in FT stage) and large-scale self-improvement를 통해 효율성 달성.
    • Single LLM capable of autoregressive search without external guidance.

Experimental Highlights

  • Datasets: OpenMathInstruct-2, NuminaMath-CoT (for training), GSM8K, MATH500, AMC2023, AIME2024, OlympiadBench (for math reasoning evaluation), FOLIO, BGQA, CRUXEval, StrategyQA, TableBench, STEM (for out-of-domain evaluation).
  • Metrics: Zero-shot pass@1 accuracy.
  • Baselines:
    • Industry-developed LLMs: o1-preview, QwQ-32B-Preview.
    • Math-specialized model: Qwen-2.5-Math-7B-Instruct (same base model as Satori).
    • Other small-scale general instruct models.
  • Results:
    • State-of-the-art performance on mathematical reasoning benchmarks.
    • Outperforms Qwen-2.5-Math-7B-Instruct (same base model) on both in-domain and out-of-domain tasks.
    • Strong generalization to out-of-domain tasks, demonstrating universal capabilities for self-reflection and self-exploration.
    • Iterative self-improvement를 통해 지속적인 성능 향상.
    • RL training을 통해 test-time scaling behavior를 보이며, progressively acquires more tokens to improve reasoning capabilities.
    • COAT reasoning이 classical CoT reasoning보다 우수함을 입증.

Limitations and Future Work

  • Limitations: 논문에 명시적인 한계점 언급은 없지만, 다음과 같은 한계점 유추 가능:
    • 7B 모델에 대한 실험 결과이므로, larger-scale LLMs에서의 효과는 추가 검증 필요.
    • COAT reasoning format의 최적화, 더 발전된 RL algorithms 적용 가능성 존재.
  • Future Work:
    • Developing reasoning formats with a broader range of meta-actions.
    • Designing more advanced RL algorithms.
    • Extending this approach to general domain.

 

 

 

 

 

Abstract

Large language models (LLMs)은 다양한 영역에서 놀라운 reasoning 능력을 보여주었습니다. 최근 연구에 따르면 test-time computation을 늘리면 LLMs의 reasoning 능력이 향상됩니다. 이는 일반적으로 inference 시간에 외부 LLM verifier의 안내에 따라 광범위한 sampling을 포함하며, 결과적으로 two-player 시스템이 됩니다. 외부 guidance에도 불구하고, 이 시스템의 효과는 복잡한 작업을 처리할 수 있는 단일 LLM의 잠재력을 보여줍니다. 따라서 우리는 새로운 연구 문제를 제기합니다: "searching 능력을 internalize하여 단일 LLM의 reasoning 능력을 근본적으로 향상시킬 수 있을까?" 본 연구는 autoregressive searching (즉, self-reflection 및 새로운 strategies의 self-exploration을 통한 확장된 reasoning 프로세스)을 위해 post-training LLMs에 초점을 맞춘 직교 방향을 탐구합니다. 이를 달성하기 위해, 우리는 Chain-of-Action-Thought (COAT) reasoning과 two-stage training paradigm을 제안합니다: 1) COAT reasoning format을 internalize하기 위한 small-scale format tuning 단계, 2) reinforcement learning을 활용하는 large-scale self-improvement 단계. 우리의 접근 방식은 open-source models 및 data에 대해 trained된 7B LLM인 Satori를 만듭니다. 광범위한 empirical evaluations는 Satori가 mathematical reasoning benchmarks에서 state-of-the-art 성능을 달성하는 동시에 out-of-domain tasks에 대한 강력한 generalization을 보여줍니다. Code, data 및 models는 완전히 open-source 될 것입니다.

 

 

1. Introduction

더보기

Large language models (LLMs)은 mathematical problems, programming, 그리고 logical reasoning을 포함한 광범위한 reasoning tasks에서 놀라운 성능을 보여주었습니다. 이러한 강력한 reasoning capabilities를 가능하게 하는 핵심 기술 중 하나는 Chain-of-Thought (CoT) prompting이며, 이는 LLMs가 일련의 intermediate reasoning steps를 generation하여 복잡한 tasks를 해결할 수 있도록 합니다. 결과적으로, 많은 초기 노력들은 human annotation 또는 더 advanced models에서 synthetic data를 distilling하여 large-scale, high-quality CoT reasoning chains를 사용하여 LLMs를 fine-tuning하는 데 집중합니다. 그러나 human annotation은 매우 노동 집약적이며, distillation은 종종 model의 reasoning capabilities를 특정 수준으로 제한합니다.

Training resources를 scaling up하는 것 외에도, 최근 연구는 test-time scaling, 즉, 더 정확한 solutions을 search하기 위해 additional inference-time compute를 할당하는 데 초점을 맞추고 있습니다. 이는 종종 multiple complete solutions을 generation하거나 multiple intermediate reasoning steps를 sampling하여 광범위한 sampling을 포함합니다. 이러한 methods는 일반적으로 search process를 guide하기 위해 external feedback이 필요하며, 대개 final solutions 또는 intermediate steps를 평가하기 위해 auxiliary reward model을 training하는 것을 통해 이루어집니다. 그러나 이러한 two-player frameworks는 더 많은 model-deployment costs를 발생시키고 search capabilities를 단일 LLM으로 internalize하지 않습니다.

위의 연구와 직교하는, 우리 연구는 LLMs에 autoregressive search capabilities, 즉 self-reflection 및 새로운 strategies의 self-exploration을 통한 확장된 reasoning process를 가능하게 하는 새로운 방향을 조사합니다. 구체적으로, 우리는 LLMs가 problem solving 중에 다양한 meta-actions을 취할 수 있도록 하는 Chain-of-Action-Thought (COAT) mechanism을 도입합니다. large-scale supervised fine-tuning (SFT) 및 reinforcement learning from human feedback (RLHF)로 구성된 conventional post-training과 달리, 우리는 새로운 two-stage training paradigm을 제안합니다: (1) COAT reasoning format을 internalize하기 위한 small-scale format tuning (FT) 단계 및 (2) "Restart and Explore" (RAE) techniques를 사용하여 reinforcement learning을 활용하는 large-scale self-improvement 단계. 우리의 접근 방식은 in-domain 및 out-of-domain tasks 모두에서 우수한 성능을 달성하는 open-source base models 및 mathematic data에 대해 trained된 7B LLM인 Satori의 개발로 이어집니다. 요약하면, 우리의 기여는 세 가지입니다.

  1. Efficiency: Satori는 external guidance 없이 autoregressive search가 가능한 단일 LLM입니다 (Section 6 및 Section A). 더욱이, 이것은 minimal supervision 및 large-scale self-improvement로 달성됩니다.
  2. Effectiveness: Satori는 in-domain mathematical reasoning tasks에서 우수한 성능을 보여주고 동일한 base model에 구축된 instruct model보다 성능이 뛰어납니다 (Section 5.1).
  3. Generalizability: math reasoning에 대한 최근 연구와 달리, Satori는 out-of-domain tasks에 대한 강력한 transferability를 보여주고 self-reflection 및 self-exploration을 위한 universal capabilities를 보여줍니다 (Section 5.2).

 

 

1. Introduction 정리 노트 (AI 연구자 대상)

핵심:

  • 기존 연구의 문제점:
    • Human annotation은 노동 집약적.
    • Distillation은 모델의 reasoning 능력을 제한.
    • Test-time scaling (e.g., extensive sampling, external reward model)은 높은 deployment cost, search capability를 단일 LLM에 내재화(internalize)하지 못함.
  • 본 연구의 차별점:
    • Autoregressive search capability를 갖춘 LLM 개발:
      • Self-reflection, self-exploration을 통해 extended reasoning process 가능.
      • Chain-of-Action-Thought (COAT) 메커니즘 도입: 문제 해결 중 다양한 meta-action 수행.
    • 새로운 Two-stage training paradigm 제안:
      • Small-scale format tuning (FT): COAT reasoning format 내재화.
      • Large-scale self-improvement: "Restart and Explore (RAE)"를 활용한 reinforcement learning.
  • Satori (7B LLM):
    • Open-source base models, mathematical data로 trained.
    • In-domain (mathematical reasoning) 및 out-of-domain tasks 모두에서 우수한 성능.
    • 장점:
      1. Efficiency: External guidance 없이 autoregressive search 가능 (minimal supervision, large-scale self-improvement).
      2. Effectiveness: In-domain에서 instruct model보다 우수.
      3. Generalizability: Out-of-domain tasks에도 잘 일반화, self-reflection/exploration 능력 보유.

쉬운 설명 :

기존의 Large Language Models (LLMs)는 reasoning 능력이 뛰어나지만, 몇 가지 문제점이 있었어요. 사람이 직접 데이터를 만들거나(human annotation), 다른 모델에서 지식을 추출하는 방식(distillation)은 효율성이 떨어졌죠. Test-time에 성능을 높이는 방법(test-time scaling)도 있었지만, 외부 모델의 도움을 받아야 해서 비용이 많이 들고, 모델 자체의 능력을 키우는 데는 한계가 있었어요.

이 논문에서는 이 문제를 해결하기 위해, LLM이 스스로 생각하고 탐색하는 능력(autoregressive search capability)을 갖도록 하는 새로운 방법을 제시합니다. 마치 사람이 문제를 풀 때 여러 가지 방법을 시도하고, 막히면 다시 처음부터 다른 방식으로 접근하는 것처럼요! 이를 위해 "Chain-of-Action-Thought (COAT)"라는 메커니즘을 도입하고, 두 단계로 구성된 학습 방법(format tuning, reinforcement learning with "Restart and Explore")을 사용합니다.

이렇게 해서 만들어진 "Satori"라는 모델은 수학 문제뿐만 아니라 다른 분야의 문제에도 잘 적응하고, 스스로 생각하고 탐색하는 능력이 뛰어나다는 것을 보여줍니다. 게다가, 외부 도움 없이 혼자서도 잘 해낸다는 장점이 있어요!

 

 

 

 

2. Related Work

더보기

이 논문의 범위와 밀접하게 관련된 문헌을 요약합니다 (자세한 내용은 Section B 참조).

Concurrent Work.

OpenAI의 o1 (OpenAI, 2024)의 영향력을 바탕으로, 연구 커뮤니티 내에서 open-source LLMs의 advanced reasoning capabilities를 향상시키기 위한 상당한 노력이 이루어졌습니다. 가장 일반적인 접근 방식은 더 강력한 teacher models에서 knowledge를 distilling하는 것입니다. 이와 대조적으로, Satori는 이 문제를 reinforcement learning (RL) 관점에서 접근하며 minimal supervision (format tuning 단계에서 10K samples만)을 필요로 합니다. 가장 관련성이 높은 concurrent work는 DeepSeek의 최근에 출시된 R1 (Guo et al., 2025)이며, 이는 small-scale cold-start SFT에 이어 large-scale RL training을 수행하는 유사한 high-level strategy를 채택합니다. 비록 두 연구 모두 이 high-level idea에서 일치하지만, 우리 연구는 data synthesis framework 및 RL algorithms을 포함한 핵심 방법론에서 R1과 다릅니다. 또한, DeepSeek-R1은 large-scale LLMs (67B) training에 중점을 두는 반면, 우리 연구는 research purpose를 위한 smaller-scale LLMs (7B) 개발에 대한 insights를 제공합니다. 마지막으로, industry-developed model로서 DeepSeek-R1 (Guo et al., 2025)의 technical details는 완전히 공개되지 않아 reproduction이 어렵지만, 우리 연구는 training data 및 training recipes를 open-source하는 것을 목표로 하는 완전히 투명한 연구입니다.

Post-training LLMs for Reasoning.

최근의 발전은 reasoning을 향상시키기 위한 광범위한 post-training에 초점을 맞추고 있습니다. 일련의 연구는 high-quality instruction-tuning datasets을 구성하는 데 중점을 두지만, 비싼 annotation costs로 어려움을 겪습니다. 보다 최근의 연구는 models가 스스로 generated한 data에 대해 trained되는 self-improvement approaches에 초점을 맞추고 있습니다. 또한, reinforcement learning methods, 특히 Proximal Policy Optimization (PPO)에 기반한 방법이 더 효과적인 것으로 입증되었으며, 이는 일반적으로 reward models를 활용하여 learning process를 guide합니다.

Enabling LLMs with Searching Abilities.

Prompting-based approaches는 LLMs가 error correction 및 alternative paths 탐색을 통해 solutions을 search하도록 guide합니다. 그러나 이러한 approaches는 LLM의 reasoning abilities를 근본적으로 향상시킬 수 없습니다. 더욱이, 최근 연구는 LLMs의 self-correction의 어려움을 지적했습니다. 최근 연구는 LLMs를 self-exploration을 위해 training하는 방향으로 전환되었습니다. 일부는 trajectory-level search (이전의 complete responses에서 errors를 반복적으로 식별하고 improved responses를 생성)를 가능하게 하는 데 초점을 맞추었습니다. 또 다른 연구는 step-level search를 탐구하여 LLMs가 더 세분화된 방식으로 mistakes를 식별하고 수정할 수 있도록 합니다. 일부는 다른 model을 사용하여 step-level feedback을 제공함으로써 이를 달성하지만, 이러한 two-player frameworks는 model deployment에 대한 높은 costs로 어려움을 겪습니다. SoS는 tree search를 flattened string으로 수행하기 위해 단일 LLM을 train하려고 시도하는 또 다른 밀접하게 관련된 작업입니다. 그러나 SoS의 effectiveness는 주로 simple symbolic tasks에서 나타났으며, 더 복잡한 problems로 generalize하는 능력은 아직 탐구되지 않았습니다.

 

 

2. Related Work 정리 노트 (AI 연구자 대상)

크게 세 가지 갈래의 관련 연구 존재:

  1. Concurrent Work (동시 진행 연구):
    • Open-source LLM reasoning 능력 향상 연구 활발 (OpenAI o1 영향).
    • 주류: Stronger teacher models에서 knowledge distillation.
    • Satori의 차별점:
      • Reinforcement learning (RL) 관점 접근.
      • Minimal supervision (format tuning 단계에서 10K samples).
    • DeepSeek-R1과의 비교:
      • 유사점: Small-scale SFT 후 large-scale RL training이라는 high-level strategy.
      • 차이점:
        • Data synthesis framework, RL algorithms 등 key methodologies.
        • DeepSeek-R1: Large-scale LLMs (67B)에 집중.
        • Satori: Smaller-scale LLMs (7B) 개발에 대한 insight 제공 (연구 목적).
        • DeepSeek-R1: 기술 세부 사항 비공개 (reproduction 어려움).
        • Satori: 완전 투명한 연구 (training data, recipes open-source 목표).
  2. Post-training LLMs for Reasoning:
    • High-quality instruction-tuning datasets 구축:
      • Annotation cost 문제.
    • Self-improvement approaches:
      • Model 스스로 생성한 data로 training.
    • Reinforcement learning (PPO 기반):
      • Reward model 활용, learning process guide.
    • Satori는 위의 방법론들을 활용하면서도, 차별점을 가짐 (뒤에서 설명).
  3. Enabling LLMs with Searching Abilities:
    • Prompting-based approaches:
      • Error correction, alternative paths 탐색 유도.
      • 근본적인 reasoning 능력 향상 어려움.
      • LLM의 self-correction 어려움 지적.
    • Self-exploration을 위한 LLM training:
      • Trajectory-level search: 이전 complete responses에서 error 식별, improved responses 생성.
      • Step-level search: 더 세분화된 mistake 식별 및 수정 (주로 external model 활용, two-player framework -> high deployment cost).
      • SoS (tree search):
        • Simple symbolic tasks에만 효과 입증.
        • Complex problems 일반화 능력 미지수.
      • Satori: Single LLM으로 autoregressive search 구현 (핵심 차별점)

쉬운 설명 :

이 섹션에서는 Satori와 비슷한 시기에 진행되었거나, Satori의 연구 배경이 되는 관련 연구들을 소개하고 있어요. 크게 세 가지 갈래로 나눌 수 있는데,

  1. 비슷한 시기에 진행된 연구들 (Concurrent Work):
    • OpenAI의 o1 발표 이후, 공개된 LLM들의 reasoning 능력을 높이려는 연구가 많았어요.
    • 대부분은 더 뛰어난 모델(teacher model)의 지식을 추출하는 방식을 사용했죠.
    • Satori는 이와 달리, reinforcement learning을 사용하고, 아주 적은 양의 데이터만으로 학습했다는 점이 달라요.
    • DeepSeek-R1이라는 모델과 비슷한 점도 있지만, 세부적인 방법론이나 연구 목적, 투명성 등에서 차이가 있어요.
  2. Reasoning 능력을 높이기 위한 LLM 후처리 (Post-training) 연구들:
    • 좋은 데이터를 만들어서 학습시키는 방법, 모델 스스로 데이터를 만들어서 학습하는 방법, reinforcement learning을 사용하는 방법 등이 있었어요.
    • Satori는 이러한 방법들을 활용하면서, 자신만의 독창적인 방식을 추가했어요.
  3. LLM에게 탐색 능력(Searching Abilities)을 부여하는 연구들:
    • 프롬프트를 사용해서 모델이 스스로 오류를 수정하거나 다른 경로를 찾도록 유도하는 방법이 있었지만, 근본적인 reasoning 능력을 높이는 데는 한계가 있었어요.
    • 모델 스스로 탐색하도록 학습시키는 연구도 있었는데, 대부분 외부 모델의 도움을 받거나, 간단한 문제에만 적용 가능했어요.
    • Satori는 외부 도움 없이, 단일 모델 안에서, 더 복잡한 문제에도 적용 가능한 autoregressive search를 구현했다는 점이 핵심적인 차이점이에요!

 

 

 

Figure 1: A High-level Overview of Satori Training Framework: Format Tuning (FT) + Self-improvement.

먼저, Satori는 small-scale demonstration trajectories에 대한 imitation learning을 통해 COAT reasoning format을 배웁니다. 다음으로, Satori는 COAT reasoning format을 활용하여 large-scale reinforcement learning을 통해 self-improve합니다.

 

 

3. Preliminaries

더보기

우리는 mathematical problem-solving을 ground truth y*와 일치하는 solution ỹ를 generation하기 위해 language model πθ를 training하여 해결합니다. 주어진 problem prompt x에 대해서 말이죠. 모든 sequences x, y, 그리고 y*는 predefined dictionary의 tokens로 구성됩니다. 우리의 approach는 math problems를 solving하기 위해 model을 training하는데 reinforcement learning (RL)을 사용하기 때문에, 아래에 주요 RL concepts를 간략하게 설명합니다.

Reinforcement Learning (RL).

RL은 agent가 environment와의 interactions를 통해 expected cumulative rewards를 maximize하기 위해 sequential decisions를 내리는 것을 포함합니다. 여기서, language model πθ는 agent의 policy 역할을 합니다. initial state z0에서 시작하여, 각 step l에서 agent는 current state z<sub>l</sub>을 observes하고, reward r<sub>l</sub>을 receives하며, πθ에 따라 action을 selects하고, next state z<sub>l+1</sub>로 transitions하며, terminal state에 도달할 때까지 계속됩니다. trajectory는 이 interaction 동안의 states와 actions의 sequence입니다. RL은 expected rewards &sum;<sub>L</sub><sup>l=1</sup> r<sub>l</sub> 를 maximize하도록 policy를 optimizes합니다. 여기서 L은 trajectory length 입니다.

 

3. Preliminaries 정리 노트 (AI 연구자 대상)

핵심:

  • 문제 정의: Mathematical problem-solving을 language model (πθ)을 사용하여 해결.
    • Input: Problem prompt (x).
    • Output: Solution (ỹ) (ground truth (y*)와 일치해야 함).
    • x, y, y*는 모두 predefined dictionary의 tokens로 구성.
  • Reinforcement Learning (RL) 사용:
    • Language model (πθ) = Agent's policy.
    • RL 과정 (일반적인 내용이지만, 이 논문의 기반):
      • Initial state (z0)에서 시작.
      • 각 step (l)마다:
        • Current state (zl) 관찰.
        • Reward (rl) 획득.
        • Policy (πθ)에 따라 action 선택.
        • Next state (zl+1)로 transition.
        • Terminal state까지 반복.
      • Trajectory: States와 actions의 sequence.
      • 목표: Expected cumulative rewards (∑L l=1 rl) 최대화 (L: trajectory length).

이 논문만의 핵심:

  • 일반적인 RL framework를 따르지만, Satori는 이 framework 안에서 "COAT reasoning"과 "two-stage training (FT + RL)"을 통해 독창적인 방식으로 문제를 해결 (다음 섹션에서 더 자세히 설명).
  • 이 섹션 자체는 일반적인 RL 개념 설명이지만, 이후 내용을 이해하기 위한 필수적인 기반 지식.

쉬운 설명 :

이 섹션에서는 논문에서 사용할 기본적인 용어와 개념을 설명하고 있어요. 마치 게임을 하기 전에 규칙을 설명해주는 것과 같아요.

  • 문제: 수학 문제(x)를 풀어서 정답(y*)과 같은 답(ỹ)을 만들어내는 것이 목표예요.
  • 해결 방법: Reinforcement Learning (RL)이라는 방법을 사용해요.
    • RL은 마치 게임처럼, 어떤 상황(state)에서 어떤 행동(action)을 해야 보상(reward)을 많이 받을 수 있는지 학습하는 방법이에요.
    • 여기서는 language model (πθ)이 게임을 하는 플레이어(agent) 역할을 하고, 문제를 풀어나가는 과정을 통해 보상을 최대화하도록 학습해요.
  • 핵심은 이 RL이라는 게임판 위에서, 이 논문에서 제시하는 특별한 방법(COAT, two-stage training)을 통해 문제를 어떻게 풀어나가는지를 보여준다는 점이에요.

 

 

 

4. Method

더보기

이 섹션에서는 reasoning의 formulation과 reasoning이 어떻게 sequential decision-making problem으로 formulated될 수 있는지 소개하는 것으로 시작합니다.

Goal: 우리는 LLMs가 final answer를 directly predicting하기보다는 multiple steps를 통해 reasoning하여 problems를 solve하도록 training하기를 원합니다. problem statement x가 주어지면, model은 reasoning steps의 sequence {y1, y2, ..., yL}을 generates합니다. 여기서 yL은 final answer를 제공합니다. 그러나 모든 intermediate steps가 도움이 되는 것은 아닙니다. errors를 반복하는 것은 accuracy를 향상시키지 않습니다. Effective reasoning은 correctness를 verifying하고, mistakes를 identifying하며, alternative solutions를 considering하는 것을 필요로 합니다. 예를 들어, x = "1 + 1 =?"이 주어지면, model은 처음에 y1 = 3을 output할 수 있고, 그런 다음 y2 (예 : "잠깐, 확인해 보겠습니다...")로 mistake를 recognize하고, 이를 y3 = 2로 correcting합니다.

Chain-of-Action-Thought reasoning (COAT).

핵심 challenge는 model이 external intervention 없이 언제 reflect, continue, 또는 explore alternatives를 할지 결정할 수 있도록 하는 것입니다. 이를 가능하게 하기 위해, 우리는 standard text generation을 넘어 model의 reasoning process를 guide하는 special meta-action tokens를 도입합니다. 이러한 tokens는 model이 계속 진행하기 전에 reasoning을 reassess할 시기를 결정하는 hint 역할을 합니다.

  • Continue Reasoning (<|continue|>) : Model이 다음 intermediate step을 generating하여 current reasoning trajectory를 build upon하도록 encourages합니다.
  • Reflect (<|reflect|>) : Model에게 prior reasoning steps의 correctness를 verify하기 위해 pause하도록 prompts합니다.
  • Explore Alternative Solution (<|explore|>) : Model에게 reasoning에서 critical flaws를 identify하고 new solution을 explore하도록 signals합니다.

각 reasoning step yl은 tokens의 sequence이며, starting token은 잠재적으로 designated meta-action tokens 중 하나일 수 있습니다. 우리는 이 formulation을 Chain-of-Action-Thought reasoning (COAT)이라고 부릅니다. 특히, typical Chain-of-Thought reasoning (CoT)은 COAT의 special case로 볼 수 있으며, CoT의 각 reasoning step은 다른 types의 meta-actions를 explicitly incorporating하지 않고 continuation으로 restricted됩니다.

Learning to Reason via RL.

우리는 reasoning을 sequential decision-making problem으로 formulate합니다. 여기서 reasoning은 step by step으로 answer를 constructing and refining하는 process입니다. 구체적으로, model πθ는 input context x (initial state z0)로 시작하여, reasoning step yl (action)을 generates하고, yl (next state zl+1 = zl ⊕ yl, 여기서 ⊕는 string concatenation을 나타냄)을 appending하여 context를 updates하고, final answer yL을 produces할 때까지 이 process를 repeats합니다. reasoning은 model이 completion을 signals할 때 (예 : EOS token 생략) terminates됩니다. 가장 simplest reward function은 1{yL = y*}일 수 있으며, final answer yL이 ground truth y*와 matches하는지 evaluates합니다. 이 formulation을 통해, 우리는 RL을 사용하여 reason하도록 model을 train할 수 있으며, expected reward를 maximize하는 reasoning steps를 generate하는 것을 목표로 합니다. 그러나 RL을 reasoning에 applying하는 것은 두 가지 key challenges를 presents합니다.

  1. Unawareness of meta-action tokens: Model은 special tokens의 purpose를 understand하지 못하고 special meta-action tokens를 encountering하는 것이 reflection 또는 proposing alternatives를 require할 수 있다는 것을 recognize하지 못합니다.
  2. Long horizon and sparse rewards: Reasoning은 rewards가 only at the end에 있는 long-term decision-making을 requires하며, 이는 learning effectiveness를 hinders합니다. Model은 rewards를 receiving하기 전에 many correct reasoning steps를 take해야 하며, failures는 initial state (즉, problem statement)에서 restart하도록 force합니다. 이는 rewards와 관련된 training data가 scarce하지만 rewards는 RL progress를 driving하는 데 essential하기 때문에 learning을 difficult하게 만듭니다.

Overview of Proposed Method.

Model의 initial unawareness of meta-action tokens를 address하기 위해, 우리는 warm-up "format-tuning" stage를 introduce합니다. 우리는 pre-trained LLM을 a few demonstrated reasoning trajectories를 featuring하는 small dataset에 fine-tune합니다 (Section 4.1). 이 step은 model이 meta-action tokens를 using하고 reacting하는 데 familiarize하도록 합니다. 둘째, long horizons 및 sparse rewards의 challenges를 tackle하기 위해, 우리는 Go-explore에서 inspired된 "restart and explore" (RAE) strategy를 propose합니다. 여기서 model은 previous reasoning attempts가 failed한 points를 포함하여 intermediate steps에서 restarts하여 scratch에서 starting하는 대신 errors를 correcting하는 데 focus할 수 있습니다. 우리는 또한 deeper reflection을 encourage하기 위해 exploration bonuses를 add하여 model이 correct answers에 arrive할 opportunities를 further increasing합니다 (Section 4.2).

4.1. Format Tuning Through Imitation Learning

Base LLM πθ를 COAT reasoning을 perform하도록 training하는 것은 significant challenge를 presents합니다. LLMs는 일반적으로 trials and errors를 incorporating하는 COAT reasoning data에 pre-trained되지 않으므로 이 capability를 inject하기 위해 post-training stage가 necessitating합니다. 이를 address하기 위해, 우리는 LLMs가 imitation learning을 통해 expert COAT trajectories를 emulate하도록 training하도록 designed된 method인 format tuning (FT)을 introduce합니다. Imitation learning techniques는 robotics domain에서 widely used되며, agents는 human experts가 provided한 demonstration trajectories를 사용하여 trained됩니다. 그러나 LLMs를 위한 high-quality demonstration trajectories를 generating하는 것은 complex tasks에 대해 prohibitively expensive합니다. Demonstration trajectory dataset Dsyn = {(x(i), ỹ(i))}N i=1을 efficiently constructing하기 위해, 우리는 세 가지 LLMs를 leverages하는 multi-agent data synthesis framework를 propose합니다.

  • Generator: Input problem이 주어지면, generator πg는 classical CoT techniques를 사용하여 주어진 input problem에 대해 multiple reasoning paths를 generates합니다.
  • Critic: Critic πc는 generator에 의해 generated된 reasoning paths의 correctness를 evaluates하여 reasoning을 refine하고 suboptimal steps를 address하기 위한 feedback을 providing합니다.
  • Reward Model: Additionally, reward model πr은 refined reasoning paths에 scores를 assigns하고 final demonstration trajectory로 most effective path를 selects합니다.

이 세 models는 high-quality demonstration trajectories를 constructing하기 위해 collaborate합니다 (trajectory synthesis에 대한 details는 Appendix C에 provided됨). 이 작업을 위해, 우리는 simplest imitation learning approach인 behavior cloning을 adopt합니다. 이는 supervised fine-tuning을 활용하여 expert COAT demonstration trajectories Dsyn에 대해 LLM policy를 train합니다. 특히, 우리는 small number (10K)의 COAT demonstration trajectories조차도 πθ가 COAT reasoning format을 effectively follow하기에 sufficient하다는 것을 observe합니다.

4.2. Self-improvement via Reinforcement Learning

Format tuning 후, LLM policy πθ는 COAT reasoning style을 adopts하지만, 특히 self-reflection을 위해 meta-actions를 using하는 데 있어 generalize하는 데 struggles합니다. 이 limitation은 format tuning 중 demonstrations의 scarcity에서 arises합니다. 더 많은 demonstrations를 collecting하는 것이 도움이 될 수 있지만, costly하고 time-consuming합니다. 대신, 우리는 model이 RL을 통해 reasoning을 self-improve할 수 있는지 explore합니다.

우리는 format-tuned LLM으로 start하고 PPO algorithm (widely used RL method)을 사용하여 training합니다. Dataset D의 problems x에 대해 training하는 것 외에도, 우리는 또한 format-tuned LLM에 의해 generated된 partial trajectories에서 reasoning을 begin하도록 model πθ를 train합니다. Reasoning errors는 일반적으로 fundamental flaws보다는 minor mistakes에서 arise하기 때문에, start에서 re-exploring하는 것은 inefficient합니다. 대신, 우리는 model이 intermediate steps에서 restart하여 errors를 correct하고 finally correct answers를 achieve할 수 있도록 합니다. Go-Explore에서 inspired되어, 우리는 Restart and Explore (RAE) strategy를 introduce합니다.

Initial States.

RAE는 model이 problem statement뿐만 아니라 past trajectories (correct 및 incorrect 모두)에서 sampled된 intermediate steps에서도 reason하도록 trains합니다. 이를 통해 redundant re-computation 없이 deeper exploration이 가능합니다. Algorithm 1에 detailed된 바와 같이, input problem x ∈ D가 주어지면, format-tuned LLM은 먼저 multiple reasoning trajectories를 generates합니다. 그런 다음 우리는 randomly T ≥ 0 steps를 backtrack하고 reflect token <|reflect|>를 append하여 model이 reasoning을 refine하도록 prompt합니다. Diverse exploration을 encourage하기 위해, correct 및 incorrect trajectories는 restart buffers (D+restart 및 D-restart)에 separately stored됩니다. 그런 다음 RL training은 original problem dataset과 함께 이러한 buffers에서 reasoning을 optimizes하여 merged dataset Drestart에서 initial states를 sampling합니다.

Reward Design.

RAE는 model에게 reasoning을 refine할 multiple opportunities를 provides하지만, effective reflection은 이러한 chances를 make use of하는 데 key입니다. Correctness를 rewards로 using하는 것 외에도, 우리는 model이 correct answers에 reach하도록 guide하는 hints로서 following bonuses rewards를 introduce합니다.

  • Rule-based Reward: Rule-based reward는 simply final answer의 correctness를 evaluates합니다.
  • rrule(ỹL, y*) = 1ỹL=y* - 1 ∈ {-1, 0}
  • Reflection Bonuses: Self-reflection을 reinforce하기 위해, 우리는 reflection bonus rbonus를 introduce합니다. Model이 negative restart buffer (D-restart)에 stored된 incorrect reasoning trajectory에서 starts하고 successfully problem을 solves하면 reward bonus를 obtains하여 past mistakes를 correct하도록 encouraging합니다. Conversely, positive restart buffer (D+restart)의 correct trajectory에서 starts하지만 problem을 solve하지 못하면 penalty를 incurs하여 이미 올바른 track에 있을 때 unnecessary revisions를 discouraging합니다. Formally, reflection bonus는 다음과 같이 defined됩니다.여기서 β는 bonus scale hyperparameter입니다.
  • rbonus(z, ỹ) = { β if z ∈ D-restart and ỹL = y*, -β if z ∈ D+restart and ỹL ≠ y*, 0 otherwise,
  • Preference Bonuses: Initial training stage에서는 correct answers가 rare하기 때문에 reward signals는 종종 effective RL training에 too sparse합니다. Reflection을 사용하더라도 model은 correct reasoning trajectories를 generate하지 못하여 sparse reward problem이 resulting할 수 있습니다. 이를 mitigate하기 위해, 우리는 Bradley-Terry (BT) preference framework를 사용하여 Outcome Reward Model (ORM)을 train합니다. ORM은 reasoning trajectories를 rates하여 correct (preferred) ones에 higher values를 assigning합니다. 각 problem x ∈ D에 대해, 우리는 πθ를 사용하여 multiple trajectories를 generate하고 correct 및 incorrect outputs를 pairing하여 preference dataset을 construct합니다. BT model은 이러한 pairs 간의 score gap을 maximize하도록 trained됩니다. ORM의 output, σ(rψ(z, ỹ)) ∈ [0, 1]은 fine-grained reward signal 역할을 하여 model이 reasoning을 further refine하도록 helping합니다. Details는 Appendix D.3를 참조하십시오.

Initial state z ∈ Drestart 및 sampled trajectory ỹ에 대해, overall reward function r(z, ỹ)는 다음과 같이 defined됩니다.

r(z, ỹ) = rrule(ỹL, y*) + σ(rψ(z, ỹ)) + rbonus(z, ỹ)

Iterative Self-improvement.

RL은 policy가 self-generated trajectories에서 self-improve할 수 있도록 하지만, policy가 local sub-optimum에 converges하고 further improve할 수 없는 vicious cycle로 lead할 수도 있습니다. (Agarwal et al., 2022; Schmitt et al., 2018)에서 inspired되어, 우리는 이 issue를 mitigate하기 위해 iterative self-improvement strategy를 propose합니다. Specifically, 각 round의 RL training 후, 우리는 current well-optimized policy의 knowledge를 supervised fine-tuning (SFT)을 통해 base model로 distill합니다. Newly fine-tuned model에서 starting하여, 우리는 또 다른 round의 RL training을 perform합니다. Intuitively, optimization perspective에서, 각 round의 distillation은 policy가 loss landscape에서 local optima를 escape할 수 있도록 하는 parameter reset mechanism으로 viewed될 수 있으며, self-improving을 continue할 수 있도록 합니다 (자세한 내용은 Section D.3에 included됨). 다음 section에서는 이 approach를 validate하기 위한 empirical evidence를 provide합니다.

 

4. Method 정리 노트 (AI 연구자 대상)

핵심:

  • Reasoning as Sequential Decision-Making:
    • 여러 단계(steps)를 거쳐 문제를 해결 (최종 정답만 예측 X).
    • 각 단계(yl)는 이전 단계에 이어붙여짐 (zl+1 = zl ⊕ yl).
    • 문제점:
      1. Meta-action token 이해 부족: LLM이 reflection, exploration 등의 meta-action token을 언제 사용해야 할지 모름.
      2. Long horizon, sparse rewards: 보상이 마지막에만 주어져 학습 어려움.
  • Chain-of-Action-Thought (COAT) Reasoning:
    • Meta-action tokens 도입:
      • <|continue|>: 다음 단계 진행.
      • <|reflect|>: 이전 단계 검토.
      • <|explore|>: 새로운 해결 방법 탐색.
    • 기존 Chain-of-Thought (CoT)는 <|continue|>만 사용하는 특수한 경우.
  • 해결 방법 (Two-Stage Training):
    1. Format Tuning (FT) - Imitation Learning:
      • 목표: LLM에게 COAT reasoning format (meta-action tokens 사용법) 학습.
      • 방법:
        • Multi-agent data synthesis:
          • Generator (πg): CoT 사용, multiple reasoning paths 생성.
          • Critic (πc): Reasoning paths 평가, feedback 제공.
          • Reward Model (πr): Reasoning paths에 점수 부여, best path 선택.
        • Behavior cloning: 생성된 demonstration trajectories (10K)로 supervised fine-tuning.
    2. Self-improvement - Reinforcement Learning (RL):
      • 목표: COAT 활용, reasoning 능력 향상 (특히 self-reflection).
      • 방법:
        • PPO algorithm 사용.
        • Restart and Explore (RAE):
          • Initial states: Problem statement뿐만 아니라, 이전 trajectories의 intermediate steps (correct/incorrect)에서도 시작.
          • <|reflect|> token 추가 -> reasoning 수정 유도.
          • Correct/incorrect trajectories 별도 저장 (restart buffers).
        • Reward Design:
          • Rule-based reward: 정답 여부 (1 or -1).
          • Reflection bonuses:
            • Incorrect trajectory에서 시작, 정답 -> bonus (+β).
            • Correct trajectory에서 시작, 오답 -> penalty (-β).
          • Preference bonuses:
            • Outcome Reward Model (ORM) 학습 (Bradley-Terry framework).
            • ORM으로 reasoning trajectories 평가, fine-grained reward 제공.
        • Iterative Self-improvement:
          • RL training 후, well-optimized policy를 base model에 distill (SFT).
          • Distilled model로 다시 RL training (반복).
          • Parameter reset mechanism -> local optima 탈출.

이 논문만의 핵심:

  • COAT reasoning: Meta-action tokens를 통해 LLM의 self-reflection, exploration 능력 부여.
  • Two-stage training (FT + RL):
    • FT: Minimal supervision (10K data)으로 COAT format 학습.
    • RL (RAE): Restart, exploration, 다양한 reward signals로 reasoning 능력 강화.
    • Iterative self-improvement: Distillation을 통해 local optima 문제 해결.

쉬운 설명 :

이 섹션에서는 Satori 모델이 어떻게 reasoning 능력을 갖게 되는지, 그 핵심 방법을 설명하고 있어요. 크게 두 가지 중요한 개념이 등장합니다.

  1. Chain-of-Action-Thought (COAT):
    • 기존의 Chain-of-Thought (CoT)는 그냥 순서대로 생각하는 거였다면, COAT는 중간에 "잠깐, 이게 맞나?" (<|reflect|>), "다른 방법은 없을까?" (<|explore|>) 하고 스스로에게 질문을 던질 수 있게 해주는 거예요.
    • 이렇게 하면 LLM이 더 유연하고 똑똑하게 문제를 해결할 수 있겠죠?
  2. Two-Stage Training (두 단계 학습):
    • 1단계 (Format Tuning): 먼저 COAT 방식으로 생각하는 "형식"을 가르쳐요. 아주 적은 양의 데이터(10K)만 사용해서, 마치 모범 답안을 보여주면서 따라 하도록 하는 거죠 (imitation learning).
    • 2단계 (Self-improvement): 이제 스스로 실력을 키울 차례예요. Reinforcement learning (RL)을 사용하는데, 여기서 핵심은 "Restart and Explore (RAE)"라는 전략이에요.
      • 문제를 풀다가 막히면, 처음부터 다시 시작하는 게 아니라, 중간 단계부터 다시 시작하거나 (<|reflect|> 사용), 다른 방법을 시도해볼 수 있게 (<|explore|> 사용) 해주는 거죠.
      • 잘못된 길로 갔다가 다시 돌아와서 정답을 찾으면 보너스를 주고, 잘 가고 있는데 엉뚱한 짓을 하면 벌점을 주는 방식으로 학습을 유도해요.
      • 그리고 주기적으로 지금까지 학습한 내용을 다시 기본 모델에 "복습"시켜서 (distillation), 더 발전할 수 있게 해줘요.

이렇게 두 단계를 거치면서 Satori는 점점 더 똑똑해지고, 복잡한 문제도 스스로 해결할 수 있는 능력을 갖추게 됩니다!