기술 요약: HERMES: 사전 학습 데이터 혼합을 위한 다중 입도 레이블 기질(Multi-Granularity Labeling Substrate)

1. 문제 정의

현재 대규모 언어 모델(LLM) 사전 학습을 위한 데이터 혼합 전략은 혼합 알고리즘 자체가 아니라 레이블 시스템에서 병목 현상에 직면해 있습니다. 기존 방식은 코퍼스를 고정된 분류 체계(출처, 주제, 형식) 또는 평면적인 임베딩 클러스터(예: KMeans)를 사용하여 그룹으로 분할합니다. 이러한 방법들은 다음과 같은 세 가지 주요 한계를 가집니다:

1. 단일 입도/축(Single Granularity/Axis): 단일한 의미적 축과 고정된 해상도에 종속됩니다. 입도(예: 광범위한 주제에서 미세한 클러스터로)를 변경하려면 전체 레이블 세트를 다시 구축해야 합니다.
2. 재계산 비용: 평면적 클러스터링 방식(KMeans 등)은 클러스터 수($K$)를 변경할 때마다 재계산이 필요하므로, 입도를 체계적으로 탐색하는 과정이 계산적으로 비용이 많이 들고 단절됩니다.
3. 유연성 부족: 선택된 그룹이 믹서(Mixer)가 표현할 수 있는 범위를 결정하며, 이는 데이터 해상도와 샘플링 전략 간의 상호작용을 테스트하는 능력을 제한합니다.

저자들은 병목 지점이 믹서 기계가 아니라 레이벨링 인프라라고 주장합니다. 따라서 재클러스터링 없이도 계층적 입도를 탐색할 수 있는 재사용 가능한 데이터 유도형 기질(Substrate)이 필요합니다.

2. 방법론: HERMES

HERMES(Hierarchical Embedding-based RVQ for Mixture Selection)는 문서를 조-중-세(coarse-to-fine) 코드로 주석 처리하도록 설계된 일회성 계층적 임베딩 기반 기질입니다. 이는 두 가지 주요 단계로 구성됩니다.

A. 학습된 의미 변환 (Learned Semantic Transform, LST)

양자화를 위해 임베딩을 준비하기 위해 HERMES는 LST를 사용합니다.

• 입력: 사전 학습된 인코더로부터 추출된 동결된 문서 임베딩 ($e_i \in \mathbb{R}^{1024}$).
• 메커니즘: 정규화가 포함된 단일 선형 레이어($h_i = \text{normalize}(W e_i + b)$)가 임베딩 공간을 회전시킵니다.
• 학습 목적 함수: 변환 행렬 $W$가 표현 붕괴를 방지하도록 근사 직교성을 강제하고, 원본과 변환된 임베딩 사이의 코사인 기하학을 유지하며, 벡터 양자화의 재구성 오차를 최소화하도록 양자화 코드북과 공동으로 학습됩니다.

B. 3단계 잔차 벡터 양자화 (3-Stage Residual Vector Quantization, RVQ)

변환된 임베딩($h_i$)은 계단식 3단계 RVQ에 의해 인코딩됩니다.

• 구조: 각각 코드북 크기가 $K=256$인 세 개의 순차적 벡터 양자화기.
• 프로세스:
  1. 1단계 (조밀/Coarse): $h_i$를 인코딩하여 가장 가까운 중심점 $c_1$을 찾고, 잔차 $r_2 = h_i - q_1$를 생성합니다.
  2. 2단계 (중간/Intermediate): $r_2$를 인코딩하여 $c_2$를 찾고, 잔차 $r_3$를 생성합니다.
  3. 3단계 (미세/Fine): $r_3$를 인코딩하여 $c_3$를 찾습니다.
• 출력: 각 문서는 계층적 코드 $(c_1, c_2, c_3)$를 부여받습니다.
• 입도 제어: 입도는 코드의 **접두사 길이(Prefix Length)**에 의해 정의됩니다:
  • L1: 코드 $(c_1)$ $\rightarrow$ 256개 버킷.
  • L12: 코드 $(c_1, c_2)$ $\rightarrow$ 약 65k 개 버킷.
  • L123: 코드 $(c_1, c_2, c_3)$ $\rightarrow$ 명목상 $1.67 \times 10^7$개 중 관찰된 약 130k 개 버킷.

C. 샘플링 분류 체계 (Sampling Taxonomy)

본 논문은 HERMES 기질 위에서 작동하는 두 단계의 구성을 샘플러로 정의합니다.

1. 1단계 (외부 가중치/Outer Weight): L1 조상(ancestors)에 가중치를 할당합니다. 논문에서는 Uniform(최대 엔트로피) 또는 DoReMi(프록시 모델에 대한 Group-DRO)를 사용합니다.
2. 2단계 (하위 버킷 샘플러/Sub-bucket Sampler): L1 조상 내에서의 적격성 및 선택을 결정합니다.
   • 질량 분포 (Mass Distribution): 균등(최대 엔트로피) vs 크기 비례(Size-proportional).
   • 문서별 적격성 (Per-Document Eligibility): 모든 문서 vs 하위 버킷 내 FineWeb-Edu 점수 기준 상위 30% 문서.

3. 주요 기여

1. HERMES 기질: 약 5천만 개의 문서를 최대 13만 개의 셀로 주석 처리하는 일회성 계층적 RVQ 시스템입니다. 이는 단일 훈련된 코드북 스택 상에서 접두사 길이를 통해 입도를 노출함으로써, 해상도를 변경할 때 재클러스터링할 필요를 제거합니다.
2. 입도-샘플러 상호작용: 샘플링 규칙(예: 품질 기반 선택)의 효과가 입도 수준에 따라 달라짐을 입증했습니다.
3. 후보 경쟁 진단 (Candidate Competition Diagnostic): 입도가 높아짐에 따라 후보 풀의 크기가 줄어들어 버킷 내 순위 결정이 불안정해지는 구조적 한계를 식별했습니다.

4. 실험 결과

실험은 1B 파라미터 규모의 LLaMA 스타일 모델을 25B 토큰 동안 학습시켰으며, 16개 태스크 역량의 매크로 평균(Avg.)으로 평가되었습니다.

A. 클러스터링 정체기 (L1=256)

• 발견: 가장 거친 입도(L1=256)에서 HERMES는 표준 클러스터링 방법(KMeans, MiniBatchKMeans, Plain RVQ)과 대등한 성능을 보였습니다.
• 지표: 내재적 지표(압축도, 질량 균형, 엔트로피) 측면에서 HERMES는 다른 256-way 클러스터러들과 함께 정체기(plateau)에 머물렀습니다.
• 다운스트림: Uniform + L1 max-entropy 샘플러와 결합했을 때, HERMES와 KMeans의 Avg. 차이는 단 0.0002에 불과했습니다.
• 함의: HERMES의 성능 향상은 더 우수한 클러스터링 알고리즘 덕분이 아니라, 바로 이 기질(Substrate) 자체 덕분입니다.

B. 발견 A: 중간 입도에서의 2단계 규칙 (L12)

• 설정: 고정된 DoReMi L1 외부 가중치, 고정된 HERMES 코드북, 입도는 L12($\sim$65k 버킷)로 설정.
• 비교: 2단계 샘플러를 max-entropy coverage(균등한 하위 버킷 질량, 모든 문서)에서 quality top-30%(크기 비례 질량, 상위 30% FineWeb-Edu 점수)로 전환.
• 결과: 이 전환을 통해 16개 태스크 Avg.가 +0.0253 (z-score +2.09) 증가했습니다.
• 해석: L12에서 하위 버킷은 충분히 크기 때문에(중앙값 풀 $\sim$2,271개 문서), 품질 순위가 전역적 품질의 안정적인 프록시 역할을 할 수 있습니다.

C. 발견 B: 미세 입도에서의 붕괴 (L123)

• 설정: 동일한 규칙을 L123($\sim$130k 버킷)에 적용.
• 관찰: 중앙값 후보 풀이 L12(2,271개)에서 L123(429개)으로 5.3배 축소되었습니다.
• 결과: 품질 top-30% 규칙의 이점이 사라졌습니다. 품질 top-30%와 로컬 랜덤 커버리지 간의 성능 격차는 +0.0002로 줄어들었으며, 이는 단일 시드 평가 하에서 통계적으로 구별 불가능한 수준입니다.
• 함의: 이는 후보 경쟁(Candidate Competition) 가설을 뒷받침합니다. 즉, 하위 버킷이 너무 작아지면 버킷 내 순위가 더 이상 전역적 품질의 신뢰할 수 있는 프록시가 되지 못하며, 이로 인해 품질 기반 선택의 이점이 상쇄됩니다.

D. 최적 구성

탐색 과정에서 가장 높은 성능을 보인 구성은 **DoReMi-L1 외부 가중치 + L12 품질 top-30% (수정된 리더)**였으며, Avg. 0.4222를 달이트했습니다.

5. 의의 및 주장

본 논문은 데이터 혼합 설계를 재정의한다는 점을 겸허히 주장합니다. 핵심적인 의의는 "고정된 레이블 세트 중 하나를 선택하는 것"에서 "재사용 가능한 데이터 유도형 입도 계층을 탐색하는 것"으로 초점을 전환했다는 데 있습니다.

• 클러스터러보다 기질이 중요: 기여점은 클러스터링 성능 그 자체가 아니라, 입도를 자유로운 변수(다이얼)로 만드는 인프라를 제공한 것입니다.
• 공동 결정 요소: 입도와 2단계 샘플러는 공동으로 결정되는 축입니다. L12에서 작동하는 규칙이 후보 풀 크기와 같은 구조적 제약으로 인해 L123에서는 실패할 수 있습니다.
• 한계: 저자들은 결론이 내부 코퍼스를 사용한 1B/25B 환경에 국한되어 있음을 인정합니다. 최적의 입도 깊이는 "후보 경쟁"에 의해 제한되며, 이는 더 큰 모델과 코퍼스에서 달라질 수 있습니다. 또한, 최적의 2단계 규칙(L12 quality top-30%)에 대한 외부 가중치 선택(DoReMi vs Uniform)이 완전히 교차 검증되지 않았음을 언급했습니다.

요약하자면, HERMES는 데이터 입도를 연속적이고 튜닝 가능한 파라미터로 취급할 수 있는 메커니즘을 제공하며, 데이터 선택 전략의 효용성이 밑바탕이 되는 레이블 시스템의 해상도와 본질적으로 연결되어 있음을 밝혀냈습니다.