TL;DR

OpenAI の Embeddings を JSTS で評価したところ、 BERT を JSTS でファインチューニングした場合には劣るものの、分散表現を得る手法としては十分な性能を持っていることが分かった。

Amazon Titan Embeddings が 2023 年 9 月 29 日に GA (一般提供) になったので追加調査したが、 JSTS スコアは OpenAI がまさった。

Cohere Embed Multilingual が Amazon Bedrock にて 2023 年 11 月 13 日に利用可能になったので追加調査したが、 JSTS スコアにおいて OpenAI (V2) を超える性能を示した。

2024 年 1 月 25 日に公開された V3 モデルを追加調査したところ、比較した埋め込み (embeddings) モデルの中で最高の性能となった。

Updates

2023/6/5: cl-tohoku/bert-base-japanese-v2 を cl-tohoku/bert-base-japanese-v3 に置き換えました。
2023/10/3: Amazon Titan Embeddings (amazon.titan-embed-text-v1) を追加しました。
2023/10/31: Amazon Titan Embeddings の評価を内積 (dot product) で行っていた問題を修正しました。お詫びいたします。
2023/11/20: Cohere Embed Multilingual (cohere.embed-multilingual-v3) を追加しました。
2024/2/6: OpenAI Embeddings の V3 を追加しました。

まえがき

ChatGPT が大いに盛り上がってますね。
実際に OpenAI の API を用いて手元で文書を生成してみた方も多いのではないでしょうか。

OpenAI はそんな ChatGPT (Chat completions) の他にも色々な API を公開していますが、その中に、 Embeddings といって、テキストをベクトルに変換する (テキストの分散表現を得る) API があります。

ベクトル化するとテキスト同士の類似度を高速で調べることができ、いろいろと応用の範囲が広いです。
(キーワード検索に対して、ベクトル検索はより意味にフォーカスした抽象的な検索ができます。一長一短ですが。)

例えば、先述の ChatGPT にはプロンプトサイズの制限 (gpt-3.5-turbo の場合で 4096 トークン = 日本語で 4000 文字程度) があるのですが、検索技術を用いてユーザの質問入力に近いコンテキストを探し、それプロンプトに含めることによって制限を超えたコンテキストからの回答出力ができるようになります。

また、本検証で利用する Embeddings 用のモデルの一つである text-embedding-ada-002 は利用料金が安く、文庫本一冊 (10 万文字程度とする) を適当な長さに切り分けてベクトル化した場合で、高々 10 円程度の出費で済みます。

OpenAI Embeddings の応用については他の記事を参照していただくことにして、表題の通り、ここではその性能を定量的に評価していきます。

評価指標

評価指標としては JGLUE に含まれる JSTS を利用します。

JSTS は二つの文がどれだけ類似しているかを 0 〜 5 でスコア付けして構築されたデータセットです。

例えば、

街中の道路を大きなバスが走っています。
道路を大きなバスが走っています。

は 4.4、

飛行場に停まっている飛行機を女の子が見ています。
女性と馬が立っている奥に建物が見えます。

は 0.4 のようにスコアがつけられています。

データセットは Hugging Face のサイトで俯瞰できます。

test セットは公開されていないので、 validation (dev) セットを使って評価を行います。

JSTS を評価に使う上での留意事項

JSTS に含まれるデータでは文書が短いので、より長い文書を扱ったときに性能の傾向が同じかどうかは分かりません。現実の分散表現の活用事例では、質問に関連する回答を探す場合など、 JSTS と異なる問題設定もよくあります。

また、あくまで日本語のベンチマークなので、 OpenAI の謳うコード検索などの用途の性能は計測されません。

ご注意ください。

結果

まず実験手法の説明を後回しにして、評価結果を示します。

総じて、 OpenAI Embeddings は、教師なし (JSTS に特化していない) かつ日本語に特化していないモデルとして十分な性能があることが分かりました。

特に V3 の large モデルは、日本語に特化した PKSHA 社の SimCSE モデルと比べても JSTS においてより高い性能を示しました。
JSTS を評価に使う上での留意事項で述べたことに加え、 PKSHA 社の SimCSE は検証 (validation) データとして JSTS の train セットを用いていることが JSTS での評価において有利に働くと思われることを考慮すると、優秀な成績です。

Cohere Embed Multilingual も JSTS スコアが高いので、強力なライバルになると考えられます。

Model	Pearson's r	Spearman's ρ
OpenAI Embeddings V3 large	0.886	0.838
OpenAI Embeddings V3 small	0.829	0.781
OpenAI Embeddings V2	0.837	0.790
- 教師なし -
PKSHA SimCSE	0.850	0.801
Word2Vec	0.625	0.607
Word2Vec (方向白色化)	0.712	0.689
fastText	0.609	0.602
Amazon Titan Embedings	0.793	0.729
Cohere Embed Multilingual	0.855	0.814
- 教師あり -
Tohoku BERT base v3	0.919 *1	0.881 *1
LUKE Japanese base	0.916 *2	0.877 *2

*1 https://github.com/cl-tohoku/bert-japanese/tree/main#model-performances から借用
*2 https://github.com/studio-ousia/luke#news から借用

実験手法

JSTS で評価する上で、教師なし学習の手法と教師あり学習の手法を用意しました。
いずれも、タスクに応じて現実的に候補に上がってくる手法であると想定しています。

用意した教師なし学習は全て、それぞれの手法で文のベクトル化を行い、コサイン類似度を計算したものをスコアとして、 JSTS のスコアとの相関係数 (Pearson's r) および順位相関係数 (Spearman's ρ) を求めています。

教師あり学習は、こちらでは再現実験を行わず、発表された相関係数および順位相関係数をそのまま利用しています。
裏取りはしていないですが、 JSTS の train セットにおけるスコアを再現するように BERT や LUKE を学習して、スコアを直接出力するモデルを作っていると考えられます。

以下、その手法を利用するための Python コードを交えながら各手法の紹介をしていきます。
掲載版の動作検証をしていないので、疑似コードとお考えください。

OpenAI Embeddings

GPT 技術を用いた文書のベクトル化サービスで、本記事の主役です。

ベクトル化した文書同士の比較にはコサイン類似度が推奨されています。

入力長が最大 8191 と長いのも特徴の一つです。

import openai

client = openai.OpenAI(api_key='(snip)')
model = 'text-embedding-ada-002'

batch = ['文書1', '文書2']
res = client.embeddings.create(model=model, input=batch)
openai_embeddings = [d.embedding for d in res.data]

ベースライン: 教師なし

PKSHA SimCSE

PKSHA 社が SimCSE という学習方法で BERT モデルをファインチューニングしたものです。 ¹
SimCSE はプリンストン大学の Gao 氏らが提案しました。
事前学習モデルとして東北大学の BERT を、学習用のデータセットとして JSNLI を用いています。

学習時にコサイン類似度を使っているため、ベクトル化した文書同士の比較にもコサイン類似度を使います。

(ところで、コサイン類似度ではなくて内積 (dot product) を使うと JSTS スコアが向上し、 0.853 / 0.805 になります。)

from sentence_transformers import SentenceTransformer

simcse = SentenceTransformer('pkshatech/simcse-ja-bert-base-clcmlp')

batch = ['文書1', '文書2']
simcse_embeddings = simcse.encode(batch)

Word2Vec

Embedding といえば Google の Mikolov 氏らの開発したコレです。
詳細は他の記事に譲ります。

Word2Vec は単語のベクトル化をする技術なので、文書をベクトル化するためには一工夫必要です。
今回は単純に全ての単語のベクトルの平均を取っています。
そんな素朴なことあるか、と思うところですが、割と有力な手段として知られています。

辞書には NEologd を用いています。

Word2Vec は今回新規に学習しました。 ²

Wikipedia の　2022年10月のダンプデータを取得
WikiExtractor でテキストを抽出
bunkai で文に切り出す
MeCab (mecab-ipadic-neologd) で分かち書き
UTF-8 (NFKC) に変換、および 10 単語未満の行を削除
gensim で Word2Vec を 300 次元で作成

import unicodedata
import numpy as np
from gensim.models import KeyedVectors
import MeCab

MECAB_DIR = '(snip)'  # mecab-config --dicdir で取得

w2v = KeyedVectors.load_word2vec_format('...', binary=True)
mecab_tagger = MeCab.Tagger(f"-O wakati -d {MECAB_DIR}/mecab-ipadic-neologd")

sentence = '文'
words = mecab_tagger.parse(unicodedata.normalize('NFKC', sentence)).split()
vector = np.sum([w2v[word] for word in words if word in w2v], axis=0)
w2v_embedding = vector

Word2Vec (方向白色化)

単に足し算する以上の一手間を加える方法はいくつかありますが、今回は東北大学の横井氏らが提案する方向白色化をやってみました。
(この他の有力な手段としては、 SIF などがあります。)

import numpy as np

def orientation_whitening(keyedvectors, word_freq):
    """KeyedVectors を方向白色化
    word_freq は単語の出現率の dict[str, float]
    dict の順番が維持されることを前提にしているので、 Python 3.7+ が必要です
    """
    X = keyedvectors.vectors

    ps = np.array([word_freq.get(w, 0) for w in keyedvectors.key_to_index)), dtype=np.float32)
    ps /= ps.sum()

    X -= np.mean([X[i] * p for i, p in zip(keyedvectors.key_to_index.values(), ps)], axis=0)
    X_norms = np.linalg.norm(X, axis=1)
    W = X * np.array([np.sqrt(p) / X_norms[i] for i, p in zip(keyedvectors.key_to_index.values(), ps)])[:, np.newaxis]
    S = W.T @ W
    L, V = np.linalg.eig(S)
    X = X @ (V @ np.diag(1 / np.sqrt(L)))

    kv = KeyedVectors(X.shape[1])
    kv.add_vectors(keyedvectors.index_to_key, X)

    return kv

white_w2v = orientation_whitening(w2v, word_freq)

前の疑似コードの w2v を white_w2v に置き換えて分散表現 (ベクトル) を取得します。

fastText

旧 Facebook の Bojanowski 氏らが提案した手法で、サブワード (単語より小さい文字のまとまり) を考慮することで、英語でいう un- や -ed などの概念を考慮できるようにした手法です。
公式で日本語のデータを配布していて簡単に実験できるので、比較に加えてみました。

一瞬調べただけだと利用した形態素解析機が分からなかったので、辞書はデフォルト設定にしています。

import unicodedata
import numpy as np
import MeCab
import fasttext
import fasttext.util

fasttext.util.download_model('ja', if_exists='ignore')
ft = fasttext.load_model('cc.ja.300.bin')

mecab_tagger = MeCab.Tagger(f"-O wakati")

sentence = '文'
words = mecab_tagger.parse(unicodedata.normalize('NFKC', sentence)).split()
vector = np.sum([ft.get_word_vector(word) for word in words], axis=0)
fasttext_embedding = vector

Amazon Titan Embeddings

2023 年 9 月 29 日に GA (一般提供) になった、 OpenAI Embeddings の類似サービスです。
AWS のネットワーク上で他のサービスと接続できるためセキュアなシステムが構築できる他、 Knowledge Bases for Amazon Bedrock で利用すればノーコードで RAG を実現できます。 ³

英語での性能は把握していないですが、 Word2Vec + 方向白色化の足し算にも劣るパフォーマンスだったので、訓練データセット中に日本語文書があまり多くないのではないかと思います。

※ 逆に、 Word2Vec もそれなりの性能ではあるので、使いものにならないということではありませんし、 JSTS に現れるものより長い文で評価するとどうなるかは未知数です。

当初、低性能な評価結果を掲載していましたが、 OpenAI Embeddings がノルム 1 に規格化されているのに対して Amazon Titan Embeddings はされておらず、コードを流用したためにコサイン類似度ではなくて内積で評価してしまっていました。お詫びして訂正いたします。

正しい評価では、 OpenAI Embeddings (V2) と Word2Vec (方向白色化) の中間くらいの性能となりました。

import json
import boto3

# NOTE: 環境変数などでアクセス可能にしているものとする
bedrock_runtime = boto3.client(
    service_name='bedrock-runtime',
    region_name='us-west-2'
)

sentence = '文'
body = json.dumps({
    'inputText': sentence,
})
response = bedrock_runtime.invoke_model(
    body=body,
    modelId='amazon.titan-embed-text-v1',
    contentType='application/json',
    accept='*/*'
)
data = json.loads(response['body'].read())
titan_embedding = data['embedding']

Cohere Embed Multilingual

Amazon Bedrock にて 2023 年 11 月 13 日に利用可能になったモデルです。
Amazon Titan Embeddings と同様、 Bedrock で利用するのに便利なはずです。

JSTS での性能は OpenAI Embeddings (V2) や PKSHA 社の SimCSE モデルを上回りましたが、 OpenAI Embeddings V3 large には及びませんでした。

import json
import boto3

# NOTE: 環境変数などでアクセス可能にしているものとする
bedrock_runtime = boto3.client(
    service_name='bedrock-runtime',
    region_name='us-west-2'
)

batch = ['文書1', '文書2']
body = json.dumps({
    'texts': batch,
    'input_type': 'search_document',
})
response = bedrock_runtime.invoke_model(
    body=body,
    modelId='cohere.embed-multilingual-v3',
    contentType='application/json',
    accept='*/*'
)
data = json.loads(response['body'].read())
cohere_embeddings = data['embeddings']

ベースライン: 教師あり

Tohoku BERT base v3

東北大学が公開している日本語 BERT モデルの定番です。

今回は手元での実験はせず、公式の JGLUE スコアを参照しました。

LUKE Japanese base

Studio Ousia 社の研究者らが提案、公開している比較的新しいモデルです。

こちらも手元での実験はせず、公式 GitHub リポジトリに掲載されているスコアを参照しました。

結論

結果の節に書きましたが、 OpenAI の Embeddings API は他の手法に対して十分な性能を持っていることが分かりました。

BERT や LUKE と比べると性能は劣りますが、これらの手法の場合、全ての文書ペアに対して都度類似度の計算をする必要があるのに対して、ベクトル化する手法の場合、一度ベクトル化してしまえば最近傍探索アルゴリズムで高速に類似文書を探すことができます (Elasitcsearch や Pinecone など)。
ベクトル化を用いて高い性能を出せていることは大きな魅力です。

JSTS においては、 V3 large モデルは比較した教師なし (埋め込み) ベースモデルの中で最高の性能となりました。

また、 SimCSE を用いる場合、おそらく自前で計算環境を準備する必要がありますが、 OpenAI では API を叩くだけで済むというのも大きな利点だと考えられます。
特に、 ChatGPT と同一の API Key を用いるため、 ChatGPT と連携する場合にシンプルな構成で済みます。

一方で、 OpenAI 社にデータを送信することになるので、プライバシーが重視される用途ではその他の手法を選択する必要があります。

Amazon Titan Embeddings 及び Cohere Embed Multilingual と比べると、Cohere の方はそれなりに性能がいいので、有力な選択肢となります。
特に、 AWS に完結したセキュアなシステムを作る場合や、 Knowledge Bases for Amazon Bedrock で使う場合などは、 Cohere Embed Multilingual を使うのがよさそうです。

総括すると、実用する上では

学習に使えるデータがあり、計算リソースが確保でき、速度も問題にならない → BERT 、 LUKE
計算リソースが準備できるが、検索に速度は必要 → SimCSE
検索に速度が必要で、プライバシーの問題はない → OpenAI Embeddings / Cohere Embed Multilingual
ChatGPT と連携する → OpenAI Embeddings
とにかく簡単な方法がいい → Word2Vec, fastText
AWS で完結したシステムを手軽に作りたい → Cohere Embed Multilingual

というような使い分けになると思います。

SimCSE には教師あり学習のものと教師なし学習のものの二種類があり、ここでは教師あり学習をしていますが、本記事の検証では、 JSTS 上では学習してないことから教師なし学習の枠で取り扱います。ややこしいですね。 ↩
日本語で学習した Word2Vec は公開されていますが、手元に学習の準備があったことと、公開されている Word2Vec が利用する MeCab の辞書のバージョンを手元と合わせることが困難なことがあって、自前で学習しました。 ↩
これらの応用を実際に試してみたことはありません。エアプです。 ↩

OpenAI の Embeddings API はイケてるのか、定量的に調べてみる [AWS の Embeddings を追加]