Translate Policy to Language: Flow Matching Generated Rewards for LLM Explanations

1. 문제 정의 (Problem Statement)

• 배경: 강화학습 (RL) 에이전트, 대규모 언어 모델 (LLM), 로봇 시스템 등 다양한 AI 에이전트와 인간이 공존하는 환경에서, 에이전트의 의사결정 과정을 자연어로 설명하는 능력은 신뢰성과 안전성을 위해 필수적입니다.
• 현황의 한계:
  • 기존 설명 생성 방법들은 특정 작업 (자율주행, 추천 시스템 등) 에 맞춰 설계되어 일반화가 어렵습니다.
  • 인간 피드백의 복잡성: 인간은 동일한 설명에 대해 다원적 (pluralistic) 이고 확률적인 판단을 내립니다. 기존 RLHF(인간 피드백을 통한 강화학습) 는 이러한 다양성을 포착하기 위해 비용이 많이 드는 인간 데이터 수집이 필요합니다.
  • AI 피드백 (RLAIF) 의 결함: 인간 피드백을 대체하기 위해 보조 LLM 을 사용한 RLAIF 가 등장했으나, 기존 방법은 단순한 점수 (scalar reward) 를 생성하거나, 보조 LLM 의 노이즈 (오류) 를 명시적으로 관리하지 못해 인간 reward 분포와의 편차를 줄이는 데 한계가 있었습니다.

2. 제안 방법론 (Methodology)

저자들은 일반적인 목적의 프레임워크를 제안하여, 에이전트의 결정 맥락 (context) 만을 입력으로 받아 해당 정책 (policy) 을 반영하는 자연어 설명을 생성하는 LLM 을 훈련시킵니다. 핵심은 **생성형 연속 정규화 흐름 (Continuous Normalizing Flows, CNF)**을 기반으로 한 보상 생성 모델을 도입한 것입니다.

A. 핵심 아키텍처: 정류 흐름 (Rectified Flow) 기반 보상 모델

• 개념: CNF 는 기본 노이즈 분포를 복잡한 분포로 매핑하는 생성 모델입니다. 저자들은 이를 보상 생성에 활용하여, 인간이 가질 수 있는 다양한 (다원적) 판단을 확률 분포로 모델링합니다.
• 작동 원리:
  1. Proxy LLM: 설명과 맥락을 입력받아 실제 에이전트 결정을 추론할 확률 (logits) 을 생성합니다. 이는 인간 피드백의 '노이즈가 포함된' 근사치로 간주됩니다.
  2. Rectified Flow (RF) 모델: Proxy LLM 이 생성한 노이즈가 있는 분포를 학습하여, 이를 '정류 (rectified)'시켜 실제 인간 reward 분포에 가깝게 복원 (denoising) 합니다.
  3. 이론적 근거: Proxy LLM 의 노이즈가 가우시안과 같은 특정 함수 형태를 따른다고 가정할 때, RF 는 이론적으로 인간 reward 분포로부터의 편차를 엄격하게 상한 (bound) 할 수 있음이 증명됩니다. 즉, Proxy LLM 의 오류를 보정하여 더 정확한 분포적 보상을 제공합니다.
• 아키텍처 세부 사항:
  • RF 모델은 LLM 의 마지막 레이어에 임베딩되어, 교차 주의 (Cross-Attention) 메커니즘을 통해 맥락 (context) 과 설명 (explanation) 의 언어적 단서를 선택적으로 주시합니다.
  • 이를 통해 설명의 각 문장마다 실제 결정이 얼마나 타당한 결과로 보이는지 확률적 변화를 계산하여 문장 단위 (sentence-level) 보상을 생성합니다.

B. 훈련 프로세스

1. 데이터 구성: 에이전트의 실제 결정 (ground-truth) 을 제외하고 맥락만 제공하여 EXPLANATION LLM 이 설명을 생성하도록 유도합니다.
2. 교대 훈련 (Iterative Training):
   • Step 1: EXPLANATION LLM 이 설명을 생성합니다.
   • Step 2: Proxy LLM 들이 이 설명을 바탕으로 결정 확률을 계산합니다.
   • Step 3: RF 모델이 Proxy LLM 의 출력을 학습하여 보정된 확률 분포 (보상) 를 생성합니다.
   • Step 4: 생성된 보상을 사용하여 PPO(Proximal Policy Optimization) 로 EXPLANATION LLM 을 미세 조정합니다.
   • 이 과정을 수회 반복하여 수렴시킵니다.

3. 주요 기여 (Key Contributions)

1. 분포적 보상 생성 (Distributional Reward Generation): 인간 피드백의 다원성과 확률적 특성을 포착하기 위해 CNF(특히 Rectified Flow) 를 보상 모델로 처음 도입했습니다.
2. 이론적 오류 상한 증명: Proxy LLM 의 노이즈를 CNF 의 초기 분포와 동일한 함수 형태로 가정할 때, 학습된 RF 모델이 실제 인간 reward 분포와의 거리를 이론적으로 보장 (bound) 한다는 것을 수학적으로 증명했습니다.
3. 맥락 인식형 보상 메커니즘: LLM 기반의 RF 모델에 교차 주의 (Cross-Attention) 를 도입하여, 설명의 언어적 단서와 의사결정 맥락을 통합적으로 분석하여 정교한 보상을 생성합니다.
4. 범용성: RL 에이전트 (SMAC, AI2-THOR) 와 LLM 태스크 (MMLU, MathQA) 등 다양한 도메인에서 작동하는 일반화된 설명 생성 프레임워크를 제시했습니다.

4. 실험 결과 (Results)

저자들은 RL 태스크 (SMAC, AI2-THOR) 와 LLM 태스크 (MMLU, MathQA) 에서 다양한 베이스라인과 비교 실험을 수행했습니다.

• 성능 향상:
  • SMAC, MMLU, MathQA: 제안된 방법은 Proxy LLM 을 직접 사용한 베이스라인보다 6.9% ~ 12.4% 높은 정확도 (ACC) 를 기록했습니다.
  • 전체 비교: SFT, RLHF (PPO, DPO, KTO), RLAIF (Skywork), 추론 기반 방법 (CoT+SFT, ReFT) 등 기존 최첨단 방법들보다 1.9% ~ 18.2% 더 우수한 성능을 보였습니다.
  • AI2-THOR: 복잡한 다중 로봇 조정 작업에서 성공률을 기존 방법 대비 최대 7.4% 향상시켰습니다.
• Human Evaluation (인간 평가):
  • 45 명의 참가자를 대상으로 한 평가에서, 제안된 방법은 **정확도 (Accuracy), 논리적 건전성 (Logical Soundness), 실행 가능성 (Actionability)**에서 가장 높은 점수를 받았으며, **인지 부하 (Cognitive Load)**는 가장 낮았습니다.
  • 특히, 설명을 통해 에이전트의 실제 결정을 추론하는 정확도가 가장 높았습니다.
• Ablation Study:
  • Cross-Attention 제거: 성능이 크게 저하되어, 맥락과 설명 간의 상호작용 학습이 필수적임을 입증했습니다.
  • Sparse Reward vs Sentence-level Reward: 문장 단위 보상을 사용한 것이 희소 보상 (Sparse Reward) 보다 학습에 더 효과적이었습니다.
  • Negative Sample 일반화: RF 모델은 Proxy LLM 이 오분류한 '거짓 부정 (False Negative)' 샘플을 올바르게 보상하여 설명을 개선하는 능력을 보여주었습니다.

5. 의의 및 결론 (Significance)

• 신뢰할 수 있는 AI 설명: 이 연구는 AI 의 의사결정을 설명하는 LLM 을 훈련할 때, 단순한 점수가 아닌 확률적 분포 기반의 보상을 사용하여 인간의 다원적 판단을 더 잘 반영할 수 있음을 보였습니다.
• 비용 효율성: 고비용의 인간 피드백 수집을 대체하면서도, 이론적으로 오류를 보정할 수 있는 AI 기반 보상 모델 (RF) 을 제안하여 RLAIF 의 한계를 극복했습니다.
• 미래 전망: 이 프레임워크는 다양한 에이전트 (RL, LLM, 로봇) 에 적용 가능한 범용 설명 생성 도구로 확장 가능하며, 고위험 분야 (의료, 금융, 자율주행) 에서 AI 의 투명성과 신뢰성을 높이는 데 기여할 것으로 기대됩니다.

요약하자면, 이 논문은 Flow Matching(정류 흐름) 기술을 활용하여 노이즈가 있는 AI 피드백을 보정하고 인간과 유사한 분포적 보상을 생성함으로써, 에이전트의 의사결정을 더 정확하고 논리적으로 설명하는 LLM 을 훈련하는 새로운 패러다임을 제시했습니다.