기술 요약: SPOKES – 다양한 사전 학습 데이터 선택을 위한 최적화

1. 문제 정의

데이터 다양성은 고정된 데이터 예산 내에서 중복을 줄이고 다운스트림 성능을 향상시키기 때문에, 사전 학습 코퍼스를 구축하는 데 있어 핵심적인 요소입니다. 그러나 다양성을 최적화하는 것은 데이터 포인트 간의 상호작용에 의존하는 **집합 수준의 속성(set-level property)**이기 때문에 본질적으로 어렵습니다. 특히 사전 학습 규모에서는 부분 집합 선택의 조합론적 특성으로 인해 다양성을 직접 최적화하는 것이 계산적으로 불가능합니다.

기존 방식들은 대개 주제, 기술 또는 비지도 클러스터링과 같은 거친 입도의 클러스터링이나 시맨틱 중복 제거(semantic deduplication)와 같은 프록시(proxy) 또는 근사치에 의존합니다. 이러한 방법들은 클러스터의 표현력을 보장하지만, 저해상도 분할에 의존하며 롱테일(long-tail)이나 과소 대표된 지식의 미세한 정보를 포착하지 못하는 경우가 많습니다. 또한, 현재의 방법들은 다양성 지표(예: G-Vendi 점수)를 최적화 목표가 아닌 사후 평가 도구로 취급하는 경-향이 있습니다. 이 연구가 다루는 핵심 과제는, 규모 확장이 어렵거나 미세한 다양성을 포착하지 못하는 휴리스틱에 의존하지 않고, 대규모 사전 학습 코퍼스에서 어떻게 신뢰할 수 있는 다양한 부분 집합을 추출할 것인가 하는 점입니다.

2. 방법론: SPOKES

저자들은 확률적 다양화 프레임워크 내에서 G-Vendi 점수를 사용하여 다양성을 직접 최적화하는 확장 가능한 프레임워크인 SPOKES(Spokes: Optimizing for Diverse Pretraining Data Selection)를 소개합니다.

2.1 핵심 목표

이 방법은 데이터 선택을 **품질(quality)**과 **다양성(diversity)**을 공동으로 최대화하는 전역 최적화 문제로 정식화합니다. $n$개의 예시가 있는 데이터셋이 주어졌을 때, 목표는 다음을 최대화하는 크기 $k$의 부분 집합 $S$를 선택하는 것입니다:
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
여기서:

• $q_i$는 예시 $i$의 품질 점수(예: FineWeb-Edu 분류기로부터의 점수)를 나타냅니다.
• $K_S$는 선택된 부분 집합에 제한된 유사도 커널(similarity kernel)입니다.
• $\alpha \in [0, 1]$는 품질과 다양성 사이의 트레이드오프를 조절하는 튜너블 파라미터입니다.

2.2 그래디언트 기반 다양성 (G-Vendi)

정적인 임베딩을 사용하는 전통적인 방법과 달리, SPOKES는 **그래디언트 공간(gradient space)**에서 작동합니다. 이 방법은 프록시 모델을 사용하여 각 데이터 포인트에 대해 모델 파라미터 $\theta$에 대한 그래디언트 $\nabla_\theta \ell(x_i)$를 계산합니다. G-Vendi 점수는 이러한 그래디언트로 구성된 유사도 행렬의 스펙트럼 엔트로피를 조사하여 다양성을 측정합니다. 이 접근 방식은 표현 벡터 간의 직교성을 장려하여 업데이트 방향의 중복을 줄이고, 더 독립적이고 정보가 풍부한 최적화 단계를 촉진합니다.

2.3 지수 그래디언트 하강법(Exponentiated Gradient Descent)을 통한 최적화

조합론적 부분 집합 선택 문제를 해결하기 위해, SPOKES는 완화된 최적화(relaxed optimization) 전략을 채택합니다.

1. 연속적 완화(Continuous Relaxation): 이산적 선택 대신, 모든 데이터 포인트에 대해 연속적인 비음수 가중치 벡터 $w$를 할당하고 심플렉스($w \in \Delta_n$) 제약 조건을 부여합니다.
2. 커널 구축: 엔트리가 $\sqrt{w_i w_j}$에 의해 스케일링되는 가중치 의존적 커널 $K(w)$를 정의합니다.
3. 최적화: 목적 함수는 지수 그래디언트 하강법을 통해 최적화됩니다. 각 반복 단계에서 가중치는 $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$를 통해 업데이트되며, 소수의 고가중치 예시로 붕괴되는 것을 방지하기 위해 정규화됩니다.
4. 이산적 복구: 수렴 후, 학습된 가중치 벡터의 상위 $k$개 항목을 선택하여 최종 부분 집합을 복구합니다.

2.4 확장성 향상

이 접근 방식이 사전 학습 규모에서도 실행 가능하도록, 저자들은 몇 가지 효율성 전략을 구현했습니다.

• 절단된 그래디언트(Truncated Gradients): 전체 모델 그래디언트를 계산하는 대신, 오직 마지막 $n$개의 트랜스포머 레이어(통상 2개 레이어)만을 사용하여 그래디언트를 근사합니다. 이는 유사도 커널의 구조를 유지하면서 계산 비용을 크게 줄여줍니다.
• 존슨-린덴슈트라우스(Johnson-Lindenstrauss, JL) 투영: 고차원 그래디언트 벡터(예: 200M 파라미터)를 처리하기 위해, 저자들은 차원을 축소하는 무작위 투영(Rademacher 행렬)을 적용합니다. 실험 결과, 낮은 차원의 투영(예: $k=1024$)만으로도 높은 부분 집합 안정성(고차원 참조값과 95.5%의 중첩)을 유지함을 보여주었습니다.
• 병렬화: 데이터셋을 독립적으로 최적화되는 서로 소의 부분 집합으로 분할하여 선형적인 속도 향상을 얻을 수 있습니다.

3. 주요 기여

1. SPOKES 프레임워크: 단순한 휴리스틱 클러스터링을 넘어, 현대적인 사전 학습 데이터셋 규모에서도 효율적으로 작동하는 다양성 최적화를 위한 확장 가능하고 원칙적인 기술을 도입했습니다.
2. 직접적인 G-Vendi 최적화: G-Vendi 점수가 단순히 사후 지표로 쓰이는 것을 넘어, 매우 다양한 부분 집합을 추출하기 위해 지수 그래디언트 하강법을 통해 직접 최적화될 수 있음을 입증했습니다.
3. 품질-다양성 공동 최적화: 품질과 다양성을 공동으로 최적화하는 것이 각각의 목적을 단독으로 최적화하거나 표준 베이스라인을 사용하는 것보다 우수한 결과를 낸다는 것을 증명했습니다.
4. 효율성 개선: 절단된 그래디언트 계산과 JL 투영을 통한 차원 축소가 막대한 계산 비용 없이도 거대 코퍼스로 확장 가능하다는 것을 검증했습니다.

4. 실험 결과

이 방법은 두 가지 대규모 사전 학습 코퍼스인 DCLM(강하게 필터링된 Dolmino 서브셋)과 FineWeb(필터링이 덜 엄격한 데이터)에서 평가되었습니다. 모델은 LLaMA-1B 아키텍처를 사용하여 1,000억 개의 토큰으로 처음부터 사전 학습되었습니다.

4.1 다양성 지표

• G-Vendi 점수 증가: 500k 샘플 부분 집합에서, SPOKES는 무작위 샘플링에 비해 +489 포인트의 G-Vendi 점수 증가를 달성했습니다. 반면, 시맨틱 중복 제거(SemDeDup)는 +7 포인트 증가에 그쳤습니다.
• 배치 다양성: 훈련 과정 전반에 걸쳐, SPOKES에 의해 선택된 배치는 무작위 샘플링에 비해 일관되게 높은 평균 다양성 점수를 유지했습니다 (DCLM에서 315에서 520으로 증가).

4.2 다운스트림 성능

평가는 OLMES 스위트의 10가지 영어 벤치마크를 통해 수행되었습니다.

• 다양성 전용(Diversity-Only): SPOKE(다양성 전용)는 무작위 샘플링 대비 평균 다운스트림 성능을 +0.4(DCLM) 및 +0.5(FineWeb) 포인트 향상시켰습니다.
• 공동 최적화(Joint Optimization): 가장 강력한 결과는 품질과 다양성을 공동으로 최적화했을 때 나타났습니다. SPOKES는 시맨틱 중복 제거 및 품질 전용 필터링을 포함한 모든 베이스라인을 능가하며, 무작위 샘플링 대비 +1.5(DCLM) 및 +1.4(FineWeb) 포인트의 이득을 달성했습니다.
• 상호 보완적 강점: 결과는 품질과 다양성이 서로 구별되면서도 보완적인 속성을 포착함을 나타냅니다. 품질 필터링이 지식 집약적 벤치마크(예: MMLU)에서 탁월한 성능을 보이는 반면, 다양성 최적화는 상식 추론(예: HellaSwag)을 개선합니다. 공동 최적화는 이 간극을 메워 모든 작업에서 경쟁력 있는 성능을 냅니다.

5. 의의 및 주장

본 논문은 다양성의 직접적인 최적화가 규모 면에서 실행 가능하며 실질적으로 유익하다고 주장합니다. 거친 입도의 근사치와 휴리스틱에서 벗어남으로써, SPOKES는 다음과 같은 점을 입증했습니다:

• 다양성은 중요한 축이다: 명시적인 품질 감독 없이도 다양성을 최적화하는 것만으로 일관된 이득을 얻을 수 있으며, 이는 데이터 분포에 대한 다양한 그래디언트 표현이 일반화 성능을 높인다는 것을 시사합니다.
• 품질 지표는 불충분하다: 표준 품질 지표는 데이터의 효용성을 완전히 포착하지 못합니다. 다양성을 공동으로 최적화하면, 선택된 부분 집합의 평균 품질 점수가 품질 전용 베이스라인보다 낮더라도 더 나은 다운스트림 성능을 이끌어내는 보완적인 훈련 신호를 제공할 수 있습니다.
• 확장성: 제안된 방법은 거대한 사전 학습 코퍼스에서 밀도가 높고 매우 다양한 부분 집합을 성공적으로 추출하며, 시맨틱 중복 제거나 무작위 샘플링에 의존하는 기존의 선택 파이프라인에 대한 유효한 대안을 제시합니다.

저자들은 SPOKES를 고품질 및 고다양성 부분 집합을 구축하기 위한 도구로 자리매김하며, 다양성의 이점이 서로 다른 필터링 체계와 노이즈 프로필을 가진 데이터셋 전반에 걸쳐 신뢰할 수 있게 나타난다는 점을 강조합니다.