원본 Figure & Table
논문의 핵심 Figure들을 원본 그대로 보존합니다.
이 논문은 LLM 사전학습 데이터 혼합을 "모델의 학습 선호도가 시시각각 바뀐다"는 관점에서 본다 — 고정된 비율로 데이터를 섞는 것은 근본적으로 한계가 있으며, 그룹 영향 함수(Group Influence)라는 효율적 신호로 모델이 지금 무엇을 '소화'하고 싶어하는지 실시간으로 측정해 혼합 비율을 동적으로 조정해야 한다.
arXiv →대규모 언어 모델(LLM)의 사전학습은 방대한 웹 데이터를 활용하며, 이 데이터는 학술 논문·책·뉴스·포럼 등 수십 개의 도메인(domain)으로 구성됩니다. 각 도메인을 어떤 비율로 섞느냐 — 이른바 데이터 혼합(data mixture) — 는 모델의 최종 성능을 결정하는 핵심 변수입니다. GPT-3 저자들이 Wikipedia를 고품질 데이터로 보고 그 비율을 늘린 사례가 대표적입니다.
그러나 현실에서 "최적의 혼합 비율"을 찾는 일은 극도로 어렵습니다. 기존 접근법들은 다음과 같은 한계를 가집니다.
"a model's learning preferences for various data domains shift dynamically throughout training. Crucially, observing these evolving preferences in a computationally efficient manner remains a significant challenge."모델의 각 데이터 도메인에 대한 학습 선호도는 훈련 전반에 걸쳐 동적으로 변화한다. 결정적으로, 이러한 진화하는 선호도를 계산 효율적인 방식으로 관찰하는 것이 중요한 도전 과제로 남아 있다.
결국 이 논문이 해결하려는 문제는 이것입니다: "모델이 현재 무엇을 더 배우고 싶어하는지를, 저렴한 비용으로 측정하고, 이를 즉각 혼합 전략에 반영할 수 있는가?"
TiKMiX의 핵심 아이디어는 영향 함수(Influence Function)를 개별 샘플이 아닌 도메인 전체 그룹에 적용하는 것입니다. 기존 영향 함수는 샘플 단위로 연산하므로 대규모 사전학습에는 비용이 너무 크지만, TiKMiX는 그라디언트 누적(gradient accumulation) 기법을 활용해 도메인 전체의 집합적 영향을 단 한 번의 헤시안-역행렬 곱(Hessian-inverse-vector product)으로 계산합니다.
Group Influence 점수를 구한 뒤, TiKMiX는 이를 최대화하는 혼합 비율 벡터 w를 찾는 최적화 문제를 설정합니다. 이를 푸는 두 가지 방법을 제안합니다.
보충 각 블록을 클릭하면 해당 모듈의 역할과 설계 이유를 확인할 수 있습니다.
먼저 기본 학습 목표를 정의합니다. 모델 파라미터 θ는 데이터셋 D = {z₁, …, zₙ}에서 경험적 손실을 최소화합니다.
| 변수 | 의미 | 비고 |
|---|---|---|
θ* | 최적 파라미터 | d차원 벡터 |
J(θ) | 경험적 위험(empirical risk) | 전체 데이터셋에 대한 평균 손실 |
L(zᵢ, θ) | i번째 샘플의 손실 | 크로스엔트로피 등 |
N | 전체 훈련 샘플 수 | 대규모 사전학습에서 수백B 토큰 |
직관적 해설
수학적 배경
그룹 S ⊆ D의 영향을 측정하기 위해, S의 모든 원소를 무한소 가중치 ε만큼 상향시킨 섭동 목적 함수를 도입합니다.
| 변수 | 의미 | 비고 |
|---|---|---|
ε | 섭동 강도 | 무한소 양수. ε=0이면 원래 θ* 복구 |
S | 영향을 측정할 데이터 그룹(도메인) | 예: "Wikipedia 도메인 전체" |
θ*_ε | 섭동된 목적 함수의 최적 파라미터 | S를 더 강조했을 때의 파라미터 |
직관적 해설
수학적 유도
최적성 1계 조건(∇θ Jε = 0)을 ε로 미분하고 ε=0에서 평가하면, 헤시안과 그룹 그라디언트만으로 파라미터 영향을 표현할 수 있습니다.
| 변수 | 의미 | 비고 |
|---|---|---|
Iₚₐᵣₐₘ(S) | S가 모델 파라미터에 미치는 영향 | d차원 벡터 |
H_{θ*} | 원래 목적 함수의 헤시안(2계 도함수 행렬) | d×d 행렬, 직접 역산은 불가 → EK-FAC 등 근사 필요 |
Σ ∇θ L(zⱼ, θ*) | 그룹 S 전체의 누적 그라디언트 | 그룹 내 샘플 수에 무관하게 gradient accumulation으로 효율 계산 |
직관적 해설
유도 과정 요약
이를 파라미터가 아닌 스칼라 함수 f(θ) (예: 검증 손실)에 대한 영향으로 변환합니다.
| 변수 | 의미 | 비고 |
|---|---|---|
If(S) | S가 검증 함수 f에 미치는 영향 (스칼라) | 양수: 검증 성능 향상, 음수: 저하 |
∇θ f(θ*) | 검증 손실의 파라미터 그라디언트 | 멀티 검증 세트 각각에 대해 계산 |
Σ ∇θ L(zⱼ, θ*) | 도메인 S의 누적 그라디언트 | gradient accumulation으로 배치 계산 |
직관적 해설
계산상의 장점
n개의 검증 태스크와 m개의 도메인에 대해 영향 행렬 S ∈ ℝⁿˣᵐ를 구성합니다. S의 원소 Sᵢⱼ는 도메인 j가 검증 태스크 i에 미치는 영향입니다. 혼합 벡터 w = [w₁, …, wₘ]ᵀ에 대해 각 태스크의 기대 영향은 P = S·w입니다. 태스크 간 스케일 차이를 제거하기 위해 정규화 영향 P̂ᵢ를 정의합니다.
| 변수 | 의미 | 비고 |
|---|---|---|
P̂ᵢ | 태스크 i의 정규화된 기대 영향 | 최대 가능 영향 대비 비율 |
max_j Sᵢⱼ | 태스크 i에 가장 큰 영향을 주는 도메인의 점수 | 태스크별 스케일 조정 기준 |
ε | 수치 안정성을 위한 소상수 | 10⁻⁸ (0 나눗셈 방지) |
왜 정규화가 필요한가
최적화 목적 함수는 세 가지 목표를 균형 있게 통합합니다.
| 항 | 의미 | 설계 이유 |
|---|---|---|
α·std(P̂) | 영향 균등성 (표준편차 최소화) | 특정 태스크만 치우치게 개선하면 다른 태스크 성능이 희생되므로, 균등 분포 장려 |
-β·Σ P̂ᵢ | 총 영향합 최대화 | 전체적인 성능 개선폭을 크게 하기 위해 총 영향의 음수를 최소화 |
-γ·H(w) | 데이터 다양성 (엔트로피 최대화) | 혼합이 특정 도메인에 과도하게 쏠리면 일반화 능력 저하 → 엔트로피 음수를 최소화해 균등 분포 장려 |
α, β, γ | 세 목표의 가중치 | 논문 실험에서 모두 1로 설정 (동등 중요도) |
직관적 해설
왜 이 세 가지 항인가
이 최적화는 다음 제약 조건 하에서 수행됩니다.
| 제약 | 의미 | 이유 |
|---|---|---|
Σwⱼ = 1 | 확률 단체(simplex) 제약 | 혼합 비율이 합이 1인 분포를 이루어야 함 |
wⱼ ≥ 0 | 비음수 제약 | 음의 혼합 비율은 물리적으로 불가 |
S·w ≥ S·w_prior | 파레토 개선 제약 | 모든 태스크에서 이전 혼합보다 나빠지면 안 됨 — 새 혼합이 전 스테이지보다 항상 개선됨을 보장 |
SLSQP 풀이
보충 아래 슬라이더로 특정 도메인(예: Wikipedia)의 비율을 조정하면, 세 목적 항이 어떻게 변하는지 체험할 수 있습니다. (실제 Group Influence 값이 아닌 개념 데모입니다.)
TiKMiX-D는 도메인 영향이 선형 가산적이라 가정하지만, 실제로는 서로 다른 도메인이 결합될 때 비선형 교차 효과가 발생할 수 있습니다. TiKMiX-M은 이를 보완하기 위해 TiKMiX-D의 결과를 출발점으로 삼아 서러게이트 회귀 모델(LightGBM)로 혼합→성능 관계를 학습합니다.
작동 과정: TiKMiX-D의 최적 혼합 w_best-D를 중심으로, 라틴 하이퍼큐브 샘플링(Latin Hypercube Sampling, LHS)으로 N개의 후보 혼합 벡터를 생성합니다. 각 후보에 대해 실제 Group Influence 점수 y를 계산하고, (w, y) 쌍으로 LightGBM을 학습합니다. 이후 Dirichlet 분포 기반 반복 탐색으로 전역 최적을 찾습니다.
| 항목 | 설정값 | 비고 |
|---|---|---|
| 데이터셋 | RefinedWeb (26 도메인) | 도메인 분류기: DeBERTaV3 기반 |
| 모델 크기 | 1B, 7B 파라미터 | 두 스케일에서 모두 검증 |
| 총 학습 토큰 | 1T 토큰 | 2 스테이지 × 500B |
| 혼합 재조정 빈도 | 200B 토큰마다 | 1T 학습 중 5회 재조정 기회 |
| LHS 후보 수 (N) | 256개 | 각 후보에 0.1B 토큰 샘플 |
| 서러게이트 모델 | LightGBM | 그라디언트 부스팅 의사결정 트리 |
| SLSQP 초기값 | 균등 분포 | α=β=γ=1 |
| Dirichlet α 범위 | α_max → α_min (로그 어닐링) | 광역 탐색→국소 수렴 |
| 헤시안 역행렬 근사 | EK-FAC (Anthropic 2023) | 50B+ 파라미터 모델 적용 가능 |
| 검증 세트 수 | 9개 벤치마크 | 5 in-domain + 4 out-of-domain |
| TiKMiX-D 비용 | REGMIX 대비 약 20% | 프록시 모델 학습 불필요 |
평가 논문은 각 모델 스케일별 정확한 GPU 종류·수량·학습 시간을 명시하지 않습니다. Group Influence 계산에 EK-FAC 근사를 사용하지만 이 근사의 정확도가 최종 성능에 미치는 민감도는 별도 분석하지 않습니다.
TiKMiX의 핵심 가설 — "Group Influence를 최대화하면 다운스트림 성능이 오른다" — 을 검증하기 위해 10가지 서로 다른 혼합 비율로 1B 모델을 500B 토큰에 학습하고, Group Influence와 실제 성능 사이의 관계를 측정했습니다.
"a strong positive correlation (i.e., Pearson correlation coefficient ρ = 0.789) between the total Group Influence and the average downstream scores"총 Group Influence와 평균 다운스트림 점수 사이에 강한 양의 상관관계(피어슨 상관계수 ρ = 0.789)가 존재한다.
이 결과는 Group Influence가 실제 성능 향상의 유효한 대리 지표(proxy)임을 실증적으로 뒷받침하며, TiKMiX-D와 TiKMiX-M의 이론적 토대를 제공합니다.
| 벤치마크 | Pile-CC | Average | DoReMi | QUAD | REGMIX | Ours TiKMiX-D | Ours TiKMiX-M |
|---|---|---|---|---|---|---|---|
| In-Domain | |||||||
| MMLU | 31.2 | 30.9 | 31.2 | 31.7 | 31.5 | 32.2 | 31.8 |
| HellaSwag | 55.6 | 55.9 | 55.3 | 56.5 | 56.0 | 57.4 | 56.6 |
| ARC Easy | 63.2 | 64.1 | 65.7 | 62.8 | 66.2 | 69.3 | 70.7 |
| ARC Challenge | 32.7 | 32.1 | 33.6 | 33.5 | 33.2 | 37.0 | 38.3 |
| TriviaQA | 16.3 | 17.3 | 15.5 | 17.6 | 15.8 | 17.7 | 17.3 |
| Out-of-Domain | |||||||
| PiQA | 69.2 | 71.5 | 73.1 | 72.4 | 73.3 | 74.1 | 74.5 |
| OpenBookQA | 37.1 | 34.6 | 36.5 | 36.6 | 37.0 | 37.4 | 37.4 |
| BoolQ | 58.7 | 58.3 | 59.2 | 60.5 | 58.9 | 61.3 | 62.2 |
| MathQA | 22.5 | 23.7 | 23.1 | 23.9 | 23.3 | 23.5 | 24.2 |
| FLOPs | 0 | 0 | 4.2e19 | 2.3e18 | 3.7e18 | 7.2e17 | 3.2e18 |
| 평균 성능 | 42.9 | 43.2 | 43.7 | 43.9 | 43.9 | 45.5 | 45.9 |
| 최고 달성 | 0/9 | 0/9 | 0/9 | 0/9 | 0/9 | 4/9 | 6/9 |
Group Influence 계산을 위해 얼마나 많은 데이터가 필요한가? 손실(perplexity) 기반 관찰과 TiKMiX의 그룹 영향 관찰을 0.1B~10B 토큰 범위에서 비교했습니다.
부록 추가 1B 모델 외에 7B 모델에서도 TiKMiX-D와 REGMIX를 비교했습니다. 7B 모델은 첫 스테이지 500B, 두 번째 스테이지 200B 토큰으로 학습했습니다.
| 벤치마크 | REGMIX (1B) | TiKMiX-D (1B) | REGMIX (7B) | TiKMiX-D (7B) |
|---|---|---|---|---|
| MMLU | 31.5 | 32.2 | 40.7 | 41.5 |
| HellaSwag | 56.0 | 57.4 | 76.6 | 76.4 |
| ARC Easy | 66.2 | 69.3 | 78.5 | 78.4 |
| ARC Challenge | 32.2 | 37.0 | 49.4 | 50.2 |
| TriviaQA | 15.8 | 17.7 | 46.4 | 45.3 |
| PiQA | 73.3 | 74.1 | 79.1 | 79.2 |
| OpenBookQA | 37.0 | 37.4 | 43.2 | 45.4 |
| MathQA | 23.2 | 23.5 | 28.8 | 29.9 |
| 평균 | 43.9 | 45.5 | 55.3 | 56.0 |
평가 7B 모델에서 HellaSwag와 ARC Easy는 REGMIX와 거의 동등합니다. 이는 7B 모델이 해당 벤치마크에서 이미 포화점에 가까워졌거나, 두 번째 스테이지의 짧은 200B 학습이 충분한 차별화를 만들기에 부족했을 가능성이 있습니다. 논문은 "7B 모델은 학습 내내 성능이 지속적으로 향상되어 더 긴 학습 데이터에서 TiKMiX의 이점이 더 두드러질 것"이라고 해석합니다.
| 얻는 것 | 잃는 것 |
|---|---|
| 동적 선호도 반영 → 성능 +2% | 학습을 여러 스테이지로 분리해야 함 |
| REGMIX 대비 80% 계산 비용 절감 | 헤시안 역행렬 근사 오차 내재 |
| 프록시 모델 학습 불필요 | Group Influence 계산용 검증 세트 필요 |
| 1B~7B 스케일 검증 완료 | 더 큰 스케일(70B+)에서의 검증 미완 |
이 논문은 세 가지 측면에서 LLM 사전학습 커뮤니티에 기여합니다.
실용적 기여: 기업 수준 LLM 사전학습 파이프라인에서 REGMIX, DoReMi를 대체하거나 보완할 수 있는 효율적인 동적 혼합 전략을 제공합니다. 특히 프록시 모델 학습 비용 없이 80% 계산을 절감하면서도 성능을 높이므로, 실용적 채택 가능성이 높습니다. 멀티스테이지 사전학습을 이미 사용하는 Llama 3, Qwen 같은 대형 모델 학습 파이프라인에 자연스럽게 통합할 수 있습니다.
이론적 기여: 영향 함수를 개별 샘플이 아닌 도메인 그룹 단위로 확장하는 Group Influence 메트릭은 그 자체로 이후 연구들이 활용할 수 있는 도구입니다. 또한 "모델의 데이터 선호도가 스케일과 학습 단계에 따라 변한다"는 현상을 1B/7B 실험으로 체계적으로 입증한 것도 중요한 경험적 기여입니다.
저자가 제안한 후속 과제: 더 큰 스케일 모델(70B+)과 더 다양한 데이터셋(코드·수학·멀티모달 포함)에서 Group Influence와 TiKMiX의 효과를 추가 검증하는 것을 명시적으로 제안합니다.
"We plan to conduct further experiments on larger-scale models and more diverse datasets to further validate the effectiveness of Group Influence and TiKMiX."더 큰 규모의 모델과 더 다양한 데이터셋에서 Group Influence와 TiKMiX의 효과를 추가 검증하기 위한 실험을 계획하고 있다.
논문의 핵심 Figure들을 원본 그대로 보존합니다.
샘플 단위 영향 함수는 각 훈련 샘플마다 헤시안 역행렬 곱을 수행해야 합니다. 대규모 LLM 사전학습에서는 수백B 토큰, 즉 수억 개의 샘플에 대해 이를 반복해야 하므로 계산적으로 완전히 불가능합니다. Group Influence는 도메인 전체의 그라디언트를 먼저 누적(gradient accumulation)한 뒤 헤시안 역행렬 곱을 단 한 번만 수행합니다. 이로써 계산 비용이 그룹 크기 |S|에 독립적이 됩니다.
Figure 3 (논문 부록 Figures 5~11)에서 시각적으로 확인할 수 있습니다. 1B와 7B 모델 모두에서 학습이 진행될수록 각 도메인의 Group Influence 점수가 변화합니다. 예를 들어 초반에 높은 영향력을 보이던 특정 도메인이 학습 후반에는 영향력이 줄어드는 패턴이 관찰됩니다. 또한 1B와 7B 모델은 동일한 데이터를 보더라도 서로 다른 선호도 패턴을 보이는데, 이는 모델 스케일에 따라 최적 혼합이 달라야 함을 의미합니다.
계산 자원이 극히 제한적이고 REGMIX 수준의 비용조차 감당하기 어렵다면 TiKMiX-D를 권장합니다 — REGMIX의 20% 비용으로 동등하거나 우수한 성능을 제공합니다. 추가 비용을 감수할 수 있고 ARC처럼 어려운 추론 태스크에서 최대 성능을 원한다면 TiKMiX-M이 낫습니다. TiKMiX-M은 최고 성능(45.9%)과 9개 태스크 중 6개 최고 달성을 보입니다.
이 제약은 "새로운 혼합이 이전 혼합보다 어떤 태스크에서도 나빠지면 안 된다"는 최소 안전 장치입니다. 단순히 총 영향합을 최대화하면 일부 태스크를 희생해 다른 태스크를 크게 향상시키는 극단적 혼합이 최적해로 나올 수 있습니다. 파레토 개선 제약은 모든 태스크에서 단조적 개선(monotonic improvement)을 보장합니다. 이전 혼합 w_prior를 기준으로 제약을 걸기 때문에 이전 스테이지에서 이미 학습된 능력이 퇴행하지 않습니다.
REGMIX는 소규모 실험 결과를 회귀 모델로 외삽하지만, 소규모→대규모 전이를 가정하고 학습 전반에 걸친 선호도 변화를 무시합니다. DoReMi는 프록시 모델을 별도로 학습해야 하는 비용이 크고, 반복적 업데이트가 필요합니다. TiKMiX는 현재 모델 자체에서 그라디언트 정보를 추출하므로 별도 프록시 모델이 불필요하고, 학습 중 어느 시점에서도 즉각 적용 가능하며, 계산 효율이 훨씬 뛰어납니다.
보충 논문이 사용하는 "under-digestion"은 데이터 소화 부족을 의미합니다. 고정 혼합 비율에서 특정 도메인의 비율이 해당 시점 모델의 학습 능력 대비 너무 낮으면, 그 도메인에서 충분한 학습이 일어나지 않습니다. 이는 모델이 해당 도메인 데이터를 "먹어도" 제대로 소화(학습)하지 못하는 상태, 즉 under-digestion입니다. 동적 혼합이 이를 방지하는 핵심 이유입니다.
평가 논문은 사전학습(pre-training) 전용이지만, Group Influence의 개념은 멀티태스크 파인튜닝(instruction tuning)의 태스크 혼합 최적화에도 유사하게 적용 가능합니다. 단, 파인튜닝에서는 태스크별 데이터 양이 훨씬 작으므로 샘플 단위 영향 함수도 실용적일 수 있어, 그룹 근사의 이점이 줄어들 수 있습니다. 이를 명시적으로 검토하는 것은 흥미로운 후속 연구가 될 것입니다.
