今回の論文
今回取り上げるのは、Aditya Kusupati、Gantavya Bhatt、Aniket Rege らによる論文「Matryoshka Representation Learning」です。公開元は arXiv で、研究分野は表現学習、埋め込み最適化、大規模検索と分類です。URL は https://doi.org/10.48550/arXiv.2205.13147 です。
この論文を選んだ理由は、埋め込みを「1回作ったら固定長で使うもの」と見なさず、「同じベクトルを状況に応じて短くも長くも使うもの」と捉え直しているからです。RAG、ベクトル検索、推薦、分類、エージェントのメモリ検索など、埋め込みを多用する実装ではそのまま応用しやすい考え方です。
どんな技術か
Matryoshka Representation Learning は、1本の高次元埋め込みの中に、低次元版から高次元版までを入れ子構造で同居させる学習手法です。
普通の埋め込みは、たとえば 768 次元なら 768 次元として固定で使います。しかし MRL では、「先頭 64 次元だけ使ってもそれなりに意味が通る」「先頭 128 次元を使えばもう少し精密になる」「全部使えば最高精度になる」というように、同じベクトルの prefix をそのまま複数の計算予算に対応させます。
イメージとしては、1つの埋め込みの中に粗い意味表現から細かい意味表現までを順番に詰め込む設計です。これにより、軽い検索では短い埋め込みを使い、重要な候補だけ長い埋め込みで再評価する、といった段階的な使い方がしやすくなります。
課題
この技術が解決しようとしているのは、固定長埋め込みが大規模運用では無駄を生みやすい、という課題です。
何が難しいのかというと、埋め込みは一度学習すると、検索でも分類でも再ランキングでも同じ次元数で使われがちだからです。ですが実際には、すべての問い合わせやデータに常に最大次元が必要とは限りません。簡単なケースなら短いベクトルでも十分ですし、難しいケースだけ細かい表現が必要です。
既存手法にも限界があります。低次元モデルを別途いくつも学習すれば多段階利用はできますが、モデル管理が増えますし、データベース全体を次元ごとに再埋め込みする負担も大きくなります。後処理で PCA や SVD などの圧縮をかける方法もありますが、重要情報がどこに残るかを学習段階で保証していないため、精度が落ちやすいです。
なぜこの課題を解く必要があるのかというと、ベクトル検索や大規模分類では「埋め込みを作るコスト」より「大量の埋め込みを保存し、比較し、探索するコスト」のほうが支配的になる場面が多いからです。特に RAG の文書インデックス、レコメンド候補集合、画像検索、商品検索では、次元数がそのままメモリ使用量や検索レイテンシに効いてきます。
実際の AI システムでは、まず粗く候補を絞ってから細かく見る二段階検索、モバイルやエッジでの軽量推論、低コストな大規模インデックス運用などで問題になります。MRL はこの「最初から最後まで同じ重さで計算する」非効率を減らす技術です。
用語解説
- 埋め込み（Embedding）
- テキスト、画像、アイテムなどをベクトルへ変換した表現です。MRL では、この埋め込みを固定長の1枚岩として扱わず、先頭部分だけでも意味を持つように学習します。
- Prefix Subvector
- 埋め込みベクトルの先頭から一部だけ切り出した部分ベクトルです。MRL の中核は、この prefix が単なる切り捨てではなく、独立した低次元表現として使えるようにする点にあります。
- Approximate Nearest Neighbor Search（ANNS）
- 巨大なベクトル集合から近い候補を高速に探す検索技術です。MRL は ANNS を置き換えるのではなく、候補絞り込み段階で使う埋め込み自体を軽くできるので、検索基盤と相性がよいです。
- Adaptive Retrieval
- 全候補に常に同じ計算量をかけず、粗い表現で候補を絞ってから精密に再評価する検索方式です。MRL はまさにこの段階的検索をやりやすくするための表現学習です。
- Contrastive / Representation Learning
- 似たデータを近く、異なるデータを遠くに配置するような表現学習の枠組みです。MRL は特定のモデル種別に閉じず、分類学習、対照学習、言語モデル学習にも組み込みやすいのが特徴です。
技術の仕組み
MRL の本質は、「重要な情報をベクトル全体に均一に散らす」のではなく、「小さい prefix ほど粗いが使える意味を持ち、次元を足すほど精密になる」ように明示的に学習することです。
基本アイデア
通常の表現学習では、最終的な 1 本の d 次元ベクトルに対して 1 つの損失だけをかけます。そのため、情報がベクトル全体へ拡散しやすく、後から先頭 64 次元だけを使っても性能が出ないことが多いです。
MRL では、1 本の d 次元ベクトル z に対して、[z_1:64]、[z_1:128]、[z_1:256]、[z_1:d] のような複数の prefix を同時に最適化します。つまり、「短い埋め込みでも仕事ができること」を学習時点で目的関数に入れます。
この設計により、同じベクトルが coarse-to-fine な多粒度表現になります。短い prefix は雑だが軽く、長い prefix は重いが精密、という役割分担が1回の学習で埋め込みの中に作られます。
モデル構造
MRL は新しい巨大モデルを提案しているわけではなく、既存の表現学習パイプラインに追加する学習方式です。論文では ResNet50、ViT、ALIGN、BERT などに適用しています。
たとえば分類設定では、エンコーダが高次元埋め込みを出したあと、その先頭 8、16、32、64 といった複数の prefix ごとに線形分類器を持たせます。そして各 prefix がそれぞれ正しく分類できるように損失を足し合わせます。
論文では、各次元専用の分類ヘッドを持つ通常版に加えて、分類器の重みを共有する Efficient MRL も扱っています。これは特にラベル空間が大きいときにヘッド側のメモリコストを抑える工夫です。
学習方法
学習のポイントは、複数の prefix に対して元のタスク損失を同時にかけることです。分類なら各 prefix に softmax cross-entropy をかけ、全体の損失として合計します。
重要なのは、MRL が低次元モデルを別々に学習しているわけではないことです。1つのエンコーダが出す最終ベクトルに対して、複数の切り出し長さを同時に良くするよう圧力をかけています。そのため追加学習コストは比較的小さく、推論時には新しい計算を増やさず prefix を切り出すだけで使えます。
論文では、次元の選び方として半分ずつ減らすような対数的な粒度が有効だと報告しています。かなり小さすぎる初期次元を入れると情報ボトルネックが強すぎて性能が落ちるため、実用上は「十分に粗いが壊れすぎていない」起点を選ぶのが大事です。
推論方法
推論時はシンプルです。学習済みエンコーダから通常どおり高次元埋め込みを1回だけ計算し、用途に応じて先頭何次元まで使うかを選びます。
たとえば検索なら、まず 64 次元で全件に粗い近傍探索をかけ、上位候補だけ 256 次元や 512 次元で再計算する運用ができます。分類なら、低次元 prefix の予測確率が十分高ければその時点で確定し、不確かな入力だけ次の次元へ進めるカスケード方式が使えます。
ここで効いているのは、低次元版を得るために別モデルの forward を回す必要がないことです。1回の埋め込み計算のあと、使う長さだけ切り替えればよいので、実運用に乗せやすいです。
データの扱い方
MRL はデータ形式に強く依存しません。論文では画像分類、画像検索、画像と言語の ALIGN、言語の BERT にまで広げています。つまり本質は「何のデータか」より「埋め込みをどう多粒度化するか」にあります。
この性質は実務でも重要です。社内文書の埋め込み、コード検索、FAQ、商品ベクトル、ユーザー表現など、ベクトル空間を使う場面ならかなり横展開しやすいです。
重要な工夫
MRL の重要な工夫は、最終ベクトルを単に圧縮しないことです。最初から「先頭側が重要情報を持つ」よう学習しているので、あとから切っても性能が落ちにくいです。
また、論文では最適化対象に含めていない中間次元でも、性能がなめらかに補間される傾向があると報告しています。つまり 64、128、256 だけを直接学習しても、その間の 96 や 192 といった次元でもある程度使える可能性があります。これは運用時のチューニング幅を広げる要素です。
実験と結果
論文では、MRL が本当に低次元でも使える表現を作れるのか、そしてそれが検索や分類の実運用効率につながるのかを検証しています。
何を検証したのか
主な検証点は 4 つあります。1 つ目は、各 prefix が独立に学習した低次元モデルに匹敵するかです。2 つ目は、段階的な分類や検索で実際に計算量を減らせるかです。3 つ目は、画像だけでなく web-scale の視覚・言語モデルや BERT にも拡張できるかです。4 つ目は、単なる圧縮や slimmable network と比べて優位かです。
どんなデータセットや評価指標を使ったのか
画像分類では ImageNet-1K、検索では ImageNet-1K と論文内で構成した ImageNet-4K を使っています。評価指標は分類精度、1-NN 精度、検索の mAP@10、Precision@k、そして検索時間です。加えて、JFT-300M で学習した ViT や ALIGN、BERT への適用も報告しています。
低次元でも独立学習モデルに近い性能
論文では、ResNet50 を使った ImageNet-1K の実験で、MRL の prefix 表現が各次元において independently trained な fixed-feature モデルと同等以上の精度を出せると示しています。特に低次元側で fixed-feature ベースラインより良いケースがあり、単なる切り詰めではなく「低次元として使えるように学習されている」ことが見えます。
さらに 1-NN 精度でも、MRL は各次元で強く、下流タスクに使う表現品質が保たれていることを確認しています。これは検索やクラスタリングへ転用しやすいことを示す結果です。
段階的分類で埋め込みサイズを大幅削減
分類では Adaptive Classification を試しており、低次元 prefix の予測が十分 confident ならそこで止め、難しい入力だけ高次元へ進めるカスケードを組んでいます。論文では、この方法により ImageNet-1K で同等精度を保ちながら、必要な期待埋め込み次元を大きく減らせたと報告しています。
要旨では、同じ精度で最大 14 倍小さい埋め込みサイズを達成したとされています。これは単にモデルを軽くしたというより、簡単な入力に毎回フル次元を使わなくてよいことを示しています。
段階的検索で大規模探索を高速化
検索では、短い prefix で候補を shortlist し、その候補だけ長い prefix で再ランキングする Adaptive Retrieval を評価しています。論文要旨では、ImageNet-1K と 4K の大規模検索で、理論 FLOPs ベースでは最大 128 倍、実測 wall-clock では最大 14 倍の高速化を報告しています。
重要なのは、ただ速いだけではなく、検索精度が single-shot の高次元検索に近いことです。つまり MRL は「最初は粗く、最後だけ細かく」という検索の定石を、追加の低次元モデルなしで実現しやすくしています。
Web-scale と他モダリティへの拡張
論文では、JFT-300M で学習した ViT、ALIGN の視覚言語設定、そして BERT にまで MRL を拡張しています。これにより、MRL が特定の画像分類専用トリックではなく、埋め込みを扱う幅広い基盤モデルに組み込みうることが示されています。
結果から何が言えるのか
結果から言えるのは、埋め込みの情報量は必ずしもベクトル全体へ均等に散っている必要はない、ということです。学習目標を工夫すれば、先頭側に coarse-to-fine に整理して詰め込めます。
もう1つ大きいのは、検索システムの効率化を index 側の工夫だけでなく、表現学習の段階で取りに行けることです。ANN、量子化、再ランキングといった既存技術と競合するというより、MRL はその前段で使う埋め込みをより運用しやすくする補完技術と見たほうが実務的です。
何に使える？
MRL は、埋め込みを大量に保存して比較するシステム全般に応用できます。特に「精度を落としすぎずにベクトルを軽くしたい」場面で効きます。
RAG の文書検索コスト削減
RAG では、全文書を高次元埋め込みで持つとストレージも検索コストも重くなります。MRL 的な埋め込みなら、まず短い prefix で候補を絞り、上位だけ長い prefix で精査する構成が取りやすいです。長い文書コーパスを抱えるナレッジ検索で特に有効です。
ベクトルDBの階層検索
ベクトル DB では ANN index を組んでも、次元数そのものはメモリ使用量に効きます。MRL を使えば、軽量 index 用の prefix と精密再評価用の full vector を1本の表現で兼用できます。別次元モデルを2系統持たなくてよいので、運用も整理しやすいです。
エージェントのメモリ検索
AI エージェントでは、会話履歴、観測ログ、ツール実行結果などをベクトル化して保存することがあります。MRL 的な表現なら、日常的な軽い想起は短い prefix、重要な再計画時だけ長い prefix、という使い分けが考えられます。長期メモリが肥大化する設計で効きそうです。
推薦や候補生成
推薦システムでも、まず大量候補から粗く引いて、その後に精密ランキングする二段階構成が一般的です。MRL はその候補生成段階を軽くしつつ、必要なら同じ表現を使って精密化できるので、候補ベクトルの再学習や二重管理を減らせます。
モバイル・エッジ向けの軽量利用
端末側で高次元ベクトルを常に扱うのが難しい場面では、短い prefix だけ送る運用も考えられます。論文の主たる検証対象ではありませんが、同じベクトルを予算に応じて切り替えられる性質は、通信量やメモリ制約のある環境とも相性があります。これは論文結果から見た実装上の応用推測です。
開発や事業へのヒント
この論文から得られるヒントは、検索や推薦の効率化を「後段の index 最適化だけの問題」にしないことです。埋め込みそのものの設計で、かなり運用性が変わります。
1本の表現で複数の予算帯をカバーする
自分で AI アプリを作るなら、低コスト版と高精度版の埋め込みモデルを別々に持つ前に、1本の表現でどこまで兼用できるかを考える価値があります。MRL の発想は、無料プランと有料プラン、オンライン処理とバッチ処理、プレビュー検索と本検索のような複数レイヤー設計に向いています。
二段階検索を前提に設計する
RAG や社内検索を作るとき、最初から「全件にフル精度をかける」設計にするとすぐ重くなります。MRL のような多粒度埋め込みを使えば、粗検索と精密検索を自然に組み合わせやすくなります。これは検索品質だけでなく、インフラ費用の抑制にもつながります。
埋め込みの prefix を API 設計に持ち込む
小規模プロダクトでも、埋め込み API を「常に full vector を返す」前提にしなくてよいかもしれません。たとえば prefix 長を引数で切り替えられる設計や、保存時は full、検索時は short を既定にする設計は、MRL の考え方と相性がよいです。
今後注目すべき方向性
今後は、Matryoshka 系の考え方がテキスト埋め込み、マルチモーダル埋め込み、推薦、メモリ圧縮へ広がっていくはずです。特に「埋め込みの次元数を固定仕様にしない」設計は、RAG やエージェントの運用コストが上がるほど重要になります。
限界
MRL にも限界はあります。まず、表現を coarse-to-fine に整理して入れる制約を学習時にかけるため、タスクによっては通常の高次元表現より最終精度がわずかに不利になる可能性があります。論文では全体として良好ですが、あらゆるデータで常に無損失とは言い切れません。
また、どの prefix 長を最適化対象にするかは設計問題です。極端に短い次元を最初から含めると情報ボトルネックが強すぎますし、粒度の取り方を誤ると使い勝手が落ちます。ここはデータ規模、検索方式、欲しい遅延帯に合わせた調整が必要です。
実装面では、既存の埋め込みモデルへそのまま差し込めるとはいえ、再学習は必要です。すでに本番運用中の埋め込み基盤に後付けで導入するなら、全文書の再埋め込み、検索閾値の再調整、index 設計の見直しまで含めて考える必要があります。
実運用では、prefix を短くしすぎると coarse な意味しか残らず、似ているが区別すべき候補を落とす危険があります。特に法務文書、コード検索、製品仕様のように細部差分が重要な検索では、shortlist 段階の recall を丁寧に評価しないと事故になります。
さらに、論文は主に画像分類と画像検索を中心に検証しており、現代の大規模テキスト埋め込み API や本番 RAG スタックでの最適値までは直接示していません。そのためテキスト検索へ適用するときは、考え方は有効でも具体的な次元設計は別途検証が必要です。
よくある質問
Q. MRL は埋め込み圧縮と何が違うのですか？
A. 大きな違いは、学習後に圧縮するのではなく、最初から短い prefix でも意味が通るよう学習する点です。PCA や SVD は後処理ですが、MRL は低次元利用を目的関数に入れているので、切り詰めたときの性能が落ちにくいです。
Q. RAG では具体的にどう使うのがよいですか？
A. まず短い prefix で広めに候補取得し、上位候補だけ長い prefix か cross-encoder で再評価する使い方が現実的です。全文書を毎回 full vector で比較しなくてよくなるため、検索レイテンシとメモリ使用量の両方に効く可能性があります。
Q. ANN や量子化があれば MRL は不要ですか？
A. 不要とは言えません。ANN や量子化は検索基盤側の最適化で、MRL は表現学習側の最適化です。組み合わせることで、短い prefix で候補探索し、必要な部分だけ高精度処理する、といった多段最適化がしやすくなります。
Q. 低次元 prefix だけ保存して full vector を捨ててもよいですか？
A. 用途次第ですが、慎重に判断すべきです。高速性だけを見るなら短い prefix だけでも動きますが、難しいクエリや細かい区別では長い次元が効く可能性があります。まずは shortlist 用と再評価用を分けて効果測定するのが安全です。
Q. この技術は小規模チームにも関係ありますか？
A. 関係あります。大規模研究に見えますが、発想自体はかなり実務向きです。特に、ベクトル DB の費用が気になる段階、RAG の検索遅延が課題の段階、同じ埋め込みを複数用途に流用したい段階で役に立ちます。
今日の学び
この論文は、固定長埋め込みをそのまま大規模検索や分類へ使うと、計算や保存の面で無駄が大きいという課題を扱いました。そこに対して、同じベクトルの先頭 prefix が低次元表現としても機能するよう同時最適化することで、多粒度な表現を作ろうとしました。
ここから得られるヒントは、検索効率化は index の工夫だけではなく、埋め込みの作り方から設計できるということです。RAG、ベクトル検索、エージェントメモリを作るときも、「1本の表現を段階的に使う」という発想はかなり応用余地があります。