토크나이저 학습하여 Vocabulary Size 획기적으로 줄이기 (Feat. Continued Pre-training)

최신 언어 모델들은 다양한 문자 체계를 포괄하기 위해서 Token Vocabulary의 수를 기하급수적으로 크게 만들고 있는 추세이다. 어느정도 인지도가 있는 모델들은 보통 대략 10만 개의 token을 포함하고 있다. 이처럼 어휘 사전에 완성된 형태의 단어 토큰이 많으면 많을수록, 모델은 문장을 더 짧고 효율적인 토큰 시퀸스 형태로 처리를 할 수 있게 된다. 그런데, 모델은 다음 단어를 예측할 때마다 사용하지 않을 모든 단어의 확률 점수를 일일이 다 계산하게 된다. 이에 따른 메모리와 컴퓨팅 자원을 소모하는 것은 굉장히 비효율적이다. 이에 따라 본 게시물에서는 Gemma 3 270M 모델을 활용하여 어휘 크기를 대략 48% 절감하는 방법에 대해 다루고자 한다.

---

 

Gemma 3 270M 모델 간단하게 알아보기 

Gemma 3:270M의 성능 (이미지 출처: Google Deepmind)

 

본 모델은 18개의 트랜스포머 블록으로 구성된 Decoder 기반의 Causal Lanuauge Model이며, 긴 문맥을 효율적으로 처리하도록 설계가 되었다. 이를 위해 대부분의 레이어에서는 계산량을 극소화 하기 위해 Sliding Window Attnetion을 적용하였으며, 주기적으로 6/12/18번째 레이어에서만 Full-Attention을 적용하였다. 이러한 방식을 통해 전체적인 계산 비용을 낮추면서도 문장 전체의 맥략을 고려할 수 있게 설계하였다. 

 

Gemma 3:270M 모델은 전체 파라미터 중 대략 1억 7천만 개를 단어 표현을 위한 임베딩에 사용이되며, 나머지 1억 개는 실제 연산을 수행하는 트랜스포머에 사용이 된다. 여기에서 가장 눈에 띄는 것은 262,144개에 달하는 사전 어휘 크기인데, 이는 모델의 크기를 고려하면 굉장히 이례적인 수준이다. 그렇다면, 왜 굳이 사용하지도 않을 토큰에 일부 컴퓨팅 자원과 메모리를 소비해야 하는 것일까? 라는 것에 대해 알아보고자 한다. 

---

언어 모델에 따른 새로운 Tokenizer 학습하기 

언어 모델의 거대한 어휘 사전을 줄이기 위해 가장 원초적인 방식은 다음과 같다고 생각한다. 단순하게 "안 쓰는 토큰들을 그냥 제외시키면 되지 않을까?" 이다. 예를 들어, 한국어 기반의 언어 모델을 사용하고 있다면 영어, 아랍어 필요하지 않는 언어들의 토큰들을 어휘에서 지우는 것이다. 이러한 방법은 기존에 모델이 학습해 둔 단어의 의미 즉 임베딩을 그대로 쓸 수 있어 굉장히 합리적으로 보인다. 

 

하지만, 이러한 방법은 모델을 망가뜨릴 수 있는 위험한 접근 방식이다. 이는 마치 젠가 블록을 빼는 것과 굉장히 유사하다. 한 두개 정도 뺄때는 굉장히 전체적인 구조가 안정하지만, 빼면 뺄수록 전체 구조가 불안정해지다가 무너지는 상황이 발생하게 된다. 이를 조금 더 직관적으로 보면 다음과 같다: 

• 단어를 쪼개는 규칙의 불안정성: 최신 토크나이저는 단순한 "단어 목록"이 아니라, 단어를 가장 효율적으로 쪼개기 위해 통계적인 규칙에 의해 생성된 집합체이다. 만약 이렇게 생성된 집합체 내에서 임의로 토큰을 제외시키게 된다면, 집합체를 이루는 통계적인 규칙이 심하게 변형이 된다. 예를 들어, "Replaying" 이라는 단어는 "Re" + "Play" + "ing"에 의해 다시 무엇을 한다는 의미를 가지게 된다. 하지만 여기에서 "Re" 토큰을 제외 시키게 된다면, LLM은 "R" + "E" + "Play" + "ing" 라는 의미를 정확하게 파악할 수 없는 문제점이 생기게 된다. 
• 할당되는 ID의 고유한 의미 값(임베딩): 일반적으로 토큰은 1, 2, 3, 4 와 같이 연속된 번호 (ID)를 통해 관리하게 된다. 일종의 "주소록"이라고 생각하면 될거 같다.  예를 들어, 서울이라는 토큰은 8000번, 부산은 8001번에 할당이 되었다고 가정을 해보자. "대한민국의 수도는?" 라는 시퀸스의 다음 단어를 예측하기 위해, 모델은 가지고 있는 전체 어휘 사전의 모든 토큰에 대해 정답일 확률을 계산하게 된다. 하지만, 만약에 7000번 부터 7999번째의 토큰이 사라졌다면,  8000번 ID를 다시 서울이라는 글자로 변환하는 Detokenization 과정에서 엉뚱한 정보를 참고하게 되는 문제점이 발생하게 된다. 

기존 어휘 사전에서 특정 토큰을 단순히 제외하는 방식은 불안정하고 특정 문제를 일으킬 가능성이 매우 높다. 이에 대한 효과적인 대안은, 처음부터 어휘 크기를 명확히 설정하고 특정 도메인이나 언어에 완벽히 최적화된 토크나이저를 새로 학습하는 것이다. 기존 모델의 지식과 겹치는 공통 토큰의 임베딩은 그대로 복사하여 재사용하고, 해당 도메인에만 존재하는 새로운 토큰에 대해서만 임베딩을 초기화하는 방식으로 진행하는 것이다. 

---

Remapping Token Embedding

새로운 토크나이저를 만든 뒤, 모델이 단어의 의미를 이해하는 임베딩 행렬 (어휘 사전)을 만드는 과정에 대해 알아보고자 한다. 조금 더 쉽게 이야기를 해보자면, 기존의 거대한 백과사전에서 필요한 부분만 가져와 나만의 전공 노트를 만드는 과정이라고 생각하면 될거 같다. 이에 대한 과정은 다음과 같다: 

• 불필요한 임베딩: 새로운 어휘 사전에 포함되지 않을 단어들의 임베딩을 삭제한다. 
• 기존 임베딩 재배치: 살아남은 토큰들의 ID는 자연스럽게 바뀌게 된다. 예를 들어, 기존 사전에서 8000번이었던 '서울' 토큰이, 더 작아진 새로운 어휘 사전에서는 150번이라는 새로운 ID를 부여받을 수 있다. 따라서 모델이 '서울'의 의미를 잃지 않도록, 기존 8000번 주소에 저장되어 있던 의미 값(임베딩)을 새로운 주소인 150번으로 정확히 복사해 옮겨주는 작업이 필수적이다. 
• 임베딩 초기화 (중요):  새로운 단어들의 임베딩을 만드는 것이 가장 중요하다. 새로운 토큰은 기존에 있던 더 작은 토큰 조각들의 조합인 경우가 매우 높다. 예를 들어, "Tokenization"이라는 새 토큰은 기존에 있던 "token"과 "ization" 이라는 토큰들로 이루어져 있을 수 있다. 이때, tokenization의 의미 값은 token의 의미 값과 ization의 의미 값을 가져와 평균을 내서 만든다. 이는 무작위로 생성된 것이 아니라, 어느정도 "Tokenization"에 대한 암묵적인 정보를 포함하고 있는 셈이다. 여기에서 새로운 토큰의 의미 값을 기존 구성 요소들의 평균으로 만드는 과정은 BPE와 같은 subword 토크나이저가 만들어낸 결과물을 활용하는 것과 비슷하다. 

def initialize_embedding_from_components(
    new_token: str,
    old_tokenizer,
    old_embeddings: torch.Tensor,
    hidden_size: int
) -> torch.Tensor:

    # 1. 입력 토큰을 기존 토크나이저가 이해할 수 있도록 정규화한다.
    text_to_tokenize = new_token
    if new_token.startswith("Ġ") or new_token.startswith(" "):
        text_to_tokenize = " " + new_token[1:]

    # 2. 기존 토크나이저를 사용해 입력 토큰을 분해한다.
    old_component_ids = old_tokenizer(text_to_tokenize, add_special_tokens=False)["input_ids"]

    # 3. 만약 분해할 수 없다면, 무작위 벡터를 생성한다
    if not old_component_ids:
        return torch.empty((hidden_size,), dtype=torch.float32).normal_(mean=0.0, std=0.02)

    # 4. 분해된 토큰들의 임베딩을 가중 평균하여 새로운 임베딩을 계산한다.
    weighted_sum_vector = torch.zeros((hidden_size,), dtype=torch.float32)
    total_weight = 0.0
    
    old_vocab_size = old_embeddings.shape[0]

    for old_id in old_component_ids:
        # 유효한 ID인지 확인
        if 0 <= old_id < old_vocab_size:
            component_token = old_tokenizer.convert_ids_to_tokens(old_id)
            
            # 가중치는 토큰 조각의 길이로 설정 (더 긴 조각이 더 중요하다고 가정)
            weight = max(1, len(component_token))
            
            # (임베딩 벡터 * 가중치)를 누적해서 더함
            weighted_sum_vector += old_embeddings[old_id] * weight
            total_weight += weight

    # 5. 유효한 조각이 하나도 없었다면, 역시 무작위 벡터를 반환한다 (Fallback)
    if total_weight == 0:
        return torch.empty((hidden_size,), dtype=torch.float32).normal_(mean=0.0, std=0.02)

    # 최종적으로 계산된 가중 평균 벡터를 반환한다.
    return weighted_sum_vector / total_weight

 

위 코드는 새로운 토큰의 초기 임베딩을 생성하는 함수이다. 먼저, 함수가 실행이 된다면, 입력된 새로운 토큰을 기존 토크나이저가 이해할 수 있도록 공백 문자 등을 정규화한 뒤, 이를 기존 토크나이저에 통과시켜 더 작은 단위의 토큰들로 분해하게 된다. 그런 다음에는 분해된 각각의 토큰들의 기존 임베딩을 가져와 문자열 길이를 기반으로 가중치를 적용한 가중 평균을 계산하게 된다. 이러한 방식은 더 긴 토큰 조각이 최종적인 의미에 더욱 더 큰 의미를 가지게 하기 위함이다. 그리고 부가적으로 새로운 토큰이 분해가 되지 않는다면, 무작위로 벡터를 생성하여 반환하는 fallback 기능도 포함되어 있다.

 

이에 대해 조금 더 구체적으로 예시를 통해 알아보고자 한다. 만약, 기존 토크나이저에 다음과 같은 토큰과 임베딩 벡터가 다음과 같다. 

• "AUTO" : [0.8, 0.1, 0.0, 0.1] 
• "MATIC" : [0.1, 0.7, 0.1, 0.1] 
• "ALLY" : [0.0, 0.1, 0.8, 0.1] 

이러한 토큰들을 가지고 기존 어휘에 존재하지 않는 "AUTOMATICALLY" 토큰에 대한 새로운 임베딩을 생성 해야하며, 이는 다음과 같이 수행될 수 있다. "< | AUTOMATICALLY | >"를 기존 토크나이저에 통과시키게 되면, ["AUTO", "MATIC", "ALLY"]와 같이 분해될 수 있다. 분해된 각 토큰의 문자열 길이를 기반으로 가중치를 다음과 같이 계산하고, 이에 대한 임베딩 벡터를 생성한다.

• "AUTO": [0.8, 0.1, 0.0, 0.1] * 4 = [3.2, 0.4, 0.0, 0.4]
• "MATIC": [0.1, 0.7, 0.1, 0.1] * 5 = [0.5, 3.5, 0.5, 0.5]
• "ALLY": [0.0, 0.1, 0.8, 0.1] * 4 = [0.0, 0.4, 3.2, 0.4]
•  최종 합산 벡터: [3.7, 4.3, 3.7, 1.3] / (4 + 5 + 4) = [0.28, 0.33, 0.28, 0.1]

이와 같은 방법을 통해 "<|Automatically|>" 라는 새로운 토큰에 대한 임베딩 값을 파악할 수 있다. 여기에서 토큰의 문자열 길이를 가중치로 사용하는 이뉴는 굉장히 단순하다. 일반적으로 하나의 단어가 여러 개의 하위 토큰으로 분리될 때, 문자열이 긴 토큰이 그 단어의 핵심적인 의미를 내포하고 있을 가능성이 높기 때문이다. 

---

새롭게 정의된 토크나이저... 그대로 사용해도 될까? 

여태까지 알아본 내용은 새로운 토크나이저와 이에 상응되는 임베딩을 교체하였다. 이때 "토크나이저만 바꿨을 뿐이니, 기존 모델을 그대로 사용해도 되지 않을까? 라는 생각을 할 수 있게 된다. 실제로 모델을 실행해 보면, 겉보기에는 그럴듯한 답변을 생성할 수도 있다. 하지만, 모델은 내부적으로 매우 불안정하다. 이는 새로운 토큰을 도입하고 임베딩 인덱스를 변경하는 과정에서 모델이 원래 학습했던 "토큰화 분포"를 완전히 바꿔버리기 때문이다. 이에 대한 간단한 예시는 다음과 같다: 

• 기존 모델: "A, B, C" 라는 토큰 조합을 기반으로 내부 파라미터가 최적화되어 있다. (예시: "AUTO" , "MATIC" , "ALLY")
• 새로운 토크나이저: 이제는 "X", Y"와 같이 완전히 다른 방식의 토큰 조합을 입력으로 받게 된다. (예시: "AUTOMATICALLY") 

따라서, 모델이 새로운 토큰 분포에 맞춰 안정적으로 실행하기 위해서는, 파라미터를 보정하는 과정, 즉 재학습이 반드시 필요하다. 하지만 여기에서 고려해야할 점이 있다. 현재 풀어내고자 하는 데이터의 특성에 따라, 어떻게 토큰 분포를 학습할지 나뉘게 된다. 

1.  지도 학습 데이터가 충분히 많은 경우 

그냥 단순하게 특정 모델에 맞게 파인튜닝을 수행하면 된다. 모델이 새로운 토크나이저의 불안정한 임베딩 값으로 특정 테스크를 수행하게 되면, 초기에는 엉뚱한 값을 도출하게 되며 이에 따른 Loss가 발생한다. 하지만, [Input]과 [Label]을 가진 지도학습 기반의 데이터라면, 역전파를 통해 해당 loss를 수정해 나아가는 방향으로 진행이 된다. 새롭게 정의된 토크나이저에서 생성되는 불안정한 임베딩 벡터라도 데이터를 통해 의미 있는 결과를 도출할 수 있도록 빠르게 보정하게 된다. 이러한 과정을 통해 모델은 새로운 토큰 분포에 적응하는 동시에 두 마리 토끼를 잡을 수 있게 된다. 

2.  지도 학습 데이터가 상대적으로 적은 경우 

대부분의 도메인에서 파인튜닝을 진행할 때 가장 큰 문제점은 [INPUT] 및 [LABEL] 쌍으로 이루어진 지도 학습 데이터 셋을 대규모로 학보하기 어렵다.  예를 들어, 특정 작업을 수행하기 위한 데이터가 1,000개 정도 존재할 수는 있지만, 그보다 훨씬 많은 데이터를 확보하는 것은 현실적으로 불가능에 가깝다. 일례로 생성형 AI 모델을 사용하여 데이터를 만드는 경우가 많아지고 있는 추세이지만, 과연 도메인 지식이 필요한 영역에서는 고품질의 데이터를 생성하는 것은 거의 불가능하다. 

 

이에 따라, 새로운 토크나이저를 정의하고 토큰 분포의 realignment가 필요한 경우, 적은 양의 데이터만으로는 한계점이 명확하게 존재한다. 지도 학습 데이터 셋이 1,000개 정도로 부족한 상황에서 곧바로 파인튜닝을 수행하게 된다면, 과적합 상태에 빠지기 굉장히 쉬워진다. 

따라서, 본 게시물에서의 가장 핵심이 되는 "지속적인 학습 (Continued Pre-training)" 관점을 소개하고, 이에 대한 문제점을 해결해보고자 한다. 

 

여기에서 Continued Pre-training 이라는 용어는 "Don’t Stop Pretraining: Adapt Language Models to Domains and Tasks" 논문에서 소개되었다. 본 논문에서 제안하는 DAPT (Domain-Adaptive Pretraing)과 TAPT (Task-Adaptive Pretraining) 방식을 제안하였다. 이는 단순하게 해당 도메인에서만 사용되는 전문적인 용어, 고유한 문체, 배경 지식 등을 모델이 학습할 뿐만 아니라, 자주 사용되는 특정한 표현 방식이나 데이터 분포에 모델을 더욱 밀접하게 맞추는 것을 의미한다. 이를 통해 DAPT/TAPT 이라는 방법을 통해 모델이 새로운 환경에 충분히 적응할 수 있게 한다면, 파인튜닝 단계에서 훨씬 더 높은 성능을 이루어낼 수 있다는 것을 증명하였다. 

 

따라서, Continued Pre-training은 두 단계로 명확하게 분리될 수 있다. 첫번 째는 임베딩 모델을 재정렬하는 것을 목표로하며, 이는 Chat Template 없이 레이블이 없는 데이터를 사용해 다음 토큰을 예측할 수 있는 과정을 수행하게 된다. 그리고 두 번째 단계는, 이렇게 CPT를 통해 새로운 토큰 분포에 정렬된 모델을 기반으로 파인튜닝을 수행한다. 

---

그렇다면, 이걸 어떻게 해야 하는걸까?  

1. 선(先) 안정화

위에서 다루었듯이, 본 단계는 단순하게 네트워크 재정렬 및 안정화이다. 이는 "Don't Stop Pretraining" 논문의 DAPT/TAPT 개념처럼 레이블 없는 대규모 텍스트로 다음 토큰 예측 학습을 수행한다. 하지만, 여기에서 [Input, Label] 구조를 완전히 무시하고 Chat Template를 제거한 채, (Input + Label)을 하나의 긴 텍스트 덩어리로 만들어 학습시킨다. 이러한 과정을 통해 모델은 "SPECIFIC TASK"를 학습하는 것이 아니라, 새로운 토크나이저의 통계적 분포에 스스로 맞추어 임베딩과 파라미터를 안정화를 거치게 되는 것이다. 

 

 

import torch
from unsloth import FastLanguageModel
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig
from transformers import AutoTokenizer
from unsloth import UnslothTrainer, UnslothTrainingArguments


# CPT용 데이터 전처리기 (챗 템플릿 없음)
def process_cpt(row, tokenizer, ~~~, ~~~):
    """CPT 모드: 챗 템플릿 없이, 텍스트를 단순 결합"""
    source = row['~'][~~~]
    target = row['~'][~~~]
    
    #'다음 토큰 예측'을 위해 레이블 없는 텍스트로 만듭니다.
    row["text"] = source + " " + target + tokenizer.eos_token
    return row

def run_cpt_stage_1(input_model_path, output_dir, pair="en-fr"):
    
    # 불안정한 모델의 토크나이저 로드
    model, tokenizer = FastLanguageModel.from_pretrained(
       model_name = input_model_path, 
       fix_tokenizer=False, # [중요] 64k 토크나이저 사용
       max_seq_length = 4096,
       full_finetuning=True,
       # ... (기타 unsloth 설정) ...
    )
    
    # CPT용 데이터 로드 및 전처리 (사용하는 데이터 셋에 따라 적용) 


    # CPT 트레이너 설정
    training_args = UnslothTrainingArguments(
          output_dir=output_dir,
          dataset_text_field='text',
          optim="adamw_8bit",
          per_device_train_batch_size=16,
          gradient_accumulation_steps=4,
          learning_rate = 5e-5, # 
          num_train_epochs=1,
          ... 
    )
    
    trainer = UnslothTrainer(
       model = model,
       train_dataset=ds_train,
       processing_class=tokenizer,
       args = training_args
    )
    
    # CPT 실행
    trainer.train()
    
    # 안정화된 모델 저장
    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    print(f"--- 1단계 CPT 완료 --- (출력: {output_dir})")

 

2. 후(後) 작업 

모델이 새로운 토큰 분포에 완벽히 적응하였으므로, 소량의 지도 학습 데이터를 활용할 차례이다. 첫번 째 단계에서 CPT로 안정화 시킨 모델을 입력으로 사용하며, 이번에는 Chat Template을 사용하여 특정 작업을 명시적으로 학습시킨다. 모델은 더 이상 "토큰 분포"에 리소스를 낭비하지 않고 오직 특정 작업 자체에만 집중할 수 있게 된다. 이는 적은 양의 데이터만으로도 과적합 위험을 최소화하고 높은 성능의 최종 목적 모델을 완성할 수 있게 된다. 

 

# SFT용 데이터 전처리기 (챗 템플릿 있음)
def process_sft(row, tokenizer,~~~ ~~~):
    """SFT 모드: 챗 템플릿을 사용하여 특정 작업 학습"""
    source = row['~'][~~~]
    target = row['~'][~~~]

    # [핵심] 사용자님의 원본 템플릿을 그대로 사용합니다.
    row["text"] = "TASK SPECIFICATION"
    return row

# 사용자님의 원본 FT 함수와 동일한 로직
def run_sft_stage_2(input_model_path, output_dir, pair="en-fr"):

    # CPT로 '안정화'된 모델 로드
    model, tokenizer = FastLanguageModel.from_pretrained(
       model_name = input_model_path, # [중요] 1단계의 CPT 완료 모델 경로
       fix_tokenizer=False,
       max_seq_length = 4096,
       full_finetuning=True,
       ... 
    )

    # 데이터 셋에 맞게 설정 
    
    # SFT 트레이너 설정
    training_args = UnslothTrainingArguments(
          output_dir=output_dir,
          dataset_text_field='text',
          ... 
    )
    
    trainer = UnslothTrainer(
       model = model,
       train_dataset=ds_train,
       # eval_dataset = ds_test,
       processing_class=tokenizer,
       args = training_args
    )

    # SFT 실행
    trainer.train()

 

여태까지의 내용을 요약해보자면, 데이터가 부족한 현실적인 도메인 환경에서 모델을 최적화할 때, 이와 같은 방법을 통해 적은 자원으로 최대의 성능을 이끌어내는 가장 합리적인 방법이지 않을까 라는 생각이 든다.