はじめての自然言語処理

第9回 Sentence BERT による類似文章検索の検証

オージス総研技術部アドバンストテクノロジセンター

鵜野和也

2020年6月23日

今回は初心に帰って類似文章検索です。連載の第1回で扱ったネタですが、 BERT を用いて再挑戦してみましょう。BERT のモデルは Hagging Face Transformers の事前学習済みモデルを用いるので、お手軽に試せるかと思います。手法としては Sentence BERT を用い、おまけとして Poor Man's BERT についても紹介します。

（本記事公開後に公開されたデータセットで再検証しています。最新情報は第18回をご覧ください。 2021.12.21 追記）

1. はじめに

本記事では Sentence BERT ¹による類似文章検索について、学習や推論のコード例と実験結果を交えてご紹介します。前々から Sentence BERT を試したいと考えていたものの、教師あり学習に必要な日本語の類似文データが用意できずにいました。その後、画像キャプションのデータセットを流用することを思いつき、試してみたところ思いのほか上手くいったのでご紹介することにしました。また、BERT のシンプルな軽量化手法である Poor Man’s BERT ² も試してみたので参考にして頂けると良いかと思います。

それでは、まず Sentence BERT について紹介していきます。

2. Sentence BERT

Sentence BERT は BERT (及びその亜種)をファインチューニングして良質な文章ベクトルを生成する手法です。話としては非常にシンプルで、「ある文章とそれに類似する文章のペアを学習データとし、似た文章から生成される文章ベクトルが似たベクトルになるように BERT をファインチューニングする」、これだけです。

ファインチューニングする際の目的関数は、用意したデータセットのラベルによって変わってきます。文章ペアに対して「含意」、「矛盾」、「どちらでもない」のようなラベルが付いていれば、それぞれの文章ベクトルを連結して、分類問題として解きますし(下図左)、ラベルが似ている度合いの数値であればコサイン類似度を算出して回帰で解く(下図右)ことになります。

図について補足しておくと、"u", “v” はそれぞれ、"Sentence A", “Sentence B” の文章に対応する文章ベクトルです。左図の “(u, v, |u-v|)” は「"u", “v”, 及び “u, v の差分ベクトルの各要素の絶対値をとったもの” を連結する」という意味です。

“pooling” についても説明しておきましょう。 BERT に文章を入力すると、出口からは入力文の各トークン毎に文脈を踏まえたベクトルが出てきます。連結するにせよ、コサイン類似度を計算するにせよ、トークン毎のベクトルをまとめて固定長の文章ベクトルにしないといけません。これが “pooling” の役割で、Sentence BERT では以下の3種類の方法が評価されています。

MEAN : 各トークンのベクトルを平均する。
MAX : 各トークンのベクトルから最大値を拾う。
CLS : “[CLS]"トークンに対応するベクトルをそのまま使う。

論文では MAX の性能が最も劣るとのことでした、本記事ではデフォルトの MEAN を用いて実験することにします。それでは、実験に使うデータセットについて考えてみましょう。

3. 学習データセット

さて、Sentence BERT の学習には類似した文章のペアが必要になります。この連載の第1回で類似文章検索をしたときは、精度評価に用いる為に手作業でデータセットを作成したのですが、学習データに十分な量を用意するとなると手作業はツライです。かといって日本語の類似文章データセットも見つかりません。

仕方がないので、他のタスク向けに作成されたデータセットを流用することにしました。利用したのは STAIR Captions ³ データセットになります。

STAIR Captions はその名のとおり、画像に対するキャプションの大規模日本語データセットであり、画像からの自動キャプション生成や文章による画像検索等に使うのが真っ当な使い方になります。公開されてる部分だけでも約12万の画像に対し、5つのキャプションが付与されています。

今回はやや乱暴ですが、「同じ画像に対するキャプションって似た文章になりそうだよね。」という仮説に基づいて Sentence BERT 用の学習データを作ろうと考えました。実際のデータの雰囲気はこんな感じです。

最後の一文はやや微妙ですが、似てますよね。「飛行機にコンテナが積み込まれている」とかよりは、よほど似ています。キャプションなので文章としては短いですが、英語の類似文章データセットもキャプションや記事タイトルから構成されてたりするので、そこは良しとします。他にないですし。

これで類似文のペアは何とかなりそうなので、次は目的関数です。

4. Triplet Loss

今回は画像キャプションのデータセットから類似文のペアを作るので、「含意」、「矛盾」のラベルもなければ、類似度の数値もありません。論文によると Sentence BERT は Triplet Loss でも良好な結果が得られているようなので、本記事の実験でも Triplet Loss を使うことにしました。

Triplet Loss は基準となるサンプル(アンカーと呼びます)とそれに類似したサンプル(ポジティブ)、類似してないサンプル(ネガティブ)の 3 者で、その関係を学習します。数式で書くと以下のようになります。

s_a, s_p, s_n はそれぞれ、アンカー、ポジティブ、ネガティブの各サンプルの文章ベクトル。||s_x - s_y|| は s_x と s_y の距離(論文ではユークリッド距離を用いています)、ε は確保するべきマージンです。アンカーとポジティブの距離(||s_a - s_p||)がアンカーとネガティブのそれ(||s_a - s_n||)よりも一定(ε)以上小さくなるように学習するものです。

わかりやすくする為に図で考えてみましょう。

上図の (a) は類似サンプルとは距離が小さく、そうでないものとは距離が大きいという望ましい結果が得られた場合です。距離差が ε 以上ですのでロスは発生しません。数式の max() で 0 が選ばれたケースです。逆に (b) は類似サンプルとそうでないもので、あまり距離に差がでず ε を確保できなかった場合です。この場合は紫の破線で示した分の長さがロスとなり、これを最小化するように学習します。

次に Poor Man’s BERT についても紹介しておきましょう。

5. Poor Man’s BERT

Poor Man’s BERT は BERT の軽量化手法の一つです。 BERT の軽量化手法は蒸留したり( DistilBERT ⁴ )、量子化したり( Q8BERT ⁵ )、いろいろしたり( ALBERT ⁶ ) と様々なアプローチが試されているのですが、Poor Man’s BERT はその中でも飛びぬけてシンプルな手法を採用しています。

第3回の BERT の記事でご紹介したとおり、 BERT は Multi Head Self-Attention と Feed Forward で構成される層が複数数珠繋ぎになった構造をしています。 BERT_BASE は 12 層なので、ちょうど下図の左端のイメージになります。

Poor Man’s BERT はこの複数層から「いくつかの層を引っこ抜いてからファインチューニングする」という手法です。Poor Man’s BERT の論文では上図に示された引き抜き方の戦略や、引く抜く層数、層を引き抜く対象( BERT, XLNet, DistilBERT ) など様々なパターンで比較実験しています。

とりあえず、BERT に関しては上位層を引き抜く "Top Layer Dropping” が良好だったようなので、本記事でも “Top Layer Dropping” で実験をしています。

では、ここからは実際に Sentence BERT でどの程度の精度がでるか実験をしてみましょう。

6. 日本語データでの実験

例によって Google Colaboratory を使用します。"ランタイム“ => "ランタイムのタイプを変更” で “ハードウェアアクセラレータ” に “GPU” を選んでおいて下さい。

環境のセットアップ

まずは、Sentence BERT をインストールします。筆者が実験した際のバージョンは 0.2.5.1 でした。

!pip install -U sentence-transformers

Sentence BERT は Hagging Face Transformers を利用して実装されています。ですので「そのまま」とはいきませんが、ちょっとした修正で、東北大学の乾・鈴木研究室の方が公開した下さった事前学習済みモデル⁷を利用することが出来ます。

トークン化で必要になる MeCab 関係をインストールします。

!apt-get install mecab mecab-ipadic-utf8 python-mecab libmecab-dev
!pip install mecab-python3

データセットの加工に使うので GiNZA もインストールします。

!pip install ginza

ここからが「ちょっとした修正」になるのですが、 Sentence BERT のコードはトークナイザとして BertTokenizer を決め打ちで使用します。Hagging Face で公開されている日本語事前学習済みモデルを利用できるように、以下のように差し替えてしまいます。

import transformers
transformers.BertTokenizer = transformers.BertJapaneseTokenizer

日本語での動作テスト

それでは、Sentence BERT で日本語の文章を固定長ベクトルに変換してみましょう。先ほどの BertTokenizer の差し替えが済んでいれば、公開されたモデルの名前(ここでは 'cl-tohoku/bert-base-japanese-whole-word-masking' )を models.BERT() に渡すだけで OK です。 “pooling” は MEAN を使用しています。

from sentence_transformers import SentenceTransformer
from sentence_transformers import models

transformer = models.BERT('cl-tohoku/bert-base-japanese-whole-word-masking')
pooling = models.Pooling(transformer.get_word_embedding_dimension(), pooling_mode_mean_tokens=True, pooling_mode_cls_token=False, pooling_mode_max_tokens=False)
model = SentenceTransformer(modules=[transformer, pooling])

sentences = ['吾輩は猫である',  '本日は晴天なり']
embeddings = model.encode(sentences)

for i, embedding in enumerate(embeddings):
  print("[%d] : %s" % (i, embedding.shape, ))

# [0] : (768,)
# [1] : (768,)

事前学習済みモデルそのままなので文章ベクトルの質は置いておいて、とりあえず固定長ベクトルに変換されることは確認できました。

それでは学習データセットの準備をしていきましょう。

学習データセットの準備

学習データセットは前述のとおり STAIR Captions を用います。約12万の画像に各5つのキャプションがついているので、以下のように加工します。

画像単位に 8 : 1 : 1 に分割
画像に付与された 5 つのキャプションから positive pair を生成
各 positive pair に対し異なる画像のキャプションからサンプリングして、negative pair を生成

とりあえず、データを取得して展開します。

!git clone https://github.com/STAIR-Lab-CIT/STAIR-captions
!tar zxvf STAIR-captions/stair_captions_v1.2.tar.gz
!ls -lh *.json

# -rw-r--r-- 1 root root 2.7K Mar 10 04:19 adc.json
# -rw-rw-r-- 1 1005 1006  67M Jan 26  2018 stair_captions_v1.2_train.json
# -rw-rw-r-- 1 1005 1006 106M Jan 26  2018 stair_captions_v1.2_train_tokenized.json
# -rw-rw-r-- 1 1005 1006  33M Jan 26  2018 stair_captions_v1.2_val.json
# -rw-rw-r-- 1 1005 1006  52M Jan 26  2018 stair_captions_v1.2_val_tokenized.json

学習データと検証データに分かれていますが両方読み込みます。

import json
with open("stair_captions_v1.2_val.json", "r") as f:
  json_data_val = json.load(f)
with open("stair_captions_v1.2_train.json", "r") as f:
  json_data_train = json.load(f)

読み込んだデータを加工しやすい形に変換します。

dataset は画像ID をキー、画像に付与された5つのキャプションの(ID, 文字列)のタプルのリストを値に持つ辞書
ids はキャプションIDのリスト
captions はキャプション文字列のリスト

dataset = {}
ids = []
captions = []

def build_dataset(dataset, json_data):
  num_samples = len(json_data['annotations'])
  for i in range(num_samples):
    anno = json_data['annotations'][i]
    image_id = anno["image_id"]
    image_captions = dataset.get(image_id, [])
    image_captions.append((anno["id"], anno["caption"]))
    ids.append(anno["id"])
    captions.append(anno["caption"])
    dataset[image_id] = image_captions

build_dataset(dataset, json_data_val)
build_dataset(dataset, json_data_train)

さらにキャプションID と ids, captions のインデックスの変換テーブルを作っておきます。

id2idx = {id:idx for idx, id in enumerate(ids)}

ここからは GiNZA でキャプション文字列を文章ベクトルに変換します。この文章ベクトルは単純に単語ベクトルの平均なのですが、画像に含まれる5つのキャプションから仲間外れをはじき出すのに用います。 (この処理はそれなりに時間がかかります。途中で画面をリロードするなどしてセッションが切れないようにして下さい。）

import numpy as np
import spacy
import pkg_resources, imp
imp.reload(pkg_resources)

nlp = spacy.load("ja_ginza")

vectors = []
for caption in captions:
  doc = nlp(caption, disable=['ner'])
  vectors.append(doc.vector)

del nlp

キャプションIDのペアでコサイン類似度を返す関数を定義し、ある程度以下の類似度になったキャプションはポジティブサンプルの対象から外すこととします。いくつか手動でチェックしてみて今回は閾値を 0.85 に設定することにしました。

from sklearn.metrics.pairwise import cosine_similarity
def similarity(id1, id2):
  return cosine_similarity([vectors[id2idx[id1]]], [vectors[id2idx[id2]]])[0][0]

以下がチェックしたサンプルの例です。"✖" がついているのが閾値(<0.85)でポジティブペアの対象外と判定したキャプションです。

夜の街にオープンしているワイナリーの入り口
 0.92 : 夜で、まだオープンしている場所の入り口
 0.75 : ✖レンジの壁の建物に黄緑色の看板がある

渋滞の中2人乗りのバイクが走っている
 0.93 : 渋滞で二人乗りのバイクが前に進めず止まっている
 0.75 : ✖黒いヘルメットを被った子供がバイクの後ろに乗っている

白い自転車に荷物を積んでライトをつけている
 0.85 : 白い自転車のライトが2つ、光っている
 0.85 : 室内で鏡の前にて折りたたみ自転車がスタンドを立てて立っている

若い男性らが、家の中でゲームをして遊んでいる
 0.88 : 数人の男女がリビングでゲームをしている
 0.60 : ✖黄色いカーテンの前に赤いシャツの女がいる

大きな時計台が繋がっている長い橋
 0.85 : 時計のある建物の近くに橋がある
 0.75 : ✖きれいな歴史的建造物が立ち並んでいる立派な町

閾値が決まったところで、アンカー、ポジティブ、ネガティブの triplet を形成します。

画像に付与されたキャプションからランダムに2つ抽出して類似度が閾値以上ならアンカーとポジティブとします。
類似度が閾値以下なら抽出しなおし、何回かやってダメならその画像はスキップします。
ネガティブサンプルはキャプション集合全体からランダムに抽出。抽出した候補が同じ画像に属していたら抽出しなおし。

import random

def make_triplets(dataset, threshold=0.85, seed=7, max_tries=25):
  triplets = []
  random.seed(seed)
  neg_candidate_indices = list(range(len(ids)))
  random.shuffle(neg_candidate_indices)
  def log(i, str):
    if i % 5000 == 0:
      print(str)

  for i, image_id in enumerate(list(dataset.keys())): 
    log(i, "### %d ###" % (image_id))

    # pickup positive pair.
    score = 0.0
    tries = 0
    while score < threshold : 
      [(id, caption), (id_pos, caption_pos)] = random.sample(dataset[image_id],2)
      score = similarity(id, id_pos)
      tries+=1
      if tries > max_tries:
        break
    if score < threshold:
      continue

    # pickup negative one.
    id_neg = id
    current_caption_ids = [id_cap[0] for id_cap in dataset[image_id]]
    while id_neg in current_caption_ids:
      idx_neg = neg_candidate_indices.pop()
      id_neg = ids[idx_neg]
    caption_neg = captions[id2idx[id_neg]]

    log(i, "  pos:  score: %4.2f [%s]:[%s]" % (score, caption, caption_pos))
    log(i, "  neg:  score: %4.2f [%s]:[%s]" % (similarity(id, id_neg), caption, caption_neg))
    triplets.append({
      "image_id": image_id,
      "id": id,
      "id_pos": id_pos,
      "id_neg": id_neg,
      "caption": caption,
      "caption_pos": caption_pos,
      "caption_neg": caption_neg  
    })
  return triplets

実際に生成するとこんな感じになります。ぱっと見た感じでは、それっぽく形成できているようです。

triplets = make_triplets(dataset)

# ### 580856 ###
#   pos:  score: 0.85 [夜の街にオープンしているワイナリーの入り口]:[踏切の近くにワイナリーが開店している]
#   neg:  score: 0.47 [夜の街にオープンしているワイナリーの入り口]:[野球選手のフォームの確認をしている]
# ### 235486 ###
#   pos:  score: 0.90 [2人の女性と1人の男性が食卓を囲んでいる]:[三人の男女がピザと飲み物を囲んで笑っている]
#   neg:  score: 0.60 [2人の女性と1人の男性が食卓を囲んでいる]:[赤い飛行機が飛行機雲をひいて飛んでいる]
# ### 133631 ###
#   pos:  score: 0.90 [1頭のゾウが柵の内側に立っている]:[土の上に立っているゾウには牙がある]
#   neg:  score: 0.84 [1頭のゾウが柵の内側に立っている]:[洗面所に花瓶に入った花が飾られている]
# ### 467843 ###
# ### 423588 ###
#   pos:  score: 0.85 [雪山の小さな小屋の屋根に、スノーボードを履いた人物が乗り、それを3人の人物が見ている]:[スキー場で数人の人がスノボーをしている]
#   neg:  score: 0.71 [雪山の小さな小屋の屋根に、スノーボードを履いた人物が乗り、それを3人の人物が見ている]:[ぞうがトラックで運ばれている]
# ### 115765 ###
#   pos:  score: 0.89 [ギターバックを持っている男性が、傘を持った2人の女性に紙を渡して話しかけている]:[傘をさした女性と、楽器をもった人]
#   neg:  score: 0.83 [ギターバックを持っている男性が、傘を持った2人の女性に紙を渡して話しかけている]:[赤いユニフォームを着た人はバッドを振っている]
...

12万件ほどの学習データができたので、8:1:1 に分割します。

len(triplets)

# 120582

from sklearn.model_selection import train_test_split
train, dev_test = train_test_split(triplets, train_size=0.8, random_state=4)
dev, test = train_test_split(dev_test, train_size=0.5, random_state=7)

TSV ファイルにして保存します。

def to_tsv(fname, triplet):
  with open(fname, "w") as f:
    lines = ["%s\t%s\t%s" % (example["caption"], example["caption_pos"], example["caption_neg"]) for example in triplet]
    f.write("\n".join(lines)+"\n")

to_tsv("triplet_train.tsv", train)
to_tsv("triplet_dev.tsv", dev)
to_tsv("triplet_test.tsv", test)

del triplets
del ids
del captions
del dataset

分割後の各ファイルの件数は以下のとおりです。

!wc -l *.tsv

#   12058 triplet_dev.tsv
#   12059 triplet_test.tsv
#   96465 triplet_train.tsv
#  120582 total

データの中身はこんな感じになります。

!head -5 triplet_train.tsv

# 森のなかに小さなテーブルと違うタイプの椅子が4つある  緑の沢山ある庭に置かれた三脚の椅子とパソコンの乗ったテーブル  緑の芝生の上に沢山のヤギがいる
# 銀の入れ物の中にハサミが入っている   テーブルの端にハサミやトングが入った容器が置いてある  緑のシャツの女が緑のミトンを両手にはめている
# 電車の窓の外に夕焼けの空が見える  乗り物の窓から夕焼け空が見える   部屋にテーブルや椅子が置かれている

今回は 1 枚の画像から 1 triplet サンプルを作っています。キャプションは5つ付与されているので複数サンプルを作ることもできたのですが、データ量を増やして学習すればするほど、「画像キャプション」ドメインに特化してしまう（汎化性能が落ちる）気がしたので、ほどほどのデータ量で試すことにしました。

学習データの準備ができたので、お待ちかねの学習ループを回していきましょう。

学習の実行

基本的なパラメータは以下のようにしました。バッチサイズやウォームアップは Sentence BERT の論文の “3.1 Training Details” の記述に従っています。

BATCH_SIZE = 16
NUM_EPOCHS = 1
EVAL_STEPS = 1000
WARMUP_STEPS = int(len(train) // BATCH_SIZE * 0.1) 
OUTPUT_PATH = "./sbert_stair"

(ここにたどり着く前にランタイムを再起動したりしてるかもしれないので) 改めて BertTokenizer を差し替えて事前学習済みモデルをロードしなおします。

import 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.evaluation import TripletEvaluator
from sentence_transformers.readers import TripletReader
from sentence_transformers.datasets import SentencesDataset
from torch.utils.data import DataLoader

transformer = models.BERT('cl-tohoku/bert-base-japanese-whole-word-masking')

pooling = models.Pooling(
    transformer.get_word_embedding_dimension(), 
    pooling_mode_mean_tokens=True, 
    pooling_mode_cls_token=False, 
    pooling_mode_max_tokens=False
)

model = SentenceTransformer(modules=[transformer, pooling])

作成したデータセットを読み込む DataLoader を定義します。

triplet_reader = TripletReader(".")
train_data = SentencesDataset(triplet_reader.get_examples('triplet_train.tsv'), model=model)
train_dataloader = DataLoader(train_data, shuffle=True, batch_size=BATCH_SIZE)

ロスには TripletLoss を使います。設定は論文の “3. Model” の記述に従いました。

train_loss = TripletLoss(model=model, distance_metric=TripletDistanceMetric.EUCLIDEAN, triplet_margin=1)

検証も Triplet で行います。

dev_data = SentencesDataset(triplet_reader.get_examples('triplet_dev.tsv'), model=model)
dev_dataloader = DataLoader(dev_data, shuffle=False, batch_size=BATCH_SIZE)
evaluator = TripletEvaluator(dev_dataloader)

学習ループを実行します。

model.fit(train_objectives=[(train_dataloader, train_loss)],
         evaluator=evaluator,
         epochs=NUM_EPOCHS,
         evaluation_steps=EVAL_STEPS,
         warmup_steps=WARMUP_STEPS,
         output_path=OUTPUT_PATH
         )

output_path に指定したディレクトリには以下のような内容が出力されています。

!find ./sbert_stair -print

# ./sbert_stair/
# ./sbert_stair/1_Pooling/
# ./sbert_stair/1_Pooling/config.json
# ./sbert_stair/triplet_evaluation_results.csv
# ./sbert_stair/config.json
# ./sbert_stair/modules.json
# ./sbert_stair/0_BERT/
# ./sbert_stair/0_BERT/pytorch_model.bin
# ./sbert_stair/0_BERT/sentence_bert_config.json
# ./sbert_stair/0_BERT/added_tokens.json
# ./sbert_stair/0_BERT/tokenizer_config.json
# ./sbert_stair/0_BERT/special_tokens_map.json
# ./sbert_stair/0_BERT/vocab.txt
# ./sbert_stair/0_BERT/config.json

学習曲線を確認してみましょう。データは triplet_evaluation_results.csv に以下のような形で出力されています。

!cat sbert_stair/triplet_evaluation_results.csv

 # epoch,steps,accuracy_cosinus,accuracy_manhatten,accuracy_euclidean
 # 0,1000,0.9897993033670592,0.989550505888207,0.989550505888207
 # 0,2000,0.9927848731132858,0.9924531431414828,0.9928678056062366
 # 0,3000,0.9940288605075469,0.9934483330568917,0.9935312655498424
 # 0,4000,0.9941947254934483,0.9940288605075469,0.9942776579863991
 # 0,5000,0.9952728479018079,0.994941117930005,0.9950240504229557
 # 0,6000,0.9953557803947587,0.9951899154088572,0.9956045778736109
 # 0,-1,0.9953557803947587,0.9951899154088572,0.9956045778736109

最後の行の steps の -1 がプロット時に邪魔なので、修正してプロットします。

import pandas as pd
eval_df = pd.read_csv("sbert_stair/triplet_evaluation_results.csv")
eval_df["steps"][6] = len(train) // BATCH_SIZE
eval_df.plot(x="steps", y=["accuracy_cosinus", "accuracy_manhatten",    "accuracy_euclidean"])

プロットされている値は、検証データの Triplet で ||s_a - s_p||、||s_a - s_n|| を文章ベクトル同士のコサイン類似度(cosinus)、マンハッタン距離(manhatten)、ユークリッド距離(euclidean)でそれぞれ計測し、s_a から見て s_p が s_n よりも近傍にあった割合を示したものです。学習量としては 1 epoch で十分そうですね。

では、ここからはテストデータでの精度を見ていきましょう。

比較対象の準備

Sentence BERT の精度だけ見ても「ふーん」になりそうなので、以下の手法と比較することにします。

GiNZA 3.0 : doc.vector で取得したベクトルを文書ベクトルとします。
fastText : 形態素解析して、各単語の原形の学習済み単語ベクトル⁸を平均して文書ベクトルとします。
TF-IDF : 学習データから作成した gensim の TF-IDF モデル⁹から文書ベクトルを生成します。

gensim と fasttext をインストールします。

!pip install gensim
!pip install fasttext

分かち書きはこんな感じです。

import MeCab
mecab = MeCab.Tagger ("-Ochasen")
mecab.parse("")

def parse_by_mecab(sentence):
  tokens = []
  node = mecab.parseToNode(sentence).next
  while node:
    feature = node.feature.split(',')
    token = feature[-3] # 標準形
    if token == '*':
      token = node.surface
    tokens.append(token)
    node = node.next
  return [token for token in tokens if len(token) > 0]

比較するベクトルの入れ物を作っておきます。

test_vectors = {"sbert":{}, "fasttext":{}, "spacy":{}, "tfidf":{}}

fasttext の学習済み単語ベクトルをダウンロードします。

!curl -O https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.ja.300.bin.gz
!gzip -d cc.ja.300.bin.gz

ダウンロードした単語ベクトルをロードし、単語ベクトルを平均して文章ベクトルとする関数を定義します。

import fasttext
import fasttext.util
ft = fasttext.load_model('cc.ja.300.bin')

import numpy as np
def fasttext_sentence_vec(sentence):
  tokens = parse_by_mecab(sentence)
  vectors = [ft.get_word_vector(token) for token in tokens]
  vectors = np.array(vectors)
  return vectors.mean(axis=0)

テストデータの各キャプションを fasttext の学習済み単語ベクトルを平均して文章ベクトルに変換します。

for example in test:
  test_vectors["fasttext"][example["id"]] = fasttext_sentence_vec(example["caption"])
  test_vectors["fasttext"][example["id_pos"]] = fasttext_sentence_vec(example["caption_pos"])
  test_vectors["fasttext"][example["id_neg"]] = fasttext_sentence_vec(example["caption_neg"])

del ft

次は TF-IDF モデルの作成と文書ベクトルを生成する関数です。

train_captions= set()
for sample in train:
  train_captions.add(sample["caption"])
  train_captions.add(sample["caption_pos"])
  train_captions.add(sample["caption_neg"])
train_captions= list(train_captions)

from gensim.corpora import Dictionary
from gensim.models import TfidfModel
NO_BELOW = 5
NO_ABOVE = 0.5

parsed_captions = [parse_by_mecab(caption) for caption in train_captions]

dictionary = Dictionary(parsed_captions)
dictionary.filter_extremes(no_below=NO_BELOW, no_above=NO_ABOVE)
corpus = [dictionary.doc2bow(tokens) for tokens in parsed_captions]
tfidf = TfidfModel(corpus)

def tfidf_sentence_vec(sentence):
  tokens = parse_by_mecab(sentence)
  bows = [dictionary.doc2bow(tokens) for tokens in [tokens]]
  vector = np.zeros(len(dictionary))
  for token_id, value in tfidf[bows[0]]:
    vector[token_id] = value
  return vector

こちらもTF-IDFの文書ベクトルを作って不要オブジェクトは削除しておきます。

for example in test:
  test_vectors["tfidf"][example["id"]] = tfidf_sentence_vec(example["caption"])
  test_vectors["tfidf"][example["id_pos"]] = tfidf_sentence_vec(example["caption_pos"])
  test_vectors["tfidf"][example["id_neg"]] = tfidf_sentence_vec(example["caption_neg"])

del dictionary
del corpus
del tfidf

同様に GiNZA でも文書ベクトルを作り、

for example in test:
  test_vectors["spacy"][example["id"]] = vectors[id2idx[example["id"]]]
  test_vectors["spacy"][example["id_pos"]] = vectors[id2idx[example["id_pos"]]]
  test_vectors["spacy"][example["id_neg"]] = vectors[id2idx[example["id_neg"]]]

del vectors

最後に Sentence BERT です。学習済みの Sentence BERT モデルは SentenceTransformer に学習時の出力先フォルダを渡してロードします。

sbert = SentenceTransformer('./sbert_stair')

こちらもロードした学習済みの Sentence BERT モデルで文章ベクトルを生成します。

for example in test:
  test_vectors["sbert"][example["id"]] = sbert.encode([example["caption"]])[0]
  test_vectors["sbert"][example["id_pos"]] = sbert.encode([example["caption_pos"]])[0]
  test_vectors["sbert"][example["id_neg"]] = sbert.encode([example["caption_neg"]])[0]

del sbert

テストデータでの評価

テストデータでの評価は検証データの評価で用いた、コサイン類似度でポジティブがネガティブより近傍にある率(“accuracy_cosinus”)、及びコサイン類似度の差分を見ていきます。変換済みの文書ベクトルを使って精度計算をする関数は以下のような感じです。

from sklearn.metrics.pairwise import cosine_similarity
def evaluate(test_set, test_vectors, model):
  correct = 0
  avg_diff = 0
  for example in test_set:
    vector_anchor = test_vectors[model][example["id"]]
    vector_pos    = test_vectors[model][example["id_pos"]]
    vector_neg    = test_vectors[model][example["id_neg"]]

    score_pos = cosine_similarity([vector_anchor], [vector_pos])[0][0]
    score_neg = cosine_similarity([vector_anchor], [vector_neg])[0][0]
    diff = score_pos - score_neg
    avg_diff += diff
    if diff > 0:
      correct += 1

  return correct/len(test_set), avg_diff/len(test)

今回用意した 4 つのモデルで精度計算してみます。

sentence_vector_models = ["sbert", "fasttext", "spacy", "tfidf"]
results = []
for model in sentence_vector_models:
  accuracy, avg_diff = evaluate(test, test_vectors, model)
  results.append([accuracy, avg_diff])
results = np.array(results)

結果をプロットします。

import matplotlib.pyplot as plt
import numpy as np
from collections import OrderedDict
from matplotlib.font_manager import FontProperties
plt.style.use('seaborn-whitegrid')
%matplotlib inline

def plot_bar(labels, scores, aspect=(6, 3), title=None, format="%5.2f", w=0.5, xlim=None):
  f, ax1 = plt.subplots(figsize=aspect, dpi=120)
  ax2 = ax1.twinx()
  fp = FontProperties(size=9)
  if xlim:
    plt.xlim(xlim)
  ax1.set_ylabel("Accuracy", fontsize=9)
  ax1.set_ylim([0.0, 1.2])
  ax2.set_ylabel("Cosine similarity avg. diff", fontsize=9)
  ax2.set_ylim([0.0, 0.50])
  plt.title(title, fontproperties=fp)

  for i in range(len(scores)):
    for j, score in enumerate(scores[i]):
      x = i-0.15+j*0.3
      y = score
      color = "C1"
      if j % 2 == 0:
        color="C0"
        y = score / 1.2 * 0.5 
      #print("x=%5.3f, y=%5.3f, score=%5.3f" % (x, y, score))
      plt.text(x=x, y=y+0.02, s=format % score, rotation=40, color=color, fontsize=9)

  ax1.bar(np.array(range(len(labels)))-0.15, np.array(scores)[:,0], color="C0", width=w, label="Accuracy")
  ax2.bar(np.array(range(len(labels)))+0.15, np.array(scores)[:,1], color="C1", width=w, label="Cosine similarity avg. diff")
  plt.xticks(np.arange(len(labels)), labels)
  ax1.grid(b=True, which='major', axis='both')
  ax2.grid(b=False)
  h1, l1 = ax1.get_legend_handles_labels()
  h2, l2 = ax2.get_legend_handles_labels()
  ax1.legend(h1+h2, l1+l2, bbox_to_anchor=(1.15, 1), loc='upper left', borderaxespad=0.3, fontsize=9)

plot_bar(sentence_vector_models, results, title="Various sentence Vectoring comparison", format="%5.3f", w=0.3, xlim=None)

期待通りに Sentence BERT が最良値となりました。ポジティブとネガティブの識別が出来ているだけでなく、類似度の差分は他のモデルより大きく出ているので、良い感じなのではないでしょうか。ただ、教師あり学習でファインチューニングしただけにドメインの異なるデータでの精度も気になります。そのあたりも確認してみました。

異なるドメインのデータを用いた比較

今回は学習データと異なるドメインのデータとして、本連載の第1回で用いた海外小説の翻訳者違いデータセットを用いて精度を出してみました(※ sbert 以外は第1回で示したデータの転載になります）。

WMD で出したスコアを 5 pt 以上上回って最高値を出ました。この辺は改めて進歩の速さを感じるところですね。

うっかり、このまま締めてしまいそうですが、 Poor Man’s BERT のことを忘れていました。もう少しだけ続きます。

Poor Man’s BERT での結果

Poor Man’s BERT は層を抜いてファインチューニングするだけです。

以下は上位4層を抜く場合の例になります。まず、Poor Man’s BERT のモデルを格納するディレクトリを作り、

!mkdir ./poor_mans_bert_base_top_4_removed

Sentence BERT は同じ名前で Tokenizer と BERTの本体をロードするので、トークナイザを作成済みのディレクトリに保存しておきます。

from transformers import BertJapaneseTokenizer, BertModel
tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
tokenizer.save_pretrained("./poor_mans_bert_base_top_4_removed")

次に、元になる事前学習済みモデルをロードして、層を抜き、層数のプロパティを合わせて保存すればOKです。

NUM_LAYERS = 8
model = BertModel.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
model.encoder.layer = model.encoder.layer[:NUM_LAYERS]
model.config.num_hidden_layers = NUM_LAYERS
model.save_pretrained("./poor_mans_bert_base_top_4_removed")

これで上位4層を抜いた Poor Man’s BERT が出来たので、Sentence BERT から以下の要領でロードします。

from sentence_transformers import models, SentenceTransformer
import transformers
transformers.BertTokenizer = transformers.BertJapaneseTokenizer
transformer = models.BERT('./poor_mans_bert_base_top_4_removed')
pooling = models.Pooling(
    transformer.get_word_embedding_dimension(), 
    pooling_mode_mean_tokens=True, 
    pooling_mode_cls_token=False, 
    pooling_mode_max_tokens=False
)
model = SentenceTransformer(modules=[transformer, pooling])

後は本記事で紹介した手順で学習ループを回すだけです。学習済みモデルは通常どおり、学習時の出力先フォルダを指定して以下のようにロードできます。

sbert = SentenceTransformer('./sbert_t4_stair')

上位4層削除、上位6層削除の2パターンでの精度は以下のとおりです。まずはテストデータセットでのスコアです。

コサイン類似度の差分平均は大分小さくなってしまいましたが、ポジティブ／ネガティブの判定結果に影響するほどではないようです。ドメインの異なるデータでのスコアも見てみましょう。

こちらも全12層を用いた時の精度を維持できていますね。

せっかくの Poor Man’s BERT なので、どの程度軽くなったのか文章分類で確認しておきましょう。

7. Poor Man’s BERT での文章分類

この連載でいつも使っている livedoor News コーパスのデータで文章分類を行い、各モデルのパラメータ数、推論時間、F1 スコアを比較してみました¹⁰。おまけで DistilBERT ¹¹も足しています。

データはこんな感じの TSV です。

!wc -l ./ldcc/*.tsv

# 1473 ./ldcc/dev.tsv
# 1473 ./ldcc/test.tsv
# 4421 ./ldcc/train.tsv
# 7367 total

!head -3 ./ldcc/train.tsv

# movie-enter   大島優子がここからどう破滅していくのか？　『闇金ウシジマくん』特報解禁“闇金”という
# movie-enter   インタビュー：クリスチャン・ベール「演じることができるのは役者だけ」公開当時、全米歴代2
# kaden-channel ブラックマジックデザイン、hyperdeck ssd レコーダーに タイムコード、dnxhd quicktime

以下は Poor Man’s BERT の上位4層削除での学習のコードです。他のパターンも model の初期化以外は同じ手順です。まず、事前学習済みモデルをロードします。

from transformers import BertJapaneseTokenizer, BertForSequenceClassification
tokenizer = BertJapaneseTokenizer.from_pretrained('cl-tohoku/bert-base-japanese-whole-word-masking')
model = BertForSequenceClassification.from_pretrained('./poor_mans_bert_base_top_4_removed', num_labels=9)

学習データのロードは面倒だったので一部 Sentence BERT のコードを借りてます。

MAX_SEQ_LEN = 512
BATCH_SIZE = 4
NUM_EPOCHS = 10

from sentence_transformers.readers import LabelSentenceReader
reader = LabelSentenceReader("./ldcc")

from torch.utils.data import Dataset
import torch

class SentenceLabelDataset(Dataset):

    def __init__(self, examples, tokenizer, max_length):
        inputs = []
        labels = []
        for index, example in enumerate(examples):       
          inputs.append(example.texts[0])
          labels.append(example.label)
        tokens = [tokenizer.encode(input, add_special_tokens=True, max_length=max_length, pad_to_max_length=True) for input in inputs]
        self.tokens = torch.tensor(tokens)
        self.labels = torch.tensor(labels, dtype=torch.long)

    def __getitem__(self, item):
        return [self.tokens[item], self.labels[item]]

    def __len__(self):
        return len(self.tokens)

train_ds = SentenceLabelDataset(reader.get_examples('train.tsv'), tokenizer, max_length=MAX_SEQ_LEN)
val_ds   = SentenceLabelDataset(reader.get_examples('dev.tsv'), tokenizer, max_length=MAX_SEQ_LEN)
test_ds  = SentenceLabelDataset(reader.get_examples('test.tsv'), tokenizer, max_length=MAX_SEQ_LEN)

from torch.utils.data import DataLoader
train_dl = DataLoader(train_ds, shuffle=True, batch_size=BATCH_SIZE)
val_dl   = DataLoader(val_ds, shuffle=False, batch_size=BATCH_SIZE)

ここから先は普通に PyTorch で学習ループを回す感じになります。

dataloaders = {"train": train_dl, "val": val_dl}

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5, weight_decay=0.00, betas=(0.9, 0.999), eps=1e-08)

import numpy as np
from sklearn.metrics import f1_score

for epoch in range(NUM_EPOCHS):
  for phase in ['train', 'val']:
    if phase == 'train':
      model.train()
    else:
      model.eval()

    epoch_loss = 0.0
    epoch_corrects = 0
    iteration = 1

    predictions = []
    ground_truths = []

    for batch in (dataloaders[phase]):
      inputs = batch[0].to(device)
      labels = batch[1].to(device)

      optimizer.zero_grad()

      with torch.set_grad_enabled(phase == 'train'):

        loss, logit = model(input_ids=inputs, labels=labels)                    
        _, preds = torch.max(logit, 1)

        predictions.append(preds.cpu().numpy())
        ground_truths.append(labels.data.cpu().numpy())

        if phase == 'train':
          loss.backward()
          optimizer.step()

          if (iteration % 200 == 0):  
            acc = (torch.sum(preds == labels.data)).double() / BATCH_SIZE
            print('itr {}, loss: {:.4f}, acc：{}'.format(iteration, loss.item(), acc))

          iteration += 1

        epoch_loss += loss.item() * BATCH_SIZE
        epoch_corrects += torch.sum(preds == labels.data)

  epoch_loss = epoch_loss / len(dataloaders[phase].dataset)
  epoch_acc = epoch_corrects.double() / len(dataloaders[phase].dataset)
  epoch_f1_score = f1_score(np.concatenate(np.array(ground_truths)), np.concatenate(np.array(predictions)), average='macro')
  print('epoch {}/{} {:^5}, loss: {:.4f}, acc: {:.4f}, f1: {:4f}'.
                  format(epoch+1, NUM_EPOCHS, phase, epoch_loss, epoch_acc, epoch_f1_score))

torch.cuda.empty_cache()

# itr 200, loss: 0.9814, acc：0.75
# itr 400, loss: 0.0654, acc：1.0
# itr 600, loss: 0.0349, acc：1.0
# itr 800, loss: 0.0197, acc：1.0
# itr 1000, loss: 0.8784, acc：0.75
# epoch 1/10  val , loss: 0.1936, acc: 0.9382, f1: 0.935539
# itr 200, loss: 0.0381, acc：1.0
# ...
# itr 600, loss: 0.0006, acc：1.0
# itr 800, loss: 0.0004, acc：1.0
# itr 1000, loss: 0.0004, acc：1.0
# epoch 10/10  val , loss: 0.1989, acc: 0.9654, f1: 0.962807

学習が終わればテストデータをロードして、

test_dl   = DataLoader(test_ds, shuffle=False, batch_size=BATCH_SIZE)
dataloaders['test'] = test_dl

Colaboratory なので %%time で推論時間を計測してます。

%%time
model.eval()
epoch_loss = 0.0
epoch_corrects = 0

predictions = []
ground_truths = []
phase = 'test'

for batch in (dataloaders[phase]):
  inputs = batch[0].to(device)
  labels = batch[1].to(device)

  loss, logit = model(input_ids=inputs, labels=labels)                    
  _, preds = torch.max(logit, 1)

  predictions.append(preds.cpu().numpy())
  ground_truths.append(labels.data.cpu().numpy())

  epoch_loss += loss.item() * BATCH_SIZE
  epoch_corrects += torch.sum(preds == labels.data)

epoch_loss = epoch_loss / len(dataloaders[phase].dataset)
epoch_acc = epoch_corrects.double() / len(dataloaders[phase].dataset)
epoch_f1_score = f1_score(np.concatenate(np.array(ground_truths)), np.concatenate(np.array(predictions)), average='macro')
print('{:^5}, loss: {:.4f}, acc: {:.4f}, f1: {:4f}'.
                format(phase, epoch_loss, epoch_acc, epoch_f1_score))

torch.cuda.empty_cache()

# test , loss: 0.3079, acc: 0.9369, f1: 0.931732
# CPU times: user 8.09 s, sys: 4.68 s, total: 12.8 s
# Wall time: 12.8 s

最後にパラメータはこんな感じでカウントしました。

def count_parameters(model):
    return sum(p.numel() for p in model.parameters())

それでは結果を確認していきましょう。"BERT" が標準の BERT_BASE です。"P_BERT_T4", “P_BERT_T6” はそれぞれ上位4層、6層を抜いた Poor Man’s BERT、"DistilBERT" はそのままですね。

まずは、パラメータ数の比較です。当たり前ですが抜いた層の数に比例して減ってる感じですね。上位6層抜いた ½ サイズのモデルで DistilBERT とほぼ同じサイズ感になります。

つづいて、F1 スコアの比較です。Poor Man’s BERT に関しては誤差レベルでほぼ横並び、精度が維持されています。DistilBERT はやや劣る結果となりました。が、ちゃんとチューニングしたら変わってくるかもしれません。

最後に推論時間です。こちらは DistilBERT が最速になりました。僅差で上位6層抜いた Poor Man’s BERT が続きます。

Poor Man’s BERT の論文には上位6層抜いた Poor Man’s BERT と DistilBERT が拮抗する性能であったと記述されているので、ほぼそのとおりの結果が確認できました。論文中で DistilBERT から層を抜く実験も行っていたので、さらに軽量のモデルが必要な場合はトライしてみてはいかがでしょうか。

8. おわりに

今回は公開された事前学習済みモデルを使って比較的お手軽に実験を行ってみました。次回は、趣向を変えて音声認識を試してみたいと思います。もう自然言語処理(NLP)じゃないような？まぁ、広い意味でとらえて下さい。そんなゆるーい連載なので。

https://arxiv.org/abs/1908.10084 ↩
https://arxiv.org/abs/2004.03844 ↩
https://stair-lab-cit.github.io/STAIR-captions-web/ ↩
https://arxiv.org/abs/1910.01108 ↩
https://arxiv.org/abs/1910.06188 ↩
https://arxiv.org/abs/1909.11942 ↩
https://github.com/cl-tohoku/bert-japanese ↩
https://dl.fbaipublicfiles.com/fasttext/vectors-crawl/cc.ja.300.bin.gz ↩
学習(訓練)データの anchor, positive, negative のキャプションを重複排除した上で、各キャプションを形態素解析、各単語を原形とし noabove= 0.5, nobellow = 5 フィルタして TF-IDF モデルを作成しました。 ↩
学習は全モデル同じハイパーパラメータで各1回ずつの結果なので F1 スコアや推論速度はそのつもりで見てください。 ↩
DistilBERT はバンダイナムコさんの公開モデルを利用させていただきました。 https://github.com/BandaiNamcoResearchInc/DistilBERT-base-jp/ ↩