[논문리뷰] (작성중)QLASS- Boosting Language Agent Inference via Q-Guided Stepwise Search

 

Abstract 

• Language agents는 복잡한 상호작용 문제를 해결하는데 좋은 성능을 보임.
• Language agents의 성공의 핵심 요소
  • agentic workflow trajectory에 대한 reward model이다.
  • 훈련 또는 추론에 중요한 가이드를 제공
• 단점. 
  • 그러나 중간 상호작용에 대한 주석이 없기 때문에, 기존 연구들은 전체 trajectory에 대한 outcome reward model을 사용하여 정책을 최적화함. 이 방법으로는 비효율적인 정책 모델이 만들어질 수 있으며 성능이 낮을 수 있다.
• 제안 
  • QLASS를 제안, Q-value를 단계별로 추정하여 자동으로 annotation을 생성하는 방법
    • 탐색 트리를 도입
    • process reward modeling을 수행하여 각 중간 단계에서 가이드를 제공, 이 중간과정을 기반으로 Q-guided generation strategy를 제안, language agents가 long-term value에 포커스하도록 한다. 그
    • 결과, 복잡한 상호작용 task에서 추론 성능이 크게 향상
      • 절반 수준의 annotation 데이터만 사용해도 높은 성능을 유지할 수 있음을 입증하였으며, QLASS가 limited supervision 상황에서도 효율적으로 작동할 수 있음을 시사

 

1. Introduction

• 일반적인 SFT 사용방법 in LLM
  • LLM이 복잡한 agent task에서 효과적인 reasoning과 planning을 수행하도록 expert trajectory를 학습한다.
  • 그러나 훈련 데이터 수집을 위해 필요한 대량의 human annotation이 어려우며, 이로인해 성능 및 데이터 확보 측면에서 제한점이 있다. real-world interactions은 본질적으로 복잡하고 다양하여 데이터 부족 문제가 심각하게 나타난다.
• 극복 방법들 : 
  • self-improvement 기법이 연구되었으며, LLM이 인간의 개입 없이 자체적으로 생성한 데이터에서 학습을 목표로 함.
  • LLM이 test-time computation을 보다 확장함으로써 성능을 향상시키는 연구가 등장
  • 이러한 연구에서 영감을 받아, language agent의 inference-time search를 개선하는 방법에 초점을 맞춤
• inference time scaling 핵심 요소
  • reward model : self-explored data의 품질을 평가
    • 기존 연구 단점
      • 기존 연구들은 ground-truth를 기반으로 단일 outcome reward를 도출하는 방식을 사용
      • 이 방식은 단순하고 직관적이지만, 복잡한 agent task에서는 효과적으로 작동하지 않음
      • outcome-reward model은 긴 trajectory에서 각 단계별 행동을 정확하게 평가할 수 없다. 예를 들어, 최종 결과가 높은 보상을 받은 trajectory라고 해서 각 단계에서 최적의 행동을 수행했다고 보장할 수는 없다. agent가 주어진 task를 성공적으로 완료했더라도, 일부 행동은 비효율적일 가능성이 존재
    • 극복 방법 : 
      • 따라서, environmental feedback으로부터 학습할 수 있는 process reward model이 필요하다. 이는 agent가 복잡한 task의 중간 단계를 충분히 이해하고 학습할 수 있도록 한다. 이를 통해, 성능 및 일반화 능력을 향상시킬 수 있다.
      • 허나, process reward model을 구축하는 것은 stepwise reward에 대한 많은 양의 인간 주석이 필요함. 즉 인간 주석 없이 self-improvement을 해야하는데, 이게 어렵다.
      • 기존 연구
        • process reward modeling에 초점을 맞춤.
          이 방법들은 비용이 많이 드는 인간 주석이나 계산량이 큰 random rollouts을 필요로 하기 때문에 self-improvement을 위한 효율적인 방법이 되기 어렵다.
    • 제안 방법 : QLASS제안, 효과적인 process reward modeling수행. agent inference 가이드 하는 방법 개발
      • 구체적으로, self-generated exploratory trajectories를 exploration tree로 변환하고, 트리 구조를 활용하여 각 tree node의 process rewards를 업데이트하는 방식을 도입
      • 미래 효용 평가 : multi-turn reasoning 과정에서, 각 단계에서 미래 효용을 보다 정확하게 포착. Bellman equation을 적용하여 Q-based process reward를 학습. 기존의 단순한 outcome-based reward와 달리, Q-value는 각 행동의 즉각적인 결과뿐만 아니라 장기적인 보상까지 고려
      • Bellman Update Rule : 미래 보상을 과거 state로 역전파하여 Q-value를 반복적으로 개선. 이를 통해 희소한 또는 delayed feedback 신호 의존도를 줄이고, full trajectory 주석 없이도 state-action pair에 대한 효과적인supervision을 제공할 수 있다
      • Q value(Value Function)를 학습하고, QNet(1992)를 활용해, 특정 partial solution의 예상 보상을 예측하는 근사 함수를 훈련함. 예측된 Q-value가 높은 action을 우선적으로 수행하게 하여 agent가 보다 효율적인 stepwise planning을 수행하도록 함.
        • 최근연구 KIMI-k1.5, DeepSeek-R1에서는 process reward modeling에 실패했다고 하나, 이 논문에서는 process reward modeling이 필수적이라고 주장함.
          • back-and-forth agent behavior은 본직적으로 단계별 비효율성을 발생시키는데, 이는 희소한 outcome reward모델만으로는 감지할 수 없다.
          • 따라서 Q-value estimation방식이 필요하고, trajectory내에 미래에대한 효용을 역전파하여, 중간에 불필요한 가지들을 가지치기하면서도 동시에 중요한 의사결정 포인트는 유지할 수 있도록 한다.
          • Q value 
            현재 행동이 미래 보상에 미치는 영향을 평가함 : 즉각적인 보상을 평가하는 reward model과는 다름.
            
            Q(s,a)로 평가하며, 상태s에서 행동a를 했을 때 기대되는 총 보상(Q-value)를 의미함.
            즉 미래에 예상되는 모든 reward의 가중합을 의미함

 

• 기여 : 
  • Process Reward Modeling with Q-Value Estimation :
    
    • Q-value estimation을 활용하 Language agent의 중간 annotation을 생성하는 새로운 전략 QLASS제안. 모델 추론 과정에서 reasoning step에서 가이드를 제공한다.
  • Q-Guided Generation Strategy
    • Q-guided generation을 통해 stepwise reward를 활용한 reasoning을 수행하도록 하여, 보다 효과적인 결정을 함
  • Superior Performance with Limited Supervision
    • WebShop, ALFWorld, SciWorld와 같은 상황에서 좋은 성능. annotation이 제한된 환경에서도 효율적으로 동작

 

2. Related Work

2.1 LLM Agent

2.2 Self-Improvement for LLM

2.3 Process Reward Modeling for LLM

• Human Annotation
• Random Rollouts
• MCTS with re-inforced training
• Step-level guidance for agent inference through training a step level value model
• Hybrid process reward modeling

기존 연구와 QLASS와 차별점

• Proprietary LLM없이 효과적인 stepwise reward modeling 수행 : closed source LLM활용 없이 preocess reward학습 가능
• 복잡한 agent환경에서 적용 가능
• 더 단순한 구조로 더 나은 process reward modeling

 

3. Preliminaries

Q-function

 

Bellman Optimality Equation of Q-function

value model Q(s_t, a_t) is updated iteratively

• Q-learning에 대해서 소개. 탐색 트리(exploration tree)에서 QLASS를 사용해 Q-value를 추출.
• Q-learning : 
  • 1992년 model-free 강화학습 알고리즘.
  • Q-value : given state에서 특정 행동을 수행할 떄 기대되는 미래 보상을 나타냄.
    
    • (1)식 : 최적의 Q함수는 주어진 상태(s)에서, 행동(a)를 수행한 이후, 장기적인 기대 보상의 최대값을 나타냄.
      • $\pi$ : 정책(모델)
      • $γ$ : 할인율
      • $r_t$ : t번째 스텝에서 받은 보상.
• Bellman Optimality Equation
  • 현 상태(t)에서 특정 행동 $a_t$를 수행한 이후 Q-valu는 즉각적인 보상 $r_t$와 미래 상태 $s_{t+1}$에서 가능한 최적 Q-value의 조합으로 계산된다.
  • 논문에서는 너무 많은 부분을 생략하였는데, 마르코프 상태라고 가정하면, 현 상태(t)와 다음 스텝(t+1)만 고려하면 된다. 결국에는 Q가 학습이 되면서 현상태의 Q-value와 t+1의 Q-value만을 고려하면 되도록 학습이 된다.
    -> Bellman equation으로 수렴이 되는 조건이 있는데, 까다로운 조건을 neural net의 유동적인 적응력으로 대체해버림.
    (너무 오래전에 공부해서 디테일한 부분이 기억이 나지 않음.)
• Q-learning 업데이트 방법
  • Value function Q(s,a)를 반복적으로 업데이트하며 학습을 진행한다. 업데이트 규칙은 (3)번식이다.
  • $\alpha$ : 학습률, A : 가능한 행동들의 집합
• Language Agent Task에서 Q-learning 적용 방법
  • 복잡한 환경에서 agent가 즉각적인 보상뿐아니라, 장기적인 여파까지 고려해야 한다.
  • 그러나 기존 Q-learning 알고리즘을 language agents에 직접 적용하는 것은 비효율적이다.
  • 이유 : 
    • 탐색 공간이 엄청 크다.
      • 일반적으로 강화 학습 환경에서는 행동 공간이 유한한 개수로 제한된다.
      • 허나, language agent task에서는 가능한 행동이 무한하다.(자연어 기반이니까 토큰단위, 문장단위)
      • 결론 : 전통적인 greedy 탐색 으로는 동작하지 않을 수 있음. 무작위 샘플링은 의미 없는 결과를 초래함.
    • 보상이 희소하다.
      • Language agent환경에서는 즉각적인 보상이 거의 없고, 대부분 trajectory의 마지막에서 보상이 주어짐.
      • 결론 : Q-learning이 학습되기 위한 자주 발생하는 보상 데이터가 부족하여 학습이 불안정해짐.
    • Action Space 너무 크다.
      • Q-value는 (상태, 행동)조합을 학습되나
      • 결론 : Language에서는 행동이 - 단어 or 문장 단위로 표현되어 exploration이 매우 비효율적임.
  • 해결방법 : 
    • 탐색 트리(exploration tree)도입하 q-value 추출을 최적화한다.
      • 완전한 trajectory를 샘플링하는 대신, 보다 효율적인 search strategy를 수행
      • 미래 보상을 backpropagation하여, 단계별 학습이 가능하도록 한다.

 

overview

4. QLASS Pipeline Details

4.1 Behavioral Cloning

• exper trajectories에 대해 SFT 수행하여 language agents의 초기 성능을 강화(reasoning alignment라고 생각함)
• 과정
  • Policy $\pi$ : 전문가 annotated 데이터 집합 $D_{expert}$를 사용하여 SFT
    • ReAct 스타일의 데이터를 활용하여 SFT수행. 각 액션을 실행하기전에 CoT 스타일의 reasoning path를 생성하도록 한다. 이 논문에서는 간결성을 위해 ReAct스타일로 생성된 전체 응답을 a로 나타냄.
    • $D_{expert}$ : = $\{{(u_i, a^i_{t}, a^i_{t})^T_{t=1}}\}^N_{i=1}$
      • $u_i$ : task 설명
      • T : 경로의 길ㄹ이
      • N : 전문가 데이터셋 내 trajectory개수
      • $o^i_t$ : t번째 step에서 행동 a^i_t를 수행한 후 환경에서 받은 observation
    • Object function : policy최적화를 위해 negative log likelihood loss사용 ( SFT했네 )
      단어 수준 SFT는 아닌것으로 보이고, reasoning단계별 SFT로 보임.

 

4.2 Constructing an Exploration Tree

SFT로 학습된 language agent는 환경을 탐색하며 대량의 trajectory를 수집할 수 있음. 그러나 language agent의 search space는 매우 커서 무작위로 trajectory를 샘플링하면 효율이 매우 낮을 수 있음.

이를 해결하기 위해 self-generation단게에서 탐색 트리를 구축하는 방식을 제안함.

• 하나의 trajectory : task설명 (symptom)을 root node로 하고
  • 이후 각 단게에서의 상태, 행동, 등 관련 정보들이 포함된 노드들로 구성된 트리로 볼 수 있다.
• 트리노드는 다음을 포함
  • 상태 : 현재 시점 t까지 축적된 context, 이전의 모든 reasoning path및 행동 데이터
    • $s_t = {u, a_1, o_1, ..., a_{t_1}, o_{t_1}}$
      • u : task 설명(symptom)
      • $a_i$ : 이전 단계에서 수행한 행동.
      • $o_i$ : 각 행동 후 반환된 관측 결과
    • $a_t$
      • 현재 노드에서 수행한 특정 행동으로 이후 상태에 영향을 끼침.
      • $a_t$는 policy에 의해 선택되며 이전 reasoning path 및 현재 상태를 기반으로 결정
    • $r_t$
      • $a_t$ : 수행후 환경으로 받는 즉각적 보상.
      • 대부분 language agent task에서 즉각적 보상은 0이거나 드뭄(최종장에서만 0~1사이 보상을 받음)
    • C : 현 노드에서 탐색된 다음 단계의 노드.
    • Q-value : 
      • s_t에서 특정행동 a_t를 수행했을 떄 기대되는 총 미래 보상

 

4.2.2 Tree Construction

• 탐색 트리는 task 루트에서 시작해. 모든 generation(branch)를 포함하는 구조를 갖음.
• 새로운 trajectory를 탐색하는 과정은, 기존 TreeNode에서 새로운 가지를 확장하는 과정으로 볼 수 있음.
  • 새로운 자식 노드 탐색 : 현재 트리에 존재하지 않은 action-state 조합을 찾아 트리에 추가
  • 불필요한 가지 정리 : 최종 보상이 0인 trajectory의 경우 탐색 중단
• 트리 가지치기
  • 탐색 트리는 매우 깊은 구조를 갖는 경우가 많아 가지치기가 필요함
    • 초기 단계에서만 탐색 트리 확장함. (3-5 step에서만 확잦ㅇ)
  • 0 보상을 받은 경로 탐색 중단
• Extracting Q-values
  • 탐색 트리가 구축된 후 각 중간 노드의 Q-value를 추정하는 과정을 수행함.
    • 과거 : 대부분은 trajectory의 마지막에서 outcome reward가 주어짐. 이것만 계산함
      • $Q(s_t, a_t) = r_T$
    • 논문 : 마지막 reward를 기반으로 중간 노드 Q-value를 역산함.
      • 예시)
        
        Q(s3, a3) = 1 (마지막 step)
        Q(s2, a2) = r2 + γ(Q(s3, a3))
                         = r2 + 0.1(1)             -> Q(s2,a2) = 0.9로 역산할 수 있음.

 

 

 

 

 

허나 Interactive 환경이 아닌 상태에서는 이 논문을 적용하는데 제한적임.

Offline RL방법론을 사용해야하나?

• 기존 데이터셋을 활용해 Exploration Tree를 구축하기
  • 과거 데이터 내에서 탐색 트리를 구축하는 방법을 사용할 수 있음.
  • 주어진 데이터셋에서 가능한 모든 action-state 경로를 트리 구조로 정리하는 방식.
  • 새로운 action을 시도할 수 없기 때문에, 탐색 트리의 깊이를 제한하거나, 과거 데이터에서 자주 등장한 action을 우선적으로 선택하는 방식이 필요.
• Process Reward를 IRL 기반으로 추정
  • Offline RL에서는 stepwise reward가 주어지지 않음 (최종 결과만 있는 경우).
  • 방법론은 찾아봐야 할 것 같음.
• Conservative Q-Learning (CQL)과 QLASS를 결합