AEM: 멀티턴 에이전트 강화학습을 위한 적응형 엔트로피 변조

LLM 에이전트가 긴 작업을 수행할 때 어떤 행동이 성공에 기여했는지 판단하는 Credit Assignment 문제는 매우 어렵다. 이 논문은 추가적인 보상 모델이나 데이터 라벨링 없이 모델 내부의 엔트로피 신호만으로 학습 효율을 극대화하는 경량화된 해결책을 제시한다.

강화학습에서 모델은 새로운 행동을 시도하는 탐색(Exploration)과 아는 것을 최적화하는 활용(Exploitation) 사이의 균형을 맞춰야 한다. 기존 LLM 에이전트 학습은 토큰 단위로 보상을 나누기 어려워 학습 초기에 특정 행동에만 매몰되는 엔트로피 붕괴 현상이 자주 발생한다.

AEM은 모델이 생성한 전체 응답의 불확실성(Entropy)이 곧 모델의 자신감을 나타낸다는 점에 착안했다. 학습 초기에는 불확실성이 높은 응답에 더 많은 가중치를 두어 다양한 시도를 장려하고, 학습 후기에는 성공 가능성이 높은 응답의 불확실성을 낮추는 방향으로 학습을 유도한다.

이 과정은 별도의 외부 데이터 없이 모델이 생성한 응답 그룹 내의 상대적인 불확실성 수치만으로 계산된다. 결과적으로 모델은 스스로의 확신 정도에 따라 학습 강도를 조절하며, 이는 마치 사람이 어려운 문제는 신중하게 검토하고 익숙한 문제는 빠르게 처리하며 학습하는 것과 유사한 효과를 낸다.

AEM은 기존 Advantage Estimator 위에 적용되는 플러그인 모듈로 설계되었다. 먼저 에이전트가 생성한 전체 응답 시퀀스를 환경과 상호작용하는 최소 단위인 Response Span으로 정의하고, 해당 구간 내 토큰들의 평균 엔트로피를 계산하여 응답 수준의 불확실성 대리 지표(Proxy)를 생성한다.

계산된 불확실성 지표는 그룹 내에서 Min-Max Normalization을 거쳐 상대적인 크기로 변환된다. [응답별 평균 엔트로피 값 입력 → 그룹 내 최솟값과 최댓값을 이용한 정규화 연산 수행 → 0에서 1 사이의 상대적 불확실성 수치 출력] 이 수치는 다시 지수 함수를 통해 변조 계수 α로 변환된다.

최종적으로 기존의 Advantage 값에 이 α를 곱하여 수정된 Advantage를 산출한다. [기존 Advantage와 변조 계수 α 입력 → 두 값의 곱셈 연산 수행 → 수정된 Advantage 출력] 이 메커니즘을 통해 부정적 보상을 받은 응답 중 불확실성이 높은 것은 엔트로피를 높이는 방향으로, 긍정적 보상을 받은 응답은 엔트로피를 낮추는 방향으로 가중치가 조정된다.

Qwen2.5-1.5B 모델을 사용한 ALFWorld 실험에서 GRPO 단독 사용 시 68.0%였던 성공률이 AEM 적용 후 76.8%로 8.8%p 상승했다. 7B 모델에서도 78.7%에서 84.4%로 유의미한 성능 향상을 보였다.

WebShop 벤치마크에서는 1.5B 모델 기준 성공률이 65.0%에서 70.6%로 증가했으며, 보상 점수(Score) 또한 83.6에서 86.4로 개선되었다. 이는 추가적인 보상 모델(PRM) 없이도 정교한 Credit Assignment가 가능함을 시사한다.

소프트웨어 공학 작업인 SWE-bench-Verified에서는 Qwen3-32B 기반의 DeepSWE 프레임워크에 통합되어 기존 42.3%의 해결률을 43.7%로 끌어올렸다. 특히 Figure 4와 5를 통해 AEM이 학습 초기 엔트로피 급락을 방지하고 안정적인 성능 우상향을 이끌어냄을 확인했다.

AEM은 Fisher-Rao Metric 기반의 정보 기하학적 분석을 통해 엔트로피 역학을 정립했다. Theorem 3.2.2에 따르면 자연 구배(Natural Gradient) 업데이트 하에서 엔트로피 드리프트는 Advantage와 Relative Surprisal의 곱에 비례한다. AEM은 이를 구현하기 위해 Doob's Decomposition을 활용하여 다루기 힘든 상태 엔트로피 대신 계산 가능한 토큰 엔트로피 합을 대리 지표로 사용한다.

구현 측면에서는 그룹 기반 강화학습(GRPO, DAPO 등)의 특성을 활용하여 동일 프롬프트에서 생성된 응답들 사이의 상대적 엔트로피를 비교하는 Group Normalization을 채택했다. 이는 배치 전체나 단일 트래직토리 내에서 정규화하는 것보다 과업별 엔트로피 편향을 효과적으로 제거한다. 또한 추가적인 모델 추론(Forward Pass)이 필요하지 않아 전체 학습 시간 중 AEM 연산이 차지하는 비중은 1.1%에 불과할 정도로 효율적이다.

AEM은 응답 공간 전체에 대한 정확한 엔트로피를 계산하는 대신 샘플링된 그룹 내의 상대적 수치를 사용하므로, 샘플링된 응답들의 품질과 다양성에 성능이 의존적일 수 있다. 또한 오픈 엔드(Open-ended) 생성 작업에서 정확한 Surprisal을 추정하는 것은 여전히 근사치에 기반한 휴리스틱이라는 한계가 있다.