[Hands-On] Unigram을 이용한 토크나이저 구현

(You can find the English version of the post at this link.)

이전 글에서 우리는 자연어 처리에서의 토크나이제이션 개념과 구체적으로 BPE, Byte-level BPE, WordPiece 및 Unigram 토크나이제이션 방법에 대해 살펴보았습니다. 이 글은 주요 토크나이제이션을 직접 구현해보는 시리즈의 세번째 글에 해당합니다.

Unigram Tokenization (Image by the author using ChatGPT)

(이 튜토리얼은 일부 핵심 기능들은 BPE tokenization 을 잘 설명해 놓은 Huggingface 글에서 일부 참조했음을 밝힙니다.)

Tokenization 방법론들에 대한 쉽고 직관적인 이해

이번 글에서는 Unigram 토크나이제이션 방법에 초점을 맞추어 기본 개념과 구현 방법에 대해 알아보겠습니다.

목표

• Unigram 토크나이제이션의 기본 개념 이해
• 텍스트 데이터 토크나이제이션부터 작은 어휘 구축까지의 전체 워크플로우 구현
• Unigram 토크나이제이션의 학습 알고리즘과 단어 분할 방법 소개
• 예제 텍스트 데이터에 대한 Unigram 토크나이제이션 수행

Unigram 토크나이제이션이란?

Unigram 토크나이제이션은 AlBERT, T5, mBART, Big Bird, XLNet 등의 모델에서 자주 사용되는 하위 단어 토크나이제이션 방법입니다. BPE나 WordPiece가 작은 어휘로 시작하여 토큰을 추가하는 것과 달리, Unigram은 큰 어휘로 시작하여 불필요한 토큰을 제거합니다.

Unigram과 BPE/WordPiece의 주요 차이점:

1. 접근 방식:
   - Unigram: “감산적” — 큰 어휘로 시작하여 불필요한 토큰 제거
   - BPE/WordPiece: “가산적” — 작은 어휘로 시작하여 필요한 토큰 추가
2. 토크나이제이션 결정:
   - Unigram: 확률사용 — 여러 가능성 중 최적의 토크나이제이션 선택
   - BPE: 결정론적 — 고정된 규칙에 따라 항상 동일한 방식으로 토크나이제이션
   - WordPiece: 준확률적 — 점수 매기기 메커니즘을 사용하지만 Unigram보다는 결정론적
3. 최적화 기준:
   - Unigram: 전반적인 언어 모델 성능 향상을 목표
   - BPE/WordPiece: 주로 빈도나 간단한 통계에 기반

이제 Unigram 토크나이제이션의 구현에 대해 자세히 알아보겠습니다. Python과 Hugging Face의 transformers 라이브러리를 사용할 것입니다.

환경 설정

먼저 필요한 라이브러리를 가져옵니다:

from transformers import AutoTokenizer
from collections import defaultdict
import math

# Load the BERT tokenizer for pre-tokenization
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

단어, 특수기호 분리 등을 위해 XLNet 토크나이저를 사용하겠습니다.

샘플 데이터 준비

Unigram 토크나이제이션 과정을 보여주기 위해 간단한 샘플 문장을 사용하겠습니다:

from transformers import AutoTokenizer
from collections import defaultdict

# Load the sample sentence
sample = "low low low low low lower lower newest newest newest newest newest newest widest widest widest"

tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")

초기 토크나이제이션 및 빈도 계산

샘플을 토크나이즈하고 각 단어의 빈도를 계산해 봅시다:

# Print frequency of sample words
word_freqs = defaultdict(int)

words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(sample)
new_words = [word for word, offset in words_with_offsets]
for word in new_words:
    word_freqs[word] += 1

word_freqs.items()

dict_items([('▁low', 5), ('▁lower', 2), ('▁newest', 6), ('▁widest', 3)])

Unigram 모델 학습

초기 어휘 생성

먼저 모든 가능한 부분 문자열의 빈도를 계산합니다:

char_freqs = defaultdict(int)
subwords_freqs = defaultdict(int)
for word, freq in word_freqs.items():
    for i in range(len(word)):
        char_freqs[word[i]] += freq
        # Loop through the subwords of length at least 2
        for j in range(i + 2, len(word) + 1):
            subwords_freqs[word[i:j]] += freq

# Vocabulary for all Chartacter
print(f"Chartacter: {char_freqs.items()}")

# Vocabulary for all possible cases
print(f"Subwords: {subwords_freqs.items()}")

Chartacter: dict_items([('▁', 16), ('l', 7), ('o', 7), ('w', 16), ('e', 17), ('r', 2), ('n', 6), ('s', 9), ('t', 9), ('i', 3), ('d', 3)])
Subwords: dict_items([('▁l', 7), ('▁lo', 7), ('▁low', 7), ('lo', 7), ('low', 7), ('ow', 7), ('▁lowe', 2), ('▁lower', 2), ('lowe', 2), ('lower', 2), ('owe', 2), ('ower', 2), ('we', 8), ('wer', 2), ('er', 2), ('▁n', 6), ('▁ne', 6), ('▁new', 6), ('▁newe', 6), ('▁newes', 6), ('▁newest', 6), ('ne', 6), ('new', 6), ('newe', 6), ('newes', 6), ('newest', 6), ('ew', 6), ('ewe', 6), ('ewes', 6), ('ewest', 6), ('wes', 6), ('west', 6), ('es', 9), ('est', 9), ('st', 9), ('▁w', 3), ('▁wi', 3), ('▁wid', 3), ('▁wide', 3), ('▁wides', 3), ('▁widest', 3), ('wi', 3), ('wid', 3), ('wide', 3), ('wides', 3), ('widest', 3), ('id', 3), ('ide', 3), ('ides', 3), ('idest', 3), ('de', 3), ('des', 3), ('dest', 3)])

문자부터 시작해서, 각 단어안에 존재할 수 있는 subword 조합들을 모두 조사해서 우선 기록해 둡니다.

subword에는 어떤것들이 가장 많이 나타나는지 조사해보죠.

# Sort subwords by frequency
sorted_subwords = sorted(subwords_freqs.items(), key=lambda x: x[1], reverse=True)
sorted_subwords[:10]

[('es', 9),
 ('est', 9),
 ('st', 9),
 ('we', 8),
 ('▁l', 7),
 ('▁lo', 7),
 ('▁low', 7),
 ('lo', 7),
 ('low', 7),
 ('ow', 7)]

이것을 character 와 합쳐서 다시 전체적이 경향을 살펴보죠.

token_freqs = list(char_freqs.items()) + sorted_subwords[: 300 - len(char_freqs)]
token_freqs = {token: freq for token, freq in token_freqs}

sorted_token_freqs = sorted(token_freqs.items(), key=lambda x: x[1], reverse=True)
sorted_token_freqs

[('e', 17),
 ('▁', 16),
 ('w', 16),
 ('s', 9),
 ('t', 9),
 ('es', 9),
 ('est', 9),
 ('st', 9),
 ('we', 8),
 ('l', 7),
 ('o', 7),
 ('▁l', 7),
 ('▁lo', 7),
 ('▁low', 7),
 ('lo', 7),
 ('low', 7),
 ....
 .... 
 ('de', 3),
 ('des', 3),
 ('dest', 3),
 ('r', 2),
 ('▁lowe', 2),
 ('▁lower', 2),
 ('lowe', 2),
 ('lower', 2),
 ('owe', 2),
 ('ower', 2),
 ('wer', 2),
 ('er', 2)]

어떤 문자는 subword 이 조합보다 많이 나타나기도 하고, 반대로 어떤 문자는 subword 의 조합보다 훨씬 적게 나타나기도 합니다. 이 커다란 어휘셋에서 어떠한 기준을 가지고 불필요한 어휘를 제거해 나가려면 분명한 기준이 필요합니다.

Unigram 기반 토크나이제이션에서는 아래와 같은 기준을 이용해서 이를 수행합니다.

• 기준 : 전체 코퍼스에 대한 손실(loss)의 증가가 가장 작은 토큰부터 제거

이 기준점에 해당하는 손실값을 알아보겠습니다.

손실계산 (Loss Calculation)

모델의 손실을 계산하는 방법에 대해 알아보겠습니다. 이 과정은 Unigram 알고리즘의 핵심이라고 할 수 있죠.

먼저, 모든 빈도의 합을 계산하여 빈도를 확률로 변환합니다. 그런데 여기서 주의할 점이 있습니다. 우리 모델에서는 확률의 로그값을 저장할 겁니다. 왜 그럴까요? 작은 숫자들을 곱하는 것보다 로그값을 더하는 것이 수치적으로 더 안정적이기 때문입니다. 이렇게 하면 모델의 손실 계산도 더 간단해집니다.

손실 함수는 다음과 같이 정의됩니다:

여기서:

• i는 각 토큰을 나타내는 인덱스입니다.
• n은 전체 토큰의 수입니다.
• f_i는 i번째 토큰의 빈도입니다.
• p_i는 i번째 토큰의 확률입니다.

이 공식은 각 토큰의 빈도와 해당 토큰의 음의 로그 확률을 곱한 값들의 총합을 계산합니다. 이렇게 하면 전체 코퍼스에 대한 모델의 성능을 하나의 숫자로 나타낼 수 있죠.

이 손실 함수를 계산하기 위해, 우리는 세 가지 중요한 구성 요소를 구현해야 합니다:

1. model: 이 딕셔너리는 어휘의 각 토큰에 대한 음의 로그 확률을 저장합니다. 키는 토큰이고, 값은 해당 토큰의 음의 로그 확률입니다. 이렇게 표현하면 알고리즘의 후속 단계에서 계산을 더 효율적이고 안정적으로 수행할 수 있습니다.
2. encode_word: 이 함수는 단어와 현재 모델을 입력으로 받습니다. 동적 프로그래밍을 사용하여 현재 어휘에서 가장 최적의 토큰 분할을 찾아냅니다. 이 과정에서 전체 손실을 최소화하는 것이 목표입니다. 함수는 최적의 토크나이제이션 결과와 그에 따른 손실을 반환합니다.
3. compute_loss: 이 함수는 현재 모델을 사용하여 전체 코퍼스에 대한 총 손실을 계산합니다. 코퍼스의 각 단어를 순회하면서 encode_word 함수를 사용해 인코딩하고, 그 손실들을 모두 더합니다. 이 결과는 현재 모델(어휘)이 전체 코퍼스를 얼마나 잘 표현하는지를 나타냅니다.

Model 준비

빈도를 기반으로 초기 모델을 생성합니다:

from math import log

total_sum = sum([freq for token, freq in token_freqs.items()])
model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

sorted_model = dict(sorted(model.items(), key=lambda item: item[1], reverse=True))
print(len(sorted_model))
print("-"*50)
sorted_model

# The 'model' dictionary in this code stores the negative log probabilities of each token.

64
--------------------------------------------------
{'r': 5.147494476813453,
 '▁lowe': 5.147494476813453,
 '▁lower': 5.147494476813453,
 'lowe': 5.147494476813453,
 'lower': 5.147494476813453,
 'owe': 5.147494476813453,
 ...
 ...
 'es': 3.643417080037179,
 'est': 3.643417080037179,
 'st': 3.643417080037179,
 '▁': 3.068052935133617,
 'w': 3.068052935133617,
 'e': 3.0074283133171824}

encode_word 함수 구현

단어를 인코딩하고 모델의 손실을 계산하는 함수를 정의합니다:

def encode_word(word, model):
    best_segmentations = [{"start": 0, "loss": 0}] + [
        {"start": None, "loss": None} for _ in range(len(word))
    ]
    for start_idx in range(len(word)):
        # This should be properly filled by the previous steps of the loop
        best_loss_at_start = best_segmentations[start_idx]["loss"]
        for end_idx in range(start_idx + 1, len(word) + 1):
            token = word[start_idx:end_idx]
            if token in model and best_loss_at_start is not None:
                loss = model[token] + best_loss_at_start
                # If we have found a better segmentation ending at end_idx, we update
                if (
                    best_segmentations[end_idx]["loss"] is None
                    or best_segmentations[end_idx]["loss"] > loss
                ):
                    best_segmentations[end_idx] = {"start": start_idx, "loss": loss}

    segmentation = best_segmentations[-1]
    if segmentation["loss"] is None:
        # We did not find a tokenization of the word -> unknown
        return ["<unk>"], None

    loss = segmentation["loss"]
    start = segmentation["start"]
    end = len(word)
    tokens = []
    while start != 0:
        tokens.insert(0, word[start:end])
        next_start = best_segmentations[start]["start"]
        end = start
        start = next_start
    tokens.insert(0, word[start:end])
    return tokens, loss

encode_word 함수는 Unigram 토크나이제이션 구현의 핵심 요소입니다. 이 함수는 동적 프로그래밍을 사용하여 주어진 단어의 최적 분할을 찾습니다. 작동 방식은 다음과 같습니다:

1. 단어의 각 위치에 대한 최적 분할 정보를 저장할 리스트를 초기화합니다.
2. 단어의 각 문자를 순회하며, 해당 문자로 시작하는 모든 가능한 부분 문자열을 고려합니다.
3. 모델에 존재하는 각 유효한 토큰(부분 문자열)에 대해 손실을 계산하고, 더 나은 분할이 발견되면 최적 분할 정보를 업데이트합니다.
4. 모든 가능성을 고려한 후, 단어의 끝에서부터 역추적하여 최적의 토크나이제이션을 재구성합니다.
5. 유효한 토크나이제이션이 없으면 <unk> (알 수 없음) 토큰을 반환합니다.

이 방식을 통해 encode_word 함수는 주어진 단어를 현재 어휘를 사용하여 가장 효율적으로 분할할 수 있는 방법을 찾아냅니다.

손실함수 계산

# Calculates the total loss for the entire corpus using the current model.
def compute_loss(model):
    loss = 0
    for word, freq in word_freqs.items():
        _, word_loss = encode_word(word, model)
        loss += freq * word_loss
    return loss

compute_loss(model)

68.28802773020526

코퍼스에 대한 모델의 손실을 계산하는 과정은 다음과 같습니다:

1. 손실 변수를 0으로 초기화합니다.
2. 코퍼스의 각 단어에 대해:
   - encode_word 함수로 최적의 토크나이제이션과 손실을 찾습니다.
   - 이 손실에 단어의 빈도를 곱하고 총 손실에 더합니다.
3. 누적된 총 손실을 반환합니다.

이 과정을 통해 현재 모델이 전체 코퍼스를 얼마나 잘 표현하는지 수치화할 수 있습니다. 위의 경우에는 68.28 정도가 계산되었네요.

그렇다면 정확히 이 수치의 의미가 무엇일까요?

68.28은 현재 모델의 전체 코퍼스에 대한 손실값입니다. 이 값의 의미는 다음과 같습니다:

1. 낮은 손실값은 모델이 코퍼스를 잘 표현함을 의미합니다.
2. 이 값은 각 단어의 로그 확률에 빈도를 곱한 합으로 계산됩니다.
3. 68.28이라는 절대적 수치보다는, 어휘 조정의 기준점으로서 중요합니다.
4. 어휘를 최적화하면서 이 손실값의 변화를 관찰하여 모델의 성능을 평가할 수 있습니다.

즉 이 손실값이 ‘커지게’ 만드는 어휘는 전에 어휘셋에서 삭제를 하면 안됩니다. 하지만 손실값이 ‘비슷하게 유지되거나 작아지게 하는’ 어휘라면 전체 어휘셋에서 삭제해도 좋다는 거죠.

정리하면:

1. 손실값을 ‘크게’ 증가시키는 어휘:
   - 이러한 어휘는 코퍼스를 효과적으로 표현하는 데 중요한 역할을 합니다.
   - 이들을 제거하면 모델의 성능이 크게 저하될 수 있으므로 어휘셋에 유지해야 합니다.
2. 손실값을 ‘비슷하게 유지하거나 작게’ 만드는 어휘:
   - 이러한 어휘는 코퍼스 표현에 큰 영향을 미치지 않습니다.
   - 이들을 제거해도 모델의 전반적인 성능에 미치는 영향이 작으므로, 어휘 크기를 줄이기 위해 제거할 수 있습니다.
3. 최적화 과정:
   - 각 어휘를 제거했을 때의 손실 변화를 계산합니다.
   - 손실 증가가 가장 작은 어휘부터 순차적으로 제거합니다.
   - 이 과정을 반복하여 원하는 어휘 크기에 도달할 때까지 진행합니다.

이 방식을 통해 Unigram 모델은 코퍼스를 효과적으로 표현하면서도 불필요한 어휘를 제거하여 효율적인 크기의 어휘셋을 만들 수 있습니다. 이는 모델의 성능과 계산 효율성 사이의 균형을 잡는 데 도움이 됩니다.

아래 함수를 통해 이를 계산해 낼 수 있습니다.

import copy
# Calculates the change in loss that would result from removing each token from the model.
def compute_losses(model):
    losses = {}
    model_loss = compute_loss(model)
    for token, _ in model.items():
        # We always keep tokens of length 1
        if len(token) == 1:
            continue
        model_without_token = copy.deepcopy(model)
        _ = model_without_token.pop(token)
        losses[token] = compute_loss(model_without_token) - model_loss
    return losses

compute_losses 함수의 losses 딕셔너리는 각 토큰이 모델의 전체 손실에 얼마나 기여하는지에 대한 정보를 담고 있습니다.

구체적으로:

1. 딕셔너리의 각 키는 현재 어휘의 토큰입니다.
2. 해당 값은 이 토큰을 어휘에서 제거할 경우 발생할 총 손실의 변화입니다.

이 정보를 통해 어떤 토큰이 모델 성능에 중요한지, 어떤 토큰을 제거해도 되는지 판단할 수 있습니다.

losses = compute_losses(model)

sorted_losses = dict(sorted(losses.items(), key=lambda item: item[1], reverse=True))
print(len(sorted_losses))
print("-"*50)
sorted_losses

53
--------------------------------------------------
{'▁newest': 18.408317610801703,
 '▁low': 15.340264675668081,
 '▁widest': 9.204158805400851,
 '▁lower': 6.136105870267244,
 'es': 0.0,
 'est': 0.0,
 'st': 0.0,
 'we': 0.0,
 ...
 ...
 '▁lowe': 0.0,
 'lowe': 0.0,
 'lower': 0.0,
 'owe': 0.0,
 'ower': 0.0,
 'wer': 0.0,
 'er': 0.0}

위 결과를 보면, 어느 어휘가 손실값에 크게 영향을 미치는지 알 수 있습니다. '▁newest', '▁low', '▁widest', '▁lower' 이 어휘들은 삭제하면 안됩니다. 이들을 제거하면 손실이 크게 증가하기 때문입니다.

하지만 나머지 ‘es’, ‘est’, ‘st’, ‘we’ 등과 그 아래에 나열된 많은 토큰들은 삭제의 대상이 될 수 있습니다. 이들을 제거해도 손실 변화가 0이기 때문입니다.

이제 이러한 삭제대상 어휘를 제거해보죠.

percent_to_remove = 0.1 # 10%
while len(model) > 30:
    losses = compute_losses(model)
    sorted_losses = sorted(losses.items(), key=lambda x: x[1])
    # Remove percent_to_remove tokens with the lowest losses.
    for i in range(int(len(model) * percent_to_remove)):
        _ = token_freqs.pop(sorted_losses[i][0])

    total_sum = sum([freq for token, freq in token_freqs.items()])
    model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

print(len(model))
print('-'*50)
sorted_model = dict(sorted(model.items(), key=lambda item: item[1], reverse=True))
sorted_model

30
--------------------------------------------------
{'r': 4.324132656254979,
 '▁lowe': 4.324132656254979,
 '▁lower': 4.324132656254979,
 'lowe': 4.324132656254979,
 'lower': 4.324132656254979,
 'owe': 4.324132656254979,
 'ower': 4.324132656254979,
 'wer': 4.324132656254979,
 'er': 4.324132656254979,
 'i': 3.9186675481468147,
 'd': 3.9186675481468147,
 '▁widest': 3.9186675481468147,
 'widest': 3.9186675481468147,
 'id': 3.9186675481468147,
 'ide': 3.9186675481468147,
 'ides': 3.9186675481468147,
 'idest': 3.9186675481468147,
 'de': 3.9186675481468147,
 'des': 3.9186675481468147,
 'dest': 3.9186675481468147,
 'n': 3.2255203675868693,
 '▁newest': 3.2255203675868693,
 'l': 3.071369687759611,
 'o': 3.071369687759611,
 '▁low': 3.071369687759611,
 's': 2.820055259478705,
 't': 2.820055259478705,
 '▁': 2.244691114575143,
 'w': 2.244691114575143,
 'e': 2.184066492758708}

위 코드에서 불필요한 토큰을 제거하는 과정은 다음과 같습니다:

1. 목표 어휘 크기(30)에 도달할 때까지 다음 과정을 반복합니다:
2. 각 반복에서:
   - 현재 모델의 각 토큰에 대한 손실을 계산합니다.
   - 손실이 가장 적은 순서로 토큰을 정렬합니다.
   - 현재 어휘 크기의 10%(percent_to_remove)에 해당하는, 가장 손실이 적은 토큰들을 제거합니다. (예: 100개 중 10개, 그 다음엔 90개 중 9개, 그 다음엔 81개 중 8개 등)
   - 남은 토큰들의 빈도를 다시 계산합니다.
   - 새로운 빈도를 바탕으로 모델을 업데이트합니다.
3. 이 과정을 통해 매 반복마다 어휘 크기가 점진적으로 줄어들며, 각 단계에서 새로운 상황을 반영하여 최적의 토큰을 유지합니다.
4. 최종적으로 30개의 토큰만 남을 때까지 이 과정을 반복합니다.

Unigram 토크나이저 사용

이제 학습된 Unigram 모델을 사용하여 텍스트를 토크나이즈할 수 있습니다:

간단하게 Tokenizer를 구현해보죠.

def tokenize(text, model):
    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)
    print("Pretokenization with Offset: \t", words_with_offsets)
    pre_tokenized_text = [word for word, offset in words_with_offsets]
    print("Pretokenization Tokenization : \t", pre_tokenized_text)

    # Returns the optimal subword segmentation for each word based on the current model
    encoded_words = [encode_word(word, model)[0] for word in pre_tokenized_text]
    return sum(encoded_words, [])

print("After Unigram Tokenization : \t", tokenize("lowest", model))

Pretokenization with Offset:   [('▁lowest', (0, 6))]
Pretokenization tokenization :   ['▁lowest']
After Unigram Tokenization :   ['▁lowe', 's', 't']

print("After Unigram Tokenization : \t", tokenize("The lowest , newest and widest !", model))

Pretokenization with Offset:   [('▁The', (0, 3)), ('▁lowest', (4, 10)), ('▁,', (11, 12)), ('▁newest', (13, 19)), ('▁and', (20, 23)), ('▁widest', (24, 30)), ('▁!', (31, 32))]
Pretokenization tokenization :   ['▁The', '▁lowest', '▁,', '▁newest', '▁and', '▁widest', '▁!']
After Unigram Tokenization :   ['<unk>', '▁lowe', 's', 't', '<unk>', '▁newest', '<unk>', '▁widest', '<unk>']

위 결과를 보면 몇몇 어휘는 [UNK] 로 되어 있는 것을 알 수 있습니다.

결론

이 글에서는 Unigram 토크나이제이션의 개념과 구현 방법을 살펴보았습니다. Unigram 방식은 큰 어휘에서 시작하여 불필요한 토큰을 제거하는 방식으로 작동하며, 이는 BPE나 WordPiece와는 다른 접근 방식입니다.

Unigram의 주요 특징:

1. 확률적 토크나이제이션: 여러 가능한 분할 중 가장 적합한 것을 선택합니다.
2. 언어 모델 성능 최적화: 단순한 빈도가 아닌 전체 모델 성능을 고려합니다.
3. 유연한 어휘 크기 조절: 원하는 어휘 크기에 맞춰 쉽게 조절할 수 있습니다.

이러한 특성으로 인해 Unigram은 특히 다국어 처리나 새로운 단어가 자주 등장하는 도메인에서 유용합니다.

다음 Colab 링크에서 직접 코드를 수행해 볼 수 있습니다:

[Hands-On] Build Tokenizer using Unigram