기술 요약: PriorZero

문제 정의

대형 언어 모델 (LLM) 을 강화 학습 (RL) 에이전트에 통합하는 것은 풍부한 세계 지식을 활용하여 장기적, 희소 보상 의사결정 과제를 해결할 수 있는 유망한 길을 제시합니다. 그러나 기존 접근법은 근본적인 **선행 지식 - 역동성 불일치 (prior-dynamics mismatch)**로 인해 방해받고 있습니다.

현재 패러다임은 다음과 같은 구체적인 구조적 병목 현상에 직면해 있습니다:

1. 순진한 정책 (P1): LLM 을 제로샷 정책으로 사용하면 언어적 확률이 환경 보수와 동일하지 않기 때문에 근시안적인 결정이 이루어져 "알고 행동하는 간극 (knowing-doing gap)"이 발생합니다.
2. 순진한 RL 파인튜닝 (P2): RL 을 통한 LLM 의 엔드투엔드 파인튜닝은 희소 보상 환경에서 특히 고분산 크레딧 할당과 최적화 불안정성에 시달립니다.
3. 역동성 모델링 (P3): LLM 을 사용하여 상태 전이를 예측하는 것은 세분성 불일치로 인해 실패합니다. LLM 은 토큰 수준의 텍스트 생성으로 사전 훈련되는 반면, 잠재 역동성 모델은 상태 - 행동 수준의 전이 예측이 필요합니다.
4. 텍스트 인코더 (P4): 고정된 LLM 특징을 정적 표현으로 사용하면 작업별 적응이 제한되어 성능 상한선이 부과됩니다.

핵심적인 과제는 LLM 의 정적, 의미론적 선행 지식과 효과적인 계획에 필요한 동적, 환경 특이적 전이 모델 간의 간극을 해소하는 것입니다.

방법론: PriorZero 프레임워크

PriorZero 는 LLM 과 세계 모델을 분리하고 분리된 롤아웃 - 훈련 설계를 통해 통합하는 통합 프레임워크를 제안합니다. 이 시스템은 LLM 과 환경 사이에 학습 가능한 세계 모델을 배치하여, LLM 이 의미론적 지침을 제공하는 동안 세계 모델이 전이 모델링, 가치 추정, 그리고 다단계 계획을 처리하도록 합니다.

1. 루트 - 선행 지식 주입 (Rollout Phase)

롤아웃 단계에서 PriorZero 는 잠재 공간에서 몬테카를로 트리 탐색 (MCTS) 을 수행하기 위해 잠재 세계 모델 (UniZero 기반) 을 활용합니다.

• 메커니즘: 새로운 루트 - 선행 지식 주입 (Root-Prior Injection) 메커니즘은 LLM 의 의미론적 행동 선행 지식을 MCTS 의 루트 노드에서만 융합합니다.
• 구현: LLM 은 현재 텍스트 컨텍스트를 기반으로 상황 분석 (Chain-of-Thought) 과 행동 분포 ($\pi_{LLM}$) 를 생성합니다. 이를 루트 노드에서 세계 모델의 정책 ($\pi_{WM}$) 과 융합합니다:
  $$P_{root}(a | z_t) = (1 - \alpha) \pi_{WM}(a | z_t) + \alpha \pi_{LLM}(a | C_t)$$
• 제약: 탐색 트리의 비루트 노드는 전적으로 세계 모델의 잠재 역동성에 의존합니다. 이는 LLM 이 상상된 잠재 상태에 대한 상태 전이를 예측하도록 강요하지 않으므로 패러다임 P3 의 "세분성 불일치"를 방지합니다.
• 이점: 이는 MCTS 가 무의미한 궤적 (예: 보상 루프) 에 빠지는 것을 방지하면서도 세계 모델의 깊은 전망 능력을 유지하면서 검색을 의미 있는 행동으로 안내합니다.

2. 교대 강화 파인튜닝 (Training Phase)

패러다임 P2 의 크레딧 할당 실패를 해결하기 위해 PriorZero 는 세계 모델 학습과 LLM 적응을 분리하는 교대 파인튜닝 (Alternating Fine-Tuning) 전략을 사용합니다.

• 일정: 훈련은 비대칭 일정 (예: LLM 업데이트 1 회당 세계 모델 업데이트 10 회) 을 따릅니다.
• 세계 모델 학습: 세계 모델은 상호작용 데이터를 기반으로 역동성, 가치, 그리고 정책 예측을 정제하기 위해 업데이트됩니다.
• LLM 적응: 고분산 온-정책 보상을 사용하는 대신, LLM 은 세계 모델의 가치 추정에서 파생된 저분산 이점 신호를 사용하여 파인튜닝됩니다.
  • 이점은 $n$-step 부트스트랩 TD 타겟을 사용하여 계산됩니다: $A(z_t, a_t) = Q_n(z_t, a_t) - v_\theta(z_t)$.
  • LLM 은 클리핑된 PPO 목적 함수를 통해 업데이트되며, 행동 토큰에는 전체 가중치를 적용하고 Chain-of-Thought 추론 토큰에는 안정성을 유지하기 위해 감소된 가중치 ($w_{cot}=0.1$) 를 적용하는 특정 가중치 체계를 따릅니다.
• 이점: 이는 세계 모델이 LLM 을 위한 안정적이고 세분화된 감독을 제공하는 폐쇄 루프를 생성하여, 엔드투엔드 최적화의 불안정성 없이 LLM 의 선행 지식이 환경 역동성과 함께 진화할 수 있게 합니다.

주요 기여

1. 패러다임 분석: 이 논문은 LLM 지식의 장기적 RL 로의 이전에서 중심적인 장애물인 선행 지식 - 역동성 불일치를 식별하고, 기존 접근법 (P1–P4) 과 그 특정 실패 모드를 분류합니다.
2. PriorZero 프레임워크: 두 가지 핵심 메커니즘을 통해 LLM 선행 지식과 세계 모델 계획을 시너지 있게 결합하는 새로운 아키텍처:
   • 루트 - 선행 지식 주입: 잠재 역동성을 오염시키지 않고 탐색을 용이하게 하기 위해 MCTS 루트에 의미론적 지침을 주입합니다.
   • 교대 강화 파인튜닝: LLM 적응을 위한 안정적이고 저분산인 크레딧 할당 신호를 구축하기 위해 세계 모델 가치 추정을 사용합니다.
3. 실증적 검증: Jericho(텍스트 기반 어드벤처 게임) 와 BabyAI(그리드 월드 지시 따르기) 에 대한 포괄적인 실험을 통해 PriorZero 가 베이스라인에 비해 탐색 효율성과 점근적 성능을 일관되게 향상시킨다는 것을 입증했습니다.

실험 결과

• Jericho 벤치마크: PriorZero 는 네 가지 다른 게임 (Detective, Acorncourt, Zork1, Omniquest) 에서 LLM 선행 지식이 없는 강력한 세계 모델 베이스라인인 UniZero 를 능가합니다.
  • Zork1(장기적, 희소 보상) 에서 PriorZero 는 더 높은 점근적 보수 달성 및 더 빠른 수렴을 보입니다.
  • Detective에서 PriorZero 는 초기 이득 이후 UniZero 를 능가하여, 선행 지식을 작업 특이적 역동성에 적응시키는 가치 입증합니다.
• BabyAI 벤치마크: BabyAI 그리드 월드의 18 개 레벨에서 PriorZero 는 더 높은 보수 플래토 (UniZero 의 0.79 대비 0.82) 에 도달하며, LLM 선행 지식이 목표를 하위 작업으로 분해하는 데 도움이 되는 구성적 작업 (예: SynthLoc) 에서 상당한 이득을 보입니다.
• Ablation 연구:
  • 교대 일정: 안정성에 결정적입니다. 비교대 변형은 노이즈가 있는 가치 신호로 인해 "지식 부패"를 겪습니다.
  • 루트 전용 주입: 전망 동안 LLM 이 세계 모델의 역동성을 덮어쓰는 것을 방지하는 데 필수적입니다.
  • CoT 손실 가중치: 균형 잡힌 가중치 ($w_{cot}=0.1$) 가 필요합니다. CoT 를 제거하면 수치적 불안정성이 발생하고, 전체 가중치는 변동적인 학습을 초래합니다.
  • MCTS: MCTS 를 비활성화하면 성능이 붕괴되어 계획 개선 타겟의 필요성을 확인시켜 줍니다.

중요성과 주장

이 논문은 PriorZero 가 선행 지식 이전의 구조적 병목 현상을 해결함으로써 LLM 기반 의사결정을 위한 유망한 프레임워크를 확립한다고 주장합니다. 그 중요성은 다음과 같습니다:

• 분리: LLM 의 역할 (의미론적 선행 지식 제공) 과 세계 모델의 역할 (역동성 학습 및 계획) 을 성공적으로 분리하여 이전 방법들의 세분성 불일치와 크레딧 할당 실패를 피합니다.
• 안정성: 세계 모델을 저분산 비판자로 활용함으로써 장기적, 희소 보상 환경에서 LLM 의 안정적인 파인튜닝을 가능하게 합니다.
• 확장성: 교대 훈련 일정을 통해 견고성을 유지하면서 LLM 용량에 따라 프레임워크가 확장됨을 (복잡한 작업에서 7B 모델로 성능 향상) 입증합니다.

저자들은 PriorZero 를 세계 모델이나 LLM 을 대체하는 것이 아니라, LLM 이 검색에 집중하고 세계 모델이 지속적인 상호작용을 통해 LLM 의 선행 지식을 정제하는 시너지적 통합으로 위치시킵니다.