BitNet Text Embeddings
Abstractの概要
BitEmbedは、エンコーダの推論コストと埋め込みのストレージコストの両方を削減することを目的とした、LLMベースのテキスト埋め込みのための極低ビットフレームワークである。事前学習済みのLLMバックボーンを、三値の重み、量子化された活性化、および学習を安定化させるための追加のサブレイヤー正規化を備えたBitNetスタイルの埋め込みエンコーダに変換する。学習パイプラインは、継続的な対照的継続事前学習と、類似度分布およびアテンション関係の蒸留を通じたフル精度の教師モデルによる教師ありファインチューニングを組み合わせている。また、この手法は複数の出力精度をサポートするように埋め込みを学習するため、1つのチェックポイントでデプロイ環境に応じたストレージと品質のトレードオフを調整できる。
新規性
本論文は、一般的な言語モデリングではなく、LLMベースのテキスト埋め込みに特化した極低ビット量子化に関する初の体系的な研究と説明されるものを提示している。その特徴的な貢献は、蒸留ベースの埋め込みフレームワーク内において、低ビットバックボーン量子化と多精度出力埋め込みの学習を統合的に扱っている点にある。
成果
MMTEB（英語、v2）において、BitEmbedは検証された両方のバックボーンでフル精度の教師モデルに近い性能を維持した（Qwen3-0.6Bで67.60対67.95、Gemma3-270Mで66.10対66.71）。また、CPUのトークンスループットをQwen3-0.6Bで364.36から830.50へ、Gemma3-270Mで1181.28から2055.47へと向上させた。多精度の出力埋め込みはストレージと性能のトレードオフを提供し、8ビットおよび4ビットの埋め込みは16ビットのブランチに対してほぼ損失ゼロであり、1ビットや2ビットの埋め込みでも実用的な性能を維持したと報告されている。
論文の注目点
- BitEmbedは、LLMの埋め込みバックボーンを三値の重みと量子化された活性化へと量子化し、サブレイヤー正規化、継続的な対照的事前学習、および教師モデル主導の蒸留を用いて学習を安定化させる。
- この手法はMMTEB（英語、v2）にてQwen3-0.6BおよびGemma3-270Mを用いて評価され、フル精度の教師モデルに近い性能を維持しつつ、CPUスループットを大幅に向上させた。
- 単一の学習済みモデルで1、2、4、8、16ビットの出力埋め込みをサポートし、ベクトルのストレージコストと下流タスクにおける埋め込み品質との間の明示的なトレードオフを可能にする。