Yi: Open Foundation Models by 01.AI

논문 개요

Yi는 중국의 AI 스타트업 01.AI가 2024년 발표한 오픈소스 대규모 언어 모델(LLM) 시리즈다. 이 보고서는 Yi-6B와 Yi-34B 모델의 학습 방법론, 아키텍처 설계, 데이터 파이프라인, 그리고 평가 결과를 상세히 기술한다.

Yi의 가장 주목할 만한 특징은 “데이터 품질이 양보다 중요하다” 는 철학이다. 대규모 데이터를 무작정 수집하는 대신, 엄격한 품질 필터링을 통해 정제된 소규모 고품질 데이터로 학습하는 전략을 채택했다. 또한 작은 모델을 확장하여 큰 모델을 만드는 depth upscaling 기법을 활용하여 학습 효율을 크게 높였다.

영어와 중국어 이중 언어 지원에 특히 강점을 보이며, 200K 토큰에 달하는 극도로 긴 컨텍스트를 처리할 수 있는 장문 컨텍스트 버전도 제공된다.

---

핵심 기여

1. 데이터 품질 우선 전략

Yi 팀은 학습 데이터 품질에 극도로 집중했다. 전체 파이프라인은 다음 단계로 구성된다:

데이터 수집 및 필터링 파이프라인:

원시 웹 데이터
    │
    ▼
언어 식별 (fastText)
    │
    ▼
품질 필터링 (휴리스틱 규칙)
    │ - 특수문자/HTML 비율 필터
    │ - 텍스트 길이 필터
    │ - 중복 제거 (MinHash LSH)
    ▼
ML 기반 품질 점수 평가
    │
    ▼
최종 학습 데이터 (~3.1T 토큰)

3.1T 토큰은 절대적 규모로는 다른 모델들(Llama 2: 2T, Qwen2.5: 18T)에 비해 작지만, 높은 품질 밀도로 인해 동등하거나 더 나은 성능을 달성했다.

2. Depth Upscaling

Yi-34B는 Yi-6B를 depth upscaling 기법으로 확장하여 초기화했다. 이 방법은 작은 모델의 가중치를 재사용하여 더 큰 모델을 초기화하는 것으로, 무작위 초기화 대비 학습 수렴 속도를 크게 높인다.

구체적인 절차:

1. Yi-6B (32 레이어) 학습 완료
2. 레이어를 반복 복사하여 Yi-34B (60 레이어) 초기화
3. Yi-34B 추가 사전학습 진행

이 접근법은 학습 비용을 절감하면서도 안정적인 학습 동학(dynamics)을 보장한다.

3. 200K 장문 컨텍스트

Yi-6B-200K와 Yi-34B-200K는 최대 200,000 토큰의 컨텍스트를 처리할 수 있다. 이는 약 40만 단어, 즉 두꺼운 소설 한 권 분량에 해당한다.

---

방법론 상세

아키텍처 설계

Yi는 LLaMA 2와 유사한 아키텍처를 기반으로 하되, 몇 가지 중요한 개선을 포함한다:

모델 구성 비교:

구성 요소	Yi-6B	Yi-34B
레이어 수	32	60
히든 차원	4096	7168
어텐션 헤드 (Q)	32	56
KV 헤드 (GQA)	4	8
컨텍스트 길이	4K (200K)	4K (200K)
어휘 크기	64,000	64,000

RoPE (Rotary Position Embedding)

Yi는 Positional Embedding으로 RoPE를 채택한다. RoPE는 쿼리와 키 벡터에 회전 행렬을 적용하여 상대 위치 정보를 인코딩한다:

$R_{\theta ,m}=\left(\begin{array}{ccccc}\cos m{\theta }_1& -\sin m{\theta }_1& & & \\ \sin m{\theta }_1& \cos m{\theta }_1& & & \\ & & \ddots & & \\ & & & \cos m{\theta }_{d/2}& -\sin m{\theta }_{d/2}\\ & & & \sin m{\theta }_{d/2}& \cos m{\theta }_{d/2}\end{array}\right)$

각 주파수는 ${\theta }_i=10000^{-2i/d}$ 로 설정된다. 내적 연산에서:

$q_m^T{k}_n=(R_{\theta ,m}q{)}^T(R_{\theta ,n}k)=q^T{R}_{\theta ,n-m}k$

즉, 어텐션 점수가 상대 위치 $n-m$에만 의존하게 된다.

GQA (Grouped Query Attention)

KV 캐시 메모리를 줄이기 위해 GQA를 사용한다. Yi-6B의 경우 32개 쿼리 헤드를 4개의 KV 헤드 그룹으로 나누어 메모리를 8배 절감한다:

$\text{메모리 절감}=\frac{n_h}{n_{kv}}=\frac{32}{4}=8\times$

SwiGLU 피드포워드 네트워크

$\text{FFN}(x)=(\text{Swish}(x{W}_1)\odot x{W}_2)W_3$

$\text{Swish}(x)=x\cdot \sigma (\beta x),\beta =1$

RMSNorm

Layer Normalization 대신 RMSNorm을 사용하여 계산 비용을 줄인다:

${\bar{x}}_i=\frac{x_i}{\text{RMS}(\mathbf{x})}{g}_i,\text{RMS}(\mathbf{x})=\sqrt{\frac{1}{n}{\sum }_{i=1}^n{x}_i^2}$

토크나이저

Yi는 64,000개의 어휘를 가진 BPE 토크나이저를 사용한다. 영어와 중국어에 최적화되어 있으며, 한국어, 일본어 등 다른 언어도 지원한다.

장문 컨텍스트 처리

200K 컨텍스트를 처리하기 위해 YaRN(Yet another RoPE extensioN) 기법을 적용한다. 기본 RoPE는 학습 시 본 컨텍스트 길이 이상으로 잘 일반화되지 않는 문제가 있는데, YaRN은 이를 보정하는 스케일링 팩터를 적용한다:

${\theta }_i^{\prime }={\theta }_i\cdot s^{-1},s=\frac{L_{\text{target}}}{L_{\text{train}}}$

이와 함께 attention sink 현상을 방지하기 위한 어텐션 마스킹 전략도 적용된다.

지시 학습 및 RLHF

Yi-Chat 모델은 사전학습 후 다음 단계의 정렬(alignment) 과정을 거친다:

1. SFT (Supervised Fine-Tuning): 다양한 지시-응답 쌍으로 미세조정
2. 보상 모델 학습: 인간 선호도 데이터로 보상 신호 학습
3. PPO: 보상 모델을 활용한 강화학습으로 응답 품질 최적화

---

실험 결과

종합 벤치마크 (Yi-34B)

벤치마크	Yi-34B	Llama 2-70B	Falcon-180B	GPT-3.5
MMLU	76.3	68.9	70.4	70.0
HellaSwag	82.6	87.3	88.9	85.5
ARC-C	65.4	67.3	70.3	71.7
TruthfulQA	56.2	44.9	45.5	47.9
GSM8K	67.9	56.8	54.3	57.1

Yi-34B는 두 배 이상 큰 Llama 2-70B 및 Falcon-180B와 비교하여 동등하거나 더 나은 성능을 보인다.

장문 컨텍스트 평가 (Needle-in-a-Haystack)

Yi-34B-200K는 200K 토큰에 달하는 긴 문서에서 특정 정보를 찾는 “Needle in a Haystack” 테스트에서 높은 정확도를 유지한다. 기존 모델들이 32K~100K 이상에서 급격한 성능 하락을 보이는 반면, Yi-200K는 전 구간에서 안정적인 성능을 보였다.

이중 언어 성능

중국어 벤치마크(C-Eval, CMMLU)에서도 최상위 수준의 성능을 달성:

벤치마크	Yi-34B	Baichuan2-13B	Qwen-14B
C-Eval	81.8	59.0	72.1
CMMLU	83.7	62.0	71.0

---

의의 및 한계

의의

데이터 품질의 중요성 재확인: Yi는 대규모 저품질 데이터보다 소규모 고품질 데이터가 더 효과적임을 실증적으로 보여주었다. 이는 단순히 더 많은 데이터를 수집하는 스케일링 전략에 대한 재고를 촉구한다.

Depth Upscaling의 실용성: 기존 소형 모델의 가중치를 재활용하여 대형 모델을 초기화하는 접근법은 학습 비용을 절감하면서도 좋은 성능을 달성하는 실용적 방법임을 입증했다.

이중 언어 LLM의 가능성: Yi는 영어와 중국어에서 동시에 뛰어난 성능을 보임으로써 다국어 LLM 개발의 가능성을 확장했다.

장문 컨텍스트의 실용화: 200K 토큰 컨텍스트는 법률 문서, 연구 논문, 코드베이스 등 실제 업무에서 요구되는 긴 입력을 처리할 수 있게 한다.

한계

수학/코딩 특화 부족: 일반 벤치마크에서는 강세를 보이지만, 수학 추론(MATH 데이터셋)과 코드 생성에서는 전문화된 모델(Qwen2.5-Math, DeepSeek-Coder)에 뒤처진다.

다국어 지원 범위: 영어와 중국어에 최적화되어 있어 다른 언어에서는 상대적으로 성능이 낮을 수 있다.

투명성 부족: 학습 데이터의 정확한 구성과 필터링 세부 사항이 완전히 공개되지 않아 완전한 재현이 어렵다.

평가 공정성: 일부 벤치마크에서의 높은 성능이 데이터 오염(data contamination)에 기인할 가능성을 배제하기 어렵다.

Yi 모델은 데이터 품질 중심 접근법과 효율적인 스케일링 전략으로 오픈소스 LLM 연구에 의미 있는 기여를 했으며, 이후 많은 모델 개발에 영향을 미쳤다.