NeurIPS 2022論文解説: Matryoshka Representation Learning

本記事は Matryoshka Representation Learning (arXiv:2205.13147) の解説記事です。

論文概要（Abstract）

Matryoshka Representation Learning（MRL）は、単一のEmbeddingモデルから複数粒度の次元ベクトルを同時に学習する手法である。著者らは、通常の学習パイプラインに最小限の変更を加えるだけで、推論コストを増やすことなく、低次元でも高次元に匹敵する精度を維持できると報告している。ImageNet-1K分類では最大14倍の次元圧縮で同等精度を達成し、大規模検索では最大14倍の実時間高速化を実現したとされる。

この記事は Zenn記事: Embedding量子化×Matryoshka次元削減の精度-コスト最適化を定量評価するの深掘りです。

情報源

会議名: NeurIPS 2022（Neural Information Processing Systems）
arXiv ID: 2205.13147
URL: https://arxiv.org/abs/2205.13147
著者: Aditya Kusupati, Gantavya Bhatt, Aniket Rege, et al.（11名）
発表年: 2022
コード: https://github.com/RAIVNLab/MRL

カンファレンス情報

NeurIPS（Neural Information Processing Systems）は機械学習分野の最高峰国際会議の1つであり、毎年数千件の投稿から厳選された論文が採択される。MRL論文は2022年の大会で発表され、その後のEmbeddingモデル設計に広く影響を与えた。2026年時点で主要なEmbeddingモデル（OpenAI text-embedding-3-large、Cohere embed-v4、Voyage 4等）のほぼすべてがMRLを採用している。

背景と動機（Background & Motivation）

従来のEmbeddingモデルは、固定次元のベクトルを出力する設計になっている。例えば1024次元のモデルであれば、検索対象が100件でも10億件でも同じ1024次元ベクトルを使用する必要があった。しかし実運用では、検索規模・レイテンシ要件・ストレージ予算に応じて最適な次元数は異なる。

この問題に対する従来のアプローチは2つあった。第一に、異なる次元数のモデルを個別に学習する方法（Fixed Feature: FF）。これは各次元数ごとに独立したモデルの学習・デプロイが必要となり、コストが線形に増加する。第二に、学習済みモデルにPCAやSVDを適用して事後的に次元削減する方法。これはモデルが低次元表現を意識して学習していないため、精度劣化が大きい。

著者らはこの課題に対し、ロシアのマトリョーシカ人形のように「大きな表現の中に小さな表現が入れ子になっている」構造を学習する手法を提案した。

主要な貢献（Key Contributions）

貢献1: 単一モデルから$O(\log d)$個の粒度でEmbeddingを同時学習する損失関数の設計。推論時のコスト増加はゼロ
貢献2: ImageNet-1Kで最大14倍の次元圧縮（2048→約37次元）で同等の分類精度を達成
貢献3: 大規模検索（4.2Mデータベース）で最大128倍の理論的高速化・14倍の実時間高速化
貢献4: Vision（ResNet, ViT）、Vision-Language（ALIGN）、Language（BERT）の3モダリティでの有効性の実証

技術的詳細（Technical Details）

MRLの損失関数

MRLの核心は、通常の学習損失を複数の次元で同時に最適化する入れ子構造にある。次元集合$\mathcal{M} = {m_1, m_2, \ldots, m_L}$（$m_1 < m_2 < \ldots < m_L = d$）に対し、損失関数は以下のように定義される。

\[\min_{\{W^{(m)}\}_{m \in \mathcal{M}}, \theta_F} \frac{1}{N} \sum_{i=1}^{N} \sum_{m \in \mathcal{M}} c_m \cdot \mathcal{L}\left(W^{(m)} \cdot F(x_i; \theta_F)_{1:m}; y_i\right)\]

ここで、

$F(x_i; \theta_F) \in \mathbb{R}^d$: 入力$x_i$に対するEmbeddingベクトル（$d$次元）
$F(x_i; \theta_F)_{1:m}$: Embeddingの先頭$m$次元（プレフィックス）
$W^{(m)} \in \mathbb{R}^{L \times m}$: 次元$m$用の線形分類器
$c_m$: 各次元の重み係数（基本実験ではすべて1）
$\mathcal{L}$: 多クラスソフトマックス交差エントロピー損失

重要な点は、Embeddingの計算は1回だけであり、その先頭$m$次元を切り出して各粒度の損失を計算する点である。推論時は必要な次元数で切り捨てるだけでよく、追加のコストは発生しない。

次元集合の設計

著者らは次元集合$\mathcal{M}$を対数的に配置している。

ResNet50 ($d = 2048$): $\mathcal{M} = {8, 16, 32, 64, 128, 256, 512, 1024, 2048}$
ViT-B/16, BERT-Base ($d = 768$): $\mathcal{M} = {12, 24, 48, 96, 192, 384, 768}$

集合サイズは$

\mathcal{M}

\leq \lfloor \log_2 d \rfloor$であり、9つ程度の粒度で十分な性能が得られると報告されている。

効率的変種: MRL-E

MRL-E（Efficient variant）では、各次元用の独立した分類器$W^{(m)}$ではなく、共有重み行列$W \in \mathbb{R}^{L \times d}$の先頭$m$列を使う重み共有（weight-tying）を導入している。

\[W^{(m)} = W_{1:m}\]

これにより分類器のパラメータ数が約50%削減される。BERT等のモデルでは入力Embeddingと出力分類器の重みが共有されるため、MRL-Eが自然な選択となる。

プレフィックス構造が機能する理由

著者らの重要な発見として、$O(\log d)$個の次元でのみ最適化しているにもかかわらず、最適化対象外の中間次元でも精度が滑らかに補間される点がある（論文Figure 5）。例えば次元32と64の間の次元50でも、精度は急激に劣化せず滑らかに遷移する。

これは逆伝播時に勾配が全次元に拡散し、低次元分類器の制約が高次元に対しても意味のある構造を強制するためと解釈される。

graph LR
    A[8次元<br/>粗い意味] --> B[16次元]
    B --> C[64次元]
    C --> D[256次元]
    D --> E[1024次元]
    E --> F[2048次元<br/>全精度]
    style A fill:#fff
    style F fill:#fff

実験結果（Results）

ImageNet-1K 分類精度

著者らはResNet50を用いたImageNet-1K（1.3Mデータベース、50Kクエリ）での分類精度を報告している。

次元	MRL	Fixed Feature (FF)	差分
8	約71%	約65%	+6%
16	約74%	約70%	+4%
64	約75.5%	約74%	+1.5%
256	約76%	約75.5%	+0.5%
2048	約76.3%	約76.3%	±0%

著者らは「MRLモデルはすべての次元でFFモデルと同等以上の精度を達成する」と主張している。低次元ほどMRLの優位性が顕著であり、8次元では約6ポイントの差が生じている。

検索性能 (mAP@10)

ImageNet-1K（1.3Mデータベース）での1-NN検索結果は以下の通りである（論文Table等より）。

次元	MRL	FF	SVD
16	約0.68	約0.60	約0.45
64	約0.75	約0.70	約0.58
256	約0.82	約0.80	約0.78
2048	約0.87	約0.86	約0.86

事後的な次元削減（SVD）と比較すると、MRLは特に低次元で大きな優位性を示している。16次元ではSVDと比べて0.23ポイント（約51%相対改善）の差がある。

Adaptive Retrieval の高速化

著者らは2段階の適応的検索パイプラインを提案している。

粗い検索: 低次元（$D_s = 16$）で候補$K = 200$件を抽出
精密な再ランキング: 高次元（$D_r = 2048$）で候補を再順位付け

この構成により、ImageNet-1Kでは理論的に約128倍、実時間で約14倍の高速化を達成しつつ、全次元検索と同等のmAP@10を維持したと報告されている。

ImageNet-4K（4.2Mデータベース）では粗い検索の次元を$D_s = 64$に上げる必要があるが、それでも理論的に約32倍、実時間で約6倍の高速化が得られている。

実装のポイント（Implementation）

MRLの実装は既存の学習パイプラインへの変更が最小限である点が特徴的である。

  
import torch
import torch.nn as nn
import torch.nn.functional as F


class MRLLoss(nn.Module):
    def __init__(
        self,
        embed_dim: int,
        num_classes: int,
        nesting_dims: list[int] | None = None,
    ):
        super().__init__()
        if nesting_dims is None:
            nesting_dims = [2**i for i in range(3, int(torch.log2(torch.tensor(embed_dim)).item()) + 1)]
            nesting_dims.append(embed_dim)
        self.nesting_dims = nesting_dims
        self.classifiers = nn.ModuleDict({
            str(m): nn.Linear(m, num_classes)
            for m in nesting_dims
        })

    def forward(
        self,
        embeddings: torch.Tensor,
        labels: torch.Tensor,
    ) -> torch.Tensor:
        total_loss = torch.tensor(0.0, device=embeddings.device)
        for m in self.nesting_dims:
            prefix = embeddings[:, :m]
            logits = self.classifiers[str(m)](prefix)
            total_loss += F.cross_entropy(logits, labels)
        return total_loss / len(self.nesting_dims)

実装上の注意点は以下の通りである。

推論時のオーバーヘッドはゼロ: 推論時は通常通りEmbeddingを計算し、必要な次元数で切り捨てるだけ
学習コスト: 論文によると約1.2倍のFLOPs増加。これは各次元の分類器を通す追加の順伝播に起因する
キャリブレーション不要: 量子化と異なり、MRLは事後的なキャリブレーションデータを必要としない
次元集合$\mathcal{M}$: 対数的に配置するのが最も効果的。均等配置よりも精度の飽和挙動が良好

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

MRLを活用した検索システムのAWS構成を、トラフィック量別に示す。MRLの特徴は「同一モデルから複数次元を取得できる」点にあり、Adaptive Retrievalの2段階パイプラインをAWS上で実現する。

規模	月間リクエスト	推奨構成	月額コスト	主要サービス
Small	~3,000 (100/日)	Serverless	$50-150	Lambda + OpenSearch Serverless
Medium	~30,000 (1,000/日)	Hybrid	$300-800	Lambda + ECS Fargate + ElastiCache
Large	300,000+ (10,000/日)	Container	$2,000-5,000	EKS + Karpenter + OpenSearch

Small構成の詳細 (月額$50-150):

Lambda: Embedding推論 + 低次元検索、1GB RAM ($20/月)
OpenSearch Serverless: ベクトルインデックス、256次元int8 ($80/月)
DynamoDB: メタデータ保存、On-Demand ($10/月)

コスト削減テクニック（2026年5月時点の概算）:

MRLによる次元削減 (1024→256) でインデックスサイズ75%削減
Adaptive Retrieval: 低次元で粗い検索 → 高次元で再ランキング、計算量を約10倍削減
OpenSearch Serverlessの自動スケーリングで夜間コスト最小化

上記は2026年5月時点のAWS ap-northeast-1料金に基づく概算値です。最新料金は AWS料金計算ツールで確認してください。

Terraformインフラコード

Small構成 (Serverless): Lambda + OpenSearch Serverless

  
module "vpc" {
  source  = "terraform-aws-modules/vpc/aws"
  version = "~> 5.0"

  name = "mrl-search-vpc"
  cidr = "10.0.0.0/16"
  azs  = ["ap-northeast-1a", "ap-northeast-1c"]
  private_subnets = ["10.0.1.0/24", "10.0.2.0/24"]

  enable_nat_gateway   = false
  enable_dns_hostnames = true
}

resource "aws_iam_role" "lambda_search" {
  name = "mrl-lambda-search-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Action    = "sts:AssumeRole"
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
    }]
  })
}

resource "aws_iam_role_policy" "opensearch_access" {
  role = aws_iam_role.lambda_search.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = ["aoss:APIAccessAll"]
      Resource = "*"
    }]
  })
}

resource "aws_lambda_function" "mrl_search" {
  filename      = "lambda.zip"
  function_name = "mrl-adaptive-search"
  role          = aws_iam_role.lambda_search.arn
  handler       = "index.handler"
  runtime       = "python3.12"
  timeout       = 30
  memory_size   = 1024

  environment {
    variables = {
      OPENSEARCH_ENDPOINT = "https://collection-endpoint.aoss.amazonaws.com"
      MRL_COARSE_DIM      = "64"
      MRL_FINE_DIM         = "1024"
      OVERSAMPLING_FACTOR  = "4"
    }
  }
}

resource "aws_dynamodb_table" "metadata" {
  name         = "mrl-document-metadata"
  billing_mode = "PAY_PER_REQUEST"
  hash_key     = "doc_id"

  attribute {
    name = "doc_id"
    type = "S"
  }

  ttl {
    attribute_name = "expire_at"
    enabled        = true
  }
}

セキュリティベストプラクティス

IAMロール: 最小権限の原則。OpenSearchへのアクセスはコレクション単位で制限
ネットワーク: Lambda は VPC 内に配置し、OpenSearch Serverless へは VPC エンドポイント経由でアクセス
暗号化: OpenSearch Serverless はデフォルトで KMS 暗号化が有効

運用・監視設定

  
import boto3

cloudwatch = boto3.client('cloudwatch')

cloudwatch.put_metric_alarm(
    AlarmName='mrl-search-latency-p99',
    ComparisonOperator='GreaterThanThreshold',
    EvaluationPeriods=2,
    MetricName='Duration',
    Namespace='AWS/Lambda',
    Period=300,
    Statistic='p99',
    Threshold=5000,
    AlarmDescription='MRL検索レイテンシP99が5秒超過'
)

  
-- CloudWatch Logs Insights: Adaptive Retrieval の段階別レイテンシ分析
fields @timestamp, coarse_search_ms, rerank_ms, total_ms
| stats avg(coarse_search_ms) as avg_coarse,
        avg(rerank_ms) as avg_rerank,
        pct(total_ms, 95) as p95_total
  by bin(5m)

コスト最適化チェックリスト

MRL次元削減でインデックスサイズを75%以上削減
Adaptive Retrieval で粗い検索次元を64以下に設定
oversampling倍率を4x以下に抑える（4x→10xでの改善はなし）
OpenSearch Serverless の OCU 最小値を設定
Lambda メモリサイズを CloudWatch Insights で最適化
開発環境は夜間停止（Lambda は自動、OpenSearch は手動スケール）

まとめと今後の展望

Matryoshka Representation Learningは、単一モデルから多粒度のEmbeddingを同時学習する手法として、2026年時点の主要Embeddingモデルに広く採用されている。著者らが報告した主要な成果は以下の通りである。

ImageNet-1K分類で14倍の次元圧縮で同等精度
大規模検索で128倍の理論的高速化（14倍の実時間高速化）
学習コストの増加は約1.2倍のみ
Vision、Vision-Language、Languageの3モダリティで有効

MRLは量子化（int8、binary）と直交する最適化手法であり、両者を組み合わせることでストレージを最大256倍圧縮できる可能性がある。SMEC（arXiv:2510.12474）等の後続研究では、MRLと量子化の共同最適化による更なる精度向上が報告されている。

参考文献

arXiv: https://arxiv.org/abs/2205.13147
Code: https://github.com/RAIVNLab/MRL
Related Zenn article: https://zenn.dev/0h_n0/articles/6d45410fe51fa1

:::message この記事はAI（Claude Code）により自動生成されました。内容の正確性については複数の情報源で検証していますが、実際の利用時は公式ドキュメントもご確認ください。 :::

📄 NeurIPS 2022論文解説: Matryoshka Representation Learning

論文概要（Abstract）

情報源

カンファレンス情報

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

MRLの損失関数

次元集合の設計

効率的変種: MRL-E

プレフィックス構造が機能する理由

実験結果（Results）

ImageNet-1K 分類精度

検索性能 (mAP@10)

Adaptive Retrieval の高速化

実装のポイント（Implementation）

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

Terraformインフラコード

セキュリティベストプラクティス

運用・監視設定

コスト最適化チェックリスト

関連研究（Related Work）

まとめと今後の展望

参考文献

📄 NeurIPS 2022論文解説: Matryoshka Representation Learning

論文概要（Abstract）

情報源

カンファレンス情報

背景と動機（Background & Motivation）

主要な貢献（Key Contributions）

技術的詳細（Technical Details）

MRLの損失関数

次元集合の設計

効率的変種: MRL-E

プレフィックス構造が機能する理由

実験結果（Results）

ImageNet-1K 分類精度

検索性能 (mAP@10)

Adaptive Retrieval の高速化

実装のポイント（Implementation）

Production Deployment Guide

AWS実装パターン（コスト最適化重視）

Terraformインフラコード

セキュリティベストプラクティス

運用・監視設定

コスト最適化チェックリスト

関連研究（Related Work）

まとめと今後の展望

参考文献

関連記事

📄 EMNLP 2025論文解説: SMEC — Matryoshka表現学習と量子化の共同最適化

📄 論文解説: Matryoshka Representation Learning — 可変次元埋め込みによる適応的表現学習

📄 ICLR 2025論文解説: NV-Embed — Latent Attentionと2段階訓練によるLLMベース汎用埋め込みモデル