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

 (github.com/dleemiller)
1 ポイント 投稿者 GN⁺ 2024-09-16 | 1件のコメント | WhatsAppで共有
  • WordLlamaは、LLMのトークン埋め込みを再利用して、あいまい重複排除、類似度計算、ランキング、クラスタリング、意味ベースのテキスト分割を行う高速で軽量なNLPツールキット
  • 推論はトークン参照と平均プーリングを中心に動作し、NumPyだけで実行可能な軽量パイプラインとCPU最適化を特徴とする
  • デフォルトモデルは256次元 16MBで、Matryoshka表現により次元を削減でき、バイナリ埋め込みはHamming類似度でさらに高速な計算をサポートする
  • MTEB表ではWL64〜WL1024がGloVe 300dやKomninosより複数の指標で高得点を示し、all-MiniLM-L6-v2よりは全体的に低いスコアとなっている
  • pip install wordllamaの後にWordLlama.load()で利用でき、.key(query)はsorted、min、maxのような標準ライブラリ関数に渡せるcallableを返す

WordLlamaがすること

  • WordLlamaは、あいまい重複排除、類似度計算、ランキング、クラスタリング、意味ベースのテキスト分割といったNLPユーティリティ作業向けの軽量ツールキット
  • LLaMA 2、LLaMA 3 70Bのような最新LLMからトークン埋め込みコードブックを抽出し、GloVe・Word2Vec・FastTextに近いコンパクトな単語表現を作る
  • 推論時の依存関係が少なく、CPUハードウェアに最適化されているため、リソース制約環境でのデプロイに適している
  • 高速で小型であることから、探索的分析、LLM出力の評価器、multi-hopやagentic workflowの前処理のようなユーティリティ用途に使える

インストールと基本的な使い方

  • インストールはpipで行う
pip install wordllama
  • デフォルトの256次元モデルはWordLlama.load()で読み込む
from wordllama import WordLlama

wl = WordLlama.load()
  • .key(query)はCallable[[str], float]を返し、候補文字列をクエリとの類似度で並べ替えたり、最大値を選んだりできる
query = "Machine learning methods"
candidates = [
    "Foundations of neural science",
    "Introduction to neural networks",
    "Cooking delicious pasta at home",
    "Introduction to philosophy: logic",
]

sim_key = wl.key(query)

sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
best_candidate = max(candidates, key=sim_key)
  • 例の結果では"Introduction to neural networks"がスコア0.3414で最も高い候補となる

主な機能

  • 埋め込み生成: 単純なトークン参照と平均プーリングでテキスト埋め込みを高速に生成する
  • 類似度計算: 2つのテキスト間のcosine similarityを計算する
  • 文書ランキング: クエリと候補文書の類似度を基準に順位付けする
  • あいまい重複排除: 類似度しきい値を基準に重複テキストを除去する
  • クラスタリング: KMeansで文書をグループ化する
  • フィルタリング: クエリとの類似度が基準以上の文書だけを残す
  • Top-K検索: クエリと最も類似したK件の文書を返す
  • 意味ベースのテキスト分割: テキストを意味的にまとまりのあるチャンクに分割する
  • バイナリ埋め込み: Hamming類似度によるさらに高速な計算をサポートする
  • Matryoshka表現: 必要に応じて埋め込み次元を切り詰め、モデルサイズと性能を調整する

モデル構造と性能

  • WordLlamaは汎用埋め込みフレームワークの中でコンテキストレスな小型モデルを学習する
  • デフォルトモデルのサイズは256次元 16MB
  • READMEのMTEB表では、WL64、WL128、WL256、WL512、WL1024をGloVe 300d、Komninos、all-MiniLM-L6-v2と比較している
    • WL256はClustering 33.25、Reranking 52.03、Classification 58.21、Pair Classification 78.22、STS 67.91、CQA DupStack 24.12、SummEval 30.99を記録
    • GloVe 300dは同じ項目でそれぞれ27.73、43.29、57.29、70.92、61.85、15.47、28.87を記録
    • all-MiniLM-L6-v2はClustering 42.35、Reranking 58.04、Classification 63.05、Pair Classification 82.37、STS 78.90、CQA DupStack 41.32、SummEval 30.81を記録
  • l2_supercatはLLaMA 2 vocabularyモデル
    • LLaMA 2 70Bやphi 3 mediumなど複数モデルのコードブックから追加のspecial tokenを除去した後に連結して学習している
    • LLaMA 2 tokenizerを使う複数モデルのコードブックをまとめて連結し、学習できる
    • LLaMA 3 70Bコードブック学習に近い性能を示しつつ、vocabularyは32k対128kで4倍小さい
  • LLaMA 3ベースのモデルとしてl3_supercatも提供されている
  • 追加の結果はResultsにある

意味ベースのテキスト分割

  • .split()は長いテキストを意味的チャンクに分割する
long_text = "Your very long text goes here... " * 100
chunks = wl.split(long_text, target_size=1536)

print(list(map(len, chunks)))

# Output: [1055, 1055, 1187]
  • target_sizeは目標サイズであり、同時に最大サイズでもある
  • 分割処理はテキストの順序、文構造、可能であれば段落構造を維持しようとする
  • WordLlama埋め込みを使って、より自然な分割インデックスを見つける
  • 出力チャンクのサイズはtarget_size以下の範囲で変動することがある
  • 推奨target sizeは512〜2048文字で、デフォルト値は1536
  • より大きなチャンクが必要なら、分割後に複数のsemantic chunkをバッチで束ねる方法が推奨される
  • 詳細はtechnical overviewにある

Model2Vecと直接推論

  • 2025-01-04アップデートでModel2Vec static embeddingsのサポートが追加された
  • WordLlama.load_m2v()でModel2Vecモデルを読み込める
wl = WordLlama.list_configs()

wl = WordLlama.load_m2v("potion_base_8m")  # 256-dim model
wl = WordLlama.load_m2v("m2v_multilingual")  # multilingual model
  • Model2VecはPCAを使ってstatic embeddingを作る別の方式
  • Model2Vec側ではmultilingual modelとgloveベースのモデルを作成しており、word similarity taskで良いスコアを出しているとされる
  • Hugging Faceのminishlabで確認できる
  • WordLlamaInferenceはローダーの代わりに、(n_vocab, dim)形式のstatic embedding配列とtokenizerを直接渡して使うこともできる
from wordllama import WordLlamaInference
from tokenizers import Tokenizer

tokenizer = Tokenizer.from_pretrained(...)
wl = WordLlamaInference(np_embeddings_ar, tokenizer)

学習と埋め込み抽出

  • バイナリ埋め込みモデルは高次元で改善幅がより顕著で、バイナリ埋め込みには512または1024次元が推奨される
  • L2 Supercatモデルは単一のA100 GPUでbatch size 512、12時間で学習された
  • LLaMAモデルからトークン埋め込みを抽出するには、ユーザー契約に同意し、Hugging Face CLIにログインする必要がある
from wordllama.extract.extract_safetensors import extract_safetensors

extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")
  • 埋め込みは通常最初のsafetensorsファイルにあるが、常にそうとは限らない
    • manifestがある場合もある
    • 自分で確認して見つける必要がある場合もある
  • 学習にはリポジトリのスクリプトを使い、既存設定をコピーまたは修正してconfiguration fileを追加する必要がある
pip install wordllama[train]
python train.py train --config your_new_config
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
  • 保存ステップではMatryoshka次元ごとにモデルを1つずつ保存する

アップデート、ロードマップ、ライセンス

  • 2025-02-01アップデートで、sorted、min、maxのような標準ライブラリ関数で使えるcallableのサポートが追加された
  • 2024-10-04アップデートでsemantic splitting inference algorithmが追加された
  • ロードマップにはDSPy evaluatorとRetrieval-Augmented Generation、すなわちRAG pipelineのサンプルノートブック追加が含まれる
  • コミュニティプロジェクトとしてGradio Demo HF SpaceとCPU-ish RAGがある
  • プロジェクトライセンスはMIT License

関連記事

1件のコメント

 
GN⁺ 2024-09-16
Hacker News のコメント

  • 小さいサイズが本当に気に入っています。すでに SBERT の最小モデルよりも利点があります。
    ただ技術的にはかなり古いアプローチに見えますし、性能とのトレードオフだという点は理解しています。それでも、意味的類似度、自然言語推論（NLI）、名詞の抽象化のような 類似度タイプの切り替えを提供できるのか気になります。
    たとえば新聞記事を「極端な環境事象」のようなカテゴリにまとめる場合、「Freezing」と「Burning」が非常に似ていると出てほしいです。MTEB/Sentence-Similarity や古典的な Word2Vec/GloVe のように動作するケースです。しかし化学記事なら、この2つはほぼ反対に出るべきですし、自然言語推論の埋め込みで2つの物事の因果関係を見たい場合もあります。
    後者2つの埋め込みタイプは2019年以降の比較的新しい方式なので、技術的な機会はより大きいと思います。古い MTEB/意味的類似度系は2014年から多くの用途に十分で、2019年に mini-lm-v2 などで大きく改善されました。
    上の3種類の埋め込みタイプは SBERT でも可能ですが、次元数が大きくモデルも大きいため、タイプごとに複数モデルをロードするとリソース負担が大きくなります。生成型埋め込みモデルや E5、自然言語推論モデルは大きく、しばしば6GB程度が必要です。

    • 良いアイデアです。いくつか実験して実現可能性を確認してみます。
      単一の類似度タイプで学習したときに性能がどうなるか見てみたいです。文脈計算なしでこれを処理する別の方法があるかは確信がありません。モデルを切り替える必要があるかもしれませんが、それ自体は大きな問題ではありません。
    • これは 17MB のモデルで、ベンチマーク上では MiniLM v2、つまり SBERT より当然低く出ます。私は23MBのモデルで ONNX 上の V3 をほぼすべてのプラットフォームで動かしています。
      けなす意図ではなく、こうした取り組みを文脈の中で理解することが重要だということです。ここでは、LLM を深く理解していくうちに LLM にも埋め込みがあることに気づき、その観点では埋め込み分野全体の現状をあらためて概観するよりも、その埋め込みを触って一歩進めるほうが自然だ、という文脈です。
    • 「ChatGPT 埋め込み」が OpenAI の埋め込みモデルを意味するなら、「burning」と「freezing」はまったく反対ではありません。text-embedding-large-3 の1024次元で試すと、コサイン類似度は0.46程度です。完全に反対の埋め込みなら、類似度は -1 であるべきです。
      反対の意味の単語は反対の埋め込みを持つ、と考えるのはよくある誤解です。実際には、反対の意味の単語同士にも共通点は多くあります。「burning」と「freezing」はどちらも温度や物理に関係し、英単語であり、動詞・名詞・形容詞としていずれも使え、綴りも正しいです。こうした特徴がすべて埋め込みに入ります。

  • 埋め込みは学習データと目的関数に応じて多くの 意味情報を含み、さまざまな有用なタスクに独立して活用できます。
    以前、CLIP モデルのテキストエンコーダの埋め込みを使って、プロンプトが対応する画像とよりよく合うように補強したことがあります。たとえばプロンプトに「building」があれば、埋め込み行列から「concrete」「underground」のような最近傍を見つけ、その単語の後に置換または追加しました。限定的な実験では、ほとんどのクエリで 再現率が上がりました。

    • その通りです。こうした ドメイン内の文脈関係を埋め込みモデルに学習させることができます。
      https://www.marqo.ai/blog/generalized-contrastive-learning-f...
    • 本当に素晴らしいアイデアです。この実装でもできそうなので、さらに考えてみます。
      wordllama でトークン埋め込みの大きさを見ると、補強すべき重要なトークンを特定するのにも役立ちそうです。ただし、この作業に合わせて選別したデータで学習すれば、はるかにうまく動く可能性があります。

  • 英語以外の言語の計画もあるのか気になります。フランス語には完璧なツールになりそうです。

    • 十分に可能です。学習用コーパスを構成する必要がありますが、フランス語にどんな資料があるのかはよく知りません。
      Mistral 系モデルで少し学習を試したことがあるので、フランス語コーパスではまずそちらを試すと思います。
      Issue を開いてくれれば、時間があるときに取り組んでみます。

  • 大きなコーパス、たとえば1万以上の文で各文をドキュメントとして扱う用途なら、TF-IDF の疎行列ベクトルを k-means でクラスタリングしても似た結果が得られます。
    ただしこのツールには、二値化のような方式で k-means 部分をさらに高速化するユーティリティがかなりあるようです。今後数週間でベンチマークしてみるつもりです。

  • 数年前に似た関数を使う 言語ゲーム集を作ったことがあります: https://github.com/Hellisotherpeople/Language-games

  • 埋め込みで Little Alchemy を解いてみようと考えた人がいるのか気になります。#sample-use

  • 良さそうです。mini-lm モデルに対して利点があるのか気になります。ほとんどの MTEB タスクでは mini-lm のほうが良さそうですが、推論速度などの面で優れている部分があるのか知りたいです。

    • Mini-lm のほうが優れた埋め込みモデルです。このモデルは アテンション計算を行わず、学習後はディープラーニングフレームワークも使いません。そのため、Transformer モデルの文脈的な利点は得られません。
      最新の最高性能モデルを目指したものでもありません。依存関係、サイズ、ハードウェア要件を下げ、速度を上げるために、かなり制約を設けたモデルです。
      単語埋め込みモデルとして見てもかなり軽量なほうです。通常その種のモデルはもっと大きな語彙を持ち、数GB規模であることが多いです。
    • モデル自体のサイズ差に見えます。より軽くて速いです。mini-lm は 80MBで、ここでの最小モデルは16MBです。

  • ゲーム制作にとても有用そうです。

  • トークン自体にどれほど多くの 意味内容が入っているかをよく示しています。

  • PostgreSQL 拡張として作れるでしょうか？