수학 문제 풀이나 코드 생성 같은 복잡한 추론 작업을 위해 LLM을 강화 학습으로 튜닝하다가, 모델의 성능 지표가 특정 구간에서 더 이상 오르지 않고 정체되는 현상을 겪어보셨을 것입니다. 분명히 GPU는 100% 가동되고 있고 막대한 비용이 지출되고 있는데, 모델이 이미 정답을 알고 있는 쉬운 문제들만 반복해서 풀며 시간을 허비하고 있다면 이는 데이터 샘플링 전략에 심각한 문제가 있다는 신호입니다. 특히 정답 여부를 명확히 확인할 수 있는 RLVR(Reinforcement Learning from Verifiable Rewards) 환경에서, 모든 프롬프트를 동일한 확률로 추출하는 기존의 방식은 학습이 진행됨에 따라 모델의 실제 성장에는 기여하지 못하는 '죽은 계산'을 양산하게 됩니다.
왜 무작위 샘플링이 LLM 강화 학습의 치명적인 병목이 되는가
강화 학습의 핵심은 모델이 새로운 시도를 통해 더 높은 보상을 얻는 방법을 배우는 것입니다. 하지만 실제 프로젝트 현장에서 사용하는 프롬프트 풀은 대개 정적이며, 모델의 현재 실력 변화를 실시간으로 반영하지 못합니다. 모델이 이미 99%의 확률로 맞히는 쉬운 문제에 계속해서 롤아웃(Rollout) 자원을 할당하는 것은 엔지니어링 관점에서 매우 비효율적입니다. 반대로 모델이 전혀 손을 댈 수 없을 정도로 어려운 문제에만 매달리는 것 또한 유의미한 보상 신호를 얻지 못해 학습을 방해합니다.
이러한 불균형은 단순히 학습 속도를 늦추는 것에 그치지 않고, 제한된 컴퓨팅 예산 내에서 모델이 도달할 수 있는 최종 성능의 고점을 낮추는 결과를 초래합니다. 개발자는 모델이 현재 '딱 한 걸음만 더 나아가면 풀 수 있는' 난이도의 문제를 집중적으로 학습하도록 유도해야 합니다. 이를 위해서는 모델의 성공률을 실시간으로 추적하고, 그 데이터에 기반해 다음에 학습할 문제를 영리하게 골라내는 동적인 메커니즘이 필수적입니다.
HeaPA가 제안하는 힙 기반 난이도 인지 샘플링의 원리
HeaPA(Difficulty-Aware Heap Sampling)는 이 문제를 해결하기 위해 데이터 구조의 '힙(Heap)' 개념을 샘플링에 도입했습니다. 모든 학습 프롬프트를 단순히 나열하는 것이 아니라, 모델의 최근 수행 능력을 바탕으로 우선순위를 부여하여 관리하는 방식입니다. 여기서 핵심은 모델이 가장 낮은 성능을 보이는 문제나 가장 높은 성능을 보이는 문제가 아니라, 성공과 실패의 경계에 있는 프롬프트에 가장 높은 가중치를 두는 것입니다.
이 메커니즘은 모델이 새로운 지식을 습득하기 가장 좋은 '근접 발달 영역'을 타격합니다. 힙 구조를 활용하면 수만 개의 프롬프트 중에서도 현재 정책(Policy) 모델이 가장 많이 배울 수 있는 데이터를 효율적으로 추출할 수 있습니다. 결과적으로 동일한 횟수의 롤아웃을 수행하더라도, 모델은 훨씬 더 밀도 높은 피드백을 받게 되며 이는 곧 가파른 성능 향상 곡선으로 이어집니다. 인프라 운영 측면에서는 불필요한 토큰 생성을 줄여 GPU 활용의 가성비를 극대화할 수 있다는 강력한 이점이 있습니다.
온폴리시 쿼리 확장을 통한 데이터 부족 문제 해결
단순히 기존 데이터를 잘 고르는 것만으로는 부족할 때가 있습니다. 모델이 특정 유형의 논리 구조에서 계속 실수를 범한다면, 그와 유사한 변형 문제를 더 많이 풀어보게 하는 것이 효과적입니다. HeaPA는 이를 위해 '온폴리시 쿼리 확장(On-policy Query Augmentation)'을 결합합니다. 이는 현재 학습 중인 모델의 상태를 반영하여, 모델이 취약한 지점을 공략할 수 있는 새로운 프롬프트를 실시간으로 생성해 학습 풀에 추가하는 기법입니다.
이 과정은 고정된 데이터셋이 가진 편향성을 극복하고 모델의 일반화 성능을 높이는 데 기여합니다. 기존 방식이 주어진 시험지만 반복해서 푸는 것이라면, 이 방식은 오답 노트를 바탕으로 유사 문제를 계속해서 만들어주는 개인 과외 선생님을 두는 것과 같습니다. 아키텍처적으로는 메인 학습 루프와 병렬로 쿼리 생성 에이전트를 운영함으로써, 학습 데이터의 신선도를 유지하고 오버피팅(Overfitting) 위험을 획기적으로 낮출 수 있습니다.
실제 구현 시 고려해야 할 아키텍처와 운영 전략
HeaPA를 실제 시스템에 이식할 때는 분산 학습 환경에서의 상태 동기화가 가장 큰 숙제입니다. 여러 GPU 노드가 동시에 샘플링을 진행할 때, 각 노드가 바라보는 힙의 상태가 일관되어야 합니다. 이를 위해 중앙화된 상태 저장소나 고성능 메시지 큐를 활용하여 모델의 성공률 업데이트를 실시간으로 공유하는 설계가 필요합니다. 또한, 샘플링 우선순위를 계산하는 로직이 전체 학습 파이프라인에 병목을 일으키지 않도록 가벼운 알고리즘으로 구현하는 것이 중요합니다.
운영 단계에서는 '난이도 인플레이션'을 경계해야 합니다. 모델이 점차 똑똑해짐에 따라 힙 상단에 배치되는 문제들의 난이도도 함께 올라가는데, 어느 순간 모델이 감당하기 힘든 수준의 문제들만 남게 되면 학습 효율이 급격히 떨어질 수 있습니다. 따라서 주기적으로 기본 난이도의 문제들을 섞어주는 '리프레시' 전략이나, 학습 초기부터 단계적으로 난이도를 높여가는 커리큘럼 스케줄링을 병행하는 것이 운영상의 노하우입니다.
도입 시 주의해야 할 함정과 해결 방안
이 기법의 가장 큰 함정은 '불가능한 문제'에 대한 집착입니다. 데이터셋 중에는 애초에 정답이 잘못되었거나, 논리적으로 모순이 있어 모델이 절대 풀 수 없는 데이터가 섞여 있을 수 있습니다. 이런 데이터에 높은 우선순위가 부여되면 모델은 무의미한 시도를 반복하며 최적화 방향을 잃게 됩니다. 이를 방지하기 위해 일정 횟수 이상 시도했음에도 불구하고 성공률이 전혀 개선되지 않는 프롬프트는 '불량 데이터'로 간주하여 힙에서 격리하거나 가중치를 강제로 낮추는 필터링 로직이 반드시 포함되어야 합니다.
또한, 특정 도메인의 문제에만 샘플링이 집중될 경우 발생하는 재앙적 망각(Catastrophic Forgetting)도 주의해야 합니다. 수학 문제는 잘 풀게 되었는데 기초적인 언어 이해 능력이 떨어지는 식의 부작용을 막으려면, 힙 기반 샘플링을 전체 학습 비중의 일정 부분으로 제한하고 나머지는 일반적인 코퍼스를 섞어주는 혼합 전략(Mixed Strategy)을 취하는 것이 안전합니다. 이는 모델의 전문성과 범용성 사이의 균형을 맞추는 핵심적인 장치가 됩니다.
3줄 요약
- 기존의 균등 샘플링은 모델이 이미 아는 문제를 반복 학습하게 하여 GPU 자원과 비용을 낭비하는 심각한 비효율을 초래한다.
- HeaPA는 힙 구조를 통해 모델이 성공과 실패의 경계에 있는 최적의 난이도 문제를 우선 샘플링하고, 실시간 쿼리 확장으로 취약점을 집중 보완한다.
- 성공적인 도입을 위해서는 불가능한 문제에 대한 필터링과 재앙적 망각을 방지하기 위한 혼합 학습 전략을 반드시 병행해야 한다.
참고: arXiv CS.LG (Machine Learning)