はじめに
こんにちは！レバレジーズデータ戦略室、データサイエンティストのJacobです。
最先端の埋め込みモデルは、精度とレイテンシの両方を意識して最適化されていることが多く、RAGシステムをはじめ様々な場面で活用されています。推論時の異なるレイテンシ要件に対応するための一般的な戦略として、学習時にMatryoshka Representation Learning（MRL）を用いる方法があります。埋め込みの先頭から段階的に大きくなるセグメントそれぞれに損失関数を適用することで、粗い特徴を前方の次元に、細かい特徴を後方の次元に学習させるという考え方です。推論時にはストレージや検索速度の要件に応じて、埋め込みを先頭k次元に切り捨てて利用できます。
MRLはsentence-transformersライブラリに実装されており、ドキュメントの例にあるように、とても簡単に使えます。
from sentence_transformers import SentenceTransformer from sentence_transformers.losses import CoSENTLoss, MatryoshkaLoss model = SentenceTransformer("microsoft/mpnet-base") base_loss = CoSENTLoss(model=model) loss = MatryoshkaLoss(model=model, loss=base_loss, matryoshka_dims=[768, 512, 256, 128, 64])
ほとんどの損失関数では、深く考えずにこれを適用しても問題ありません。この記事では、これがうまくいかない可能性のあるケースを検証してみたいと思います。
AnglE損失
AnglE損失は、文同士の微妙な違いをモデルに学習させるためのランキング損失です。
CoSENT損失と似ていますが、埋め込みを実部と虚部に分割することで最適化を容易にしています。難しく聞こえるかもしれませんが、核となるアイデアはシンプルです。コサイン類似度は2つの埋め込みがほぼ平行または正反対のときに飽和してしまうため、細かい違いを区別したい場面で勾配が消失してしまいます。複素空間での角度距離は同じようには飽和しないため、最適化がしやすくなります。
sentence-transformersの実装では、埋め込みの前半が実部、後半が虚部になるように分割されます。では、この実装をそのままMRLと組み合わせるとどうなるでしょうか？例えば、128次元のモデルをMRLサイズ（スケール）[32, 64, 128]でファインチューニングするとします。最初のセグメントでは、先頭16次元が実部、16〜32次元が虚部になります。しかし2番目のセグメントでは、先頭32次元が実部、32〜64次元が虚部になります。つまり16〜32次元は虚部になったり実部になったりと役割が入れ替わってしまい、奇妙な帰納バイアスが生まれます。
さらに、この入れ替わりには偏りがあることにも注目してください。16〜32次元が虚部になるのは32次元の損失項だけで、64次元と128次元の損失項では実部として扱われます。つまり学習シグナルの大半は、これらの次元を実部として扱っています。一般に、各プレフィックスの後半はちょうど1つの損失項でのみ虚部となり、それより大きいすべての損失項では実部となります。したがって前方の次元はほぼ実部として学習される一方、末尾の次元（この例では64〜128次元）は虚部としてしか登場しません。
簡単な修正方法であり、より筋の良いアプローチでもあるのが、偶数インデックスを実部、奇数インデックスを虚部とするインターリーブ方式です。上記の実装のうち2行を書き換えるだけなので、実装の変更は最小限で済みます。
AnglEはこのように実装した方が良いのでしょうか？それとも、細かいことを気にしすぎでしょうか？このような帰納バイアスがあっても、モデルはうまく損失地形を進んでいけるのかもしれません。迷ったときは、いつも通り試してみましょう！
実験
実験にはJSTSデータセット（JGLUEの一部）を使用します。これは、人手で0〜5の段階的な類似度が付与された文ペアのデータセットです。JGLUEのGitHubリポジトリで公開されているv1.3リリースを使用しており、こちらではテストセットを含む3つのスプリットすべてが公開されています。典型的なRAGスタイルのクエリとパッセージのペアとは少し違いますが、単純な語彙の重なりを超えた言語理解が求められるため、AnglE損失を使う良い題材になります。
モデルにはruri-v3-130mを使用し、MRL次元は[32, 64, 128, 256, 512]としました。512はこのモデルの埋め込みのフルサイズなので、最大のMRL次元では切り捨てが発生しないことに注意してください。各AnglEバリアントをMatryoshkaLossでラップし、標準的なハイパーパラメータ（バッチサイズ32、学習率2e-5、3エポック）で、各構成につき、シードを変えて5モデルずつ学習します。評価には各モデルの学習終了時点の状態を使うため、両バリアントの学習量は完全に同じです。評価では、埋め込みを各MRL次元に切り捨てて再正規化し、文ペアごとのコサイン類似度を計算します。
結果
テストセットでの結果を以下にプロットしました。評価指標は、モデルの予測と人手ラベルの間のスピアマン相関とピアソン相関（高いほど良い）で、切り捨て後の次元数（MRL次元）ごとに計算しています。各シードを小さい点で表し、平均を線で、標準偏差を帯で示しています。
インターリーブは128・256・512次元でわずかに高いスコアを示しました。差は小さい（スピアマン相関で約0.0035）ものの、一貫しています。同じシードでは両バリアントの初期化とデータ順序が同一なので、この比較は対応のある比較になっています。そして同じシード同士で比べると、この3つのサイズすべてで5シードすべてインターリーブが勝っています。64次元では差はごくわずかで、32次元ではむしろチャンク分割の方がピアソン相関でやや良い結果になりました（こちらも5シードすべて）。ただしスピアマン相関では実質的な差はありません。
この差の出方は、もう少し詳しく見る価値があります。役割の入れ替わりはプレフィックスが小さいほど大きな割合の次元に影響するため、差は最小の切り捨てレベルで最大になると予想したくなります。ところが実際には32次元と64次元では差がなく、128次元以上でほぼ一定の差が現れました。特に注目すべきは、何も切り捨てられない512次元でも差が同じ大きさで残っている点です。これは、インターリーブが切り捨て後に残るプレフィックスを守っているだけではなく、埋め込み全体の学習ダイナミクスを変えていることを示しています。
低次元側については、先ほど述べた役割割り当ての偏りがひとつの解釈を与えてくれます。16〜32次元は5つの損失項のうち4つで実部として学習されるため、チャンク分割の小さいプレフィックスはほぼ実部だけのベクトルとして学習されます。また、32次元でのチャンク分割の優位はピアソン相関にしか現れません。スピアマン相関はペアの順位のみで決まる一方、ピアソン相関は類似度の値そのものにも影響されます。つまりチャンク分割は順位付けが優れているわけではなく、類似度の値と人手スコアの関係がより線形に近いということです。この仮説の検証には類似度スコアではなく埋め込みの幾何そのものを見る必要があるため、今後の課題として面白そうです。
MRLの目的は様々な切り捨てレベルで良い性能を維持することであり、チャンク分割とインターリーブで計算量に違いはないため、この結果からはインターリーブ方式を使うのが妥当と言えそうです。もちろんこれは単一の実験であり、データセット、モデル、ハイパーパラメータによって結果は変わり得ます。今後は、実運用でMRLの恩恵を受けやすいRAGスタイルのデータセットでも検証してみたいところです。
まとめ
この記事では、MRLでAnglE損失を使う2つの方法を検証しました。2行の実装変更だけで、小さいながらも一貫した性能向上が得られました。損失関数をやみくもに試すのではなく、実装の詳細をきちんと確認することが大切だと分かります。
