기술 요약: 매니폴드 밴딧(Manifold Bandits) – 대규모 언어 모델의 잠재 기하학에 대한 베이지안 커리큘럼 학습

1. 문제 정의

강화 학습(RL)은 특히 GRPO, DAPO, GSPO와 같은 그룹 상대적 최적화 알고리즘을 통해 대규모 언어 모델(LLM)의 추론 능력을 향와시키는 주요 방법론이 되었습니다. 이러한 방법들은 그룹 내 롤아웃(rollout) 간의 **보상 분산(reward variance)**에 의존하여 정책 그래디언트를 생성합니다. 만약 모델이 특정 프롬프트를 지속적으로 해결하거나 지속적으로 실패한다면, 보상 분산은 0이 되어 학습 신호가 발생하지 않으며 계산 자원이 낭비됩니다.

기존의 적응형 커리큘럼 학습 방법들은 이를 해결하기 위해 중간 난이도(비제로 분산을 생성할 가능성이 높은 프롬프트)를 우선시하려고 시도합니다. 그러나 이러한 접근 방식은 일반적으로 문제 선택을 독립적인 팔(arm)을 가진 표준 멀티 암드 밴딧(multi-armed bandit) 문제로 프레이밍합니다. 이러한 프레이밍은 LLM 학습의 두 가지 핵심 역학을 간과합니다:

1. 내생적 비정상성(Endogenous Non-Stationarity): 보상 분포는 외부 요인뿐만 아니라, 모델 자체의 업데이트와 샘플링 결정의 직접적인 결과로서 진화합니다.
2. 구조적 이질성(Structured Heterogeneity): 문제들은 독립적이지 않습니다. 유사한 유형의 프롬프트들은 공통된 추론 패턴과 잠재적 특징을 공유합니다. 이러한 관계를 무시하는 것은 비효적인 탐색으로 이어집니다.

나아가, 현재의 전략들은 종종 생산성(즉각적인 학습 신호 극대화)에만 집중하며 다양성(태스크 공간의 커버리지)과 유용성(다운스트림 평가 벤치마크에 대한 관련성)을 간과합니다. 이는 좁은 커버리지와 하위 최적의 일반화로 이어질 수 있습니다.

2. 방법론

저자들은 문제를 샘플링하는 것을 내생적 비정상성을 가진 매니폴드 구조의 밴딧 문제로 재정의하는 프레임워크인 **베이지안 매니폴드 커리큘럼(Bayesian Manifold Curriculum, BMC)**을 제안합니다. 이 방법론은 세 가지 핵심 구성 요소로 이루어집니다:

A. 잠재 태스크 트리 (Latent Task Trees)

BMC는 프롬프트를 독립적인 팔로 취급하는 대신, 정책 모델의 내부 표현으로부터 직접 계층적인 잠재 태스크 트리를 구축합니다.

• 구축: 트리는 정책 모델의 중간 레이어 임베딩을 사용하여 재귀적으로 구축됩니다. 각 노드에서 임베딩은 PCA를 통해 차원이 축소되고, UMAP을 통해 투영되며, HDBSCAN을 사용하여 클러스터링됩니다.
• 중단 기준: 재귀는 로컬 매니폴드가 근사적으로 매끄럽고 일관적임(낮은 고유 차원 및 연결성)을 검증하는 "차트 테스트(Chart Test)"를 통과할 때까지 계속됩니다.
• 목적: 이 계층 구조는 모델의 내부적 태스크 관계를 반영하는 다중 척도 영역으로 프롬프트를 조직함으로써, 모델의 잠재적 태스크 매니폴드를 근사합니다. 이는 외부 주석이나 고정된 분류 체계에 의존하지 않습니다.

B. 베이지안 매니폴드 커리큘럼 (BMC)

BMC는 다음 세 단계의 사이클을 통해 잠재 태스크 트리 상에서 구조화된 의사결정을 수행합니다:

1. 하향식 선택 (계층적 톰슨 샘플링): 학습 배치를 형성하기 위해 에이전트는 루트에서 리프까지 트리를 순회합니다. 각 노드에서 에이전트는 가우시안 신념 분포($\mathcal{N}(\mu, \sigma^2)$)로부터 자식 노드의 보상 추정치를 샘플링하고, 가장 높은 샘플을 가진 자식을 선택합니다. 이는 다양한 추상화 수준에서 탐색(다양성)과 착취(생산성)의 균형을 맞춥니다.
2. 비정상 신념 모델링 (베이지안 필터링): 리프 노드(개별 프롬프트)는 기대 학습 신호(보상 분산)에 대한 신념을 유지합니다. 이러한 신념은 Normal–Inverse-Gamma 모델에서 영감을 얻은 경량 베이지안 필터를 사용하여 온라인으로 업데이트됩니다. 업데이트 메커니즘은 모델의 진화하는 능력에 적응하기 위해 놀라움(surprise)(표준화된 예측 오차)과 노후도(staleness)(마지막 샘플링 이후 경과 시간)를 통합합니다.
3. 상향식 신념 전파 (경험적 베이즈): 업데이트 후, 정보는 상위로 전파됩니다. 부모 노드의 신념은 자식 노드들의 신념을 정밀도 가중 평균(precision-weighted average)으로 추정하며, 이는 **서브트리 이질성(subtree heterogeneity)**에 의해 조정됩니다. 이를 통해 한 프롬프트에서의 관측이 동일한 영역 내의 관련 프롬프트에 대한 신념에 정보를 제공하도록 하여, 잠재 구조를 활용합니다.

C. 유용성 인식 확장 (BMC-T)

학습 생산성과 다운스트림 평가 성능 사이의 불일치를 해결하기 위해, 저자들은 BMC-T를 도입합니다. 이 변형은 잠재 공간에서 타겟 분포(예: 특정 평가 벤치마크)와 겹치는 훈련 분포의 영역으로 샘플링을 편향시킵니다. 이는 잠재 태스크 트리를 근사적 개입 모델로 취급하며, 인접한 태스크를 학습하는 것이 더 강한 전이 효과를 가진다고 가정합니다.

3. 주요 기여

1. 잠재 태스크 트리: 외부 라벨 없이 모델 임베딩으로부터 태스크 매니폴드의 구조를 직접 근사하는 계층적 방법으로, 다중 척도 프롬프트 관계 표현을 제공합니다.
2. 베이지안 매니폴드 커리큘럼 (BMC): 이러한 트리 상에서 베이지안 의사결정을 수행하여, 내생적이고 모델 유도적인 비정상성을 고려하면서 효율적인 탐색을 가능하게 하는 구조 인식 적응형 커리큘럼 프레임워크입니다.
3. 생산성–다양성–유용성 트레이드오프: 저자들은 적응형 커리큘럼이 다운스트림 성능만으로는 특징지어질 수 없음을 보여주는 진단 분석을 제시합니다. 이들은 세 가지 축을 구분합니다:
   • 생산성: 학습 신호 (보상 분산).
   • 다양성: 태스크 매니폴드의 커버리지.
   • 유용성: 특정 평가 벤치마크에 대한 관련성.
     그들은 이 축들이 서로 교환 가능하지 않음을 보여주며, 하나를 최적화한다고 해서 다른 것들이 반드시 최적화되지는 않는다는 점을 명시합니다.

4. 실험 결과

저자들은 수학적 추론 태스크(DAP_Math-17K 데이터셋)와 Qwen3 모델(4B 및 8B)을 사용하여 GSPO 및 GRPO 환경에서 BMC를 평가했습니다.

• 학습 효율성 (생산성): BMC와 그 "난이도 전용(Difficulty Only)" 어블레이션은 다이내믹 샘플링(분산을 확보하기 위해 배치를 재생성함)과 유사한 학습 속도를 달면서도, 실제 실행 시간(wall-clock time) 측면에서는 **균등 샘플링(Uniform Sampling)**의 효율성에 가깝게 유지하며 훨씬 낮은 시간을 기록했습니다.
• 커버리지 및 정보 공유 (다양성): BMC는 Tree Only(커버리지를 강조)의 다양성과 Difficulty Only(생산성을 강조)의 생산성 사이에서 성공적으로 균형을 잡았습니다. 지표에 따르면 BMC는 희소 영역(희소성 가중 노출)에 대해 더 나은 커버리지를 제공했으며, 잠재 태스크 트리가 학습 신호의 상당한 변동을 설명함을 보여주었습니다. 이는 학습 신호가 잠재 구조와 상관관계가 있음을 확인시켜 줍니다.
• 평가 성능 (유용성):
  • 다이내믹 샘플링은 희귀 프롬프트 유형을 암묵적으로 필터링하기 때문에, 덜 대표적인 벤치마크(예: 중국어 수학, OOD 과학)에서 낮은 성능을 보였습니다.
  • BMC는 OOD 벤치마크인 GPQA-Diamond에서 강력한 성능을 보였으며, 이는 생산성과 다양성을 동시에 최적화하는 것이 표준 인-디스트리뷰션 수학 평가에서 포착되지 않는 능력을 촉진한다는 것을 시사합니다.
  • BMC-T는 유용성을 조절할 수 있음을 입증했습니다. 특정 벤치마크(예: AIME2024)를 타겟팅하면 전체적인 생산성을 유지하면서도 해당 벤치마크에 대한 성능을 개선할 수 있었으며, 이는 서로 다른 문제 유형이 서로 다른 평가 관련성을 가진다는 것을 확인해 줍니다.

5. 의의 및 주장

본 논문은 난이도만으로는 충분하지 않다고 주장합니다. 학습 신호를 극대화하는 것이 목표 태스크로의 일반화를 이끄는 유형의 문제들에 모델을 노출시키지 못한다면, 더 나은 결과를 보장하지 않습니다.

저자들은 문제 선택을 정책의 잠재 기하학 위에서의 매니폴드 구조 밴딧 문제로 프레이밍해야 한다고 주장합니다. 인간이 정의한 분류 체계 대신 데이터와 컴퓨팅을 사용하여 구조를 학습하는 "비터 레슨(Bitter Lesson)"을 준수하며, 모델 자체의 학습된 표현으로부터 커리큘럼 구조를 도출함으로써 BMC는 다음과 같은 기능을 제공합니다:

• 생산성, 다양성, 유용성 간의 트레이드오프를 탐색합니다.
• 샘플링 결정이 미래의 보상 지형을 형성하는 내생적 비정상성을 활용합니다.
• 외부의 텍스토노미(taxonomy)나 외부 임베딩 모델 없이도 도메인, 언어, 모달리티를 가로질러 일반화합니다.

이 연구는 "평가 데드존(evaluation deadzone)" 현상을 강조합니다: 모델이 특정 문제 세트에 대해 생산적으로 학습할 수 있지만, 그 문제들이 평가 분포와 일치하지 않는다면 다운스트림 성능은 개선되지 않을 수 있습니다. BMC-T는 타겟 관련 잠재 매니폴드 영역으로 학습 노력을 유도함으로써 이를 해결할 수 있는 메커니즘을 제공합니다.