【ITニュース解説】Sentence embeddings in JAX, after Transformers v5 dropped it

Hugging Faceの機械学習ライブラリ「Transformers」のバージョン5で、Googleが開発した数値計算ライブラリJAXのサポートが打ち切られ、主にPyTorchに注力する方針が示された。これにより、これまでJAX環境でテキストを数値に変換する「文の埋め込み」を計算していた多くの開発者は、既存のコードが使えなくなり、新たな方法を模索する必要に迫られていた。

このような状況の中、「eqx-zoo」というオープンソースライブラリが登場し、この問題の解決策を提供している。eqx-zooは、Hugging Faceで公開されている学習済みのモデル（チェックポイントと呼ぶ）を、JAX上で動作するEquinoxというライブラリのモジュールとして直接読み込むことを可能にする。これにより、PyTorch環境に依存することなく、JAXで文の埋め込みを計算できるようになった。eqx-zooは、モデルのデータ形式である「safetensors」ファイルを直接読み込むため、PyTorchをインストールする必要がない点が特徴的だ。

「文の埋め込み」とは、人間が使う言葉の文章を、コンピュータが扱える数値のベクトル（向きと大きさを持つ数の列）に変換する技術を指す。このベクトルは、文章の意味的な特徴を捉えており、似た意味の文章であれば、そのベクトルも空間上で近い位置に配置されるように設計されている。これにより、コンピュータは文章の意味を比較したり、検索したりできるようになる。

eqx-zooを使い始めるには、pip install eqx-zoo tokenizersというコマンドで必要なライブラリをインストールする。ここでtokenizersは、Hugging Faceが提供する高速なトークナイザーライブラリであり、テキストデータをコンピュータが理解できる数値IDの羅列に変換する役割を担う。例えば、「all-MiniLM-L6-v2」のような英語モデルを使って最初の埋め込みを生成する場合、まずモデルのリポジトリ（保存場所）を指定し、トークナイザーと埋め込みモデルをそれぞれ読み込む。その後、入力したい文章群をトークナイザーで数値IDと注意マスク（パディングされた部分を無視するための情報）に変換し、モデルのembedメソッドを使って埋め込みベクトルを得る。この埋め込みベクトルは単位長に正規化されているため、それらの内積（ドット積）を計算することで、文章間のコサイン類似度、つまり意味の近さを直接測ることができる。例として「The cat sits on the mat.」と「A feline rests on a rug.」という二つの文章は高い類似度を示し、猫とは関係ない「Stock markets fell today.」とは低い類似度を示す。

埋め込みモデルは、単にTransformerと呼ばれる基盤モデルだけでなく、トークンごとの出力を最終的な一つのベクトルに集約する「プーリング」という処理方法など、モデルごとに異なる「レシピ」を持っている。eqx-zooは、Hugging Faceのチェックポイントに含まれるこれらの設定（modules.jsonやプーリング設定）を正確に読み取り、モデルが学習された通りの方法で埋め込みを生成する。もしeqx-zooが未対応のプーリング方法などが指定された場合、明確なエラーメッセージを出力することで、誤った埋め込みが計算されるのを防ぐ仕組みになっている。

この文の埋め込みは、セマンティック検索のような実用的な応用にも活用できる。セマンティック検索とは、キーワードの一致だけでなく、文章の意味内容を基に情報を検索する技術のことだ。eqx-zooを使って文書の埋め込みを事前に計算しておき、クエリ（検索したい言葉）の埋め込みと比較することで、意味的に関連性の高い文書を効率的に見つけ出すことができる。JAXにはeqx.filter_jitという機能があり、これを使うと埋め込み計算のステップを事前にコンパイル（JITコンパイル）しておくことで、高速な処理が可能になる。ただし、JITコンパイルは入力データの形状（例えば文章の長さ）が変わると再実行されるため、一定の処理速度を保ちたい場合は、入力文章を固定の長さにパディングするなどの工夫が必要となる。

eqx-zooは、英語モデルだけでなく、約100言語に対応する「multilingual-e5」のような多言語モデルも同じAPIで利用できる。多言語モデルを使用する際には、「query:」や「passage:」といったプレフィックス（接頭辞）をテキストの先頭に付与することが重要だ。これは、モデルがこれらのプレフィックス付きで学習されているため、正確な埋め込みを得るために必要となる。また、Qwen3-Embeddingのような比較的新しいタイプの埋め込みモデルもサポートしている。Qwen3-Embeddingは、従来のBERTスタイルのエンコーダーベースのモデルとは異なり、チャットモデルに似たデコーダーベースの構造を持ち、因果的アテンション（各トークンが自分より前のトークンしか見ない）を使用する。そのため、埋め込みは最後のトークンの隠れ状態となる。こちらも、クエリには特定のインストラクションプロンプトを付与するという利用上の慣習がある。

eqx-zooは、計算される埋め込みが正確であることを保証するために、厳密な検証プロセスを経ている。各レイヤーの出力がHugging Face Transformersの参照実装とfloat32（単精度浮動小数点数）で一致するか、最終的な埋め込みがsentence-transformersライブラリの結果とfloat32の丸め誤差範囲内で一致するか、さらにbfloat16（半精度浮動小数点数）での動作も検証されている。また、LayerNormのような精度に影響しやすい部品についても個別にテストを行い、小さなバグも見逃さないようにしている。これにより、eqx-zooが提供する埋め込みは、高い信頼性を持って利用できる。

現在のところ、eqx-zooはまだ若いプロジェクトであり、いくつかの制約がある。テストスイートは主にCPU上で実行されており、GPUやTPUでの数値チェックは体系的に行われていない。また、MPNetベースのモデルなど、一部のアーキテクチャにはまだ対応していない。テキストのトークン化はユーザーの責任で行う必要があるが、tokenizersライブラリを使えば、sentence-transformersの動作と一致させることが可能だ。しかし、このプロジェクトは活発に開発されており、GitHubリポジトリでバグ報告や機能提案が歓迎されている。