modernbert-ja-310m-sae

English | 日本語

A Top-K sparse autoencoder (SAE) trained on the final-layer hidden states of sbintuitions/modernbert-ja-310m. It decomposes the 768-dimensional hidden states into 65,536 features, only 16 of which are active per token, and reconstructs the original hidden states from them.

This SAE was used to initialize the output layer of the learned sparse retrieval model mahiyama/sae-splade-ja-310m. In sae-splade-ja-310m, the base model was trained for retrieval while the weights of this SAE were kept frozen. For retrieval, use sae-splade-ja-310m.

Model Overview

Property	Value
Base model	sbintuitions/modernbert-ja-310m (revision 77675fc, weights unchanged)
SAE input	Final-layer hidden states of the base model (768 dimensions)
Number of features	65,536
Active features per token (k)	16
Parameters	About 415M (base model about 315M, SAE encoder and decoder about 50M each)
Maximum input length	1,024 tokens (the length used to train the SAE)
Framework	SparseEncoder in sentence-transformers (tested with 6.2.0)

The SAE weights are in 1_SAETopKPooling/model.safetensors.

Name	Shape	Description
encoder.weight	[65536, 768]	Encoder weights $W_{\mathrm{e}\mathrm{n}\mathrm{c}}W\_\{\backslash mathrm\{enc\}\}$
encoder.bias	[65536]	$b_{\mathrm{e}\mathrm{n}\mathrm{c}}-W_{\mathrm{e}\mathrm{n}\mathrm{c}}{b}_{\mathrm{d}\mathrm{e}\mathrm{c}}b\_\{\backslash mathrm\{enc\}\}\; -\; W\_\{\backslash mathrm\{enc\}\}\; b\_\{\backslash mathrm\{dec\}\}$
decoder.weight	[768, 65536]	Decoder weights $W_{\mathrm{d}\mathrm{e}\mathrm{c}}W\_\{\backslash mathrm\{dec\}\}$
decoder.bias	[768]	Decoder bias $b_{\mathrm{d}\mathrm{e}\mathrm{c}}b\_\{\backslash mathrm\{dec\}\}$

The features $zz$ of the hidden state $hh$ and its reconstruction $\widehat{h}\backslash hat\{h\}$ are computed by the following formulas. ${\mathrm{T}\mathrm{o}\mathrm{p}\mathrm{K}}_{16}\backslash mathrm\{TopK\}\_\{16\}$ is an operation that keeps only the 16 largest values and sets the rest to 0.

$$z={\mathrm{T}\mathrm{o}\mathrm{p}\mathrm{K}}_{16}\left(\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(W_{\mathrm{e}\mathrm{n}\mathrm{c}}(h-b_{\mathrm{d}\mathrm{e}\mathrm{c}})+b_{\mathrm{e}\mathrm{n}\mathrm{c}})\right)z\; =\; \backslash mathrm\{TopK\}\_\{16\}\backslash left(\backslash mathrm\{ReLU\}\backslash left(W\_\{\backslash mathrm\{enc\}\}\; \backslash left(h\; -\; b\_\{\backslash mathrm\{dec\}\}\backslash right)\; +\; b\_\{\backslash mathrm\{enc\}\}\backslash right)\backslash right)$$

$$\widehat{h}=W_{\mathrm{d}\mathrm{e}\mathrm{c}}z+b_{\mathrm{d}\mathrm{e}\mathrm{c}}\backslash hat\{h\}\; =\; W\_\{\backslash mathrm\{dec\}\}\; z\; +\; b\_\{\backslash mathrm\{dec\}\}$$

The step of subtracting $b_{\mathrm{d}\mathrm{e}\mathrm{c}}b\_\{\backslash mathrm\{dec\}\}$ from the input is folded into encoder.bias in advance. Therefore, $zz$ is computed from encoder.weight and encoder.bias as ${\mathrm{T}\mathrm{o}\mathrm{p}\mathrm{K}}_{16}\left(\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(W_{\mathrm{e}\mathrm{n}\mathrm{c}}h+\text{encoder.bias})\right)\backslash mathrm\{TopK\}\_\{16\}\backslash left(\backslash mathrm\{ReLU\}\backslash left(W\_\{\backslash mathrm\{enc\}\}\; h\; +\; \backslash text\{encoder.bias\}\backslash right)\backslash right)$.

Usage

Decomposing and Reconstructing Hidden States

The base model can be loaded with AutoModel in transformers, and the SAE weights can be loaded as a safetensors file.

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModel, AutoTokenizer

repo = "mahiyama/modernbert-ja-310m-sae"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModel.from_pretrained(repo).eval()
sae = load_file(hf_hub_download(repo, "1_SAETopKPooling/model.safetensors"))

# Final-layer hidden states of the base model [tokens, 768]
inputs = tokenizer("東京は日本の首都であり、政治と経済の中心地である。", return_tensors="pt")
with torch.no_grad():
    h = model(**inputs).last_hidden_state[0]

# Decomposition into features: keep only the 16 features with the largest values for each token [tokens, 65536]
pre = torch.relu(h @ sae["encoder.weight"].T + sae["encoder.bias"])
values, indices = pre.topk(16, dim=-1)
z = torch.zeros_like(pre).scatter(-1, indices, values)

# Reconstruction of hidden states [tokens, 768]
h_hat = z @ sae["decoder.weight"].T + sae["decoder.bias"]

# Indices of active features for each token
tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
for token, ids in zip(tokens, indices.tolist()):
    print(token, sorted(ids))

Using as a SparseEncoder

When loaded as a SparseEncoder, it encodes text into a 65,536-dimensional sparse vector. Feature values for each token are transformed with log1p, and the maximum of each feature is taken over the whole text, excluding special tokens. This corresponds to the state of sae-splade-ja-310m before training for retrieval. The decoder is not used for this computation.

from sentence_transformers import SparseEncoder

model = SparseEncoder("mahiyama/modernbert-ja-310m-sae", trust_remote_code=True)

embeddings = model.encode(["東京は日本の首都であり、政治と経済の中心地である。"])
print(embeddings.shape)  # torch.Size([1, 65536])

Because the output layer is a custom module, trust_remote_code=True is required.

Training

A Top-K SAE (Gao et al., 2024) was trained using the final-layer hidden states of the base model as input. It was trained to reconstruct hidden states using only the 16 features with the largest values for each token, and the auxiliary loss (AuxK) from the same paper was added so that no features are left unused. During training, the weights of the base model were frozen.

Training Data

The documents (including negatives) and queries from the retrieval training datasets were deduplicated, and the first 4,000 characters of Japanese Wikipedia articles with 2,000 characters or more were added. The total is 846 million tokens, consisting of 70.1% documents, 14.9% queries, and 14.9% Wikipedia articles.

Dataset	Domain	Tokens	Share
mahiyama/mqa-ja	General web	425,378,366	50.3%
mahiyama/auto-wiki-qa	Wikipedia QA	244,642,338	28.9%
wikimedia/wikipedia (20231101.ja)	Wikipedia articles	126,374,330	14.9%
mahiyama/mmarco-ja	General web	27,871,485	3.3%
Non-public government FAQ data	Public sector	8,799,958	1.0%
sbintuitions/JMTEB (jaqket-query train)	Knowledge QA	5,865,936	0.7%
mahiyama/quiz-no-mori	Wikipedia QA	2,955,264	0.3%
mahiyama/quiz-works	Wikipedia QA	2,110,801	0.2%
mahiyama/anlp-meeting-retrieval	Academic	1,315,666	0.2%
mahiyama/miracl-retrieval	Wikipedia QA	354,226	Less than 0.1%
mahiyama/mrtydi	Wikipedia QA	332,480	Less than 0.1%
mahiyama/amagasaki-qna	Public sector	207,276	Less than 0.1%
AmazonScience/mintaka (ja train)	Knowledge QA	55,467	Less than 0.1%
mahiyama/kosodate-faq-pairs-ja	Public sector	8,291	Less than 0.1%
Total		846,271,884	100%

The queries and relevant documents used in the JMTEB v2 lite evaluation and all documents from the evaluation corpora with few documents (JaGovFaqs, NLP Journal, and MultiLongDocRetrieval) were removed from this corpus. In addition, 5% of the total was held out for evaluation and was not used for training.

Training Settings

Setting	Value
Auxiliary loss (AuxK)	aux_k 256, coefficient 1/32
Optimizer	Adam (β 0.9, 0.999)
Learning rate	5e-4 (warmup over the first 5%, followed by cosine decay)
Batch size	8,192 tokens
Max sequence length	1,024 tokens
Training tokens	1 billion (about 1.18 passes over the corpus, 122,071 steps)
GPU	NVIDIA RTX PRO 6000 Blackwell Server Edition × 1 (Google Colab)
Training time	About 4.9 hours

SAE Performance After Training

FVU was measured on 2 million tokens of the held-out documents, and the MLM loss recovery rate was measured on 100 held-out documents.

Metric	Value	Description
FVU	0.155	Fraction of variance unexplained by reconstruction. Lower is better
MLM loss recovery rate	0.944	How much of the MLM loss was recovered when replacing the final-layer hidden states with their reconstructions, compared to replacing them with the mean vector. 1 indicates full recovery
Share of unused features	0%	Share of features that were never active during the last 10 million tokens of training

License

The model weights are released under the MIT License. The base model, sbintuitions/modernbert-ja-310m, is also released under the MIT License.

However, the training data includes the following data that require attention:

• Data released for non-commercial use: mmarco-ja (a Japanese translation of MS MARCO), quiz-no-mori, and JAQKET
• Non-public government FAQ data: the author's non-public data based on FAQs from national government agencies and local governments

If you use this model, especially for commercial purposes, please take these points into account and make your own judgment at your own responsibility.