바이트 레벨 인터페이스를 통한 교차 토크나이저 LLM 증류

서로 다른 토크나이저를 사용하는 모델 간의 지식 증류는 어휘 사전 불일치 문제로 인해 매우 까다로운 과제였다. 이 논문은 모든 토크나이저의 공통 분모인 바이트 레벨에서 증류를 수행하는 단순하고 효과적인 방법론을 제시하여, 도메인 특화 모델 학습이나 이종 모델 간의 지식 융합 가능성을 열어준다.

기존의 지식 증류는 교사 모델과 학생 모델이 동일한 어휘 사전(Vocabulary)을 공유한다는 가정하에, 각 토큰에 대한 확률 분포(Logit)를 일치시키는 방식으로 이루어진다. 하지만 토크나이저가 다르면 교사 모델의 50,000개 토큰에 대한 확률을 학생 모델의 32,000개 토큰 확률과 직접 비교할 수 없게 되어 정보 손실이 발생하거나 복잡한 매핑 전략이 필요해진다.

이 논문은 모든 텍스트 데이터의 근본 단위인 바이트(Byte)에 주목한다. 어떤 토크나이저를 사용하든 결국 데이터는 바이트 시퀀스로 표현될 수 있다는 점을 이용해, 교사 모델의 토큰 확률을 바이트 단위의 확률 분포로 변환한다. 학생 모델에는 바이트를 예측할 수 있는 작은 출력층을 추가하여, 교사가 생성한 바이트 확률을 직접 따라 하도록 학습시킨다.

결과적으로 학생 모델은 자신의 고유한 토크나이저를 유지하면서도 교사 모델의 고차원적인 추론 패턴을 바이트라는 공통 언어를 통해 전수받게 된다. 증류가 끝난 후에는 추가했던 바이트 헤드만 제거하면 일반적인 토큰 기반 모델로 즉시 활용이 가능하다.

BLD는 크게 두 단계로 구성된다. 첫 번째 단계는 학생 모델에 바이트 레벨 인터페이스를 추가하는 것이다. 기존 Transformer 아키텍처의 마지막 은닉 상태(Hidden State)를 입력으로 받아 256개의 바이트 값과 특수 토큰들을 예측하는 경량 선형 레이어(Linear Layer)를 기존 토큰 레벨 헤드와 병렬로 배치한다.

두 번째 단계는 실제 증류 과정이다. 교사 모델의 토큰 확률 분포를 바이트 레벨 확률로 변환하기 위해 Vieira 등의 근사 알고리즘을 사용한다. 이 알고리즘은 빔 서치(Beam Search)를 통해 특정 바이트 시퀀스에 대응하는 토큰 조합들의 확률을 합산하여 바이트별 조건부 확률을 계산한다. [토큰 확률 분포 입력 → 빔 서치 기반 바이트 시퀀스 매핑 → 바이트별 조건부 확률 출력]

최종 손실 함수(Loss Function)는 세 가지 요소의 합으로 정의된다. 학생 모델의 원래 토큰 예측에 대한 Cross-Entropy, 학생 모델의 바이트 헤드가 정답 바이트를 맞추도록 하는 Cross-Entropy, 그리고 교사의 바이트 확률 분포와 학생의 바이트 확률 분포 사이의 KL Divergence이다. [교사/학생 바이트 확률 입력 → KL Divergence 연산 → 지식 전이량 수치화]

BPE 토크나이저 간 전이 실험(Llama 3.2 3B → Qwen2 토크나이저)에서 BLD는 PiQA(75.68)와 AGI-ZH(35.97)에서 가장 높은 점수를 기록하며 기존 SOTA 기법인 ALM+SFT와 대등한 성능을 보였다. 다만 지시 이행 능력을 측정하는 IFEval에서는 30.58점으로 다른 기법들에 비해 낮은 수치를 기록하여 구조적 출력 유지에는 한계가 있음을 시사했다.

교차 모델 증류 실험(OpenMath2-Llama3.1-8B → Gemma2 2B)에서는 수학 벤치마크인 GSM8K에서 62.55점을 기록하여 SFT(59.29)나 ALM+SFT(61.56)보다 우수한 성적을 거두었다. 이는 바이트 레벨 증류가 복잡한 추론 능력을 전이하는 데 효과적임을 입증한다.

바이트 레벨 확률 근사 품질 분석 결과, 빔 폭(K)을 10으로 설정하고 프루닝 임계값(epsilon)을 0.01로 설정했을 때 Jensen-Shannon Divergence가 0.0045 수준으로 낮게 유지되면서도 계산 효율성을 확보할 수 있음을 확인했다.

BLD의 핵심은 토큰 기반 언어 모델을 바이트 레벨 확률 생성기로 취급하는 것이다. 이를 위해 $P(b_i | b_{<i})$를 계산해야 하는데, 이는 해당 바이트 접두사를 포함하는 모든 가능한 토큰 시퀀스들의 확률 합으로 정의된다. 본 연구에서는 이를 효율적으로 계산하기 위해 트라이(Trie) 구조와 빔 서치를 결합한 근사 방식을 채택했다.

학생 모델에 추가되는 바이트 디코더 $O_b$는 단순한 선형 투영 레이어로 구현되었다. 실험 결과, 복잡한 자기회귀(Autoregressive) 구조 없이도 10개의 병렬 선형 레이어만으로 토큰당 최대 10바이트까지 동시에 예측하도록 구성했을 때 충분한 성능이 나왔다. 이는 바이트 레벨 인터페이스가 매우 가벼워도 지식 전송 통로로서 기능할 수 있음을 보여준다.

학습 시에는 LoRA(Rank=64)를 사용하여 학생 모델의 백본을 미세 조정하며, 임베딩 레이어와 LM 헤드는 새로운 토크나이저에 맞춰 Fast Vocabulary Transfer(FVT) 방식으로 초기화한다. 이러한 설정은 파라미터 효율성을 높이면서도 새로운 어휘 체계에 빠르게 적응할 수 있도록 돕는다.

본 연구는 8B 이하의 모델 규모에서만 실험이 진행되어, 수십억 개 이상의 파라미터를 가진 초대형 모델에서의 동작 특성은 아직 검증되지 않았다. 또한, 바이트 레벨 증류가 일반적인 지식 전이에는 효과적이지만 IFEval과 같은 지시 이행 벤치마크에서는 성능 저하가 관찰되어, 정교한 출력 구조를 유지하는 능력에는 개선이 필요하다.