MLエンジニアへの道 #23 - トークナイゼーション & 分散表現

このブログ記事では、NLPにおけるトークナイゼーションや分散表現について紹介します。

NLP、すなわち自然言語処理は、翻訳、感情分析、質問応答など、自然言語に関するさまざまなタスクを含みます。特に、大規模言語モデル（LLM）の台頭によって、驚くべき成果と社会的影響を目にしてきた分野です。本記事では、自然言語がどのように文字列や文字の列から、アルゴリズムが処理しやすい数値に変換されるかについて説明します。

テキスト前処理

文字列を数値に変換する前に、テキストを観察し、クリーニングする必要があります。例えば、データに「1」や「①」、そして「😘」などが含まれている場合、これらをどのように処理するかを決定しなければなりません。「①」などを削除してすべて「1」に統一するか、「😘」を「（投げキッスの顔）」と変換するか、あるいはデータから削除することも考えられます。また、データ品質に関する多くの課題、例えば空白の不整合や文法の誤り、スペルミスなどがあり、これらは研究者たちが常に対処している問題です。

トークナイゼーション

前処理が完了すると、文字列を数値、すなわちトークンに変換する準備が整います。このプロセスをトークナイゼーションと呼びます。トークナイゼーションにはいくつかのアプローチがありますが、その一つが規則に基づいた方法で、空白や句読点でテキストを単語に分割し、初めて登場する単語に一意のトークンIDを割り当てます。

この規則に基づいたアプローチはシンプルで、場合によっては十分に効果的です。特に、ステミングやレンマ化と組み合わせると、単語のバリエーションを減らすことができます。ステミングは単語の末尾を機械的に削除する手法で、「ing」や「able」といった接尾辞を削ります。一方、レンマ化は意味が通るように接尾辞を削除する点で、ステミングよりも精緻です。しかし、このアプローチでは、LLMに微妙なニュアンスや文法的に正確な文章を理解させることは困難です。

このアプローチがうまくいかない理由を考えるために、「fishing」という単語を例にとりましょう。「fishing」は、ステミングやレンマ化によって「fish」に変換されますが、これは魚を捕まえる行為が「魚」という名詞そのものに変わってしまうということです。ステミングやレンマ化をしない場合、「fishing」と「fish」は全く別の単語として扱われ、両者の関連性をモデルが認識するのにより多くの労力が必要となります。

サブワードトークナイゼーション：BPE

この問題に対する直感的なアプローチは、「fish」と「ing」の両方にトークンを作成し、「fishing」の各側面を捉えることです。このアプローチはサブワードトークナイゼーションと呼ばれ、**Byte Pair Encoding（BPE）**のような統計的手法を使用して行われます。まず、文字を分割してバイトとして表現します。その後、アルゴリズムは最も頻繁に出現するバイトのペアを再帰的に探し、一つのエンティティまたはサブワードにまとめます。理想的には、データに多くの「n」と「g」ペアや「i」と「ng」ペアがあり、それによって「ing」が一つのサブワードとしてエンコードされることになります。

aabaabfc: aa appears the most, X = aa
-> XbXbfc: Xb appears the most, Y = Xb or aab
-> YYfc

テキストを再帰的に処理した後、サブワードには新しいトークンIDが割り当てられ、それを使用して文字列をトークンの列に変換します。たとえば、「fishing」がペアリングプロセスでデータに含まれていない場合でも、アルゴリズムが「fish」と「ing」を複数回見てそれらをサブワードとして認識した場合、「fishing」は「fish」と「ing」の2つのトークンで表されます。変換プロセスで未知のサブワードが存在する場合、例えば「buffling」のようなものは、アルゴリズムが「[UNK]」タグ用の特別なトークンを使用し、「[UNK]」トークンと「ing」トークンに変換されます。（サブワードトークナイゼーションにはWordPieceなど他にも多くのアプローチがあります。興味があれば調べてみてください。公開されているトークナイザーも多いので、自分のニーズに合ったものを活用できます。）

Word2Vec

これで、単語やサブワードにトークンが割り当てられましたが、単一のトークンIDには単語の意味に関する情報が含まれていません。理想的には、単語を多次元ベクトルで表現し、ベクトルの各要素が単語の特性を表し、似た文脈で登場する単語が似たベクトルで表現されることが望まれます。これらのベクトル表現を分散表現（word embeddings）と呼び、シンプルなニューラルネットワークであるWord2Vecを使って自己教師あり学習によって取得できます。

Word2Vecはトークンのワンホットエンコードベクトルを入力し、それを任意の数のニューロンを持つ全結合層に通します。（ニューロン数が分散表現の次元数になります。）Word2Vecのトレーニングには、 CBOW（Continuous Bag of Words）とSkip-gramの2つの一般的なアプローチがあります。 CBOWは文脈にある単語からマスクされた単語を予測し、Skip-gramは単語から文脈を予測します。モデルが訓練されるにつれて、各単語に対応する全結合層の重みが調整され、単語が使用される文脈を捉えられるようになり、結果的に単語の意味を表すことができるようになります。訓練後、トークンをワンホットエンコードベクトルに変換し、対応する重みのセットを使用して分散表現に変換できます。

サブサンプリングとネガティブサンプリング

しかし、Word2Vecのトレーニングは簡単ではありません。大規模なデータセットを扱う際には、数百万、数十億のトークンが存在することがあり、ワンホットエンコードベクトルには何百万、何十億もの要素が含まれ、ソフトマックス層の出力も同様に大規模になります。このようなモデルを訓練することは現実的ではないため、いくつかの技術を用いる必要があります。その一つがサブサンプリングで、頻度の高い単語（「the」や「a」など）や頻度の低い単語（「gobbledygook」など）は、情報量が少ないため文脈から無視します。もう一つの技術がネガティブサンプリングで、トレーニングステップごとに10〜20個の負例サンプルを使用し、全ての語彙を評価するのではなく、効率的にモデルを更新します。（階層的ソフトマックスもありますが、これはこの記事の範囲外です。）

数学的定式化

これまでの説明から、Word2Vecの目的関数を最終的に定式化できます。ここでは、サブサンプリングとネガティブサンプリングを使用したSkipgramモデルを考えます。

J(\theta) = -log(\sigma(u_o^T v_c)) + \sum_{j \sim P(w)} log(\sigma(u_j^T v_c)) \\ P(w) = \frac{U(w)^{3/4}}{Z}

入力単語分散表現は $v_c$ で表され、正しい文脈単語の分散表現 $u_o$ と、ネガティブサンプルとしての文脈単語の分散表現 $u_j$ が使われます。シグモイド関数 $\sigma$ を用いて確率をシミュレーションし、単調増加関数 $\log$ を使って対数確率を表現します。Skipgramのネガティブサンプリングでは、 $v_c$ と $u_o$ の組み合わせによる対数確率を最大化し、 $v_c$ とネガティブサンプル $u_j$ の対数確率を最小化することを目指します。

ネガティブサンプルは $P(w)$ の分布から抽出されます。これは $3/4$ の乗数をかけたユニグラム分布で表されます。ユニグラム分布は、単語が出現する確率を単純にカウントすることで求めることができ、サブサンプリングにおいて高頻度単語を少なくサンプルするために $3/4$ を累乗します。

コードの実装

プロセスを理解したところで、Pythonで実装を行いましょう。デモンストレーションのために、NLTKライブラリからReuters Corpusを使用します。

import nltk
nltk.download('reuters')
from nltk.corpus import reuters

Reuters Corpusは上記のようにダウンロードできます。このコーパスには、ファイルIDを取得するためのfileids、生テキストを取得するためのraw、単語リストを表示するためのwordsなど、多くの便利なメソッドが備わっています。

fileids = reuters.fileids()
words = list(reuters.words())
corpus = reuters.raw()

前処理

テキストの前処理フェーズでは、改行記号\nと不要な空白のみを削除します。

def text_preprocessing(text):
    text = text.replace('\n', ' ')
    text = re.sub(r'\s+', ' ', text).strip()
    return text
 
corpus = text_preprocessing(corpus)

トークナイゼーション (BPE)

Byte Pair Encodingアルゴリズムを実装することから始めましょう。まず、単語をスペースで区切り、終了タグ</w>を追加し、それらを辞書dictに入れます。

def get_vocab(words):
    vocab = defaultdict(int)
    for word in words:
        vocab[' '.join(list(word)) + ' </w>'] += 1
    return vocab

その後、各バイトペア（英語では文字ペア）の頻度をカウントし、辞書に保存します。

def get_pairs(vocab):
    pairs = defaultdict(int)
    for word, freq in vocab.items():
        symbols = word.split()
        for i in range(len(symbols)-1):
            pairs[symbols[i],symbols[i+1]] += freq
    return pairs

この辞書を使用して、最も頻繁に出現するペアを見つけ、スペースを削除して組合せます。

def merge_vocab(pair, v_in):
    v_out = {}
    bigram = re.escape(' '.join(pair))
    p = re.compile(r'(?<!\S)' + bigram + r'(?!\S)')
    for word in v_in:
        w_out = p.sub(''.join(pair), word)
        v_out[w_out] = v_in[word]
    return v_out

これらを組み合わせて、BPE関数を作成できます。

def byte_pair_encoding(words, n):
    vocab = get_vocab(words)
    for i in range(n):
        pairs = get_pairs(vocab)
        max_pair = max(pairs, key=pairs.get)
        vocab = merge_vocab(max_pair, vocab)
    return vocab
 
vocab = byte_pair_encoding(words, 10000)
vocab = dict(sorted(vocab.items(), key=lambda x:x[1], reverse=True))

nはハイパーパラメータで、何組のペアを統合するかを決定します。vocab辞書を使用して、サブワードトークンペアの辞書を作成します。

def build_token_dictionary(vocab):
    token_to_id = {}
    id_to_token = {}
    id_frequency = {}
    current_id = 0
 
    for word, frequency in vocab.items():
        tokens = word.split()  # Split the word based on spaces
        for token in tokens:
            if token not in token_to_id:
                token_to_id[token] = current_id
                id_to_token[current_id] = token
                id_frequency[current_id] = frequency
                current_id += 1
            else:
                id_frequency[token_to_id[token]] += frequency
 
    token_to_id['[UNK]'] = current_id
    id_to_token[current_id] = '[UNK]'
    id_frequency[current_id] = 0
 
    return token_to_id, id_to_token, id_frequency
 
tokens, reverse_tokens, token_frequency = build_token_dictionary(vocab)

サブワードトークンペアを格納する辞書を使用して、各単語に対して最も長いサブワードをチェックし、適切にトークンを割り当てるトークン化関数を作成します。

def tokenize(text, token_dict, verbose=False):
    tokens = []
    words = text.split()
 
    for word in words:
        word_with_marker = word + '</w>'
        i = 0
        while i < len(word_with_marker):
            for j in range(len(word_with_marker), i, -1):
                subword = word_with_marker[i:j]
                if subword in token_dict:
                    print(subword) if verbose else None
                    tokens.append(token_dict[subword])
                    i = j
                    break
            else:
                tokens.append(token_dict['[UNK]'])
                i += 1
 
    return tokens
 
tokenized_corpus = tokenize(corpus, tokens)

Word2Vec: Skip-gram

Skip-gram Word2Vecをトレーニングする際には、ネガティブサンプルを準備する必要があります。頻度テーブルを使用して確率テーブルを作成し、それをネガティブサンプリングに使用します。

def transform_frequencies(id_frequency):
    transformed = {key: value**(3/4) for key, value in id_frequency.items()}
    
    mean = np.sum(list(transformed.values()))
    
    normalized = {key: value / mean for key, value in transformed.items()}
    
    return normalized
 
normalized_frequency = transform_frequencies(token_frequency)
 
def negative_sampling(normalized_frequency, num_samples, avoid_ids=[]):
    # Create a list of ids based on their normalized frequency
    ids = list(normalized_frequency.keys())
    probabilities = list(normalized_frequency.values())
    
    # Filter out avoided IDs
    filtered_ids = []
    filtered_probabilities = []
    
    for i, id in enumerate(ids):
        if id not in avoid_ids:
            filtered_ids.append(id)
            filtered_probabilities.append(probabilities[i])
 
    # Normalize the filtered probabilities
    filtered_probabilities = np.array(filtered_probabilities)
    filtered_probabilities /= np.sum(filtered_probabilities)
    
    # Sample `num_samples` IDs based on the filtered probabilities
    negative_samples = np.random.choice(filtered_ids, size=num_samples, p=filtered_probabilities)
 
    return list(negative_samples)

上記に基づいて、次の関数を使ってトレーニングデータを生成できます。

import tqdm
def generate_skipgram_training_data(tokens, token_frequency, window_size=2):    
    center_words = []
    contexts = []
    labels = []
 
    for center_idx in tqdm.tqdm(range(window_size, len(tokens) - window_size)):
        center_word = tokens[center_idx]
        context_words = [tokens[i] for i in range(center_idx - window_size, center_idx + window_size + 1) if i != center_idx]
        negative_samples = negative_sampling(token_frequency, 2*window_size)
        context = context_words + negative_samples
        label = [1 for _ in context_words] + [0 for _ in negative_samples]
 
        center_words.append(center_word)
        contexts.append(context)
        labels.append(label)
    
    return center_words, contexts, labels
 
center_words, contexts, labels = generate_skipgram_training_data(tokenized_corpus, token_frequency, window_size=5)
center_words, contexts, labels = np.array(center_words), np.array(contexts), np.array(labels)

numpy配列から、TensorFlowやPyTorch用のデータセットを作成できます。

BATCH_SIZE = 1024
BUFFER_SIZE = 10000
 
# TensorFlow
import tensorflow as tf
dataset = tf.data.Dataset.from_tensor_slices(((center_words, contexts), labels))
 
dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
dataset = dataset.cache().prefetch(buffer_size=tf.data.AUTOTUNE)
 
# PyTorch
import torch
import torch.nn as nn
 
center_words, contexts, labels = map(lambda X: torch.tensor(X, dtype=torch.float32), (center_words, contexts, labels))
dataset = torch.utils.data.TensorDataset(center_words, contexts, labels)
data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=BATCH_SIZE, shuffle=True)

最後に、TensorFlowやPyTorchでモデルを定義し、トレーニングを行います。以下は、公式のTensorFlow実装を基に、重みの初期化子とL1正則化を追加してトレーニングの安定化を図った修正版です。

from tensorflow.keras import layers
class Word2Vec(tf.keras.Model):
  def __init__(self, vocab_size, embedding_dim):
    super(Word2Vec, self).__init__()
    self.word_embedding = layers.Embedding(vocab_size,
                                      embedding_dim,
                                      embeddings_initializer="he_normal",
                                      embeddings_regularizer="l1",
                                      name="w2v_embedding")
    self.context_embedding = layers.Embedding(vocab_size,
                                       embedding_dim,
                                       embeddings_initializer="he_normal",
                                       embeddings_regularizer="l1",)
 
  def call(self, pair):
    word, context = pair
    # word: (batch, dummy?)  # The dummy axis doesn't exist in TF2.7+
    # context: (batch, context)
    if len(word.shape) == 2:
      word = tf.squeeze(word, axis=1)
    # word: (batch,)
    word_emb = self.word_embedding(word)
    # word_emb: (batch, embed)
    context_emb = self.context_embedding(context)
    # context_emb: (batch, context, embed)
    dots = tf.einsum('be,bce->bc', word_emb, context_emb)
    # dots: (batch, context)
    return dots
 
def custom_loss(x_logit, y_true):
      return tf.nn.sigmoid_cross_entropy_with_logits(logits=x_logit, labels=y_true)
 
embedding_dim = 1024
vocab_size = len(tokens)
word2vec = Word2Vec(vocab_size, embedding_dim)
word2vec.compile(optimizer='adam',
                 loss=custom_loss,
                 metrics=['accuracy'])
word2vec.fit(dataset, epochs=20)
 
## Accessing to word embedding
# word2vec.word_embedding(center_words[0])

PyTorchでも同じ手法を適用し、Word2Vecモデルを構築して単語の分散表現を取得して、実践的な理解を深めてみてください。

結論

前処理、トークナイゼーション、そして単語埋め込みを通じてテキストを数値表現に変換することは、機械学習モデルが言語を処理できるようにするための重要な要素です。 Byte Pair Encodingのような技術やWord2Vecのようなモデルは、単語の意味を捉えながら構造を維持するのに役立ちます。しかし、分散表現はWord2Vecだけしか作成できないわけではありません。次の記事では、分散表現を作成するための別のアプローチを探ります。

リソース

Thakur, A. 2021. Subword Tokenization: Byte Pair Encoding. YouTube.
The Semicolon. 2019. Word2Vec - Skipgram and CBOW. YouTube.