Show HN: Wordllama – LLMのトークン埋め込みでできること

(github.com/dleemiller)

1 ポイント投稿者 GN⁺ 2024-09-16 | 1件のコメント | WhatsAppで共有

WordLlama

WordLlamaは、高速で軽量なNLPツールで、あいまい重複排除、類似性、ランキング付けのようなタスクを、推論時依存を最小限に抑えて処理し、CPUハードウェア向けに最適化されている。

クイックスタート

インストール:
```
pip install wordllama
```

256次元モデルをロード:

from wordllama import WordLlama
wl = WordLlama.load()

2つの文の間の類似性を計算:

similarity_score = wl.similarity("i went to the car", "i went to the pawn shop")
print(similarity_score)  # output: 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)
# output:
# [
#  ('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）の構成要素を再利用して、効率的でコンパクトな単語表現を生成するNLPおよび単語埋め込みモデル。GloVe、Word2Vec、FastTextのようなモデルに近い。

Matryoshka Representations: 必要に応じて埋め込み次元を削減可能
低いリソース要件: 平均プーリングを使った単純なトークン参照により、CPU上で高速に動作可能
二値化: ストレートスルー推定器を使って学習したモデルは、小さな整数配列にパック可能（近日対応予定）
Numpy専用推論: 軽量でシンプル

WordLlamaはさまざまなNLPタスクに適しており、高速で持ち運びやすいサイズのため、探索的分析やユーティリティアプリケーションに有用。

MTEBの結果

Metric	WL64	WL128	WL256 (X)	WL512	WL1024	GloVe 300d	Komninos	all-MiniLM-L6-v2
Clustering	30.27	32.20	33.25	33.40	33.62	27.73	26.57	42.35
Reranking	50.38	51.52	52.03	52.32	52.39	43.29	44.75	58.04
Classification	53.14	56.25	58.21	59.13	59.50	57.29	57.65	63.05
Pair Classification	75.80	77.59	78.22	78.50	78.60	70.92	72.94	82.37
STS	66.24	67.53	67.91	68.22	68.27	61.85	62.46	78.90
CQA DupStack	18.76	22.54	24.12	24.59	24.83	15.47	16.79	41.32
SummEval	30.79	29.99	30.99	29.56	29.39	28.87	30.49	30.81

テキスト埋め込み

事前学習済み埋め込みをロードしてテキストを埋め込む方法:

from wordllama import WordLlama
wl = WordLlama.load(trunc_dim=64)
embeddings = wl.embed(["the quick brown fox jumps over the lazy dog", "and all that jazz"])
print(embeddings.shape)  # (2, 64)

二値埋め込みモデルの使用例:

wl = WordLlama.load(trunc_dim=64, binary=True)
wl.embed("I went to the car")  # output: array([[3029168427562626]], dtype=uint64)
wl = WordLlama.load(dim=1024, binary=True)
similarity_score = wl.similarity("i went to the car", "i went to the pawn shop")
print(similarity_score)  # output: 0.57421875
ranked_docs = wl.rank("i went to the car", ["van", "truck"])
wl.binary = False  # ハミング類似度の代わりにコサイン類似度を使用
wl = WordLlama.load(config="l3_supercat", dim=1024)

学習ノート

二値埋め込みモデルは高次元でより明確な改善を示し、512または1024次元が推奨される。L2 Supercatは単一のA100で12時間、バッチサイズ512で学習された。

ロードマップ

推論機能を追加作業中:
- 意味的テキスト分割
- サンプルノートブックの追加
- DSPy evaluator
- RAGパイプライン

トークン埋め込みの抽出

モデルからトークン埋め込みを抽出するには、ユーザー同意書に同意し、Hugging Face CLIでログインする必要がある。その後、次のスニペットを使用できる:

from wordllama.extract import extract_safetensors
extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")

引用

WordLlamaを研究やプロジェクトで使う場合は、次のように引用してほしい:

@software{miller2024wordllama,
  author = {Miller, D. Lee},
  title = {WordLlama: Recycled Token Embeddings from Large Language Models},
  year = {2024},
  url = {https://github.com/dleemiller/wordllama},
  version = {0.2.6}
}

ライセンス

このプロジェクトはMITライセンスの下で提供されている。

GN⁺の要約

WordLlamaは、大規模言語モデルの構成要素を再利用して、効率的でコンパクトな単語表現を生成するNLPツール。
CPU上で高速に動作し、さまざまなNLPタスクに適した「スイスアーミーナイフ」的ユーティリティとして活用できる。
二値埋め込みモデルは高次元でより明確な改善を示し、512または1024次元が推奨される。
高速で持ち運びやすいサイズのため、探索的分析やユーティリティアプリケーションに有用。

1件のコメント

GN⁺ 2024-09-16

Hacker Newsのコメント

サイズが小さいのが気に入った。SBERTの最小モデルより有利
- 技術的には古いが、性能のためのトレードオフ
- さまざまな類似性のタイプ（例: 意味的、NLI、名詞-抽象）を切り替えられる方法がほしい
- たとえば、"Freezing" と "Burning" を新聞記事の分類では似ているものとして、化学記事では逆のものとして扱いたい
- NLI埋め込みを使って因果関係を把握したい
- SBERTはサイズが大きく、複数のモデルをロードする必要があるため、リソース消費が大きい
埋め込みは多くの意味情報を捉えており、単独でも有用なタスクに使える
- CLIPモデルのテキストエンコーダ埋め込みを使ってプロンプトを強化している
- たとえば、"building" という単語が与えられたら、埋め込み行列から "concrete"、"underground" などを見つけて置き換えたり追加したりする
- 限られた実験では、ほとんどのクエリに対して高い再現率が得られた
英語以外の他言語への対応予定があるかという質問
- フランス語向けには完璧なツールになりそう
トークン自体に多くの意味内容が含まれていることを示している
埋め込みを使って Little Alchemy を解くことについての考え
ゲーム制作にとても役立つ。ありがとう
すごく良さそう。mini-lmモデルの利点について質問
- ほとんどのmtebタスクではより良いようだが、推論などでも優れているのか気になる
数年前に似た機能を使った「言語ゲーム」を書いた
- Language-games