과거는 과거가 아니다: 메모리 강화 동적 보상 셰이핑

LLM을 강화학습으로 훈련할 때 모델이 특정 오류 패턴에 갇혀 비슷한 오답을 반복 생성하는 '정책 붕괴' 현상을 해결합니다. 과거의 실패 경로를 기억하고 동적으로 페널티를 부여함으로써 모델이 더 다양한 해결책을 탐색하도록 유도하여 수학적 추론 성능을 크게 높였습니다.

강화학습에서 에이전트는 보상을 최대화하는 방향으로 정책을 업데이트한다. 하지만 LLM 훈련 시 특정 오답 경로가 높은 확률을 갖게 되면, 모델은 엔트로피 정규화(Entropy Regularization)를 적용하더라도 단순히 단어 선택만 바꾼 채 동일한 논리적 오류를 반복하는 '에러 베이슨(Error Basin)'에 빠지게 된다. 이는 기존의 무작위성 주입만으로는 거대한 행동 공간에서 유의미한 탐색을 이끌어내기 어렵기 때문이다.

MEDS는 인간이 과거의 실수를 기억하고 같은 실수를 반복할 때 더 강한 심리적 가책을 느끼는 원리에서 착안했다. 모델이 추론 과정에서 생성하는 중간 레이어들의 로짓(Logits)은 모델의 내부 논리 구조를 반영하는 임베딩 역할을 한다. 이를 벡터화하여 저장함으로써 모델이 현재 어떤 '생각의 길'을 걷고 있는지 기록한다.

훈련 중 새로운 샘플이 생성되면 과거의 실패 기록들과 비교하여 얼마나 유사한지 측정한다. 만약 현재의 오답이 이미 수차례 반복된 논리적 오류 클러스터에 속한다면, 단순한 오답 감점 외에 '반복 실수 페널티'를 추가로 부여한다. 결과적으로 모델은 보상을 극대화하기 위해 이미 실패했던 논리 구조를 버리고 새로운 추론 경로를 탐색하게 되어 성능이 향상된다.

MEDS는 크게 세 단계로 구성된다. 첫째, 논리 특징 추출(Logic Feature Extraction) 단계에서는 모델이 최종 답변의 첫 번째 토큰을 생성할 때 발생하는 마지막 14개 Transformer 레이어의 로짓을 수집한다. 각 레이어의 로짓 l*(n)을 입력으로 받아 이를 연결(concat)하여 특징 벡터 f(y)를 생성하며, 이는 모델의 암시적인 추론 궤적을 가볍고 효과적으로 표현한다.

둘째, 메모리 기반 클러스터링(Memory-based Clustering) 단계에서는 각 프롬프트별로 과거의 실패 샘플 특징들을 저장하는 에러 메모리 Gx를 유지한다. 여기에 HDBSCAN 알고리즘을 적용하여 유사한 논리적 오류들을 그룹화한다. HDBSCAN은 데이터의 밀도를 계산하여 클러스터의 크기 |Ck|를 결정하며, 이는 해당 오류 패턴이 얼마나 자주 반복되었는지를 나타내는 지표가 된다.

셋째, 보상 셰이핑(Reward Shaping) 단계에서는 클러스터 크기에 비례하는 페널티를 계산한다. penalty = min(α log(|Ck| + 1), β) 수식을 통해 [클러스터 크기 입력 → 로그 연산 및 하이퍼파라미터 적용 → 페널티 값 출력] 과정을 거친다. 최종 보상은 원래의 과업 보상 r(x, y)에서 이 페널티를 뺀 값으로 조정되어 정책 업데이트에 사용된다.

Qwen3-1.7B, Qwen3-8B, Qwen2.5-Math-7B 모델을 대상으로 AIME24, MATH500 등 5개 수학 벤치마크에서 평가를 진행했다. MEDS는 모든 모델과 데이터셋에서 기존의 GRPO, DAPO 등 강력한 베이스라인을 일관되게 능가했다. 특히 Qwen3-8B 모델의 경우 평균 pass@1 점수가 53.89점에서 66.72점으로 크게 상승했다.

탐색 다양성 분석 결과, MEDS는 훈련이 진행됨에 따라 감소하는 경향이 있는 '단계 내 다양성(Within-Step Diversity)'과 '단계 간 다양성(Across-Step Diversity)'을 베이스라인보다 높게 유지했다. 또한 Top-1 Eigen Ratio 지표를 통해 확인한 결과, 로짓 공간에서 특정 방향으로의 쏠림 현상이 줄어들어 더 고른 탐색이 이루어짐이 확인됐다.

Ablation Study에서는 마지막 14개 레이어의 로짓을 직접 사용하는 방식이 레이어 간 차이(diff)를 사용하는 방식보다 성능이 좋았으며, 클러스터링 품질과 하류 작업 성능 사이에 강한 양의 상관관계가 있음이 밝혀졌다. 이는 로짓 기반 클러스터링이 실제 LLM의 의미론적 오류 카테고리를 잘 대변하고 있음을 시사한다.

MEDS는 RLVR(Reinforcement Learning with Verifiable Rewards) 환경에서 발생하는 정책 붕괴(Policy Collapse) 문제를 해결하기 위해 도입됐다. 기존의 엔트로피 정규화는 현재 정책 분포의 무작위성만 높일 뿐, 누적된 훈련 이력에서의 행동 중복을 방지하지 못한다는 한계가 있다. MEDS는 이를 '동적 보상 셰이핑' 문제로 정의하고 이론적으로 페널티 부여가 기대 수익을 감소시키지 않음을 증명했다.

아키텍처 측면에서 가장 큰 특징은 별도의 외부 모델(Reward Model 등) 없이 대상 모델의 내부 로짓을 특징값으로 활용한다는 점이다. 이는 계산 비용을 최소화하면서도 텍스트 수준에서는 보이지 않는 모델의 '생각의 흐름'을 캡처할 수 있게 한다. 특히 초기 레이어보다 복잡한 의미 정보를 담고 있는 후반부 레이어의 로짓이 클러스터링에 더 유효하다는 것을 실험적으로 확인했다.

구현 시 HDBSCAN을 사용하여 클러스터 수를 동적으로 결정함으로써 다양한 오류 유형에 유연하게 대응한다. 또한 페널티 함수에 로그 스케일을 적용하고 상한선(β)을 두어 특정 오류에 대한 과도한 페널티가 전체 학습의 안정성을 해치지 않도록 설계했다. 이러한 메모리 기반 접근법은 향후 더 긴 시계열의 행동 이력을 관리하는 방향으로 확장이 가능하다.

현재 연구에서 사용된 로짓 활용 방식은 비교적 단순하며, 더 정교한 집계 함수(Aggregation Function)를 도입할 경우 성능이 더 개선될 여지가 있다. 또한 복잡한 클러스터링 전략이 성능에 미치는 영향에 대해 추가적인 탐구가 필요하다.