바이트 레벨 모델을 위한 토큰 기반 LLM 디스틸레이션 혁신

초록

본 논문은 기존 토큰 기반 대형 언어 모델(LLM)을 그대로 활용해 바이트 레벨 언어 모델(BLM)을 효율적으로 만들기 위한 두 단계 디스틸레이션 파이프라인을 제안한다. 첫 단계에서는 바이트 임베딩을 토큰 임베딩에 정렬하고, 교차 엔트로피와 KL 손실을 통해 의미 정보를 전달한다. 두 번째 단계에서는 바이트‑전용 슈퍼바이즈드 파인튜닝을 수행해 완전한 바이트 생성 능력을 부여한다. Llama, Qwen, OLMo 등 다양한 모델군에 적용했으며, 약 125 B 바이트(≈1.25 억 토큰) 학습 데이터만으로도 원본 토큰 모델의 90 % 이상 성능을 유지한다.

상세 분석

이 연구는 바이트 레벨 모델이 토큰화에 의존하지 않아 언어 다양성 및 희소 언어 처리에서 장점을 가질 수 있다는 점을 출발점으로 삼는다. 그러나 기존 BLM은 수조 바이트 규모의 데이터와 수백 GPU·수주에 달하는 학습 비용이 필요했으며, 토큰 기반 LLM을 그대로 옮겨오려면 토큰‑바이트 경계 불일치라는 근본적인 구조적 장애물을 극복해야 했다. 논문은 이를 해결하기 위해 두 단계 커리큘럼을 설계한다.

1️⃣ Progressive Knowledge Distillation (PKD)

• Embedding Alignment (L_align): 바이트 시퀀스 중 토큰 경계에 해당하는 바이트 위치(ˆe_k)의 인코더 출력과 교사 모델의 토큰 임베딩(e_k)을 L2 거리로 최소화한다. 이는 바이트 수준 표현이 토큰 의미 공간에 매핑되도록 하는 기초 단계다.
• Joint Distillation (L_distill): 토큰 경계에 맞춰 바이트 시퀀스를 재샘플링한 뒤, 교사와 학생 모델의 출력 분포에 KL 발산을 적용한다. 여기서 교사의 조건부 확률 P_T(·|t<k)와 학생의 P_S(·|x<pos(k))를 비교함으로써 토큰‑바이트 길이 차이를 보정한다.
• Boundary Learning (L_boundary): 라우팅 모듈에 One‑Byte Lookahead Router를 도입해 바이트 간 유사도 기반으로 경계 확률을 예측한다. 이때 바이너리 교차 엔트로피 손실을 사용해 교사의 토크나이저 경계(y_i)를 학습한다.

세 손실을 순차적으로 최적화하는 커리큘럼은 각각의 모듈이 안정된 기반 위에서 학습되게 하여, 복합 손실을 한 번에 적용할 때 발생하는 불안정성을 크게 감소시킨다.

2️⃣ Byte‑Level Supervised Fine‑Tuning (SFT)
PKD 단계가 토큰 공간과 정렬된 바이트 모델을 만든 뒤, 전체 파이프라인을 바이트‑전용으로 전환한다. 여기서는 Dechunk 모듈과 Decoder를 활용해 바이트 수준에서 직접 다음 바이트와 경계 정보를 예측한다. 두 가지 경계 예측 전략을 제시한다.

• Joint Boundary Prediction (JBP): 출력 어휘를 바이트 × 2 로 확장해 바이트와 경계 플래그를 동시에 예측한다.
• Multi‑Byte Prediction (MBP): 보조 헤드를 두어 t+2 위치의 바이트를 예측하고, 이를 라우팅 모듈에 입력해 t+1의 경계를 추정한다. 이는 인퍼런스 시 미래 정보를 요구하지 않으면서도 라우팅 로직과 일관성을 유지한다.

또한 Shifted‑Upsampling 기법을 도입해 청크‑레벨 표현을 바이트 레벨로 복원할 때, 청크 마지막 바이트에만 해당 청크 정보를 노출시켜 자동 회귀성을 보존한다.

실험 결과는 Llama‑3.2‑3B, Qwen‑3‑4B, OLMo‑1.5B 등 다양한 토큰 모델을 학생 모델로 변환했을 때, MMLU, GSM‑8K, BBH 등 벤치마크에서 원본 모델 대비 92 %~96 % 수준의 정확도를 달성했음을 보여준다. 특히 전체 학습 데이터가 125 B 바이트에 불과함에도 불구하고, 토큰‑기반 BLM을 처음부터 학습한 경우(수조 바이트 요구)와 비교해 동일하거나 더 나은 효율성을 입증한다.

기여는 다음과 같다.

• 토큰‑바이트 경계 불일치를 해결하는 두 단계 디스틸레이션 프레임워크 제시.
• One‑Byte Lookahead Router와 Shifted‑Upsampling을 통한 경계 학습 및 자동 회귀 보장.
• 다양한 모델군에 대한 광범위한 실증과 On‑Policy Distillation 등 추가적인 사후 강화 기법 탐색.
• 코드·체크포인트 공개를 통한 연구 재현성 및 커뮤니티 확장성 확보.

이러한 접근은 BLM 연구의 진입 장벽을 크게 낮추고, 토큰화에 얽매이지 않은 언어 모델을 빠르게 구축할 수 있는 실용적인 길을 제시한다.