WordLlama: 超軽量の単語埋め込みモデル

WordLlamaとは?

WordLlamaは、大規模言語モデル(LLM)の部品を再利用して作られた、効率的でコンパクトな単語表現モデルです。LLMから抽出した単語の埋め込み情報を基に、様々な自然言語処理タスクで活用できる単語のベクトル表現を生成します。従来の単語埋め込みモデル(GloVe、Word2Vecなど)よりも高性能でありながら、モデルサイズが非常に小さい点が特徴です。

WordLlamaの特長

  • 高性能: さまざまな自然言語処理タスクにおいて、従来のモデルよりも高い精度を実現。
  • コンパクト: モデルサイズが非常に小さいため、メモリ使用量を抑え、高速な処理が可能。
  • 柔軟性: 異なるタスクに合わせて、モデルのサイズや次元数を調整可能。
  • 高速性: CPUでも高速に動作するため、様々な環境で利用可能。

WordLlamaは、単語の類似度計算、文章のクラスタリング、テキスト分類など、幅広い自然言語処理タスクに適用できます。また、LLMの評価や、複雑な自然言語処理パイプラインの構築にも役立ちます。

WordLlamaの技術的な背景

WordLlamaは、LLMのトークン埋め込み情報を再利用し、シンプルなニューラルネットワークで訓練することで実現されています。マトリョーシカ表現と呼ばれる手法により、モデルのサイズを柔軟に調整することが可能です。また、バイナリ化により、モデルサイズをさらに縮小し、高速な検索を可能にする研究も進められています。

MTEBベンチマーク結果

WordLlamaは、MTEBベンチマークにおいて、従来のモデルを大きく上回る性能を示しています。特に、クラスタリング、ランキング、分類タスクにおいて高い精度が確認されています。

MetricWL64WL128WL256 (X)WL512WL1024GloVe 300dKomninosall-MiniLM-L6-v2
Clustering30.2732.2033.2533.4033.6227.7326.5742.35
Reranking50.3851.5252.0352.3252.3943.2944.7558.04
Classification53.1456.2558.2159.1359.5057.2957.6563.05
Pair Classification75.8077.5978.2278.5078.6070.9272.9482.37
STS66.2467.5367.9168.2268.2761.8562.4678.90
CQA DupStack18.7622.5424.1224.5924.8315.4716.7941.32
SummEval30.7929.9930.9929.5629.3928.8730.4930.81

WordLlamaを使ってみる

WordLlama の実装については、GitHubのリポジトリを確認してください。ここでは、簡単な使い方だけ紹介します。

まずは、pip を使ってWordLlama をインストールします。

$ pip install wordllama

インストール後は、以下のサンプルコードのように、様々な機能を簡単使用することができます。

from wordllama import WordLlama

# デフォルトのWordLlamaモデルをロード
wl = WordLlama.load()

# 2つの文の類似度を計算
similarity_score = wl.similarity("i went to the car", "i went to the pawn shop")
print(similarity_score)  # 出力: 0.06641249096796882

# クエリに対する類似度に基づいてドキュメントをランク付け
query = "i went to the car"
candidates = ["i went to the park", "i went to the shop", "i went to the truck", "i went to the vehicle"]
ranked_docs = wl.rank(query, candidates)
print(ranked_docs)
# 出力:
# [
#    ('i went to the vehicle', 0.7441646856486314),
#    ('i went to the truck', 0.2832691551894259),
#    ('i went to the shop', 0.19732814982305436),
#    ('i went to the park', 0.15101404519322253)
# ]

# その他、推論メソッド
wl.deduplicate(candidates, threshold=0.8)  # あいまいな重複を排除
wl.cluster(docs, k=5, max_iterations=100, tolerance=1e-4)  # kmeans/kmeans++初期化を使用してラベル付け
wl.filter(query, candidates, threshold=0.3)  # クエリに基づいて候補をフィルタリング
wl.topk(query, candidates, k=3)  # クエリに基づいて上位k個の文字列を返す

WordLlamaは、LLMの技術を応用して開発された、強力かつコンパクトな単語埋め込みモデルです。自然言語処理の様々な分野で活用されることが期待されます。