[논문 리뷰] Scaling of Search and Learning: A Roadmap to Reproduce o1 from Reinforcement Learning Perspective

Zhiyuan Zeng, Qinyuan Cheng외 7명|arXiv (Cornell University)|2024. 12. 18.

Auction Theory and Applications인용 수 4

한 줄 요약

이 논문은 OpenAI o1의 전문가 수준의 추론 능력을 재현하기 위한 강화학습 기반 로드맵을 제시하며, 인간과 유사한 추론을 위한 정책 초기화, 결과 및 과정 신호를 포함한 보상 설계, MCTS 및 Best-of-N과 같은 검색 전략, 그리고 검색에 의해 생성된 데이터를 통한 피니팅을 통한 확장 가능한 학습을 포함한 네 핵심 요소를 강조한다. 주요 기여는 기존 오픈소스 o1 재현 노력들을 통합하고 확장 가능한 성능에 이를 수 있는 검색과 학습을 통합한 체계적이고 모듈화된 프레임워크를 제공하는 것이다.

ABSTRACT

OpenAI o1 represents a significant milestone in Artificial Inteiligence, which achieves expert-level performances on many challanging tasks that require strong reasoning ability.OpenAI has claimed that the main techinique behinds o1 is the reinforcement learining. Recent works use alternative approaches like knowledge distillation to imitate o1's reasoning style, but their effectiveness is limited by the capability ceiling of the teacher model. Therefore, this paper analyzes the roadmap to achieving o1 from the perspective of reinforcement learning, focusing on four key components: policy initialization, reward design, search, and learning. Policy initialization enables models to develop human-like reasoning behaviors, equipping them with the ability to effectively explore solution spaces for complex problems. Reward design provides dense and effective signals via reward shaping or reward modeling, which is the guidance for both search and learning. Search plays a crucial role in generating high-quality solutions during both training and testing phases, which can produce better solutions with more computation. Learning utilizes the data generated by search for improving policy, which can achieve the better performance with more parameters and more searched data. Existing open-source projects that attempt to reproduce o1 can be seem as a part or a variant of our roadmap. Collectively, these components underscore how learning and search drive o1's advancement, making meaningful contributions to the development of LLM.

연구 동기 및 목표

• OpenAI o1의 전문가 수준의 추론 성능을 가능하게 하는 핵심 요소를 규명하고 체계화하기.
• 강화학습 원리를 활용한 o1 재현을 위한 종합적이고 모듈화된 프레임워크 제공하기.
• 기존 오픈소스 o1 재현 프로젝트들이 이 로드맵과 어떻게 일치하거나 다름을 밝혀내기.
• 검색과 학습의 체계적 통합을 통해 언어 모델의 추론 능력 확장에 도전 과제 해결하기.

제안 방법

• 언어 이해, 세계 지식, 기본 추론 능력에 대한 사전 훈련을 통해 정책 초기화를 수행한 후, 목표 명확화 및 작업 분해와 같은 인간과 유사한 추론 행동을 심어주기 위해 지시어에 대한 피니팅 수행.
• 결과 기반 및 과정 기반 보상 모두를 사용한 보상 설계로, 보상 형태 조정, 학습된 보상 모델, 앙상블 방법을 통해 신호 밀도 및 일반화 능력을 향상.
• 자기 평가, 가치 함수, 히우리스틱 규칙 등 내부 및 외부 가이드라인을 기반으로 한 검색 전략과 함께 MCTS, Best-of-N, 비트 서치 등의 계획 알고리즘을 조합하여 고품질의 해답 후보 생성.
• 정책 그래เดียน트 방법(예: PPO, DPO)과 행동 클로닝을 통한 학습으로, 검색에 의해 생성된 데이터를 기반으로 정책을 피니팅하여 성능을 반복적으로 향상.
• 세계 모델을 통합하여 시뮬레이션 환경에서의 계획 및 검색을 가능하게 하여 실시간 환경 상호작용에 대한 의존도 감소.
• 다중 모odal 정보(예: 이미지)를 연속적 표현으로 통합하여 체인 오브 톰 추론에 통합함으로써 지연 시간을 최소화하고 일치도를 향상.

Figure 1: The overview of this roadmap including policy initialization, reward design, search and learning.

실험 결과

연구 질문

• RQ1 강화학습을 통해 o1의 추론 능력을 재현하기 위해 필요한 필수 요소는 무엇인가?
• RQ2 정책 초기화가 언어 모델에 인간과 유사한 추론 행동을 효과적으로 심어줄 수 있는 방법은 무엇인가?
• RQ3 복잡한 추론 과제에서 신호 밀도와 일반화 능력을 극대화하기 위한 보상 설계 전략은 무엇인가?
• RQ4 MCTS 및 Best-of-N과 같은 검색 전략이 계산 자원 확장 조건에서 해답 품질에 기여하는 방식은 무엇인가?
• RQ5 세계 모델은 실세계 에이전트 배포를 위한 확장 가능하고 안전하며 취소 가능한 계획 수립에 어떤 역할을 하는가?

주요 결과

• Marco-o1 및 o1-coder와 같은 기존 오픈소스 o1 재현 프로젝트들은 제안된 로드맵의 특정 구현 사례로 확인되어 그 일반성 검증.
• 거부 샘플링 또는 MCTS에서 생성된 데이터에 대해 DPO 및 SFT를 적용한 강화학습 방법은 추론 성능을 크게 향상시키며, 특히 디스틸리케이션과 결합할 경우 추가 성능 향상 기대.
• MCTS 중 자기 반성은 검색 효율성과 해답 품질을 향상시키며, 계획 과정에서 내부 피드백의 가치를 입증.
• 표준 답안이 존재하는 추론 과제에서는 결과 기반 보상 모델이 효과적이지만, 비추론 또는 일치도 관련 과제에서는 과정 기반 및 학습된 보상 모델이 필수적.
• 세계 모델은 실세계 환경에서 행동하기 전에 상호작용을 시뮬레이션할 수 있도록 해주어 안전하고 취소 가능한 계획 수립을 가능하게 하며, 실세계 배포에 매우 중요.
• 체인 오브 톰 추론에 다중 모달 데이터(예: 이미지)를 통합할 경우 추론 길이와 지연 시간이 증가하므로, 연속적이고 압축된 표현 방식의 필요성이 제기됨.

Figure 3: The visualization of the interaction between agent and einvironment in reinforcement learning for LLMs. Left: traditional reinforcement learning. Right: reinforcement learning for LLMs. The figure only visualizes the step-level action for simplicity. In fact, the action of LLM can be eithe