splade-ja-310m

ModernBERT-Ja-310M をベースに、日本語 SPLADE (Learned Sparse Retrieval) モデルとしてトレーニングしたモデルです。

入力テキストを語彙次元 (102,400 次元) のスパースベクトルに変換し、関連語への自然な拡張を伴ったまま転置インデックスや疎行列演算で高速に検索できます。句読点・記号・サブワード断片・装飾プレースホルダ等のノイズトークンを訓練段階で抑止する仕組みを組み込んでおり、推論側で trust_remote_code や追加モジュールは不要です。素の SparseEncoder.from_pretrained() でそのまま読み込めます。

モデルの概要

項目	内容
ベースモデル	sbintuitions/modernbert-ja-310m
アーキテクチャ	Transformer (fill-mask) + SpladePooling (max)
パラメータ数	約 315M
最大入力長	8,192 tokens (訓練は 512 tokens で実施)
語彙サイズ (= 出力次元数)	102,400
トークナイザ	SentencePiece (ModernBERT-Ja の語彙をそのまま使用)
主要フレームワーク	sentence-transformers >= 5.4 (SparseEncoder)
損失関数	CachedSpladeLoss(SparseMultipleNegativesRankingLoss)
混合精度	bf16

SPLADE v2 は MLM head の logits を語彙次元ごとに max pooling して得られる、文ごとに非ゼロ次元が数十〜数百しかないスパース埋め込みです。検索時は内積でスコアリングできるので、転置インデックス (Elasticsearch, OpenSearch, Vespa, Qdrant 等) との親和性が高いのが特徴です。

使い方

from sentence_transformers import SparseEncoder

model = SparseEncoder("mahiyama/splade-ja-310m")

queries = [
    "日本の首都はどこですか？",
    "機械学習とは何ですか",
]
documents = [
    "東京は日本の首都であり、政治・経済の中心地である。",
    "京都は古都として知られ、神社仏閣が多くある。",
    "機械学習は人工知能の一分野で、データからパターンを学習する手法である。",
]

q_emb = model.encode_query(queries)
d_emb = model.encode_document(documents)

# 内積で類似度スコアを計算
scores = model.similarity(q_emb, d_emb)
print(scores)

# 拡張トークンを覗く (top_k=15 の語彙とその重み)
for token, weight in model.decode(q_emb[0], top_k=15):
    print(f"  {token:>20}  {weight:.3f}")

クエリと文書で同じ encoder を使う設計です (asymmetric ではなく symmetric)。スコアは内積で計算してください (cosine ではなく dot)。

語彙拡張サンプル

訓練済みモデルが各クエリをどのような語彙次元に拡張するかを 3 例示します。{token: weight} 形式の JSON で、上位のトークンほど重みが大きく、各サンプルとも上位 10 トークンのみ抜粋しています。

例 1: 短い名詞句「機械学習」

{
  "query": "機械学習",
  "active_dims": 307,
  "top_tokens": {
    "機械学習":     2.828,
    "学習":         1.906,
    "分類":         1.734,
    "人工知能":     1.711,
    "機械":         1.672,
    "アルゴリズム": 1.578,
    "用語":         1.508,
    "コンピューター": 1.492,
    "サイエンス":   1.461,
    "科学":         1.453
  }
}

「機械学習」というクエリ自身に加え、学習・人工知能・アルゴリズム・コンピューター・科学といった上位概念や関連分野の語に綺麗に拡張されています。BM25 のような表層一致では拾えない「機械学習について書かれているが文中に "機械学習" という語が現れない文書」もこの拡張でヒットしやすくなります。

例 2: 質問文「日本の首都はどこですか？」

{
  "query": "日本の首都はどこですか？",
  "active_dims": 11,
  "top_tokens": {
    "首都":         2.562,
    "日本の":       2.156,
    "の首都":       1.859,
    "日本":         1.273,
    "首相":         1.219,
    "都":           0.902,
    "日本で":       0.801,
    "都市":         0.734,
    "はどこですか": 0.391,
    "日本における": 0.289
  }
}

質問文形式でも active_dims が 11 と非常に疎で、検索の核となる「首都」「日本」「都市」に重みが集中しています。「は」「どこ」「ですか」のような疑問文の機能語に重みが浪費されておらず、検索シグナルとしてノイズの少ない疎ベクトルになっています。

例 3: 多義語「アップル」

{
  "query": "アップル",
  "active_dims": 533,
  "top_tokens": {
    "アップル": 2.781,
    "別名":     2.062,
    "リンゴ":   1.992,
    "Apple":    1.984,
    "apple":    1.922,
    "▁Apple":   1.820,
    "▁apple":   1.781,
    "英語":     1.773,
    "りんご":   1.711,
    "APPLE":    1.680
  }
}

果物のリンゴ・りんご、企業名の Apple・apple・APPLE、英日表記揺れ (▁Apple は SentencePiece の語頭マーカ付き) の両方が同時に高い重みで活性化しており、文脈なしクエリに対する語義カバレッジが広いことがわかります。一方で句読点・記号・絵文字・特殊トークン ([CLS] 等) は上位に一切含まれていません。

想定される用途

日本語 Web 文書 / Wikipedia 派生コンテンツの検索
BM25 補強 (BM25 + SPLADE のスコア融合)
Dense retrieval とのハイブリッド検索 (dense + sparse の rank fusion)
転置インデックス (Elasticsearch, OpenSearch, Vespa, Qdrant Sparse 等) でのオフライン索引

実験手法

ベース ModernBERT-Ja-310M を出発点に、6 段階の段階的訓練で仕上げています。各段階で前段階の best checkpoint を起点に、データ量・学習率・正則化を少しずつ変えて少しずつ収束させていく方針です。

段階	起点モデル	訓練データ件数	epoch	batch_size	learning_rate	query_reg	document_reg
第一訓練	ベース MLM	200,000	1	64	1.5e-5	2e-5	1.5e-4
第二訓練	第一訓練 best	200,000	1	64	5e-6	1.5e-5	1e-4
第三訓練	第二訓練 best	500,000	1	64	3e-6	1.5e-5	1e-4
第四訓練	第三訓練 best	1,000,000	1	64	2e-6	1.5e-5	1e-4
第五訓練	第四訓練 best	200,000	1	64	1e-6	1.5e-5	1e-4
第六訓練 (最終)	第五訓練 best	200,000	1	64	5e-7	1e-5	5e-5

すべての段階で共通する設定:

訓練データ: mahiyama/auto-wiki-qa の triplets サブセット (anchor / positive / negative) を件数制限してシャッフル
損失: CachedSpladeLoss(SparseMultipleNegativesRankingLoss), mini_batch_size=32
正則化スケジューラ: SpladeRegularizerWeightSchedulerCallback (quadratic, warmup_ratio=0.4)
学習率スケジューラ: linear warmup (warmup_ratio=0.1 第一訓練, 以降 0.05)
混合精度: bf16
ノイズトークンの抑止 (詳細は後述)

各段階のねらい

第一訓練: ベース MLM から SPLADE 化する初期収束。比較的高めの学習率でスパース化と関連語拡張の方向性を作る。
第二訓練: メタトークン (○○, わからない, 質問等) を抑止対象に追加し、正則化を緩めて active_dims を target レンジ内に乗せる。
第三訓練: 訓練データ量を 2.5 倍 (500K) に拡大し、自然な語彙拡張を維持したまま事実知識のカバレッジを広げる。
第四訓練: さらに 2 倍 (1M) のデータ量で訓練し、事実知識のカバレッジを拡張する。
第五訓練: データ量を 200K に戻し、低い学習率で精緻化。active_dims プロファイルを安定化。
第六訓練 (最終): 正則化重みをさらに緩めて query_active_dims を target レンジ (40〜80) の下端に近づけ、corpus_active_dims を 200〜400 の領域に収める最終仕上げ。これが本モデルとして公開されています。

段階を進めるごとに学習率を約 1/3 ずつ下げ (1.5e-5 → 5e-6 → 3e-6 → 2e-6 → 1e-6 → 5e-7)、収束済みモデルを過剰に動かさずに各段階の目的だけを取り込むよう調整しています。

ノイズトークンの抑止

自然な語彙拡張を実現するため、検索シグナルとして無価値なトークンを訓練段階で抑止する仕組みを組み込んでいます。

抑止対象のトークン群

SPLADE は MLM head の logits をそのまま語彙次元として使うため、放っておくと検索シグナルとして無価値なトークンが拡張結果の上位に紛れ込みがちです。具体的には、句読点や記号、SPM の語頭マーカ ▁ 単独、バイトフォールバック (<0xHH>)、装飾プレースホルダ (○○, ※, ■, → 等)、特殊トークン ([CLS], [SEP], [PAD], [MASK])、Wikipedia コーパス由来のメタトークン (○○, わからない, 質問など) といったものです。これらは転置インデックス上で高頻度次元となり、メモリと計算時間の両方を浪費します。

抑止対象は、語彙全体に対して以下のルールで自動判定したものに、コーパス固有のメタトークンを手動で加えたものです:

Unicode カテゴリが Punctuation / Symbol / Separator / Other / Mark のみで構成されるトークン
バイトフォールバック (<0xHH> 形式)
SPM 語頭マーカ ▁ 単独
同一文字の繰り返し (ーー, ・・, !! など)
modifier letter (Unicode カテゴリ Lm) のみで構成されるトークン
ASCII 1 文字の英字・数字
装飾プレースホルダのリテラルリスト (○○, ××, △△, ■, ※, ▼, ◆, …, ★, ☆, ♪, →, ・, ーなど)
特殊トークン (CLS / SEP / PAD / MASK / UNK / BOS / EOS)

FLOPS 正則化との役割分担

SPLADE は損失関数に組み込まれた FLOPS 正則化によって、高頻度かつ低貢献の次元が自然に抑えられる設計です。したがって句読点や記号は理屈上は徐々に重みを失っていくはずですが、実運用では以下の理由から自然消滅だけに頼るのは難しいと判断しました:

FLOPS 正則化は data-agnostic な頻度ペナルティで、対照損失上で正解スコアに偶発的に寄与してしまう次元 (たとえばクエリと正例の両方によく出てくる句読点) は罰されずに残ることがある
FLOPS 重みを強くしすぎると、ノイズ次元と一緒に関連語拡張側の有用な低頻度次元まで巻き添えで死ぬため、思い切って上げ切れない
訓練途中の途中状態では一時的にノイズ次元が大きく出ることがあり、best checkpoint がその近傍で選ばれる可能性もある

そのため本モデルでは、構造的にノイズと判定できるトークンはハード制約で潰しておき、FLOPS 正則化には「関連語拡張の中での疎性コントロール」という本来の役割に集中してもらう方針を取っています。これにより検索品質と語彙拡張の自然さが安定して両立し、active_dims のチューニングも見通しよくなります。

抑止の実装 (2 段階)

以下の 2 段階で抑止を恒久化しています:

訓練時の forward pre-hook: SpladePooling の直前で MLM logits の抑止対象次元へ -1e4 を加算し、勾配がほぼ立たないようにする。bf16 環境で NaN を生まないよう -inf ではなく -1e4 を使用 (relu(log(1+exp(-1e4))) ≈ 0 で実効ゼロ)。
訓練後の bias 焼込: MLM head の bias パラメータの抑止対象次元へ -1e4 を恒久的に書き込む。hook を外したあとも推論時にノイズ次元が復活せず、保存モデルは素の SparseEncoder.from_pretrained() で読み込めます (trust_remote_code や追加モジュール不要)。

段階間の自動フィードバック

第二訓練以降では、各 eval ステップでサンプルクエリの語彙拡張を JSON ダンプし、上位に出てしまったノイズ語を抽出して次の段階の抑止リストへ伝搬する自動フィードバックを回しています。これによりコーパス固有のメタトークンも段階的に潰し込んでいけます。

Lessons Learned

本モデルの開発を通して得られた、SPLADE 訓練に関する知見を記録しておきます。

起点モデルの語彙拡張プロファイルは固定的

一度収束した SPLADE モデルからの追加訓練では、正則化重み (query_reg / document_reg) を多少動かしても active_dims が大きく変わらない傾向が見られました。たとえばベース MLM から訓練した最初の段階で query_active_dims が 20 程度に収束した場合、その後の段階で query_reg を 1/10 まで弱めても 30 前後までしか上がりません。target レンジを大きく変えたい場合は、ベース MLM から起点を取り直す必要があります。

コーパス偏りによる構造ノイズは FLOPS 正則化では除けない

Wikipedia 系コーパス (auto-wiki-qa) のみで訓練すると、「別名」「英語版」「学名」「主な」「タグ」など、コーパス由来の構造的な語が拡張上位に頻出します。FLOPS 正則化は data-agnostic な頻度ペナルティで「使われない高頻度次元」しか抑制しないため、対照損失上で正解スコア向上に寄与してしまうこの種の語は罰されず残ります。本質的な対処はデータミックス (異なるドメイン/レジスタのコーパスを混ぜる) で、単一コーパスでの訓練だけでは限界があります。

自然な語彙拡張には構造判定 + 手動メタトークン追加の併用が必要

Unicode カテゴリベースの自動判定 (Punctuation / Symbol / Modifier / 繰り返し等) は句読点・記号・絵文字・SPM マーカ・バイトトークンを綺麗に拾えますが、「○○」「わからない」「質問」といったコーパス固有のメタトークンは構造的特徴を持たないため取りこぼします。各段階の語彙拡張ログを目視で確認し、上位に出るメタトークンを次段階の抑止リストへ手動追加するフローを組み込んだことで、Top-30 への punctuation 混入率と subword 断片混入率を最終的に 0% に抑えられました。

評価ベンチマーク選定が段階間の差別化を左右する

今回、訓練中の評価データセットとして使用した JaCWIR Retrieval ベンチマークの少量版は、queries 1,000 / corpus 10,000 という比較的小規模な構成のため、訓練ループに組み込んで頻繁に評価できる軽量さが利点です。一方で、本モデルの第一訓練段階から nDCG@10 ≈ 0.97 と高水準に達してしまい、段階間の改善幅が 0.005 程度と小さく、ベンチマーク単体では訓練判断の指標として機能しにくいことがわかりました。本モデルの開発では active_dims プロファイル (q/c の target レンジ適合度) と語彙拡張サンプルの自然さを副次的評価軸として併用しています。本格的な評価には MIRACL-ja / JSQuAD のような難しめのベンチマークの併用が望ましいです。

SparseInformationRetrievalEvaluator は max_active_dims が事実上必須

sentence-transformers の SparseInformationRetrievalEvaluator は corpus 側の疎ベクトルに対する CPU 上の sparse matmul で類似度を計算します。訓練初期は active_dims が数千次元に膨らむため、queries=1,000 / corpus=10,000 規模でも eval 1 回が 10 分以上かかることがあります。max_active_dims=512 程度に capped 評価を入れることで eval ループを実用的な時間 (1〜2 分) に収めつつ、最終評価のみ無制限で素のスパース性を測る運用が安定しています。

制限事項

訓練データは Wikipedia 由来 QA (auto-wiki-qa) が中心です。法律実務・医療・専門技術ドメイン等での運用は別途 fine-tune を推奨します。
Wikipedia コーパスの構造由来語 (「別名」「英語版」「学名」「主な」など) が一部のクエリで上位に出ることがあります。
ベース ModernBERT-Ja は最大 8,192 tokens の長文に対応していますが、本モデルの訓練自体は 512 tokens で実施しています。長文での挙動は未検証なので、必要に応じてチャンク分割を検討してください。
関連語拡張により表層一致では拾えない文書も検索できる反面、active_dims が大きくなるクエリでは BM25 と比べてインデックス膨張・クエリ時間に注意してください。

ライセンス

MIT License

Downloads last month: 125

Safetensors

Model size

0.3B params

Tensor type

F32

Model tree for mahiyama/splade-ja-310m

Base model

sbintuitions/modernbert-ja-310m

Finetuned

(21)

this model

Quantizations

1 model

mahiyama
/

splade-ja-310m