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

If you would like to know more details:

SSE Technical Article

(a) Retrieval performance (nDCG@10) across NanoBEIR Japanese tasks. (b) Mean nDCG@10 vs. inference speed (QPS: queries per second) measured on Miracl using an Intel® Core™ Ultra 7 265K (3.90 GHz) with batch size 32.

🩵 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以下です。

---

🩵 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 最適化 — エッジデバイスや限られたハードウェアでも高速に動作

---

🩵 モデル詳細 🩵

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) $$

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

$$ \frac{\partial y_k}{\partial x_k} = c_k a_k , \mathrm{sech}^2(a_k x_k + b_k). $$

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

$$ |a_i x_i + b_i| \gg 1 $$

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

$$ \mathrm{sech}^2(z) \sim 4e^{-2|z|} $$

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

$$ \partial y_i / \partial x_i \to 0 $$

となります。

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

$$ \mathrm{sech}^2(z) \approx 1 $$

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

$$ \partial y_j / \partial x_j \approx c_j a_j $$

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

---

🩵 評価結果 (NanoBEIR_ja) 🩵

---

🩵 使い方 🩵

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個のデータセットを使用しました。

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}
}