💬 NLP

TRACE: A Unified Rollout Budget Allocation Framework for Efficient Agentic Reinforcement Learning

이 논문은 상호작용을 트리 구조의 노드로 모델링하고, 보상 대비를 높이고 정책 학습 효율을 향방하기 위해 혼합된 터미널 보상을 사용하여 프롬프트 루트와 중간 접두사를 모두 동적으로 타겟팅함으로써 멀티 턴 에이전틱 강화 학습에서 롤아웃 예산 할당을 최적화하는 통합 프레임워크인 TRACE를 소개한다.

원저자: Heming Zou, Qi Wang, Yun Qu, Yuhang Jiang, Lizhou Cai, Yixiu Mao, Ru Peng, Xin Xu, Weijie Liu, Kai Yang, Saiyong Yang, Xiangyang Ji

게시일 2026-06-10

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Heming Zou, Qi Wang, Yun Qu, Yuhang Jiang, Lizhou Cai, Yixiu Mao, Ru Peng, Xin Xu, Weijie Liu, Kai Yang, Saiyong Yang, Xiangyang Ji

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 매우 똑똑하지만 경험이 부족한 학생(AI)에게 복잡한 퍼즐(수학 문제나 거대한 도서관에서 답을 찾는 일 등)을 푸는 법을 가르치려는 선생님이라고 상상해 보세요. 이 학생은 무언가를 시도하고, 실수를 하고, 마지막에 "예" 또는 "아니오"라는 단순한 결과를 얻으며 학습합니다.

문제는 선생님에게 학생이 연습할 수 있는 제한된 시간(예산)이 있다는 점입니다. 만약 선생님이 학생에게 똑같이 쉬운 퍼즐을 100번 풀게 하거나, 똑같이 불가능한 퍼즐을 100번 풀게 한다면, 학생은 아무것도 배우지 못합니다. 학생에게는 성공할 수도 있고 실패할 수도 있는, 즉 '적절한 난이도'의 퍼즐을 제공해야 합니다.

이 논문은 이 문제를 해결하기 위한 새로운 방법인 TRACE를 소개합니다. 이해를 돕기 위해 간단한 비유를 들어 설명하겠습니다.

1. 문제점: "지루한" 경로에 시간을 낭비하는 것

과거에는 AI를 가르칠 때, 연구자들이 퍼즐 하나를 골라 AI가 처음부터 끝까지 풀어보도록 했습니다.

문제점: 퍼즐이 너무 쉬우면 AI는 항상 이깁니다. 반대로 너무 어려우면 AI는 항상 집니다. 두 경우 모두, 마지막에 나오는 "예/아니오"라는 대답은 AI에게 어디서 틀렸는지를 알려주지 못합니다.
기존 방식: 연구자들은 더 나은 시작 퍼즐을 고르려고 노력했지만, 일단 AI가 풀기 시작하면 끝까지 그냥 내버려 두었습니다. AI가 중간에 막히고 있는지 확인하기 위해 멈추지는 않았습니다.

2. 해결책: "가지가 뻗어 나가는 길" (TRACE)

TRACE는 AI의 시도를 단순한 직선이 아니라, 많은 가지가 뻗어 나온 나무로 취급함으로써 게임의 판도를 바꿉니다.

AI가 보물을 찾기 위해 산을 오르는 등산객이라고 상상해 보세요.

뿌리 (시작점): 먼저, TRACE는 시작점(퍼즐)을 살펴봅니다. 그리고 예측합니다: "이 퍼즐이 성공과 실패가 섞여 있을 가능성이 높은가?" 만약 너무 쉽거나 너무 어렵다면, 이를 건너뜁니다. 만약 적절한 "학습용" 퍼즐이라면, AI를 산 위로 보냅니다.
가지 (중간 과정): 이것이 핵심적인 부분입니다. AI가 산을 오르다 보면 길이 갈라지는 지점(결정을 내려야 하는 '선택의 순간')에 도달합니다. 이때 TRACE는 멈춰서 질문합니다: "만약 AI가 이 특정 경로를 택한다면, 그것이 승리로 이어질까 아니면 패배로 이어질까?"
- 만약 그 경로가 확실히 승리로 이어지거나 확실히 패배로 이어질 것으로 보인다면, TRACE는 그 경로를 더 탐색하며 시간을 낭비하지 않습니다.
- 만약 그 경로가 불확실하다면(성공할 확률과 실패할 확률이 반반이라면), TRACE는 이렇게 말합니다: "이 특정 경로로 더 많은 등산객을 보내서 어떤 일이 일어나는지 확인해 보자!"

3. "수정구슬" (예측 모델)

TRACE는 어떤 경로가 불확실한지 어떻게 알까요? 바로 "수정구슬"(예측 모델)을 사용합니다.

이 수정구슬은 지금까지의 등산 기록(AI가 취한 생각과 행동들)을 살펴봅니다.
그리고 성공 확률을 추정합니다.
만약 수정구슬이 "여기서 성공할 확률이 50%입니다"라고 말한다면, 그곳이 바로 더 많은 시간을 투자해야 할 완벽한 지점입니다. 이는 AI가 승리하는 경로와 패배하는 경로를 비교하며 배울 수 있는 '학습 구역'에 와 있음을 의미합니다.

4. 결과: 적은 노력으로 더 똑똑하게 배우기

제한된 시간을 오직 "불확실한" 부분에만 집중함으로써, TRACE는 풍부한 **대조(contrast)**의 지도를 만들어냅니다.

AI는 단순히 "실패했다"라고 아는 것에 그치지 않고, "내가 갈림길에서 왼쪽 길을 택했기 때문에 실패했구나, 하지만 오른쪽 길을 택했다면 성공했을 거야"라고 배우게 됩니다.
이는 AI가 학습할 수 있는 훨씬 강력한 신호를 만들어내며, 이 모든 과정은 이전 방식과 동일한 양의 연습 시간(예산)을 사용하면서도 가능합니다.

요약하자면

TRACE를 단순히 운동장을 무작위로 뛰는 선수가 아니라, 영리한 코치라고 생각해보세요.

적절한 경기를 선택합니다: 도전적이면서도 승산이 있는 경기를 고릅니다.
까다로운 코스에서 멈춥니다: 선수가 달리는 것을 지켜봅니다. 만약 선수가 미끄러질 수도 있고 혹은 잘 버틸 수도 있는 까다로운 코너에 도착하면, 코치는 그 정확한 코너를 더 많은 선수들이 시도하게 하여, 미끄러지는 것과 잘 버티는 것의 차이를 확인하게 합니다.
시간을 절약합니다: 쉬운 직선 구간이나 불가능한 절벽 구간은 무시합니다.

이 논문은 이 방법을 사용함으로써 AI 모델(특히 Q-wen3)이 기존 방식과 동일한 양의 컴퓨팅 자원을 사용하면서도 수학, 다단계 질문, 도구 사용 능력에서 더 발전했음을 보여줍니다. TRACE는 평범하고 지루한 연습 세션을, 매 단계마다 새로운 것을 가르쳐주는 역동적이고 가지가 뻗어 나가는 탐험으로 바꾸어 놓았습니다.

기술 요약: TRACE - 효율적인 에이전트 강화 학습을 위한 통합 롤아웃 예산 할당 프레임워크

1. 문제 정의

검증 가능한 보상(Verifiable Rewards, RLVR)을 활용한 강화 학습은 대규모 언어 모델(LLM)의 추론 및 에이전트 능력을 향상시키는 핵심 메커니즘으로 부상했습니다. 그러나 표준 RLVR 파이프라인은 긴 사고 사슬(Chain-of-Thought, CoT) 궤적 생성과 상호작용의 간헐적 발생으로 인한 높은 계산 비용 때문에 상당한 계산 비효률성에 직면해 있습니다.

핵심 과제는 **보상 대비(reward contrast)**에 있습니다. 결과 기반 보상 설정(종료 시 단일 이진 보상이 할당되는 방식)에서는 학습 신호가 희소하고 불균형하게 나타납니다:

프롬프트 수준: 지나치게 단순하거나 지나치게 복잡한 프롬프트는 낮은 분산의 피드백(예: 모두 성공하거나 모두 실패함)을 생성하여, 정책 업데이트를 위한 신호를 거의 제공하지 못합니다.
궤적 수준: 긴 다회차 롤아웃에 단일 최종 보상을 할당하면 로컬 신용 할당(local credit assignment)이 모호해집니다. 중간 단계의 피드백 없이는 모델이 어떤 특정 턴이나 결정이 최종 결과에 기여했는지 구분할 수 없습니다.

기존 방법들은 유망한 프롬프트에 롤아웃 예산을 할당(프롬프트 선택)하거나 프롬프트당 롤아웃 횟수를 늘리는(롤아웃 횟수 할당) 방식으로 이를 완화하려 시도했습니다. 그러나 이러한 접근 방식은 롤아웃을 **원자적 궤적(atomic trajectories)**으로 취급합니다. 일단 프롬프트가 선택되면 전체 롤아웃이 선형적으로 생성되며, 결과가 불확실한 중간 단계(접두사, prefixes)에서 분기할 수 있는 가능성을 간과합니다. 이는 고정된 샘플링 예산 내에서 정보 밀도를 극대화하지 못하는 "평면적인" 롤아웃 수집 결과로 이어집니다.

2. 방법론: TRACE

저자들은 롤아웃 할당을 트리 구조의 프로세스로 재구상하는 통합 프레임워크인 **TRACE (Tree Rollout Allocation for Contrastive Exploration)**를 제안합니다. 독립적인 선형 궤적을 생성하는 대신, TRACE는 동일한 조건부 맥락 내에서 혼합된 최종 보상(성공과 실패 모두)을 관찰할 확률을 극대화하기 위해 **프롬프트 루트(root)**와 중간 접두사(prefix) 모두에 예산을 동적으로 할당합니다.

주요 기술 구성 요소:

트리 구조 롤아웃:
- 본 프레임워크는 ReAct 스타일의 상호작용(생각-행동-관찰)을 트리의 노드로 모델링합니다.
- 루트(Roots): 초기 프롬프트 ( $x$ ).
- 접두사(Prefixes): $t$ 번의 턴 후에 형성된 중간 히스토리 ( $H_t$ ).
- 분기(Branching): 방문된 모든 접두사에서, 시스템은 "사실적(factual)" 브랜치로부터 벗어나는 "반사실적(counterfactual)" 연속형(새로운 접미사)을 생성하기 위해 추가 예산을 할당할 수 있습니다.
통합 할당 목적 함수 (혼합 보상 대비):
- 목표는 조건부 성공 확률( $V^\pi_t$ )이 0이나 1이 아닌, 중간 단계인 앵커(루트 또는 접두사)에 예산을 할당하는 것입니다.
- 이론적 근점: 논문은 롤아웃 접미사를 따르는 조건부 성공 예측의 기대 이차 변동(expected quadratic variation)이 $V^\pi_t(1 - V^\pi_t)$ 와 같음을 입증합니다. 이 값은 "대비 잠재력(contrast potential)"을 나타냅니다. 대비 잠재력이 높은 앵커에 예산을 할당함으로써, 동일한 조건부 맥락 내에서 서로 반대되는 최종 결과를 가진 형제 브랜치를 관찰할 확률을 극대화하며, 이를 통해 최적화 도구를 위한 암시적 쌍체 선호(pairwise preference) 신호를 생성합니다.
2단계 할당 프로세스:
1. 글로벌 루트 할당: 샘플링 전, 예측기가 후보 프롬프트 풀에 대한 성공 확률을 추정합니다. 결정론적인 성공/실패를 보이는 경우를 제외하고, 혼합된 결과를 낼 가능성이 높은 프롬프트에 예산을 할당합니다. 예측 확률이 0 또는 1에 가까운 프롬프트는 예산을 받지 못합니다.
2. 로컬 접두사 확장: 활성화된 프롬프트에 대해, 시스템은 방문된 비종료 접두사를 식별합니다. 이후 이 접두사들로부터 추가적인 연속형을 생성하기 위해 로컬 예산을 할당합니다. 할당은 관찰된 사실적 결과가 예측된 확률과 다른 접두사(즉, 반사실적 브랜치가 결과를 뒤집을 가능성이 있는 곳)를 우선시합니다.
공유 가능한 일반화 예측기:
- 경량 예측기( $\tilde{V}_\psi$ )는 접두사 히스토리로부터 조건부 성공 확률을 추정합니다.
- 이는 수집된 롤아웃 트리에서 유도된 재귀적 타겟(최종 보상의 하향식 경험적 평균)을 통한 평균 제곱 회귀를 통해 훈련됩니다.
- 이 예측기는 루트 선택과 접두사 분기 결정을 모두 가이드합니다.
트리 인식 정책 최적화:
- 결과물인 트리 구조는 암시적인 단계별 선호 쌍(예: 동일한 접두사로부터의 성공적인 연속형 vs 실패한 연속형)을 제공합니다.
- 이는 표준 결과 기반 업데이트보다 더 조밀한 그룹 상대적 또는 트리 인식 그래디언트 신호를 활용하여 정책을 업데이트하는 데 사용됩니다.

3. 주요 기여

통합 공식화: 본 연구는 혼합 보상 대비 할당을 샘플 효율적인 RLVR의 단일 관점으로 공식화하여, 프롬프트 필터링, 롤아웃 횟수 할당, 그리고 접두사 수준의 분기를 롤아웃 트리 앵커에 대한 전략적 예산 결정으로 통합했습니다.
이론적 유도: 저자들은 공유된 자손 집합(descendant-set) 목적 함수를 바탕으로 할당 유틸리티를 유도하였으며, 접두사 수준의 정보가 그룹 난이도 예측을 엄격하게 개선하고, 불확실성을 가진 중간 지점에 예산을 할당하는 것이 대비 잠재력을 극대화함을 증명했습니다.
TRACE 프레임워크: 이 관점을 예측기 가이드형, 예산 제한형 트리 롤아웃 프레임워크로 구현하여, 수작업으로 제작된 프로세스 보상 없이도 암시적 신용 신호를 강화하며 글로벌 루트 할당과 로컬 접후사 확장을 통합했습니다.

4. 실험 결과

저자들은 Qwen3-8B 및 Qwen3-14B 백본을 사용하여 세 가지 다회차 에이전트 벤치마크(수학적 추론: DeepScaler, AIME 등; 다단계 QA: HotpotQA, 2Wiki 등; 함수 호출: BFCL v4)에서 TRACE를 평가했습니다.

성능 향상: 동일한 샘플링 예산 하에서, TRACE는 표준 GRPO, 프롬프트 선택 방법(PCL), 무작위 트리 기반 방법(TreePO)을 포함한 베이스라인들을 일관되게 능가했습니다.
- 예시: Multi-Hop QA에서, TRACE는 Qwen3-14B의 평균 정확도를 경쟁 베이스라인 대비 2.8 포인트 향상시켰습니다.
- 예시: 수학적 추론에서, TRACE는 Qwen3-8B의 정확도를 70.0%(GRPO)에서 71.1%로 개선했습니다.
효율성 및 대비: TRACE는 유효 비율(Effective Ratio)(성공과 실패 결과를 모두 포함하는 프롬프트/롤아웃의 비율)을 크게 높였습니다.
- 수학적 추론(Qwen3-8B)에서 유효 비율은 26.8%(GRPO)에서 60.6%(TRACE)로 상승했습니다.
- 이는 TRACE가 포화된(결정론적인) 앵커로부터 대비가 풍부한 영역으로 예산을 성공적으로 재배치했음을 나타냅니다.
절제 연구(Ablation Studies):
- 루트 할당과 접두사 확장 단계 모두 성능 향상에 독립적으로 기여했습니다.
- 프레임워크는 다양한 예산 규모에 대해 견고함을 보였으며, 예산의 "형태"(루트 커버리지의 너비 vs 연속형의 깊이)가 결정적임을 입으로 증명했습니다.

5. 의의 및 주장

본 논문은 TRACE가 에이전트 RLVR 샘플링을 "얼마나 많은 독립적인 롤아웃을 뽑을 것인가"의 문제에서 "롤아웃 트리 내의 어디에서 분기할 것인가"의 문제로 재정의한다고 주장합니다.

더 조밀한 학습 신호: 확장 가능한 결과 보상을 더 조밀한 혼합 보상 대비 및 암시적 단계별 선호 쌍으로 변환함으로써, TRACE는 계산 비용을 증가시키지 않고도 정책 업데이트 신호를 증폭합니다.
일반화 가능성: 이 접근 방식은 루트와 접두사 모두에 대한 공유 예측기에 의존하며, 이는 접두사 히스토리가 분기 결정을 가이드할 수 있는 학습 가능한 로컬 불확실성 개념을 보유하고 있음을 보여줍니다.
한계: 저자들은 TRACE가 명확한 최종 검증이 있는 결과 기반 RLVR을 위해 설계되었음을 언급했습니다. 명확한 최종 보상이 없는 작업은 대비 목적 함수를 재검토해야 할 수 있습니다. 또한, 현재의 예측기 구현은 바닐라 설정이며, 특정 에이전트 작업을 위한 더 강력한 예측기 개발이 향리 연구 분야임을 밝혔습니다.

요약하자면, TRACE는 동적으로 트리 구조 롤아웃을 구축하여 불확실성을 타겟팅하고 보상 대비를 극대화함으로써, 다회차 에이전트 RL에서 고정된 롤아웃 예산의 정보 가치를 극대화하는 체계적인 방법을 제공합니다.