はじめての自然言語処理

今回は文章のベクトル化を扱います。文章のベクトル化は 第9回 で扱っていますが、当時に比べてデータセット、事前学習モデル、ライブラリ等でいろいろと状況が好転しているので、改めて扱ってみることにしました。最近は大規模データセットを用いた事前学習が公開されているので、作り比べてみます。

1. はじめに

今回は sentence-transformers¹ で文章のベクトル化にチャレンジしてみます。文章をベクトル（埋め込み表現）化することで、文章間の意味合い的な比較が可能になり、類似文章検索やクラスタリングなどが可能になります。

このライブラリは 第9回 で紹介済みですが、当時のバージョンは 0.2.5.1 であり、その後に損失関数が追加されていたり、サンプルコードが充実したりとかなりの更新が入って執筆時点で 2.1.0 になっています。ついでに言うと 第9回 は結構アクセス数があるみたいなので、食いつきイイんならもう一回やりましょうかと。

また、 第9回 では学習データを画像のキャプションデータから無理やり作ったのですが、現在は日本語SNLI(JSNLI)データセット² が公開されています。今回のタスクにピッタリのデータセットなのでより良いモデルが作れそうです。

日本語の事前学習モデルも充実してきました。 第9回 当時は日本語の事前学習済みモデルというと日本語 Wikipedia で学習したものが定番だったのですが、最近は CC-100³ や mC4⁴ などより大規模な日本語データセットを用いた事前学習済みモデルが登場してきているので、それらを利用して文章ベクトル化モデルを作り比べ見ることにしましょう。

sentence-transformers は Sentence BERT の論文⁵ （＋α）を実装したものになるわけですが、Sentence BERT 関しては 第9回 で簡単に説明しているので、一読してから戻ってきてもらえると、より理解しやすくなるかと思います。

それでは今回使用する JSNLI データセットから見ていきましょう。

2. JSNLI データセット

JSNLI データセット²は SNLI(Standord Natural Language Inference) を日本語化したデータセットになります。 データとしては自動翻訳の後にクラウドソーシングや計算機でフィルタリングしたもの/してないものの二種類が提供されていますが、 今回はフィルタリングされたバージョンを使用しました。

データは以下のようにラベル、前提、仮説の TSV 形式になっており、前提と仮説は形態素毎に “ "（空白）で分割されています。

entailment      自転車 で ２ 人 の 男性 が レース で 競い ます 。       人々 は 自転車 に 乗って います 。

ラベルは前提と仮説の文の関係性について、entailment(含意), contradiction(矛盾), neutral (中立)の3種類が割り当てられています。 学習データで 533,005 件、評価データで 3,916 件と、かなりの量があります。第9回 で画像キャプションから作ったデータは 96,465 件だったので期待できそうです。

このデータを元に [anchor(基準となる文), posivive(基準と類似した文), negative(基準と類似してない文)] の triplet に加工して学習に投入するのですが、 今回は損失関数もより精度のでるものを使います。

3. Multiple Negatives Ranking Loss

第9回は損失関数に Triplet Loss を使いました。少し復習しましょう。 anchor, positive, negative を S_a, S_p, S_n として εを閾値とすると以下のような感じでしたね。

triplet の 1 サンプルだけ見て anchor と positive, anchor と negative の距離の差が ε 以上になるように学習していました。

今回使用する Multiple Negatives Ranking Loss では triplet の 1 サンプルだけではなくバッチ全体を考慮します。 バッチサイズを 3 として anchor, positive, negative をそれぞれ a_n, p_n, n_n と表現するとイメージ的には以下のような感じになります。

a₁ に注目すると自分に似た p₁ との距離を縮めて、それ以外の p_2,3, n_2,3 との距離を広げようとする動きになります。a₂, a₃ に関しても同じことをします。

計算としては図中の矢印間のコサイン類似度を計算して緑矢印を正解として Cross Entropy で分類問題を解く感じになります。

anchor と positive だけでも良いのですが、今回はデータセットに contradict のラベルが付いたペアがあります。これを negative サンプルとして使用することで単語の並びが似ていても意味合いが異なるもの⁶は類似度が低い埋め込み表現が生成されるようにしているわけです。

また、バッチサイズを大きくすることで性能向上が期待できます。正解を選び出す選択肢の数が増える訳ですから、感覚的にもわかる気がしますよね。

学習結果の検証方法も変更するので、今度はそっちを確認してみましょう。

4. ParaphraseMiningEvaluator

第9回の検証は triplet 1 サンプルで「 anchor に対して positive が negative よりも近かったら正解」という判定で accuracy を評価していました。正直、すぐに正解率 99 % とかになってしまって扱いにくいなぁと思っていた訳で。。。

今回は ParaphraseMiningEvaluator を使用します。簡単に説明すると、文章集合 C に存在する類似文章ペア集合 G を文章ベクトル化モデルを用いた類似度評価でどの程度拾い出せたか？という評価です。処理としては以下のようになります。

• 文章集合 C に含まれる全ての文章ペアで類似度を計算します。これを P とします。
• P に含まれれる文章ペアをスコアが高いペアから順に G に含まれていれば正解と評価し Average Precision と 最良のF1 スコア及びその閾値を求めます。

イメージ的には以下のような感じですね。

今回は以前やった内容の焼き直し的な感じですから、あまり説明することもないですね。 それでは、実際に文章をベクトル化するモデルを作っていきましょう。

5. 文章ベクトル化モデルの学習

学習は sentence-transformers のサンプルコード⁷ をベースに事前学習モデルやデータセットのロード処理を修正したものになります。

いつものように、記事内のコードスニペットは、特に断りがない場合は Google Colaboratory (以下、Colab)で動かす想定にしています。ノートブックを開き、アクセラレータは GPU を選んで下さい。

セットアップ

まずは sentence-transformers と MeCab 関係をインストールします。

!pip install sentence-transformers==2.0.0
!apt-get install mecab mecab-ipadic-utf8 python-mecab libmecab-dev
!pip install mecab-python3 fugashi ipadic

今回は事前学習モデルの一つとして Megagon Labs さんの ELECTRA を使います。このモデルはトークナイザに sudachitra を使用するので、 インストールしておきます。

!pip install sudachitra

ここで一旦、ランタイムを再起動しておいてください。

JSNLI データセットの準備

次に JSNLI データセットをダウンロードして展開します。

!wget https://nlp.ist.i.kyoto-u.ac.jp/DLcounter/lime.cgi?down=https://nlp.ist.i.kyoto-u.ac.jp/nl-resource/JSNLI/jsnli_1.1.zip&name=JSNLI.zip
!mv *zip* jsnli_1.1.zip
!unzip jsnli_1.1.zip

事前学習モデルのロード関数

今回使用する事前学習済みモデルは以下の4つです。

• cl-tohoku/bert-base-japanese-whole-word-masking :
  東北大学さんが日本語 Wikipedia で学習した BERT です⁸。
• rinna/japanese-roberta-base :
  りんなさんが CC-100 と日本語 Wikipedia で学習した RoBERTa です⁹。
• megagonlabs/transformers-ud-japanese-electra-base-discriminator :
  Megagon Labs さんが mC4 で学習した ELECTRA です¹⁰。
• gs://somewhere/pretrained/electra/cc100/pytorch :
  僭越ながら筆者が 第12回 の記事をベースにして CC-100 で学習した ELECTRRA です。モデルは公開してないので GCS のパスは架空のものです¹¹。

という訳で今回は自前で作った ELECTRA で参戦するので GCS の認証を通しておきます。

from google.colab import auth
auth.authenticate_user()

JSNLI データセットでの学習の起点となる事前学習モデルのロード関数です¹²。 Megagon さんの ELECTRA が ElectraSudachipyTokenizer を使用する都合上、 sentence-transformers から綺麗に読み込めなかったので、 一旦ローカルに落としてロードした上で tokenizer を強引に差しかえる小細工を入れています。

import os
import json
import tensorflow as tf
from sentence_transformers import models, SentenceTransformer 
from transformers import ElectraModel, AutoTokenizer, PretrainedConfig, BertJapaneseTokenizer
from sudachitra import ElectraSudachipyTokenizer

def load_model(model_name, max_seq_length=75):

  def load_from_local_dir(dirname):
    with open (os.path.join(dirname, "tokenizer_config.json"), "r") as f:
      tokenizer_config = json.load(f)
    tokenizer_args = {"config": PretrainedConfig(**tokenizer_config)}
    word_embedding_model = models.Transformer(dirname, tokenizer_args=tokenizer_args, max_seq_length=max_seq_length)
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='mean')
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
    return model

  # Megagon's Electra
  if model_name == "megagonlabs/transformers-ud-japanese-electra-base-discriminator":
    TEMP_MODEL_DIR = "./megagon_electra"
    DUMMY_TOKENIZER_NAME = "cl-tohoku/bert-base-japanese-whole-word-masking"
    # Save pretrained model with dummy tokenizer.
    if not os.path.exists(TEMP_MODEL_DIR):
      electra = ElectraModel.from_pretrained(model_name)
      dummy_tokenizer = AutoTokenizer.from_pretrained(DUMMY_TOKENIZER_NAME)
      electra.save_pretrained(TEMP_MODEL_DIR)
      dummy_tokenizer.save_pretrained(TEMP_MODEL_DIR)
    model = load_from_local_dir(TEMP_MODEL_DIR)
    # Replace tokenizer with correct one.
    tokenizer = ElectraSudachipyTokenizer.from_pretrained(model_name)
    model[0].tokenizer = tokenizer

  # Ours Electra
  elif model_name == "gs://somewhere/pretrained/electra/cc100/pytorch" :
    TEMP_MODEL_DIR = "./our_electra"
    TOKENIZER_NAME = "cl-tohoku/bert-base-japanese-whole-word-masking"
    print("Downloading pretrained_model from GCS to {}...".format(TEMP_MODEL_DIR))
    if not os.path.exists(TEMP_MODEL_DIR):
      tf.io.gfile.mkdir(TEMP_MODEL_DIR)
      for fname in ["config.json", "pytorch_model.bin"]:
        tf.io.gfile.copy(os.path.join(model_name, fname),
            os.path.join(TEMP_MODEL_DIR, fname), overwrite=True)
    print("Save tokenizer data files...")
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    tokenizer.save_pretrained(TEMP_MODEL_DIR)
    model = load_from_local_dir(TEMP_MODEL_DIR)

  # cl-tohoku's Bert and Rinna's RoBERTa and others..
  else:
    word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)
    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='mean')
    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
  return model

学習の実行

まずは必要なモジュールをインポートしておきます。

import math
from sentence_transformers import models, losses, datasets
from sentence_transformers import LoggingHandler, SentenceTransformer, util, InputExample
from sentence_transformers.evaluation import ParaphraseMiningEvaluator
import logging
from datetime import datetime
import sys
import os
import gzip
import csv
import random

logging.basicConfig(format='%(asctime)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    level=logging.INFO,
                    handlers=[LoggingHandler()])

学習データをロードします。

def add_to_samples(data, sent1, sent2, label):
  if sent1 not in data:
    data[sent1] = {'contradiction': set(), 'entailment': set(), 'neutral': set()}
  data[sent1][label].add(sent2)

def load_data(filename):
  data = {}
  with open(filename, "r") as f:
    lines = f.readlines()
    lines = [line.strip().split("\t") for line in lines]
    rows = [[line[0], line[1].replace(" ", ""), line[2].replace(" ", "")] for line in lines]
    for row in rows:
      label = row[0] 
      sent1 = row[1]
      sent2 = row[2]
      add_to_samples(data, sent1, sent2, label)
      add_to_samples(data, sent2, sent1, label) 

  samples = []
  for sent1, others in data.items():
    if len(others['entailment']) > 0 and len(others['contradiction']) > 0:
      samples.append(InputExample(texts=[sent1, random.choice(list(others['entailment'])), random.choice(list(others['contradiction']))]))
      samples.append(InputExample(texts=[random.choice(list(others['entailment'])), sent1, random.choice(list(others['contradiction']))]))

  dedup = {}
  for sample in samples:
    key = "".join(sample.texts)
    dedup[key] = sample 
  return list(dedup.values())

train_samples = load_data("jsnli_1.1/train_w_filtering.tsv")
len(train_samples)
# 294577

ロードしたデータは以下のような感じになります。

for sample in train_samples[:3]:
  print(sample.texts)

# ['ガレージで、壁にナイフを投げる男。', 'ガレージに男がいます。', '男が台所のテーブルで本を読んでいます。']
# ['ガレージに男がいます。', 'ガレージで、壁にナイフを投げる男。', '男が台所のテーブルで本を読んでいます。']
# ['ラップトップコンピューターを使用して机に座っている若い白人男。', '人は椅子に座っています。', '黒人はデスクトップコンピューターを使用します。']  

つづいて評価データのロード処理です。

def load_data_for_paraphrase_mining(filename):
  sentences_map = {} # id -> sent
  sentences_reverse_map = {} # sent -> id
  duplicates_list = [] # (id1, id2)

  def register(sent):
    if sent not in sentences_reverse_map:
      id = str(len(sentences_reverse_map))
      sentences_reverse_map[sent] = id
      sentences_map[id] = sent
      return id
    else:
      return sentences_reverse_map[sent]

  with open(filename, "r") as f:
    lines = f.readlines()
    lines = [line.strip().split("\t") for line in lines]
    rows = [[line[0], line[1].replace(" ", ""), line[2].replace(" ", "")] for line in lines]
    for row in rows:
      label = row[0] 
      sent1 = row[1]
      sent2 = row[2]
      ids = [register(sent) for sent in [sent1, sent2]]
      if label == "entailment":
        duplicates_list.append(tuple(ids))
  return sentences_map, duplicates_list

sentences_map, duplicates_list = load_data_for_paraphrase_mining("jsnli_1.1/dev.tsv")
len(sentences_map)
# 5809

評価データの文章数は 5809 件ですね。 そして、類似文ペアの数は以下のようになります。

len(duplicates_list)
# 1432

では、まずはりんなさんの RoBERTa で試してみましょう。他の事前学習モデルで試すときは以下のセルの # を付け替えて下さい。

#model_name = "cl-tohoku/bert-base-japanese-whole-word-masking"
#model_name = "megagonlabs/transformers-ud-japanese-electra-base-discriminator"
#model_name = "gs://somewhere/pretrained/electra/cc100/pytorch"
model_name = "rinna/japanese-roberta-base"

model_save_path = "./strf_{}".format(model_name.replace("gs://","").replace("/","_"))
model_save_path
# './strf_rinna_japanese-roberta-base'

ハイパーパラメータは以下のとおりです。

train_batch_size = 48
max_seq_length = 75
num_epochs = 1

最近の Colab は GPU アクセラレータとして割り当てられるのがもっぱら Tesla K80 ですね。今回の MultipleNegativesRankingLoss バッチサイズが大きい方が有利なのですが、 K80 ではバッチサイズ = 64 で OOM になったので、とりあえず 48 で我慢することにしました。

つづいて、事前学習モデルをロードします。

model = load_model(model_name, max_seq_length=max_seq_length)

りんなさんの RoBERTa はトークナイザが T5Tokenizer なのですが、大丈夫そうですね。

model.tokenizer.__class__.__name__
# 'T5TokenizerFast'

続いてデータローダを定義し、

train_dataloader = datasets.NoDuplicatesDataLoader(train_samples, batch_size=train_batch_size)

損失関数には前述の MultipleNegativesRankingLoss を使います。

train_loss = losses.MultipleNegativesRankingLoss(model)

検証に使うのは ParaphraseMiningEvaluator です。

dev_evaluator = ParaphraseMiningEvaluator(sentences_map, duplicates_list, name="paramin-jsnli-dev")

最初の 10 % はウォームアップです。

warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1) #10% of train data for warm-up
logging.info("Warmup-steps: {}".format(warmup_steps))
# 2021-09-28 06:16:56 - Warmup-steps: 614

あとは学習ループを回すだけです。

model.fit(train_objectives=[(train_dataloader, train_loss)],
          evaluator=dev_evaluator,
          epochs=num_epochs,
          evaluation_steps=int(len(train_dataloader)*0.1),
          warmup_steps=warmup_steps,
          output_path=model_save_path,
          use_amp=False
          )

# ...
# 2021-09-28 06:47:33 - Paraphrase Mining Evaluation on paramin-jsnli-dev dataset after epoch 0:
# 2021-09-28 06:47:42 - Number of candidate pairs: 308134
# 2021-09-28 06:47:43 - Average Precision: 10.97
# 2021-09-28 06:47:43 - Optimal threshold: 0.8148
# 2021-09-28 06:47:43 - Precision: 16.50
# 2021-09-28 06:47:43 - Recall: 31.01
# 2021-09-28 06:47:43 - F1: 21.54

同様の手順で事前学習モデルを切り替えながら4回学習してみました。どういう結果になったか見てみましょう。 あと、 Colab の人は4回学習する前に、 GPU ランタイムの寿命が尽きると思うので GCS に学習済みモデルを退避するなどして下さい。

6. 結果の比較

以下が比較結果になります。複数回試行して平均とかはしていない一発勝負になります。

結果としては、東北大さんの BERT が AP = 12.87、F1(best) = 23.61 で最良になりました。

今回比較した事前学習モデルはサイズ的には基本的に同格で事前学習コーパスのサイズがモノを言う結果になると想像していたので個人的には意外でした。 二番手がりんなさんの RoBERTa で、筆者と Megagon Labs さんの ELECTRA が少し離されて 三、四番手です。

ELECTRA の（Discriminatorの）事前学習は真偽分類¹³であって、BERT や RoBERTa のような Masked LM ではないので、今回のタスクとは相性が良くないのかもしれませんね。

りんなさんの RoBERTa が東北大さんの BERT を下回った理由はなんでしょうね。トークナイザが T5Tokenzier(Sentencepiece) なので、 pooling_mode が cls だと（トークナイズ結果の先頭に”[CLS]“が入らないので)問題ありそうですが、今回は mean なので問題ないかと思うんですけどね。あるいは東北大さんの BERT の whole-word-masking が功を奏しているのか、ちょっとそこまではわかりません。

赤の水平線は第9回で作った Sentence BERT のモデルを今回の ParaphraseMiningEvaluator で評価したスコアです。 今回の 4 モデルは学習データと検証データのドメインが一致しているので元から有利ではありますが、かなりスコアが上がりましたね。

さて Sentence Transformer のコードには色々とサンプルがついているので、もう少し試してみることにしましょう。

7. モデルの蒸留

ここからは、こちらのサンプルコード¹⁴をベースとして先ほど作ったモデルを蒸留して軽量化してみましょう。 蒸留というのは深層学習におけるモデルの軽量化手法の一つで、学習済みモデル(teacher)の振る舞いをより小規模なモデル(student)に写し取ることを言います。

具体的な学習方法はケースバイケースでアレコレあるのですが、今回は単純に「同じ文章を入力したときの teacher と student の出力ベクトルの差が小さくなる」ように student のパラメータを更新しています。

ここにたどり着くころには Colab の GPU ランタイムの寿命が尽きていると思うので、その場合はもう一度セットアップして学習データをロードするくらいのところまで動かしておいてください。あと、GCS に学習済みモデルを退避したのであれば、それも取得しておいて下さい。

それでは改めて以下をインポートします。

import torch
from torch.utils.data import DataLoader
from sentence_transformers.datasets import ParallelSentencesDataset
from sentence_transformers import models, losses, evaluation, SentenceTransformer
from sentence_transformers.evaluation import ParaphraseMiningEvaluator

teacher には先ほど成績の良かった東北大さんの BERT を起点にしたモデルを使いましょう。

teacher_model_name = "./strf_cl-tohoku_bert-base-japanese-whole-word-masking"
teacher_model = SentenceTransformer(teacher_model_name)

今回の出力先です。

output_path = "./strf_distilled_cl-tohoku_bert-base-japanese-whole-word-masking"

student を準備します。teacher より規模が小さいモデルを用意する訳ですが、 このサンプルではまっさらな初期状態から学習するのではなくて、 teacher の層を歯抜け状態にしたものを使っています。

teacher は BERT_BASE なので 12 層(0～11)のモデルですが、そこから 1, 4, 7, 10 の4層を拾いだしています¹⁵。 これでパラメータ数はほぼ 1/3 になりました。

student_model = SentenceTransformer(teacher_model_name)
auto_model = student_model._first_module().auto_model
layers_to_keep = [1, 4, 7, 10]    

new_layers = torch.nn.ModuleList([layer_module for i, layer_module in enumerate(auto_model.encoder.layer) if i in layers_to_keep])
auto_model.encoder.layer = new_layers
auto_model.config.num_hidden_layers = len(layers_to_keep)

バッチサイズは 48 にしました。

inference_batch_size = 48
train_batch_size = 48

JSNLI データセットのロード関数です。今回は単純にデータセットに含まれる全ての文章を重複排除してリストにするだけですね。

def load_sentences(filename):
  data = []
  with open(filename, "r") as f:
    lines = f.readlines()
    lines = [line.strip().split("\t") for line in lines]
    rows = [[line[0], line[1].replace(" ", ""), line[2].replace(" ", "")] for line in lines]
    for row in rows:
      label = row[0] 
      sent1 = row[1]
      sent2 = row[2]
      data.append(sent1)
      data.append(sent2)
    return list(set(data))  

train_sentences = load_sentences("jsnli_1.1/train_w_filtering.tsv")
len(train_sentences)
# 584921

データセットには ParallelSentencesDataset を使います。

train_data = ParallelSentencesDataset(student_model=student_model, teacher_model=teacher_model, batch_size=inference_batch_size, use_embedding_cache=False)
train_data.add_dataset([[sent] for sent in train_sentences], max_sentence_length=75)

データローダと損失関数です。student の出力ベクトルを teacher に近づけるだけなので MSE (Mean Squared Error) を使います。

train_dataloader = DataLoader(train_data, shuffle=True, batch_size=train_batch_size)
train_loss = losses.MSELoss(model=student_model)

検証は MSE も併用します（※以下を動かす前にこちらのセルを再実行しておいて下さい）。

dev_sentences = load_sentences("jsnli_1.1/dev.tsv")
dev_evaluator_mse = evaluation.MSEEvaluator(dev_sentences, dev_sentences, teacher_model=teacher_model)
dev_evaluator = ParaphraseMiningEvaluator(sentences_map, duplicates_list, name="paramin-jsnli-dev")

後は学習を回すだけですね。

student_model.fit(train_objectives=[(train_dataloader, train_loss)],
                  evaluator=evaluation.SequentialEvaluator([dev_evaluator, dev_evaluator_mse]),
                  epochs=1,
                  warmup_steps=1000,
                  evaluation_steps=5000,
                  output_path=output_path,
                  save_best_model=True,
                  optimizer_params={'lr': 1e-4, 'eps': 1e-6, 'correct_bias': False},
                  use_amp=False)

抜き取る層を変えて何パターンか試したところ、精度と実行速度は以下のようになりました。

精度の下落幅が sentence_transformer で示されている値¹⁶よりもかなり大きいですが、 検証に用いたタスクが異なるのでその影響もあるかもしれませんね。

次回でも使う予定なので ”./strf_distilled_cl-tohoku_bert-base-japanese-whole-word-masking" フォルダは 丸ごと GCS かどこかにバックアップしておいてもらえると良いかと思います。

あと、次元削減のサンプルも見つけたので、そっちもやっちゃいましょう。

8. 文章ベクトルの次元削減

引き続き、こちら¹⁷のサンプルコードを参考にして、文章ベクトルの次元削減をやってみます。

もちろん文章ベクトルの次元数を落とせば精度的にはマイナスなのですが、検索等で大量のベクトルを処理する必要がある場合は計算量や必要なメモリサイズを抑えるという意味で有効かと思います。

では、動かしてみましょう。次元削減には PCA が使われています。

from sklearn.decomposition import PCA
import numpy as np
import torch
from sentence_transformers import models, losses, evaluation, SentenceTransformer

次元削減の元ネタには先程、蒸留したモデル([1, 4, 7, 10] の 4層）を使いました。

model = SentenceTransformer("./strf_distilled_cl-tohoku_bert-base-japanese-whole-word-masking")

削減後の次元数は 128 にしてみましょう。

new_dimension = 128

まずは学習データの文章を 768 次元の文章ベクトルにして、

train_embeddings = model.encode(train_sentences, convert_to_numpy=True)

次に PCA を使って元の 768 次元の情報をできるだけ残しつつ 128 次元にする変換パラメータを求めます。

pca = PCA(n_components=new_dimension)
pca.fit(train_embeddings)
pca_comp = np.asarray(pca.components_)

最後に PCA で求めたパラメータを用いて全結合層(Dence)を作り、先ほどの蒸留済みモデルの末尾に追加します。 4 層の BERT とプーリング層で生成された 768 次元の文章ベクトルが、最後にこの全結合層を通ることで 128 次元になるわけですね。

dense = models.Dense(in_features=model.get_sentence_embedding_dimension(), out_features=new_dimension, bias=False, activation_function=torch.nn.Identity())
dense.linear.weight = torch.nn.Parameter(torch.tensor(pca_comp))
model.add_module('dense', dense)

もう一度精度を確認してみましょう。

dev_evaluator(model)

# 2021-10-21 08:00:14 - Paraphrase Mining Evaluation on paramin-jsnli-dev dataset:
# 
# /usr/local/lib/python3.6/dist-packages/transformers/tokenization_utils_base.py:2227: UserWarning: `max_length` is ignored when `padding`=`True`.
#   warnings.warn("`max_length` is ignored when `padding`=`True`.")
# 
# 2021-10-21 08:00:20 - Number of candidate pairs: 308373
# 2021-10-21 08:00:21 - Average Precision: 10.86
# 2021-10-21 08:00:21 - Optimal threshold: 0.8367
# 2021-10-21 08:00:21 - Precision: 16.27
# 2021-10-21 08:00:21 - Recall: 31.70
# 2021-10-21 08:00:21 - F1: 21.50
# 
# 0.10859622212309042

AP = 10.86, F1(best) = 21.50 になりました。元々の 12 層で 768 次元のモデルは AP = 12.87 だったので、-16 % 程の劣化になります。 まぁ、パラメータ数で 1/3 、文章ベクトルの次元数で 1/6 まで削っているので、こんなモノかもしれませんね。 このあたりはユースケースに応じてバランスを見ながら調整することになるでしょう。

さて、最後にもう一つおまけです。

9. おまけ

今回の MultipleNegativesRankingLoss は前述したとおり、バッチサイズが大きい方が精度的に有利になります。 前章までは Colab の GPU ランタイムで動かせることを優先して、泣く泣くバッチサイズを 48 にしていましたが、 元ネタのサンプルではバッチサイズが 128 な訳ですよ。

どうしても試してみたくなり、GCE で Tesla T4 のインスタンスを立てたんですが AMP を有効¹⁸にしても 80 くらいが一杯一杯です。 V100 とか使っても良いのですが高そうなので Tesla T4 を二枚刺しのインスタンスを使いました。

そんな訳なので以下のコードは Colab では動かせません

ただ、今回使用したバージョンの Sentence Transformer のコードは複数 GPU に対応していないっぽかったので、 モデルのロード関数を以下のように書き換えて誤魔化しました。

import os
import torch
import json
import tensorflow as tf
from typing import List, Dict, Tuple, Iterable, Type, Union, Callable, Optional
from sentence_transformers import models, SentenceTransformer 
from transformers import ElectraModel, AutoTokenizer, PretrainedConfig, BertJapaneseTokenizer

import types
import logging
import torch
import transformers
from sentence_transformers.models import Transformer

logger = logging.getLogger(__name__)
__version__ = "2.0.0.multi-gpu"

def load_model(model_name, max_seq_length=75):
  word_embedding_model = models.Transformer(model_name, max_seq_length=max_seq_length)
  pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='mean')
  sequential = torch.nn.Sequential(word_embedding_model, pooling_model)
  parallel = torch.nn.DataParallel(sequential) 
  parallel.tokenizer = word_embedding_model.tokenizer
  parallel.tokenize = word_embedding_model.tokenize
  parallel.max_seq_length = word_embedding_model.max_seq_length  
  model = SentenceTransformer(modules=[parallel])

  # Replace save method
  # copied from https://github.com/UKPLab/sentence-transformers/blob/afee883a17ab039120783fd0cffe09ea979233cf/
  #  sentence_transformers/SentenceTransformer.py#L331-L374
  def save(self, path: str, model_name: Optional[str] = None, create_model_card: bool = True):
    if path is None:
      return

    os.makedirs(path, exist_ok=True)

    logger.info("Save model to {}".format(path))
    modules_config = []

    if '__version__' not in self._model_config:
      self._model_config['__version__'] = {
        'sentence_transformers': __version__,
        'transformers': transformers.__version__,
        'pytorch': torch.__version__,
      }

    with open(os.path.join(path, 'config_sentence_transformers.json'), 'w') as fOut:
      json.dump(self._model_config, fOut, indent=2)

    for idx, name in enumerate(sequential._modules):
      module = sequential._modules[name]
      if idx == 0 and isinstance(module, Transformer):    #Save transformer model in the main folder
        model_path = path + "/"
      else:
        model_path = os.path.join(path, str(idx)+"_"+type(module).__name__)

      os.makedirs(model_path, exist_ok=True)
      module.save(model_path)
      modules_config.append({'idx': idx, 'name': name, 'path': os.path.basename(model_path), 'type': type(module).__module__})

      with open(os.path.join(path, 'modules.json'), 'w') as fOut:
        json.dump(modules_config, fOut, indent=2)

      if create_model_card:
        self._create_model_card(path, model_name)
  model.save = types.MethodType(save, model)
  return model

• 複数 GPU が使いたいなら torch.nn.DataParallel でラップすればよい。
• だけど、SentenceTransformer の fit() は使いたい。

という方針で修正を始め、修正過程で発生したエラーを泥縄的に抑え込んだら上記のようになりました。

上記の関数を使って東北大さんの BERT を事前学習済みモデルとしてロードするとこんな感じになります。

model_name = "cl-tohoku/bert-base-japanese-whole-word-masking"
model = load_model(model_name)
model
# SentenceTransformer(
#   (0): DataParallel(
#     (module): Sequential(
#       (0): Transformer({'max_seq_length': 75, 'do_lower_case': False}) with Transformer model: BertModel 
#       (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
#     )
#   )
# )

DataParallel が挟まってるので、これで複数 GPU 使えそうです（というか使えました）。

この状態で AMP を有効にしてバッチサイズ = 160 で学習したところ、バッチサイズ = 48 の場合に比べ、 AP = 14.45(+12%)、F1(best) = 25.05(+6%) となりました。確かにバッチサイズを拡大した効果がありましたね¹⁹。

save() をかなり適当に書き換えたのでロードできないんじゃないかと思っていましたが、以下のようにして学習済みモデルを普通にロードできました。

model = SentenceTransformer("./strf_2gpu_cl-tohoku_bert-base-japanese-whole-word-masking")
model

# SentenceTransformer(
#  (0): Transformer({'max_seq_length': 75, 'do_lower_case': False}) with Transformer model: BertModel 
#  (1): Pooling({'word_embedding_dimension': 768, 'pooling_mode_cls_token': False, 'pooling_mode_mean_tokens': True, 'pooling_mode_max_tokens': False, 'pooling_mode_mean_sqrt_len_tokens': False})
# )

いろいろモデルを作ったので最後にスコアを整理しておきます。

bs はバッチサイズ、dim は文章ベクトルの次元数です。層数と次元数が明記されてないものは 12 層、768 次元です。

10. おわりに

今回は、文章のベクトル化を題材にいろいろやってみました。じつは Augmented SBERT²⁰という手法もあって、そのサンプルコード²¹も付属しています。 その中の Elastic Search を使うやつを参考にして、

• JSNLI データセットで分類モデルをつくり、
• その分類モデルを使って STAIR Captions データセットの文章にラベル付け、
• ラベル付けしたデータと元々のJSNLI データセットを合わせて Stentence Transformer の学習!

って感じでやってみたんですが、精度は落ちちゃいました。。。元々のサンプルは分類モデルではなくって STS データセットで回帰モデルになっていると思うので、そのあたりに原因があるのかもしれません。分類を Soft ラベルにすると良いかとも思ったのですが、いろいろと改造しないとイケなさそうだったので諦めちゃいました。

次回はまだちゃんと考えていないのですが、今回せっかく文章をベクトル化するモデルを作ったので、これを活用して CLIP²²風なことでもやろうかなーと思っています。 Colab で動かすのであくまで“風"ですけどねー。