以下文章转载自文章 100倍速で実用的な文章ベクトルを作れる、日本語 StaticEmbedding モデルを公開。

static-embedding-japanese

文本的密集向量可用于信息检索、文本分类、相似文本提取等多种用途。然而，即使是最先进的 Transformer 模型，即便是较小的模型，特别是在 CPU 环境下处理速度也较慢，因此在实际应用中往往不太实用。

为解决这一问题，作为一种新方法，最近发布的非 Transformer 模型 StaticEmbedding 模型，例如在与 intfloat/multilingual-e5-small（以下简称 mE5-small）的基准测试比较中，达到了 85% 的分数，这是一个相当不错的性能，更重要的是在 CPU 上运行时，可以以 126 倍的速度生成文本向量，这个速度令人惊叹。

因此，我立即创建并发布了用日语（和英语）训练的模型 sentence-embedding-japanese。

• https://huggingface.co/hotchpotch/static-embedding-japanese

以下是用于评估日语文本向量性能的 JMTEB 的结果。虽然综合得分略低于 mE5-small，但在某些任务上表现更好，而且达到了比其他日语 base 大小的 BERT 模型得分还要高的程度，至少达到了实用可行的性能水平。在亲自训练之前，我对此半信半疑，但结果确实令人惊讶。

此外，关于 StaticEmbedding 日语模型训练等技术细节写在文章后半部分，感兴趣的读者请继续阅读。

使用方法

使用很简单，可以通过 SentenceTransformer 用通常的方法生成文本向量。这次我们不使用 GPU，而是在 CPU 上运行。SentenceTransformer 版本为 3.3.1。

pip install \"sentence-transformers>=3.3.1\"

from sentence_transformers import SentenceTransformer

model_name = \"hotchpotch/static-embedding-japanese\"
model = SentenceTransformer(model_name, device=\"cpu\")

query = \"美味しいラーメン屋に行きたい\"
docs = [
    \"素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。\",
    \"新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。\",
    \"あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。\",
    \"おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。\",
]

embeddings = model.encode([query] + docs)
print(embeddings.shape)
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f\"{similarity:.04f}: {docs[i]}\")

(5, 1024)
0.1040: 素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。
0.2521: 新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。
0.4835: あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。
0.3199: おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。

如上所示，可以计算出与查询匹配的文本得分较高。在这个例句中，例如使用 BM25，由于文章中没有出现查询中包含的「ラーメン」这样的直接词汇，因此很难进行匹配。

接下来是相似文本任务的示例。

sentences = [
    \"明日の午後から雨が降るみたいです。\",
    \"来週の日曜日は天気が良いそうだ。\",
    \"あしたの昼過ぎから傘が必要になりそう。\",
    \"週末は晴れるという予報が出ています。\",
]

embeddings = model.encode(sentences)
similarities = model.similarity(embeddings, embeddings)

print(similarities)

# 显示第一句话与其他句子的相似度
for i, similarity in enumerate(similarities[0].tolist()):
    print(f\"{similarity:.04f}: {sentences[i]}\")

tensor([[1.0000, 0.2814, 0.3620, 0.2818],
        [0.2814, 1.0000, 0.2007, 0.5372],
        [0.3620, 0.2007, 1.0000, 0.1299],
        [0.2818, 0.5372, 0.1299, 1.0000]])
1.0000: 明日の午後から雨が降るみたいです。
0.2814: 来週の日曜日は天気が良いそうだ。
0.3620: あしたの昼過ぎから傘が必要になりそう。
0.2818: 週末は晴れるという予報が出ています。

同样，相似文本也获得了较高的分数。

另外，使用 Transformer 模型在 CPU 上生成文本向量时，即使文本量很少也会花费相当长的时间，有过这种经验的人应该不少。而 StaticEmbedding 模型在 CPU 速度还不错的情况下应该能瞬间完成。不愧是 100 倍速。

减小输出维度

标准生成的文本向量维度是 1024，但这可以进一步减小维度。例如让我们试试指定 128。

# truncate_dim 可以从 32, 64, 128, 256, 512, 1024 中指定
model = SentenceTransformer(model_name, device=\"cpu\", truncate_dim=128)

query = \"美味しいラーメン屋に行きたい\"
docs = [
    \"素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。\",
    \"新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。\",
    \"あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。\",
    \"おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。\",
]

embeddings = model.encode([query] + docs)
print(embeddings.shape)
similarities = model.similarity(embeddings[0], embeddings[1:])
for i, similarity in enumerate(similarities[0].tolist()):
    print(f\"{similarity:.04f}: {docs[i]}\")

(5, 128)
0.1464: 素敵なカフェが近所にあるよ。落ち着いた雰囲気でゆっくりできるし、窓際の席からは公園の景色も見えるんだ。
0.3094: 新鮮な魚介を提供する店です。地元の漁師から直接仕入れているので鮮度は抜群ですし、料理人の腕も確かです。
0.5923: あそこは行きにくいけど、隠れた豚骨の名店だよ。スープが最高だし、麺の硬さも好み。
0.3405: おすすめの中華そばの店を教えてあげる。とりわけチャーシューが手作りで柔らかくてジューシーなんだ。

变成了 128 维的向量，结果的分数也略有变化。由于维度减小，性能有所下降（后半部分有基准测试）。不过，从 1024 维减少到 128 维，可以减少存储空间，以及检索时使用的相似度计算成本大约降低 8 倍等，根据用途不同，较小的维度可能更有优势。

为什么 CPU 推理速度这么快？

StaticEmbedding 不是 Transformer 模型。也就是说，完全没有 Transformer 特有的 "Attention Is All You Need" 的注意力计算。只是将文本中出现的单词标记存储在 1024 维的表中，生成文本向量时只取其平均值。另外，由于没有注意力机制，所以不进行上下文理解等。

此外，在内部实现中，使用 PyTorch 的 nn.EmbeddingBag，通过传递所有连接的标记和偏移量进行处理，从而实现 PyTorch 优化的高速 CPU 并行处理和内存访问。

根据原文章的速度评估结果，在 CPU 上比 mE5-small 快 126 倍。

评估结果

JMTEB 的所有评估结果都记录在这个 JSON 文件中。与 JMTEB Leaderboard 中的其他模型进行比较，可以看到相对差异。考虑到模型大小，JMTEB 的整体评估结果非常出色。另外，JMTEB 的 mr-tidy 任务需要对 700 万篇文章进行向量化，处理时间相当长（根据模型不同，RTX4090 需要 1~4 小时）。这在 StaticEmbeddings 中也非常快，RTX4090 只需要大约 4 分钟就能完成处理。

信息检索中能否替代 BM25？

让我们看看 JMTEB 中信息检索任务的 Retrieval 结果。在 StaticEmbedding 中，mr-tidy 的项目明显较差。mr-tidy 比其他任务的文章数量多得多（700 万篇文章），也就是说在大量文章检索的任务中，结果可能会比较差。由于只是忽略上下文的简单标记平均，文章越多，平均相似的句子就越多，结果也可能如此。

因此，对于大量文章的情况，性能可能比 BM25 差很多。不过，对于文章较少且精确词汇匹配较少的情况，结果往往比 BM25 更好。

另外，信息检索任务 jaqket 的结果相对于其他模型异常好，虽然可能是因为学习了包含 jaqket 问题的 JQaRa（dev，未使用），但也高得有些奇怪。我认为没有泄露 test 信息...

聚类结果较差

这一点也没有详细追踪，但得分比其他模型差很多。分类任务并不差，所以很奇怪。可能是因为嵌入空间是通过套娃表示学习创建的影响吧。

JQaRA, JaCWIR 的重排序任务评估

JQaRA 的结果如下。

JaCWIR 的结果如下。

JQaRA 评估结果略好于 BM25，略低于 mE5-small，JaCWIR 结果明显低于 BM25 和 mE5。

JaCWIR 是从查询中查找的文本，是 Web 文章的标题和摘要，所以很多情况下不是所谓的"干净"文本。Transformer 模型对噪声有较强的抵抗力，而简单标记平均的 StaticEmbedding 在得分上会有差距，这也可以理解。BM25 匹配出现特征词汇的文章，因此在 JaCWIR 中，作为噪声的文章上的词汇本来就不会匹配查询，所以留下了与 Transformer 模型有竞争力的良好结果。

从这个结果来看，与 Transformer / BM25 相比，StaticEmbedding 在包含大量噪声的文本情况下得分可能较差。

输出维度缩减

StaticEmbedding 输出的维度根据训练方式而定，这次创建的是 1024 维，大小适中。维度数大，推理后的任务（聚类、信息检索等）的计算成本会增加。然而，由于训练时使用了套娃表示学习（Matryoshka Representation Learning(MRL)），因此可以轻松地将 1024 维进一步减小到更小的维度。

MRL 在训练时将更重要的维度放在向量的前面，例如即使是 1024 维，只使用前面的 32、64、128、256... 维，截断后面的部分，也能显示出相当良好的结果。

根据这个图表参考来源的 StaticEmbedding 文章，128 维保持 91.87% 的性能，256 维保持 95.79%，512 维保持 98.53% 的性能。如果对精度要求不是很严格，但想降低后续计算成本，可以直接进行维度缩减使用。

StaticEmbedding 日语模型的维度缩减结果

在 JMTEB 中，由于可以在输出时控制模型参数，因此通过传递 truncate_dim 选项，可以轻松测量维度缩减后的基准测试结果。这真是太棒了。因此，我们也对 StaticEmbedding 日语模型进行了维度缩减后的基准测试。

观察得分变化，当维度缩减到 512 维时，Retrieval、Classification、Reranking 的性能会变得很差。反而缩减到 256 维结果更好。在 256 维时，得分是维度缩减前模型的 98.93%，但这是因为聚类结果莫名其妙地比 1024 维更好。

修正了 512 维的得分计算错误。套娃表示学习很好地反映了，随着维度数的减少，得分略有下降，但由于维度数减少，后续成本可以降低。

在聚类任务中，即使维度缩减到 128 维，得分也比 1024 维更高，这是一个本应保留更多信息量得分更好的结果，但只有聚类任务出现了相反的得分上升的有趣结果...。在套娃表示学习中，前面的维度更好地考虑了整体特征，因此对于聚类用途（虽然也取决于聚类算法），可能只使用前面的特征维度而不使用后面的维度，会得到更好的结果。

因此，在 static-embedding-japanese 模型中进行维度缩减时，512、256、128 维左右在性能和维度缩减之间取得了较好的平衡。

创建 StaticEmbedding 模型后的感想

说实话，我对于简单的标记嵌入平均值能否产生如此好的性能半信半疑，但实际训练后，我对这种简单架构却能产生高性能感到惊讶。在这个 Transformer 盛行的时代，这种利用传统词嵌入的模型能够在现实世界中应用，我无法掩饰我的惊讶。

在 CPU 上推理速度快的文本向量生成模型，不仅在本地 CPU 环境中用于大量文本转换，还可以在边缘设备、网络速度慢（无法调用远程推理服务器）的环境等各种场景中使用。

StaticEmbedding 日语模型训练技术笔记

为什么能够很好地训练

StaticEmbedding 非常简单，只是将文本标记化后的 ID 从存储单词嵌入向量的 EmbeddingBag 表中获取 N 维（这次是 1024 维）向量，然后取其平均值。

到目前为止，说到单词嵌入向量，像 word2vec 和 GloVe 那样使用 Skip-gram 或 CBOW 来学习单词周围的方式。但是，StaticEmbedding 使用整个文本进行训练。此外，使用对比学习对大量各种文本进行大批量训练，成功学习了良好的单词嵌入表示。

对比学习基本上将正例以外的所有内容作为负例进行学习，例如，如果批量大小为 2048，则对 1 个正例和 2047 个负例进行 2048 种组合，即 2048x2047 约 400 万次比较学习。因此，可以一边对原始单词空间进行适当的权重更新，一边推进训练。

训练数据集

在日语模型训练中，我创建并使用了以下可用于对比学习的数据集。

• hotchpotch/sentence_transformer_japanese
  • 整理为 SentenceTransformer 易于训练的列名和结构。
    • 结构为 (anchor, positive), (anchor, positive, negative), (anchor, positive, negative_1, ..., negative_n)。
  • 基于以下数据集创建了 hotchpotch/sentence_transformer_japanese。再次感谢数据集的作者们，特别是 hpprc 氏。
    • https://huggingface.co/datasets/hpprc/emb
      • 使用 https://huggingface.co/datasets/hotchpotch/hpprc_emb-scores 的重排序得分，进行了 positive(>=0.7) / negative(<=0.3) 的过滤。
    • https://huggingface.co/datasets/hpprc/llmjp-kaken
    • https://huggingface.co/datasets/hpprc/msmarco-ja
      • 使用 https://huggingface.co/datasets/hotchpotch/msmarco-ja-hard-negatives 的重排序得分，进行了 positive(>=0.7) / negative(<=0.3) 的过滤。
    • https://huggingface.co/datasets/hpprc/mqa-ja
    • https://huggingface.co/datasets/hpprc/llmjp-warp-html
• 在上述创建的数据集中，使用了以下数据。另外，由于想要加强信息检索，因此通过增强增加了适合信息检索的数据集的数据量。
  • httprc_auto-wiki-nli-triplet
  • httprc_auto-wiki-qa
  • httprc_auto-wiki-qa-nemotron
  • httprc_auto-wiki-qa-pair
  • httprc_baobab-wiki-retrieval
  • httprc_janli-triplet
  • httprc_jaquad
  • httprc_jqara
  • httprc_jsnli-triplet
  • httprc_jsquad
  • httprc_miracl
  • httprc_mkqa
  • httprc_mkqa-triplet
  • httprc_mr-tydi
  • httprc_nu-mnli-triplet
  • httprc_nu-snli-triplet
  • httprc_quiz-no-mori
  • httprc_quiz-works
  • httprc_snow-triplet
  • httprc_llmjp-kaken
  • httprc_llmjp_warp_html
  • httprc_mqa_ja
  • httprc_msmarco_ja
• 英语数据集使用了以下数据集。
  • sentence-transformers/msmarco-co-condenser-margin-mse-sym-mnrl-mean-v1
  • sentence-transformers/squad
  • sentence-transformers/all-nli
  • sentence-transformers/trivia-qa
  • nthakur/swim-ir-monolingual
  • sentence-transformers/miracl
  • sentence-transformers/mr-tydi

日语分词器

为了训练 StaticEmbedding，使用 HuggingFace 分词器库的 tokenizer.json 格式处理似乎比较简单，因此我创建了 hotchpotch/xlm-roberta-japanese-tokenizer 分词器。词汇量为 32,768。

这个分词器使用 unidic 分割 wikipedia 日语~~、wikipedia 英语（采样）、cc-100（日语，采样）~~（更正：确认创建代码后发现只使用了 wikipedia 日语）的数据，并用 sentencepiece unigram 训练。它也可以作为 XLM-Roberta 格式的日语分词器使用。这次使用了这个分词器。

超参数

与原始训练代码的更改点和笔记如下。

• batch_size 从原始的 2048 设置为 6072。
  • 在对比学习中处理大批量时，如果同一批次中包含正例和负例，可能会对学习产生负面影响。为了防止这种情况，有 BatchSamplers.NO_DUPLICATES 选项。但是，当批量大小很大时，为了不包含在同一批次中而进行的采样处理可能需要时间。
  • 这次指定了 BatchSamplers.NO_DUPLICATES，设置为 RTX4090 24GB 可容纳的 6072。批量大小更大可能结果更好。
• epoch 数从 1 改为 2
  • 2 比 1 结果更好。不过，如果数据大小更大，1 可能更好。
• 调度器
  • 从标准的 linear 改为经验上感觉更好的 cosine。
• 优化器
  • 保持标准的 AdamW。改为 adafactor 后，收敛变差。
• learning_rate
  • 保持 2e-1。虽然怀疑值是否太大，但降低后结果变差。
• dataloader_prefetch_factor=4
• dataloader_num_workers=15
  • 由于标记化和批量采样器的采样需要时间，设置得较大。

训练资源

• CPU
  • Ryzen9 7950X
• GPU
  • RTX4090
• 内存
  • 64GB

使用这台机器资源，完全从头开始训练大约需要 4 小时。GPU 核心负载非常小，与其他 transformer 模型训练时负载在 90% 左右不同，StaticEmbedding 几乎为 0%。这可能是因为大部分时间都花在将大批量传输到 GPU 内存上。因此，如果 GPU 内存带宽更快，训练速度可能会进一步提高。

进一步提高性能

这次使用的分词器不是专门为 StaticEmbedding 设计的，如果使用更合适的分词器，性能可能会提高。通过进一步增大批量大小，可以提高训练稳定性，有望提高性能。

此外，通过使用各种领域和合成数据集，将更广泛的文本资源纳入训练，可以期待进一步提高性能。

原始训练代码

训练使用的代码以 MIT 许可证在以下公开。运行脚本应该可以重现...!

• https://huggingface.co/hotchpotch/static-embedding-japanese/blob/main/trainer.py

许可证

static-embedding-japanese 以 MIT 许可证公开模型权重和训练代码。