AI で文章をかしこく比較！ Sentence-Transformers のご紹介

文章のベクトル化とは？

冒頭でも文章をベクトル化すると述べましたが、まずはその意味について説明したいと思います。そもそもベクトルとは「向きを持った量」と説明されますが、数字の並びで表現される配列のようなものという認識で問題ありません。n 次元のベクトルであれば n 個の数字の並びとなります。機械学習における特徴量はもっぱらベクトルの形で表現されます。したがって文章のベクトル化とは、文章を数字の並びで表現することになります。ただ注意点として、文章を文字コードに変換してベクトル化するような操作ではありません。文章の特徴をうまく表現するようにベクトル化する操作になります。具体的なアルゴリズムは省略しますが、基本的には先に単語をベクトルで表現してから、それらを集約することで文章をベクトル化します。文章ベクトルを利用することで文章間の意味的な類似度を定量的に測れるようになります。

Sentence-Transformers とは？

Sentence-Transformers は、その名前の通り文章を対象とした Transformer ベースのモデルが利用できる Python のライブラリです。実装されているモデルを利用することで、前章で説明したような文章ベクトルを簡単に作成することができます。この文章ベクトルを利用すると、文章の分類や類似文章の検索といったタスクを高い精度と効率で実行することが可能になります。

• Sentence-Transformers : https://www.sbert.net/

本ライブラリは pip で簡単にインストールすることができます。

pip install -U sentence-transformers

Pytorch と Transformers がベースとなっており、Transformers で利用したモデルをそのままロードして使用することができます。

実装されているモデルは Sentence-BERT がベースとなっており、これは BERT に対して Pooling やLoss を工夫することで、より質の良い文章ベクトルを生成するモデルです。具体的にどのような Pooling や Loss が実装されているかは次の章で紹介します。

Pooling と Loss について

機械学習をはじめると Pooling や Loss という単語をよく目にすると思います。Sentence-Transformers でも文章ベクトルを得るために様々な種類の Pooling や Loss が用意されており、目的に応じて選択できるようになっています。

Pooling について

Pooling は日本語でも「プーリング」とそのまま呼ばれることが多く、あえて訳すなら集約のような操作となります。平均をとるような操作がイメージしやすいと思います。文章ベクトルを作成する場合においては、BERT などから出力される単語ベクトルを集約する操作に対応します。Sentence-Transformers では以下の 3 種類の Pooling が用意されています。

• Mean-Pooling：文中の各単語ベクトルの平均を文章ベクトルとする方法です。デフォルトの Pooling となります。
• Max-Pooling：全単語を通して単語ベクトルの各次元の最大値を集約したものを文章ベクトルとする方法です。
• CLS：BERT の CLS トークンを文章ベクトルとする方法です。厳密には Pooling ではありませんが、よく利用される方法です。

Loss について

Loss は日本語では「損失」と表現されるのですが、「誤差」と言った方がイメージはつきやすいかもしれません。通常の教師あり機械学習では、推定したラベルと真のラベルの Loss を小さくするように学習することが多いと思います。このように「ラベルの推定」のみに興味がある場合は、通常の教師あり機械学習のアプローチが採用されます。一方で、文章ベクトルのように「特徴量の表現」自体に興味がある状況では、特徴量間の距離から Loss を算出して学習する方法があります。このような方法による学習はメトリック学習と呼ばれ、Sentence-Transformers による文章ベクトルの作成にも、メトリック学習のアプローチが採用されています。

Sentence-Transformers では 14 種類の Loss が実装されており、ほとんどがメトリック学習に使われるものとなっています。この中から 3 つだけピックアップして紹介したいと思います。

• Contrastive Loss：ラベルが付与された文章 A と 文章 B の組を利用する方法です。2 つの文章のラベルが等しければそれぞれの文章ベクトルを近づけ、ラベルが異なれば文章ベクトルを遠ざけるように学習されます。
• Triplet Loss：文章 A、文章 A とラベルが同じ文章 P、文章 A とラベルが異なる 文章 N から成る 3 つ組（Triplet）の文章を利用する方法です。 文章 A が文章 P と文章 N のどちらに近いかという相対的なメトリック学習が可能な方法となります。
• Batch All Triplet Loss：機械学習実行時のバッチ (一回で学習するデータの単位) に含まれる全ての Triplet に対して Loss を計算する方法です。3 つ以上ラベルがある場合は、同時に複数のラベルに対してメトリック学習が可能な方法となります。

他の Loss に関しても、公式のリファレンスにて説明されています。興味がある方はぜひご覧ください。

• Loss のリファレンス : https://www.sbert.net/docs/package_reference/losses.html

Triplet Loss を利用して文章ベクトルを作成する

本記事では BERT の事前学習済みモデルと Triplet Loss を利用して、文章ベクトルの作成とその可視化をしてみようと思います。

文章ベクトルの作成にあたって、以前の記事と同様の事前学習済みモデルを利用していきます。

• BERT モデル : https://huggingface.co/cl-tohoku/bert-Base-japanese-whole-word-masking

データセットも前回の記事と同様に Livedoor ニュースコーパスを利用します。Livedoor ニュースコーパスは、ライブドアニュースの記事が 9 種類のカテゴリに分類されているものです。今回はカテゴリのラベルを推定するのではなく、Triplet Loss を利用し、同じラベルの文章ベクトルを近づけて、異なるラベルの文章ベクトルを遠ざけるようにモデルを Fine-Tuning していきます。

• Livedoor ニュースコーパス : https://www.rondhuit.com/download.html

Triplet Loss を利用するにあたって、上記のようなデータから下記のような Anchor, Positive, Negative から成る Triplet データを用意する必要があります。Anchor を基準となる文章として Positive は Anchor と同じラベルの文章、Negative は Anchor と異なるラベルの文章をそれぞれランダムに抽出して Triplet とします。実装によっては学習時に逐一 Triplet をサンプリングする方法もありますが、今回はあらかじめ用意する方法で実装しています。

さっそく Sentence-Transformers を利用してモデルを実装してみましょう。Sentence-Transformers をインストール後、以下のコードを Python で実行するだけで、トークナイザと事前学習済みモデルを用意することができます。全体の流れとしては、

1. ライブラリをインポートします。日本語の Transformers のモデルを利用する際には、この時点で トークナイザを日本語のものに上書きしておきます。
2. モデルを読み込んで構築します。ここで Pooling の設定を行っています。
3. データセットの読み込みとパラメータを設定します。用意した Triplet のファイルを読み込み、Triplet Loss の設定を行います。

# 1. ライブラリのインポート
import transformers
import sentence_transformers
transformers.BertTokenizer = transformers.BertJapaneseTokenizer
from sentence_transformers import SentenceTransformer
from sentence_transformers import models
from sentence_transformers.losses import TripletDistanceMetric, TripletLoss
from sentence_transformers.readers import TripletReader
from sentence_transformers.datasets import SentencesDataset
from torch.utils.data import DataLoader

# 2. モデルの定義
bert = models.Transformer("cl-tohoku/bert-base-japanese-whole-word-masking")
pooling = models.Pooling(
        bert.get_word_embedding_dimension(),
        pooling_mode_mean_tokens=True,
)
model = SentenceTransformer(modules=[bert, pooling])

# 3. データの読み込み、パラメータの設定
triplet_reader = TripletReader(".")
train_dataset = SentencesDataset(
        triplet_reader.get_examples("triplet.tsv"),
        model=model,
)

BATCH_SIZE = 8
NUM_EPOCH = 15
EVAL_STEPS = 100
WARMUP_STEPS = int(len(train_dataset) // BATCH_SIZE * 0.1)

train_dataloader = DataLoader(
        train_dataset,
        shuffle=False,
        batch_size=BATCH_SIZE,
)
train_loss = TripletLoss(
        model=model, 
        distance_metric=TripletDistanceMetric.EUCLIDEAN,
        triplet_margin=1,
)

モデルの用意が完了したら、以下を実行するだけで学習が可能です。

model.fit(
        train_objectives=[(train_dataloader, train_loss)],
        epochs=NUM_EPOCH,
        evaluation_steps=EVAL_STEPS,
        warmup_steps=WARMUP_STEPS,
        output_path="./sbert",
        )

モデルの利用も非常に簡単です。テスト用に分割したデータを作成したモデルに入力することで、簡単に文章ベクトルを得ることができます。

sbert = SentenceTransformer('./sbert')
vector = sbert.encode('ベクトル化したい文章')

文章ベクトルを可視化する

前章の手順で作成した文章ベクトルを可視化することで、その性質を確認します。今回作成した文章ベクトルは BERT の出力と同じ 768 次元あります。これを t-SNE というアルゴリズムで 2 次元へと次元削減を行い可視化します。比較のため以下の文章ベクトルについても可視化を行います。

1. Fine-Tuning 前の文章ベクトル（事前学習済み BERT の出力を Mean-Pooling したもの）
2. Triplet Loss を利用して Fine-Tuning した文章ベクトル
3. 文書分類タスクとして Fine-Tuning した際の文章ベクトル

以下に結果を示します。