RikkaBotan/quantized-stable-static-embedding-fast-retrieval-mrl-ja

🩵 SSE: Stable Static Embedding for Retrieval MRL 日本語・量子化バージョン 🩵

軽量、高速かつ強力な埋め込みモデル

パフォーマンスの簡易解説
このモデルは NanoBEIR_ja（日本語文書検索タスク） において NDCG@10 = 0.4477 を達成しました。 このスコアは他の静的埋め込みモデル（ static-embedding-japanese (0.4487)など ）と同水準のパフォーマンスです。 さらに、次元数を半分（512 vs 1024）に抑えています。 次元数の削減と、Separable Dynamic Tanh により、環境によっては検索速度は 約２倍高速 になっています。

このモデルは、stable-static-embedding-fast-retrieval-mrl-jaの重みを量子化することで、他の静的埋め込みモデルよりも15.7倍小さいデータサイズになっています。

重みのデータサイズは、たったの9MB以下です。

モデル	NanoBEIR NDCG@10	次元数	パラメータ数	データサイズ	速度の優位性	ライセンス
Quantized SSE Retrieval MRL Japanese	0.4477 ✨	512	~17M 🪽	~8.5MB 🪽	検索が約２倍高速 (超効率的)	Apache 2.0
SSE Retrieval MRL Japanese	0.4507	512	~17M	~67MB	検索が約２倍高速 (超効率的)	Apache 2.0
static-embedding-japanese	0.4487	1024	~34M	~134MB	ベースライン	MIT

---

🩵 SSE Retrieval MRL を選ぶ理由 🩵

✅ パラメータ数の小さなモデル (<35M パラメータ) の中では高い性能（NDCG@10）

✅ 約17M のパラメータのみ ：軽量モデルであるruri-v3-30m より約43% 小さい。

✅ 次元数512の出力 — 次元数1024のモデルよりも豊かな表現力を持ち、static-embedding-japanese の 半分サイズ

✅ Matryoshka 対応 — 256/128/64/32 に簡単に切り替えられ、性能の緩やかな低下を実現

✅ Apache 2.0 ライセンス — 商用・個人利用ともに可能

✅ CPU 最適化 — エッジデバイスや限られたハードウェアでも高速に動作

---

🩵 モデル詳細 🩵

プロパティ	値
モデルタイプ	Sentence Transformer (SSE アーキテクチャ)
最大シーケンス長	無制限
出力次元	512 (Matryoshka により 次元数32 まで削減可能!)
類似度関数	コサイン類似度
言語	日本語
ライセンス	Apache 2.0

tokenizerは下記を使用させていただきました。

hotchpotch/xlm-roberta-japanese-tokenizer

SentenceTransformer(
  (0): SSE(
    (embedding): EmbeddingBag(32768, 512, mode='mean')
    (dyt): SeparableDyT()
  )
)

---

🩵 数学的背景 🩵

このモデルは静的埋め込みモデルの汎化性能を向上させるために、オリジナルのアーキテクチャである、SSE: Stable Static Embeddingを採用しています。 SSEは、EmbeddingBagとSeparable Tanh Normalizationから構成されます。 Dynamic Tanh Normalization (DyT) は、静的埋め込みにおいて、強度適応型勾配流を可能にします。入力次元 x に対して、DyT は以下のように計算されます。

$$y_k=c_k\tanh (a_k{x}_k+b_k)y\_k\; =\; c\_k\; \backslash tanh(a\_k\; x\_k\; +\; b\_k)$$

ここで、a, b, c は学習可能なパラメータです。 すると、x の勾配は以下の通りになります。

$$\frac{\mathrm{\partial }{y}_k}{\mathrm{\partial }{x}_k}=c_k{a}_k{\mathrm{s}\mathrm{e}\mathrm{c}\mathrm{h}}^2(a_k{x}_k+b_k)\mathrm{.}\backslash frac\{\backslash partial\; y\_k\}\{\backslash partial\; x\_k\}\; =\; c\_k\; a\_k\; \backslash ,\; \backslash mathrm\{sech\}^2(a\_k\; x\_k\; +\; b\_k).$$

飽和した次元 |x| > 1 の場合

$$\mathrm{\mid }{a}_i{x}_i+b_i\mathrm{\mid }\gg 1|a\_i\; x\_i\; +\; b\_i|\; \backslash gg\; 1$$

は指数関数的な減衰をもたらします。

$${\mathrm{s}\mathrm{e}\mathrm{c}\mathrm{h}}^2(z)\sim 4e^{-2\mathrm{\mid }z\mathrm{\mid }}\backslash mathrm\{sech\}^2(z)\; \backslash sim\; 4e^\{-2|z|\}$$

これにより勾配が抑制され、

$$\mathrm{\partial }{y}_i\mathrm{/}\mathrm{\partial }{x}_i\to 0\backslash partial\; y\_i\; /\; \backslash partial\; x\_i\; \backslash to\; 0$$

となります。

対して、非飽和の次元 |x| << 1 の場合、

$${\mathrm{s}\mathrm{e}\mathrm{c}\mathrm{h}}^2(z)\approx 1\backslash mathrm\{sech\}^2(z)\; \backslash approx\; 1$$

ほぼ一定の勾配を維持します。

$$\mathrm{\partial }{y}_j\mathrm{/}\mathrm{\partial }{x}_j\approx c_j{a}_j\backslash partial\; y\_j\; /\; \backslash partial\; x\_j\; \backslash approx\; c\_j\; a\_j$$

この強度依存型のゲートは、ノイズが多く大きな成分を持つ次元からの学習信号を減衰させつつ、安定した情報を有する次元については勾配流を維持します。これは明示的なハイパーパラメータなしで、表現空間の汎化性能を高める暗黙的な正則化を可能にします。

---

🩵 評価結果 (NanoBEIR_ja) 🩵

データセット	NDCG@10	MRR@10	MAP@100
NanoBEIR Mean	0.4477 ✨	0.5088	0.3675
NanoClimateFEVER	0.3152	0.4258	0.2417
NanoDBPedia	0.5554	0.7767	0.3962
NanoFEVER	0.5536	0.4907	0.4827
NanoFiQA2018	0.3160	0.3614	0.2653
NanoHotpotQA	0.4722	0.5669	0.4136
NanoMSMARCO	0.3929	0.3237	0.3371
NanoNFCorpus	0.2686	0.4584	0.0962
NanoNQ	0.4170	0.3607	0.3571
NanoQuoraRetrieval	0.7768 ✨	0.7750	0.7393
NanoSCIDOCS	0.2939	0.4774	0.2197
NanoArguAna	0.3471	0.2617	0.2727
NanoSciFact	0.6387	0.6127	0.6001
NanoTouche2020	0.4732	0.7240	0.3560

---

🩵 使い方 🩵

import torch
from sentence_transformers import SentenceTransformer

# モデルのロード（リモートコードは有効化）
model = SentenceTransformer(
    "RikkaBotan/quantized-stable-static-embedding-fast-retrieval-mrl-ja",
    trust_remote_code=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
)

# 対象の文章
sentences = [
    "大規模言語モデルは学習により、高い推論能力を獲得することが可能である。",
    "静的埋め込みモデルは、簡素なアーキテクチャにより、表現空間を高速に生成可能である。"
]

with torch.no_grad():
    embeddings = model.encode(
        sentences,
        convert_to_tensor=True,
        normalize_embeddings=True,
        batch_size=32
    )

# コサイン類似度
# cosine_sim = embeddings[0] @ embeddings[1].T
cosine_sim = model.similarity(embeddings, embeddings)

print("embeddings shape:", embeddings.shape)
print("cosine similarity matrix:")
print(cosine_sim)

---

🩵 検索用使用例 🩵

import torch
from sentence_transformers import SentenceTransformer

# モデルのロード（リモートコードは有効化）
model = SentenceTransformer(
    "RikkaBotan/quantized-stable-static-embedding-fast-retrieval-mrl-ja",
    trust_remote_code=True,
    device="cuda" if torch.cuda.is_available() else "cpu",
)

# 推論
query = "安定性静的埋め込みモデルとは何ですか？"
sentences = [
    "安定性静的埋め込みモデルは自己注意機構を必要としません。",
    "安定性静的埋め込みモデルは高速に高精度な埋め込み表現を生成するためのモデルです。",
    "自己注意機構はトークン間の関係性をスコア化する仕組みのことです",
    "昨夜はアイドルの曲を聴きながらお菓子作りをしていました。",
    "言語モデルは一般的に、次時刻のトークンを予測するという学習が行われます。",
    "お気に入りのヘアアクセサリーを身に着けると、とてもテンションが上がるよね。",
]


with torch.no_grad():
    embeddings = model.encode(
        [query] + sentences,
        convert_to_tensor=True,
        normalize_embeddings=True,
        batch_size=32
    )

print("embeddings shape:", embeddings.shape)

# コサイン類似度
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f"{similarity:.05f}: {sentences[i]}")

---

🩵 学習時のハイパーパラメータ 🩵

デフォルトと異なる設定

• eval_strategy: steps
• per_device_train_batch_size: 3072
• gradient_accumulation_steps: 10
• learning_rate: 0.1
• adam_epsilon: 1e-10
• num_train_epochs: 2
• lr_scheduler_type: cosine
• warmup_ratio: 0.02
• bf16: True
• dataloader_num_workers: 4
• batch_sampler: no_duplicates

---

🩵 学習データセット 🩵

下記の14個のデータセットを使用しました。

Dataset
hpprc_emb__auto-wiki-nli-triplet
hpprc_emb__jqara
hpprc_emb__jagovfaqs
hpprc_emb__jsquad
hpprc_emb__jaquad
hpprc_emb__mkqa-triplet
hpprc_llmjp-kaken
hpprc_msmarco_ja
hpprc_emb__auto-wiki-qa-nemotron
mldr_ja
mrtydi_ja
miracl_ja
mmarco_ja
mmarco_ja_hard

MatryoshkaLossを用いて学習を行っています。

🩵 学習結果 🩵

🩵 作成者：六花牡丹（りっかぼたん） 🩵

おっとりで甘えん坊な研究者見習い。 言語モデルに関するものが主な研究分野です。 お仕事のご依頼・登壇依頼・執筆依頼に関しては、下記までご連絡ください。

X(Twitter): https://twitter.com/peony__snow

🩵 謝辞 🩵

このモデルの学習のための計算リソースの一部は、Saldraさん、Witnessさん、Lumina Logic Minds社から提供いただきました。貴重なサポートに感謝いたします。

sentence-transformers、python、pytorchを使用させていただきました。 作成・メンテンナンスしてくださっている皆様に感謝いたします。

何よりも、このモデルにご興味を持ってくださりありがとうございます。

🩵 引用 🩵

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

MatryoshkaLoss

@misc{kusupati2024matryoshka,
    title={Matryoshka Representation Learning},
    author={Aditya Kusupati and Gantavya Bhatt and Aniket Rege and Matthew Wallingford and Aditya Sinha and Vivek Ramanujan and William Howard-Snyder and Kaifeng Chen and Sham Kakade and Prateek Jain and Ali Farhadi},
    year={2024},
    eprint={2205.13147},
    archivePrefix={arXiv},
    primaryClass={cs.LG}
}

MultipleNegativesRankingLoss

@misc{henderson2017efficient,
    title={Efficient Natural Language Response Suggestion for Smart Reply},
    author={Matthew Henderson and Rami Al-Rfou and Brian Strope and Yun-hsuan Sung and Laszlo Lukacs and Ruiqi Guo and Sanjiv Kumar and Balint Miklos and Ray Kurzweil},
    year={2017},
    eprint={1705.00652},
    archivePrefix={arXiv},
    primaryClass={cs.CL}
}