Speculate While You Reason: Teaching Agents to Predict Their Next Tool Call via Joint Agent-Speculator RL

기술 요약: 추론하는 동안 추측하라 (Speculate While You Reason)

문제 정의

대규모 언어 모델(LLM) 에이전트는 외부 도구 호출 결과(예: 웹 검색, 데이터베이스 쿼리, API 호출)를 기다리는 동안 상당한 실시간 지연 시간(wall-clock latency)을 초래하곤 한다. 유망한 완화 전략은 **도구 호출 추측(tool-call speculation)**으로, 시스템이 에이전트의 추론과 병렬로 다음 도구 호출을 예측하고 미리 실행하는 것이다. 만약 예측이 에이전트가 최종적으로 수행할 호출과 정확히 일치한다면, 미리 실행된 결과를 재사용하여 지연 시간을 효과적으로 숨길 수 있다.

그러나 기존의 추측 방법들은 추측기-에이전트 간극(speculator–agent gap) 문제를 겪는다. 현재의 접근 방식들은 다음 사항들에 의존한다:

1. 외부 드래프트 모델(External draft models): 일반적인 도구 사용 어시스턴트를 모사하는 더 작은 규모의 LLM이지만, 배포된 에이전트의 구체적인 동작과 정렬(align)되는 데 실패한다.
2. 캐싱된 트레이스(Cached traces): 서로 다른 사용자 쿼리, 추론 상태 또는 오래된 환경에 대응할 수 있는 과거의 도구 호출 로그들이다.

이러한 방법들은 시스템 오버헤드(추가적인 모델 가중치, 별도의 KV 캐시)를 발생시키며, 타겟 에이전트가 수행할 특정한 다음 구조화된 호출을 예측하는 데 실패하는 경우가 많아 추측의 효용성을 제한한다.

방법론

본 논문은 별도의 파라미터나 KV 캐시를 요구하지 않고 두 가지 뚜렷한 모드로 작동하도록 훈련된 단일 모델인 **자기 추측 에이전트(Self-Speculating Agent)**를 제안한다:

1. 에이전트 모드(Agent Mode): 모델은 추론, 도구 호출, 관찰, 그리고 최종 답변을 교차하며 과업을 해결한다.
2. 추측기 모드(Speculator Mode): 중간 단계의 궤적(trajectory)과 고정된 "추측 접미사(speculation suffix)"가 주어지면, 모델은 다음 구조화된 도구 호출(도구 이름 및 인자)을 직접 예측한다.

공동 에이전트-추측기 강화 학습 (Joint Agent–Speculator RL)

두 목적을 위해 성능 저하 없이 단일 정책($\pi_\theta$)을 훈련하기 위해, 저자들은 공동 에이전트-추측기 RL 프레임워크를 제안한다:

• 온-폴리시 데이터 생성(On-Policy Data Generation): 추측 대상은 에이전트 자신의 현재 롤아웃(rollout)에서 직접 유도된다. 샘플링된 궤적의 모든 도구 호출 턴 $t$에 대해, 접두사 $h_t$를 사용하여 추측기 입력을 구성하며, 에이전트의 실제 다음 호출 $a_t$가 정답 타겟 역할을 한다.
• 보상 산정(Reward Formulation):
  • 에이전트 보상 ($R_{ag}$): 최종 과업 성공 여부(예: 답변의 정확성)에 기반한다.
  • 추측 보상 ($R_{sp}$): 이름 조건부 인자 보상(name-conditioned argument reward)이다. 이는 도구 이름의 정확한 일치($R_{name}$)를 요구하며, 인자 딕셔너리에 대한 매크로 토큰-F1 점수($R_{args}$)를 계산한다. 최종 보상은 $R_{sp} = R_{name} \cdot R_{args}$이다.
• 교대 최적화(Alternating Optimization): 훈련은 에이전트 업데이트와 추측기 업데이트 사이를 교대하며 진행된다(예: 4:8 스케줄). 이는 두 모드가 서로의 최적화 지형을 방해하는 것을 방지한다.
• 안정화 기술(Stabilization Techniques):
  1. SFT 웜업(SFT Warmup): 성공적인 궤적과 추측 예시를 혼합하여 사전 지도 미세 조정(SFT) 단계를 거침으로써, 두 모드 모두에서 잘 형성된 동작을 확립한다.
  2. 옵티마이저 리셋(Optimizer Resets): 훈련 스케줄이 에이전트 모드에서 추측기 모드로 전환될 때마다 옵티마이저 상태를 리셋하여, 한 목적 함수에서 다른 목적 함수로 모멘텀이 전이되는 것을 방지한다.
  3. 반복에 대한 페널티(Penalty for Repetition): 에이전트 보상에는 생산적이지 않은 루프에 빠지는 것을 방지하기 위해 반복적이고 무의미한 도구 호출에 대한 페널티가 포함된다.

주요 기여

1. 추측기-에이전트 간극 식별: 저자들은 기성(off-the-shelf)의 작은 드래프트 모델들이 타겟 에이전트의 다음 호출 분포와 일치하지 못하는 반면, 타겟 에이전트 자체가 자신의 행동을 예측하는 데 더 강력한 예측기임을 입증하였다.
2. 자기 추측 에이전트 아키텍처: 과업 해결과 다음 호출 예측 모드 사이를 전환하며, 접두사 KV 캐시를 완전히 재사용하고 보조 모델이나 이력 트레이스 캐시를 제거하는 통합된 모델 설계를 제공한다.
3. 공동 에이전트-추측기 RL: 에이전트 자신의 롤아웃으로부터 추측 타겟을 유도하고, 과업 해결 성능을 유지하면서 추측 정확도를 높이기 위해 업데이트를 교대하는 훈련 방법론을 제시한다.

실험 결과

본 방법론은 Qwen3-4B 및 Qwen3.5-4B 백본을 사용하여 Agentic SearchQA(HotpotQA, MuSiQue, BrowseComp-Plus) 및 Conversational Tool-Use(τ-bench Airline 및 Retail) 과업에서 평가되었다.

• 추측 성능: 공동 RL 방법은 Hit@1 점수(도구 이름과 인자의 정확한 일치)를 크게 향상시켰다.
  • Qwen3-4B: 44.1(SFT 베이스라인)에서 61.2로 향상.
  • Qwen3.5-4B: 48.9(SFT 베이스라인)에서 66.3으로 향상.
  • 이 결과들은 외부 추측기(동일 계열의 더 작은 모델) 및 자신의 호출을 추측하도록 프롬프트된 기본 에이전트보다 우수한 성능을 보였다.
• 과업 성능: 결정적으로, 다운스트림 과업 성공률은 안정적이거나 약간 향상되었다(예: Qwen3-4B 과업 성공률이 26.6에서 27.7로 이동). 이는 이중 모드 훈련이 에이전트의 주요 과업 해결 능력을 저하시키지 않았음을 확인시켜 준다.
• 효율성: 자기 추측은 두 번째 모델의 가중치와 별도의 KV 캐시를 필요로 하지 않음으로써 서빙 오버헤드를 줄이는 동시에, 접두사 캐시 재사용을 유지하였다.

의의 및 주장

본 논문은 도구 호출 추측이 에이전트 자체에 직접 통합될 수 있음을 주장한다. 에이전트와 추측기를 단일 모델 내에 통합하고 온-폴리시 RL을 통해 공동으로 훈련함으로써, 시스템은 외부 드래프트 모델과 관련된 시스템 오버헤드나 정렬 문제를 피하면서 높은 정확도의 다음 호출 예측을 달实现한다.

저자들은 이 연구를 읽기 전용 도구(검색, 추출, 데이터베이스 조회)에 특히 적합한 **추측적 도구 사용(speculative tool use)**을 향한 첫걸음으로 규정한다. 이러한 도구들은 잘못된 사전 실행이 발생하더라도 안전하게 폐기될 수 있기 때문이다. 또한, 외부 상태를 변경하는 도구(예: 주문 넣기)의 경우 드라이 런(dry-run) 모드나 트랜잭션 롤백과 같은 추가적인 안전장치가 필요함을 인정하며, 현재 평가는 4B 규모의 모델과 특정 과업군으로 제한되어 있음을 명시하였다.