NL-336, ReTok: Replacing Tokenizer to Enhance Representation Efficiency in Large Language Model, Preprint 2024

◼ Comment

  • 오픈소스 모델에 중국어 vocab expansion하여 실험한거다
  • 기본 백본에다가 중국어 코퍼스로 vocab을 만들어서 merge한다.
    • 이때 원래있던 vocab은 그대로 유지하고
    • 새로운 vocab은 기존 토크나이저로 쪼개서 각 임베딩을 평균낸다
    • 이렇게 하면, 기존에 백본이 어떻게 토큰을 이해하는지 반영할 수 있기 때문에 좋은 초기화라고 생각하는듯
    • 실험에서도 이런 초기화가 더 좋다고 보여줌
  • 학습은
    • 백본은 고정하고
    • embedding / LM head layer만 튜닝하는것
    • 학습데이터는 영어/중국어/코드 혼합을 사용했고 (자세한 혼합 비율은 안나오는듯)
    • 성능은 기존 백본에서 조금만 떨어진다는 것
  • 이 방법은 중국어에서 기존 백본의 성능은 유지하면서
    • 토크나이저는 새롭게 제안해서, 압축이 잘되니 속도나 이런점에서 강점이 있다고 주장하는 것

Abstract

Tokenizer는 대규모 언어 모델(LLM)의 핵심 구성 요소이며, 높은 압축률을 가진 토크나이저는 모델의 표현 효율성과 처리 효율성을 향상시킬 수 있다. 그러나 토크나이저는 모든 상황에서 높은 압축률을 보장하지 못하며, 입력과 출력의 평균 길이가 증가하면 모델의 학습 및 추론 비용도 증가하게 된다. 따라서 모델의 성능을 유지하면서 최소한의 비용으로 효율성을 향상시키는 방법을 찾는 것이 중요하다.

본 연구에서는 LLM의 토크나이저를 교체함으로써 모델의 표현 및 처리 효율성을 개선하는 방법을 제안한다. 우리는 기존 모델의 입력층과 출력층의 파라미터를 기반으로 이를 교체 및 재초기화하고, 다른 파라미터는 고정한 채 해당 층들만 학습하는 방식을 제안한다.

다양한 LLM에 대해 실험을 수행한 결과, 제안한 방법은 토크나이저를 교체한 이후에도 모델의 성능을 유지하면서, 특히 긴 텍스트에 대해 디코딩 속도를 크게 향상시킬 수 있음을 보였다. 

1 Introduction

Tokenizer는 대규모 언어 모델(LLM)의 기본 구성 요소로서 (Brown et al., 2020; OpenAI, 2023; Touvron et al., 2023), 텍스트를 토큰 시퀀스로 변환하는 전처리 과정에 사용된다. 이 과정은 텍스트를 더 작은 단위로 나누어 모델이 보다 효율적으로 처리하고 분석할 수 있도록 한다. 현재 서로 다른 LLM들은 일반적으로 서로 다른 토크나이저를 사용하며, 이러한 토크나이저는 각자의 학습 데이터에 기반하여 BPE, WordPiece, Unigram 등의 다양한 방법으로 학습된다.

일반적으로 이러한 방법들은 학습 데이터에서 등장 빈도에 따라 토큰의 vocabulary를 구성한다. 이를 통해 학습 데이터에 대해서는 비교적 높은 압축률(즉, 더 짧은 시퀀스 길이)을 달성할 수 있다. 그러나 학습 데이터와 테스트 데이터 간 분포 차이가 클 경우, 토크나이저의 압축률이 낮아지는 문제가 발생한다 (Tran, 2020; Liu et al., 2023; Dobler and de Melo, 2023). 예를 들어, 일반 도메인 코퍼스로 학습된 토크나이저는 코드나 수학과 같은 특정 도메인에서 낮은 압축률을 보인다.

압축률이 낮은 토크나이저를 사용할 경우, 추론 시 모델이 더 긴 시퀀스를 생성하게 되어 계산 비용과 추론 시간이 증가한다 (Ahia et al., 2023). 또한 특정 언어나 태스크에 대해 supervised fine-tuning(SFT)을 수행할 때도 더 많은 연산이 필요하게 되어 학습 효율이 저하되고, 최종 성능에도 영향을 줄 수 있다 (Dagan et al., 2024).

더 나아가 실제 응용에서 여러 LLM을 함께 사용하는 경우(예: 모델 간 협업, Xiong et al., 2023), 서로 다른 토크나이저에 맞춰 최적화 및 적응이 필요해지며, 이는 개발 및 유지 비용을 증가시킨다. 따라서 더 높은 압축률을 제공하는 임의의 토크나이저로 사전학습된 모델의 기존 토크나이저를 대체할 수 있다면, LLM 사용 비용을 크게 줄일 수 있다.

본 연구에서는 LLM의 토크나이저를 교체(Replace Tokenizer, ReTok)하여 모델의 표현 및 처리 효율을 향상시키면서, 성능은 최대한 유지하는 방법을 제안한다. 먼저, 영어와 코드에서 높은 압축률을 보이는 Llama3 토크나이저를 예로 들어, vocabulary를 확장하여 중국어에서도 높은 압축률을 달성하도록 한다. 이후 새로운 토크나이저에 맞추어 모델의 입력층과 출력층을 교체하고, 새로운 토크나이저의 vocabulary와 기존 토크나이저 간의 대응 관계를 기반으로 파라미터를 초기화한다.

마지막으로, 새로 추가된 파라미터만 학습하고 나머지 파라미터는 고정한 채 모델을 학습한다. 실험은 Qwen1.5-0.5B, Aquila2-7B, Llama3-8B 모델에서 수행되었으며, 그 결과 토크나이저를 교체한 이후에도 모델 성능을 유지하면서 특정 도메인에서 추론 시간을 크게 줄일 수 있음을 확인하였다. 

2 Related Work

토크나이저에 관한 연구는 일반적으로 그 목적에 따라 두 가지 범주로 나눌 수 있다.

첫 번째 범주는 모델이 보다 효과적으로 학습할 수 있도록 최적의 토크나이저를 설계하는 연구이다. Singh and Strouse (2024)는 수학 문제에서 숫자 토크나이제이션 방식이 미치는 영향을 분석하였다. Clark et al. (2022)는 토크나이저를 제거하고 문자 단위로 직접 입력을 표현하는 방식을 탐구하였다. Zheng et al. (2023)는 코드 작업에서 모델 성능을 향상시키기 위해 Syntax-Aware Tokenization 방법을 제안하였다. Rajaraman et al. (2024)는 다양한 토크나이제이션 알고리즘을 비교·분석하기 위한 이론적 프레임워크를 제시하고, 토크나이제이션 성능을 개선하기 위한 여러 제안을 제공하였다. 또한 아랍어(Alyafeai et al., 2023)나 일본어(Tolmachev et al., 2018)와 같은 특정 언어에서의 토크나이제이션 성능을 개선하는 연구들도 존재한다.

두 번째 범주는 사전학습된 모델의 기존 토크나이저를 수정하는 데 초점을 둔 연구이다. Xue et al. (2022)는 최소한의 수정으로 사전학습 모델이 바이트 시퀀스를 처리할 수 있도록 하는 방법을 연구하였다. Minixhofer et al. (2024)는 토크나이저를 입력으로 받아 해당 임베딩을 예측하는 하이퍼네트워크를 학습하는 방법을 제안하였다. Gee et al. (2024)와 Dagan et al. (2024)은 사전학습 모델의 토크나이저를 적응시켜 downstream 태스크에서의 효율성과 성능을 향상시키는 방법을 연구하였다.

이러한 연구들 중에서 본 연구는 두 번째 범주에 더 가깝다. 즉, 사전학습된 모델을 기반으로 토크나이저를 개선한다는 점에서 유사하다. 다만 본 연구는 모델의 전체 성능을 유지하면서 가능한 한 적은 비용으로 표현 효율을 향상시키는 데 보다 초점을 맞추고 있다. 

3 Methods

제안하는 방법은 세 가지 뚜렷한 단계로 구성된다. 즉, 높은 압축률의 토크나이저를 얻는 단계, 모델의 입력층과 출력층을 교체하고 초기화하는 단계, 그리고 모델을 학습하는 단계이다. 

우리 방법의 개요는 그림 1에 제시되어 있다.

3.1 Vocabulary Expansion

먼저, 높은 압축률의 토크나이저를 확보해야 한다. 이는 처음부터 새로 학습하여 얻을 수도 있고, 기존 토크나이저를 확장하여 얻을 수도 있다. 

본 연구에서는 영어와 코드에서 우수한 압축 성능을 보이는 Llama3 토크나이저를 예시로 사용한다. 

  • 여기에 중국어 어휘를 vocabulary에 통합함으로써 중국어 텍스트에 대한 압축 효율을 향상시켜, 보다 효율적인 표현을 얻는다. 
  • 먼저, 원래 토크나이저와 동일한 pre-tokenization 방식을 사용하여 중국어 학습 코퍼스를 서로 다른 조각들로 분할한다. 
  • 그 다음, 이 코퍼스에 대해 byte-level BPE 방법을 사용해 새로운 vocabulary를 학습하고, 이를 기존 vocabulary와 병합하여 새로운 vocabulary와 토크나이저를 얻는다. 
  • 우리는 중국어, 영어, 코드, 수학 코퍼스 전반에서 서로 다른 토크나이저들의 압축률을 비교했으며, 압축률은 다음과 같이 계산된다:

여기서 byte(s)는 문장을 바이트로 표현했을 때의 시퀀스 길이를 의미하고, tokenize(s)는 토크나이징 이후의 시퀀스 길이를 의미한다. 결과는 표 1에 제시되어 있으며, 새롭게 확장한 토크나이저(ReTok)가 가장 높은 전체 평균 압축률을 달성했음을 보여준다.

근데 특별한건 없는거 같은데..? 중국어만 가지고 따로 vocab 만들고 이를 합쳤다는거 느낌

3.2 Layer Replacement & Initialization

이 단계에서는 원래 모델의 입력층(즉, embedding layer)과 출력층(즉, LM head layer)을 교체하고 재초기화한다. 기존 모델의 지식을 최대한 활용하고 새로운 파라미터의 수렴을 가속하기 위해, 우리는 새로 추가된 파라미터를 기존 모델의 층을 기반으로 초기화한다.

구체적으로, 새로운 vocabulary에 포함된 각 토큰에 대해, 해당 토큰이 기존 vocabulary에도 존재하는 경우, embedding layer와 LM head layer에서의 해당 파라미터를 기존 모델의 파라미터로 초기화한다. 

  • 기존 vocabulary에 존재하지 않는 토큰의 경우에는, 원래 토크나이저를 사용하여 해당 토큰을 기존 vocabulary의 ID 시퀀스로 변환한다. 
  • 이후 해당 ID들에 대응하는 embedding layer와 LM head layer의 파라미터를 가져와 평균을 계산하고, 이 평균값을 새로운 파라미터의 초기값으로 사용한다.
  • 중국어 vocab -> 기존 토크나이저로 쪼개서 해당하는 임베딩, head 레이어 평균으로 초기화
  • 새로운 vocab중 기존 vocab과 겹치는 경우는 기존 파라미터를 그냥 사용

3.3 Model Training

그 다음으로, 우리는 모델의 재초기화된 embedding layer와 LM head layer를 학습하고, 모델의 주요 파라미터들—즉 attention layer, FFN layer, 그리고 layer norm 파라미터—은 catastrophic forgetting을 방지하기 위해 고정한다.

학습 데이터 T={t0,…,tn}T = \{t_0, \dots, t_n\}가 주어졌을 때, 우리는 표준 language modeling objective를 사용하여 새로 추가된 파라미터를 학습한다:

여기서 조건부 확률 PP는 파라미터 Θ\Theta를 갖는 신경망으로 모델링된다. 모델이 수렴한 이후에는, 모델의 주요 파라미터들을 unfreeze하고 모든 파라미터에 대해 공동 학습을 수행할 수 있으며, 이는 모델 성능을 추가적으로 향상시킬 수 있다. 

4 Experiments

4.1 Experimental Setups

Data
재초기화된 파라미터를 효과적으로 학습하기 위해서는 포괄적인 학습 데이터셋을 구성하는 것이 중요하다. 이상적으로는 사전학습 모델의 원래 데이터를 사용하는 것이 모델 성능을 복원하는 가장 좋은 방법이다. 그러나 실제로는 사전학습 데이터 전체에 접근하는 것이 불가능한 경우가 많다. 따라서 본 연구에서는 가능한 한 공개 데이터셋을 활용하여 실험을 진행하였다.

수집된 학습 데이터셋은 세 부분으로 나뉜다: 중국어, 영어, 코드.
영어 데이터는 주로 RedPajama, Pile, C4, Falcon 등의 데이터에서 수집되었으며, 중국어 데이터는 Wudao와 ChineseWebText에서 수집되었다. 코드 데이터는 StarCoder 등의 데이터에서 수집되었다.

Training Details
실험은 Qwen1.5-0.5B, Aquila2-7B, Llama3-8B 모델을 기반으로 수행되었다.
학습 시 global batch size는 384K tokens로 설정하였다.
learning rate scheduler는 cosine 스케줄을 사용했으며, 최대 learning rate는 1e-4, warm-up은 1000 step으로 설정하였다.
학습은 NVIDIA A100 GPU 32개를 사용하여 수행되었다. 

Evaluation
평가 태스크는 두 가지 범주로 나뉜다: 이해(comprehension) 태스크와 생성(generative) 태스크이다.
이해 태스크에는 PIQA (5-shot), ARC-Challenge (25-shot), HellaSwag (10-shot), MMLU (5-shot), CMMLU (5-shot)가 포함된다.
생성 태스크에는 AGIEval (0-shot), Big-Bench Hard (BBH) (3-shot), Human-Eval (0-shot), Gsm8k (5-shot)가 포함된다.

이러한 태스크들은 중국어, 영어, 코드, 수학에 대한 모델의 능력을 평가한다.
모든 태스크는 lm-evaluation-harness와 OpenCompass 프레임워크를 사용하여 평가되었으며, 모든 결과는 단일 실행(single run)에서 얻어진 것이다. 

4.2 Main Results

전체 실험 결과는 Table 2에 제시되어 있다. 

대부분의 테스트셋에서, 우리의 방법(ReTok)은 원래 모델과 유사한 성능을 달성하여 제안한 방법의 효과를 입증하였다.

그러나 일부 테스트셋에서는 성능 변화가 두드러지게 나타났는데, 예를 들어 HumanEval과 CMMLU에서 그러한 변화가 관찰되었다. 이는 우리의 학습 데이터에 코드 데이터가 많이 포함된 반면, 중국어 데이터는 상대적으로 적었기 때문일 수 있다.

이러한 현상은 사전학습 데이터의 분포가 모델 성능에 영향을 줄 수 있음을 시사하며, 향후 연구에서는 다양한 데이터가 모델 성능에 미치는 영향을 추가적으로 분석할 예정이다. 

4.3 Convergence Speed Analysis

이 절에서는 제안한 방법의 수렴 속도를 분석하였다. 우리는 학습 데이터셋에서 일부 데이터를 무작위로 선택한 뒤, 철저한 중복 제거를 수행하여 이를 validation set으로 사용하였다.

Llama3-8B 모델을 기반으로, 학습 과정에서 우리의 방법을 적용했을 때 validation loss의 변화를 측정하였다. 또한, 모델의 입력층과 출력층을 무작위로 초기화한 경우의 validation loss 변화와 비교하였다.

결과는 Figure 2에 제시되어 있다. 

결과에 따르면, 일반적으로 사용되는 random initialization 방식과 비교했을 때, 우리의 방법은 더 빠르게 수렴하며 더 나은 성능을 보인다.

4.4 Inference Speed Comparison

이 실험에서는 Llama3-8B 모델을 기반으로, 토크나이저 교체 전후의 모델 디코딩 속도 변화를 비교하였다. 우리는 제안한 방법(ReTok)이 주로 중국어의 압축률을 개선하기 때문에, 중국어 validation set에서 모델의 디코딩 속도를 측정하였다.

우리는 바이트 기반 시퀀스 길이에 따라 중국어 데이터를 그룹화하고, 각 길이 그룹에서 모델의 속도를 측정하였다. 이후 force decoding을 사용하여 해당 데이터를 처리함으로써, 서로 다른 모델 간 입력과 출력 텍스트의 길이는 동일하게 유지하되(단, 토크나이즈된 이후의 시퀀스 길이는 동일하지 않음), 비교가 가능하도록 하였다.

서로 다른 모델의 디코딩 속도를 비교한 결과는 Figure 3에 제시되어 있다. 

결과는 특히 긴 텍스트에서 제안한 방법이 디코딩 속도를 향상시킬 수 있음을 보여준다. 

5 Conclusion

본 연구에서는 LLM의 토크나이저를 교체하여 모델의 표현 및 처리 효율을 어떻게 향상시킬 수 있는지 탐구하였다. 우리는 모델의 입력층과 출력층을 초기화하기 위해 기존 파라미터를 사용하는 방법을 제안하여, 수렴 속도를 가속하였다.

여러 테스트셋에서 실험을 수행한 결과, 우리의 방법은 모델의 원래 성능을 유지하면서 디코딩 속도를 크게 향상시킬 수 있음을 보여주었다.

Limitations

본 연구에는 두 가지 주요 한계가 있다. 첫째, 우리가 사용한 데이터는 여러 출처에서 수집되었으며, 각 데이터 소스와 서로 다른 데이터 비율이 서로 다른 모델에 미치는 영향을 분석하지 않았다. 둘째, 우리의 테스트셋은 주로 중국어, 영어, 코드, 수학에 대한 모델의 능력에 초점을 맞추고 있으며, 다국어 능력과 같은 다른 능력은 평가하지 않았다. 이러한 부분은 향후 연구 과제로 남겨둔다. 

Reference

댓글

댓글 쓰기