SIF/uSIFを使ってRustで簡単高速文埋め込み

本記事は、情報検索・検索技術 Advent Calendar 2023 9日目の記事です。

SIF/uSIFという文埋め込み手法と、そのRust実装であるsif-embeddingを紹介します。最後にちょこっとベクトル検索もします。

はじめに

自然言語文の密ベクトル表現を文埋め込みと呼びます。文埋め込み同士のコサイン類似度などを使って、文同士の意味的な類似度が計算できるので、自然言語処理や情報検索などで重宝します。特に最近では、今年のAdvent Calendarでも既にいくつかの記事で書かれているように、Retrieval Augmented Generation (RAG) の検索で使用されているのを目にする機会が多いです。

このような応用範囲の拡大に伴い、品質の良い文埋め込みを得ることは重要なタスクです。近年では、Transformerを用いた手法の開発が盛んです。特に、SimCSEを代表する対照学習を用いた手法は、複数の評価タスクで高い性能を示しています。Open AIのEmbeddings APIを叩くことでも、LLMにより推論された文埋め込みが簡単に得られます。素晴らしいです。

ですがこの記事では敢えて、それ以前の主流であった静的単語埋め込みを用いた文埋め込み手法を振り返ります。

静的単語埋め込みとは、あるコーパスから既に学習された単語埋め込みを指します。学習方法としては、Word2VecやFastTextが有名です。これら静的単語埋め込み（以下、単語埋め込み）を使った最も簡単な文埋め込みの計算方法は、文に現れる各単語の単語埋め込みを平均することです。非常に単純ですが、この方法は意外と効果的に機能することが知られています。1

また、単語の出現頻度などに基づいた重み付け平均を使用することで、性能が改善することが知られています。その代表的な手法が、この記事で紹介するSIFです。

SIFは単語に着目したシンプルなモデリングに基づき設計されているため、Transformerなどの系列を扱えるモデルと比べると、文に含まれる複雑な言語的特徴を捉えるのは難しいでしょう。しかし、そのシンプルさ故に以下のような利点もあります。

  • CPUのみで動作
  • 計算が非常に高速

これらの利点は、文埋め込み手法の技術選定をする際のベースラインとして価値があると考えます。例えば、Transformerを使用するためにGPUなどの計算基盤は必須ですが、「その基盤は本当に必要か？」「妥当なレイテンシか？」などの議論を可能にします。Open AIなどのAPIを使用する場合にも「それは妥当な課金か？」といった議論ができます。

また、SIFの計算アルゴリズムには理論的な裏付けと導出があります。それらが結果の解釈や使用すべき用途の判断の手助けになるかもしれません。2

上記のような理由から、SIFは現在でも手段のひとつとして備えておくと利益のある手法ではないでしょうか。前置きが長くなりましたが、このような動機からこの記事ではSIF、及びその改善であるuSIFを紹介します。また、Rustでライブラリsif-embeddingを作成したので、そちらの紹介もします。

本記事で記述する内容は以下の通りです。

  • SIFの計算アルゴリズムと使用上の注意点
  • uSIFの概要説明
  • Rustライブラリsif-embeddingの使用方法と性能評価
  • sif-embeddingとQdrantを使ったベクトル検索

一方で、以下の内容は記述しません。

  • SIF/uSIFのモデリングと導出
  • sif-embeddingの仕様説明
  • Qdrantの使用方法

SIF

SIFはICLR 2017で発表された文埋め込み手法です。Smoothed Inverse Frequencyの頭文字を取ってSIFです。タイトルに「Simple but Tough-to-Beat Baseline」とあるように、ベースラインとしての活躍が期待されます。

openreview.net

前置きとして、この論文ではSIF-weightingとCommon Component Removal (CCR)という2つの手法を提案しています。ここで注意したいのが、これら2つの手法を一括りに"SIF"と呼ばれることもあれば、SIF-weightingだけを指して"SIF"と呼ばれることもあるという点です。

この記事では以降、一括りに2つの手法を指すときは"SIF"という用語を用いて、それぞれの手法を指すときは"SIF-weighting"と"CCR"という用語を用います。

SIF-weighting

簡単に言えば、単語埋め込みと単語の発生確率をGivenとして、文に現れる単語の埋め込みについてTFIDFのような重み付け平均を計算する手法です。

具体的には、以下のような計算式で文 $s$ の埋め込み $v_s \in \mathbb{R}^d$ を計算します。

$$ v_s := \frac{1}{|s|} \sum_{w \in s} \frac{a}{p(w) + a} v_w $$

where

  • $d \in \mathbb{N}$: 埋め込みの次元数
  • $w$: 単語
  • $s = (w_1,w_2,\dots,w_{|s|})$: 文 or 単語の系列
  • $v_w \in \mathbb{R}^d $: 単語 $w$ の埋め込み
  • $p(w) \in [0,1]$: 単語 $w$ の発生確率
  • $a \in \mathbb{R}_{> 0}$: 重みを調整するハイパーパラメータ

改めて式を見てみると、文 $s$ に現れる各単語 $w$ について $\frac{a}{p(w) + a}$ で重み付けをして、平均を計算していることがわかります。重みについて、単語の発生確率 $p(w)$ が分母に含まれていることから、観測されやすい単語ほど小さな重みとなることがわかります。すなわち、TFIDFのような重み付けをしています。

SIF-weightingはシンプルな重み付け平均です。"the"や"and"などの文脈に関係せず常に頻出する単語に対処した結果として、単語の発生確率が導入されています。そのモデリングや導出の過程は論文を参照してください。もしくは、持橋先生のスライドが非常に参考になります。

$a$ はハイパーパラメータです。論文では、実験を通して $a \in [10^{-4},10^{-3}]$ を設定するのが良いと報告しています。

Common Component Removal

Common Component Removal (CCR) は、文埋め込みから文法に関係する成分を取り除く手法です。SIF-weightingが適用された文埋め込みに対して、後処理として適用されます。

文の集合 $S = \{ s_1, s_2, \dots, s_n \}$ を入力として、各 $s_i$ についてSIF-weightingの結果として得られた文埋め込みを $v_{s_i}$ と表記します。CCRでは以下のように $v_{s_i}$ を更新します。

$$ v'_{s_i} := v_{s_i} - u u^{\top} v_{s_i} $$

ここで、$u$ は行列 $ [v_{s_1} v_{s_2} \dots v_{s_n}] $ の第一主成分です。一番目の右特異ベクトルと一致するので、Truncated SVDなどを使って計算できます。3

論文では実験を通して、このベクトル $u$ は"just", "when", "even", "one"などの構文情報に関係する単語の埋め込みとコサイン類似度で近かったことを報告しています。このような成分を、文の意味には寄与しないノイズとして取り除くのがCCRです。

アルゴリズム

結果として、SIFによる文埋め込みの計算は以下のようなアルゴリズムになります。

論文から引用

重そうな処理は、右特異ベクトルを取り出す行列計算くらいです。こちらも予め何かのコーパスから学習して置けば良いので、推論時にはスキップできます。高速に文埋め込みが計算できそうです。

使用上の注意

この記事では、SIFのモデリングと導出については解説しませんが、その文埋め込みが何を計算しているかを知っておくことは、誤用を防ぐためにも大切です。

簡単にですが、SIFのモデリングについて自分が注意しておきたい部分を紹介します。以下の2点です。

  • 単語埋め込みと文埋め込みの内積は、単語と文の相関関係を捉えているという仮定
  • 文に含まれる文脈（or 話題）は一定という仮定

1つ目は、単語埋め込みから文埋め込みを計算しているので当然のようにも見えますが、気をつけるべき点です。特に専門的な領域の文の埋め込みを計算する場合は、使用する単語埋め込みがそれに適したものかを注意して下さい。

2つ目も強い仮定なので注意です。話題が途中で変化するような長い文を埋め込む場合は、上手く機能しない可能性があります。

このようにSIFを使用する場合、一度論文に目を通し、その"お気持ち"を理解しておくことは大切でしょう。（SIFに限った話ではありませんが。。。）

uSIF

SIFの論文では、SIFを教師なしと位置づけています。しかし、ハイパーパラメータ $a$ を含んでいるので、例えばこちらの論文では弱教師ありに分類されています。

そこで、コーパスや語彙から適切なパラメータ $a$ を見積もる Unsupervised SIF (uSIF) が以下で提案されています。

aclanthology.org

この記事では具体的には解説しませんが、uSIFは以下のような利点を持ちます。

  • 適切な $a$ を見積もり、ハイパーパラメータを排除
  • 単語埋め込みの長さに起因する不都合なケースへの対処
  • CCRで第5主成分まで使用するように修正（Piecewise CCR, PCCR)4

経験的に、uSIFはSIFよりも品質の良い文埋め込みを生成します。しかし、計算が増えるため速度は低下します。

sif-embedding

sif-embeddingは、SIF/uSIFを使った文埋め込みライブラリです。Rustでササッと文埋め込みが計算できるライブラリが欲しかったので作りました。

github.com

この節では、以下のチュートリアルから重要な部分を抽出してsif-embedding (v0.6.1) の使用方法を紹介します。

sif-embedding/tutorial at v0.6.1 · kampersanda/sif-embedding · GitHub

準備

SIF/uSIFで文埋め込みを計算するためには、以下の2つの入力が必要です。

幸いにもこれらは一般に配布されているので、ダウンロードして使用することができます。sif-embeddingでは、何か特定のモデルに依存する設計にはしていませんが、有名なモデルはプラグインとして使えるように実装しています。

単語埋め込みの準備

単語埋め込みは、GloVeやFastTextなどで配布されている、訓練済みモデルをダウンロードして使用するのが簡単です。Pythonでは、Gensimなどを使って読み込むことが多いかと思います。

Rustの場合はどうするのでしょうか？調べてみたところ、以下のfinalfusionというライブラリが良さそうでした。

github.com

finalfusionは、GloVeやFastTextなどの異なるモデルを統一的に扱うためのプロジェクトのようです。GloVeやFastText形式からfinalfusion形式へのモデル変換は、以下で提供するコマンドラインツールが使用できます。READMEに従って変換してください。5

sif-embedding/finalfusion-tools at v0.6.1 · kampersanda/sif-embedding · GitHub

ここでは、glove.42B.300d.fifu というfinalfusion形式の単語埋め込みモデルが手元に準備できたとします。

Manifestファイルには、以下のように依存関係を記述します。

[dependencies]
finalfusion = "0.17.2"

そしたら、以下のように単語埋め込みを読み込むことができます。

use std::fs::File;
use std::io::BufReader;

use finalfusion::prelude::*;

let mut reader = BufReader::new(File::open("glove.42B.300d.fifu")?);
let word_embeddings = Embeddings::<VocabWrap, StorageWrap>::mmap_embeddings(&mut reader)?;

この例では、mmapを使ってモデルを読み込んでいます。単語埋め込みのモデルサイズは大きいことも多いので、mmapに対応してくれているのはありがたいです。

ユニグラム言語モデルの準備

ユニグラム言語モデルは、コーパスから単語の出現頻度を数えれば簡単に構築できますが、コーパスの準備や表記揺れなどへの対処が少々煩わしいです。

そういうとき、Pythonではwordfreqという選択肢があります。このライブラリには、多種多様なコーパスから学習された様々な言語のユニグラム言語モデルが同梱されており、関数の引数に en や ja などを指定するだけで簡単に単語の発生確率が得られます。また、数値の処理や言語依存の標準化も備わっています。

これは羨ましいです。というわけでご用意致しました。Rust移植版です。

github.com

導入は非常に簡単です。例えば、英単語のモデルが必要な場合は、features に large-en を指定すればモデルを自動でダウンロードしバイナリに埋め込んでくれます。

[dependencies]
wordfreq-model = { version = "0.2.3", features = ["large-en"] }

モデルはコンパイル時にバイナリに埋め込まれるため、以下のように引数で指定するだけで簡単にユニグラム言語モデルのインスタンスを生成できます。

use wordfreq_model::ModelKind;

let word_probs = wordfreq_model::load_wordfreq(ModelKind::LargeEn)?;

Let's 文埋め込み

単語埋め込みとユニグラム言語モデルが準備できたら、実際に文埋め込みを計算してみましょう。

sif-embeddingの features に finalfusion と wordfreq を指定することで、これらをプラグインとして使用することができます。

[dependencies.sif-embedding]
version = "0.6"
features = ["finalfusion", "wordfreq"]
default-features = false

実際に文埋め込みを計算するコードは以下になります。

use sif_embedding::SentenceEmbedder;
use sif_embedding::Sif;

// コーパス
let sentences = vec![
    "This is a sentence.",
    "This is another sentence.",
    "This is a third sentence.",
];
// 1) インスタンスの生成
let model = Sif::new(&word_embeddings, &word_probs);
// 2) SIFモデルの学習
let model = model.fit(&sentences)?;
// 3) 文埋め込みの計算
let sent_embeddings = model.embeddings(sentences)?;
println!("{:?}", sent_embeddings);
  1. 先ほど準備した単語埋め込みとユニグラム言語モデルから Sif インスタンスを生成します。
  2. コーパスから Sif モデルを学習します。ここでは、CCRで導入したベクトル $u$ を計算しています。
  3. fit で計算した $u$ を使って文埋め込みを計算します。ちなみに、fit で使用したものとは異なる文を入力しても機能します。

結果として、例えば以下のような出力が得られます。簡単ですね。

[[0.001560542, -0.009350948, 0.004585041, -0.0023774628, -0.005104985, ..., -0.0021531796, -0.00494717, -0.004704632, 0.00046236068, 0.007420417],
 [0.00095811766, 0.007716676, -0.0046481024, 0.0014976687, 0.0034232885, ..., -0.0023950879, -0.0035641952, -0.00756339, -0.0003676638, -0.0021527726],
 [-0.0028154051, -0.00013797171, 0.0011601038, 0.0005516759, 0.000922475, ..., 0.005240795, 0.009591073, 0.014395955, -1.0702759e-5, -0.004908829]], shape=[3, 300], strides=[300, 1], layout=Cc (0x5), const ndim=2

今回は入力した単語埋め込みの次元数が300だったので、出力の文埋め込みの次元数も300です。

基本的な流れは以上です。fit と embeddings の振る舞いは SentenceEmbedder で定義されており、uSIFの実装 USif でも同じように文埋め込みが計算できます。

性能評価

この節では、sif-embeddingを以下について評価します。

  1. 計算速度
  2. 得られる文埋め込みが自然言語文の意味を上手く捉えられているか

速度性能

英語Wikipedia記事からランダムに抽出して得られた100万文を使用して、文埋め込みの速度性能を評価しました。コーパスはこちらで配布されています。一文当たりの平均単語数は22.8です。使用したマシンはM2 MacBook Air (24 GB RAM) です。全てシングルコアで計測しています。

実際に使用したコードは以下にあります。

sif-embedding/benchmarks/wiki1m at v0.6.1 · kampersanda/sif-embedding · GitHub

以下の4種類の手法を評価しました。

  • SIF-weighting
  • SIF-weighting + CCR
  • uSIF-weighting
  • uSIF-weighting + PCCR

100万文を埋め込む処理を5回試行し平均しました。1秒当たりに処理できた文量についての結果を以下に示します。

Method Sentences per second
SIF-weighting 81,972 ± 44.3
SIF-weighting + CCR 69,080 ± 231.2
uSIF-weighting 28,452 ± 42.9
uSIF-weighting + PCCR 26,603 ± 167.4

当然ですが、SIF-weightingのみが最速です。1秒当たり8万文近く処理できているので、十分に高速なんじゃないでしょうか。CCRが加わると、少し低速化しますが、それでも7万文近く処理できています。uSIFは計算が増えるので一回り遅くなります。それでも3万文近く処理できているので、実用には問題ない速度でしょう。

本当は他の手法とも比較したいところなのですが、今後の予定とさせてください。

評価用データセットを使ったベンチマーク

評価用データセットを使い、「得られる文埋め込みが自然言語文の意味を上手く捉えられているか？」を評価します。英語と日本語でそれぞれ評価します。

英語

SentEval STS/SICK Taskを用いて評価します。いくつかの評価セットが存在しますが、この記事ではSTS-Benchmarkを使った結果を報告します。

詳細な実験設定と評価に使用したコード、全ての実験結果は以下にあります。

sif-embedding/evaluations/senteval at v0.6.1 · kampersanda/sif-embedding · GitHub

STSタスクでは、二文の意味的関連性を [0,5] のスコアで人手でアノテーションされたデータが与えられます。そして、文埋め込み同士のコサイン類似度とそのスコアの相関係数を、その文埋め込み手法の性能とします。

SIFとuSIFに加え、princeton-nlpで配布されている2種類のSimCSEモデルを比較対象とします。STS-Benchmarkは {train, dev, test} の3種類のデータを含みますが、今回は訓練する必要は無いので全て評価に使用します。

使用した相関係数はスピアマンの順位相関係数（×100）です。結果を以下に示します。

Model train dev test Avg.
sif_embedding::Sif 65.2 75.3 63.6 68.0
sif_embedding::USif 68.0 78.2 66.3 70.8
princeton-nlp/unsup-simcse-bert-base-uncased 76.9 81.7 76.5 78.4
princeton-nlp/sup-simcse-bert-base-uncased 83.3 86.2 84.3 84.6

やはりSimCSEは強く、SIF/uSIFを上回る結果です。10ポイント以上の差が見られます。これを大きいと捉えるか小さいと捉えるかはおまかせします。

日本語

JGLUE JSTSとJSICKタスクを用いて評価します。詳細な実験設定と評価に使用したコードは以下にあります。

sif-embedding/evaluations/japanese at v0.6.1 · kampersanda/sif-embedding · GitHub

この実験では、cl-nagoyaで配布されている日本語SimCSEモデルを比較対象とします。評価セットや評価方法の説明は英語の場合と同じです。スピアマンの順位相関係数（×100）の結果を以下に示します。

Model JSICK (test) JSTS (train) JSTS (val) Avg.
sif_embedding::Sif 79.7 67.6 74.6 74.0
sif_embedding::USif 79.7 69.3 76.0 75.0
cl-nagoya/unsup-simcse-ja-base 79.0 74.5 79.0 77.5
cl-nagoya/unsup-simcse-ja-large 79.6 77.8 81.4 79.6
cl-nagoya/sup-simcse-ja-base 82.8 77.9 80.9 80.5
cl-nagoya/sup-simcse-ja-large 83.1 79.6 83.1 81.9

やはりSimCSEは強いです。ですが英語の場合と比べて、SIF/uSIFが接近してるように見えます。英語と対応関係のある評価セットを使っている訳では無いので差を比較しても仕方ないのですが、それでもSIF/uSIFが現代でも意外と健闘する姿が見られたような気がします。

ベクトル検索

情報検索・検索技術 Advent Calendarですので、最後にQdrantを使って実際にベクトル検索を試してみようと思います。

qdrant/rust-clientとsif-embeddingを使って、Qdrantに文埋め込みを登録し、セマンティック検索をしてみます。といってもQdrandやrust-clientの使用方法の解説は、この記事の趣旨では無いのでしません。使用例のコードを以下に置いてありますので、興味のある方は見てみてください。

sif-embedding/qdrant-examples at v0.6.1 · kampersanda/sif-embedding · GitHub

ここでは、速度性能の評価に使用した英語Wikipedia記事データセットを使用します。100万文をインデックスに登録し、試しに以下の文をクエリとして検索します。

UNIX is basically a simple operating system, but you have to be a genius to understand the simplicity.

結果として得られたTop5が以下です。Scoreはクエリとのコサイン類似度です。

Rank Score Entry
1 0.732 The Monitor Call API was very much ahead of its time, like most of the operating system, and made system programming on DECsystem-10s simple and powerful.
2 0.699 It was a simple, efficient system, very effective primarily because of its simplicity.
3 0.683 Erzya has a simple five-vowel system.
4 0.680 True BASIC.“ Upon Kemeny's advice, True BASIC was not limited to a single OS or computer system.
5 0.679 An open-loop controller is often used in simple processes because of its simplicity and low cost, especially in systems where feedback is not critical.

実のところの適合率や再現率はわかりませんが、一見それっぽい結果が返って来ているように見えます。Computing SystemやSimplicityに関連した結果が得られています。

他の文埋め込みやBM25の結果とも比較したいところですが、Future Workです。

おわりに

長くなりましたが終わりです。もう少し色々と実験したいところではありますが、時間と体力切れです。ベクトル検索も申し訳程度ですみません。

もしSIF/uSIFに興味を持って貰えたなら、そしてあなたがRustユーザなら、是非sif-embeddingを使ってみて下さい。

ちなみに、PythonではfstというパッケージがSIF/uSIFを提供しています。よく整備されてるようなので、こちらも使ってみると良いと思います。

github.com

最後に、最近同僚からBERTとSIFの単語重み付けの関連性に関する研究を教えて貰いました。とても面白かったので共有します。

arxiv.org

  1. John Wieting, Mohit Bansal, Kevin Gimpel, and Karen Livescu. Towards Universal Paraphrastic Sentence Embeddings. ICLR 2016.↩
  2. 一方で、SIFのモデリングや導出が妥当なものかについても注意する必要があります。SIFに対する批評も存在しますので併せて紹介しておきます。Aidana Karipbayeva, Alena Sorokina, Zhenisbek Assylbekov. A Critique of the Smooth Inverse Frequency Sentence Embeddings. AAAI 2020.↩
  3. 行列の第k主成分とk番目の右特異ベクトルが一致する理由: Singular Value Decomposition and its applications in Principal Component Analysis | by Mohamed Afham | Towards Data Science↩
  4. ところで、CCR or PCCRで第何主成分までを使用するのか、というのはハイパーパラメータでは無いのでしょうか？この部分は、論文における平滑化項 $\beta$ に対応すると思います。自分はこの値をハイパラとして扱わなくて良い理由が実のところよくわかっていないです。↩
  5. プロジェクトページでfinalfusion形式のモデルを配布してくれているようなのですが、2023-12-04時点ではリンク切れとなっています。↩