Matryoshka Embeddings:让同一个语义向量在多个维度下都能工作

普通稠密嵌入模型通常固定输出 768 或 1024 维向量。分类、聚类、语义搜索以及后续排序都使用这一完整表示。如果数据库有数千万条向量,维度就会直接影响存储、传输和相似度计算成本。Matryoshka Representation Learning(MRL,套娃表征学习)改变的是训练目标:让向量的前一部分也包含有用信息,从而允许部署时选择不同长度的向量。

本文译编自 Sentence Transformers 官方文档,原文作者为 Sentence Transformers 文档贡献者,实验模型由 Tom Aarsen 发布。核对日期为 2026 年 10 月 5 日;同时核对了原文链接的 NLI 训练脚本和 STSBenchmark 评估脚本。下文所有性能数字均来自原文,未在本次编辑中重新训练或复测。

同一条768维嵌入依次取512、256、128与64维前缀,每个前缀在训练时分别计算损失,推理时按检索预算选择维度。
原创技术示意图:多个前缀共用一次完整编码结果;图形长度用于表达包含关系,不表示实测速度。

先分清节省发生在哪里

套娃向量可以用于“两阶段”检索:先用较短向量筛出候选,再用完整向量对候选重新排序。这是一种 shortlisting and reranking 方法;这里的重新排序仍可基于向量相似度,并不意味着自动引入 cross-encoder。另一个用途是按预算选择存储维度,在检索质量、向量库容量和计算量之间取得合适平衡。

截断输出不会自动缩小模型权重,也不会减少编码器的层数或前向计算。原文明确指出,单纯使用较短 Matryoshka 向量,不会使模型训练和编码推理更快、更省内存或更小。节省主要发生在生成向量之后的存储与处理阶段。如果候选生成只保存短向量,而精排需要完整向量,系统仍须另外保存、缓存或重新生成完整表示。

原文实验说明了什么

原文比较两个以 microsoft/mpnet-base 为基础的模型:mpnet-base-nli-matryoshka 使用 MatryoshkaLoss 包装 MultipleNegativesRankingLoss;mpnet-base-nli 只使用后者。两者都用 AllNLI 训练,AllNLI 由 SNLI 和 MultiNLI 合并而来;之后在 STSBenchmark 测试集上,分别使用不同向量长度评估 Spearman 相关性。

该实验中,Matryoshka 模型在各测试维度上的相关性都较高,随维度减少的下降也较缓。保留原始维度的约 8.3% 时,套娃模型保留了自身最高表现的 98.37%,普通模型为 96.46%。这两个比例的分母是各模型在该实验中的最大表现,而非一个通用的检索准确率。它们支持“短向量可能在此任务上保留较多质量”的判断,不能直接推导出任何语料、语言或向量数据库都会获得相同收益。

把同一损失应用于多个前缀

MRL 的训练方法很直接。假设模型输出 768 维向量,可对 768、512、256、128、64 等前缀分别计算同一个损失,再把损失加总;也可以为不同维度指定权重。原文文字还举了 32 维这个候选,但实际代码列表的最小维度是 64,二者不应混为同一次配置。

from sentence_transformers import SentenceTransformer
from sentence_transformers.sentence_transformer.losses import CoSENTLoss, MatryoshkaLoss

model = SentenceTransformer(
    "microsoft/mpnet-base",
    model_kwargs={"torch_dtype": "float32"},
)
base_loss = CoSENTLoss(model=model)
loss = MatryoshkaLoss(
    model=model,
    loss=base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],
)

这段代码只定义模型和损失,尚未准备数据或启动训练。训练显存允许时,原文倾向于用 FP32 加载模型;实际混合精度选项需与硬件匹配。模型最大输出维度也必须覆盖列表里的各维度。不能把这个列表照搬到一个只有 384 维输出的模型上。

如果同时希望减少编码器层数,可结合 AdaptiveLayerLoss。Sentence Transformers 提供 Matryoshka2dLoss 作为层数与输出维度联合训练的简写:

from sentence_transformers.sentence_transformer.losses import Matryoshka2dLoss

loss = Matryoshka2dLoss(
    model=model,
    loss=base_loss,
    matryoshka_dims=[768, 512, 256, 128, 64],
)

减少层数的模型另有训练与评估要求,应参见官方 Adaptive Layers 文档。它与仅截断输出向量是两件事,不能把前者的潜在编码收益归给后者。

完整 NLI 训练脚本怎样串起流程

原文链接的 matryoshka_nli.py 才是完整训练入口。其默认模型是 distilbert/distilroberta-base,也接受命令行传入模型名;原文性能比较使用的是另行指定的 MPNet,而非这个默认值。如果基础模型还不是 Sentence Transformer,加载过程会建立带 mean pooling 的表示模型。

脚本读取 sentence-transformers/all-nli 的 triplet 配置,分别取 train 和 dev。NLI 的蕴含关系形成正样本,矛盾句形成困难负样本。MultipleNegativesRankingLoss 利用批内负样本,因此脚本设置 BatchSamplers.NO_DUPLICATES,避免同一批中的重复样本破坏负样本假设。它再用 MatryoshkaLoss 同时训练五个前缀维度。

from datasets import load_dataset
from sentence_transformers.sentence_transformer.losses import (
    MatryoshkaLoss, MultipleNegativesRankingLoss,
)

train_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="train")
eval_dataset = load_dataset("sentence-transformers/all-nli", "triplet", split="dev")
dims = [768, 512, 256, 128, 64]
train_loss = MatryoshkaLoss(model, MultipleNegativesRankingLoss(model), matryoshka_dims=dims)

训练之外,脚本给 STSBenchmark 验证集的每个维度分别建立 EmbeddingSimilarityEvaluator,设置余弦相似度与对应的 truncate_dim,再用 SequentialEvaluator 组合。这样既能观察训练损失,也能观察每种维度的语义相似度质量。脚本采用 1 个 epoch、每设备 batch size 128、warmup_steps=0.1、每 100 步评估/保存/记录日志、最多保留 2 个 checkpoint,并启用 fp16=True、bf16=False。这些是示例参数;FP16 不受支持时不能原样使用,增加 batch size 也会增加显存需求。

注意脚本头部注释写着“每 10% 训练步数评估”,实际参数却是固定 eval_steps=100。复现时应以代码配置为准并记录总步数,不能同时声称执行了两种计划。训练结束后脚本对 STSBenchmark 的 test 切分逐维评估,再保存到输出目录的 final 子目录。测试集应只用于最终报告,不能反复据此挑选维度、超参数或模型后仍声称是独立测试。

静态审核发现:官方训练脚本最后的 push_to_hub() 虽标作“可选”,代码本身仍会尝试执行,只在失败时记录异常。如果机器已经登录,它可能对外上传模型。本文的训练流程止于 model.save(final_output_dir),不包含该调用;这是明确的安全改编。脚本还可能在安装了 W&B 时启用相应日志集成,训练私有数据前应明确配置报告目标。模型、数据集和依赖都应固定可追溯版本,不让远程默认分支变化冒充同一实验。

推理时选择维度

只有经过相应训练的模型,才有依据在指定前缀长度下保持语义质量。对任意现成模型直接裁掉尾部,不等价于 Matryoshka 训练。原文使用 nomic-ai/nomic-embed-text-v1.5 演示 64 维推理,同时分别添加查询前缀 search_query: 与文档前缀 search_document:。

from sentence_transformers import SentenceTransformer

matryoshka_dim = 64
model = SentenceTransformer(
    "nomic-ai/nomic-embed-text-v1.5",
    trust_remote_code=True,
    truncate_dim=matryoshka_dim,
)
embeddings = model.encode([
    "search_query: What is TSNE?",
    "search_document: t-distributed stochastic neighbor embedding (t-SNE) is a statistical method for visualizing high-dimensional data by giving each datapoint a location in a two or three-dimensional map.",
    "search_document: Amelia Mary Earhart was an American aviation pioneer and writer.",
])
assert embeddings.shape[-1] == matryoshka_dim
similarities = model.similarity(embeddings[0], embeddings[1:])

原文给出的两个相似度是 0.7839 与 0.4933,对应相关文档和无关文档。本次未运行该示例,不能把这两个数当成本机测试输出。不同模型 revision、依赖、精度和处理方式都可能改变数值。原文导入了未使用的 torch.nn.functional,此处已删除,仅属清理。

trust_remote_code=True 允许执行模型仓库里的 Python 代码,是实质性信任决定。上面的原文示例为便于核对而保留;正式运行前需审阅仓库代码并把 revision 固定到已审核的提交 SHA,而不是直接信任可变化的最新代码。本文没有编造一个“已审核 SHA”。查询与文档的维度、提示前缀和归一化策略必须一致;若手工切片后使用点积替代余弦,还需重新确认归一化假设。

逐维评估时的兼容性检查

独立评估脚本 默认比较前述两个 MPNet 模型,维度为 768、512、256、128、64。它读取 mteb/stsbenchmark-sts 测试切分,将原始分数除以 5,使用余弦相关性评估,然后生成绝对 Spearman 曲线与“相对各自最大表现”的比例图。

静态检查还发现两个运行前需要处理的细节。其一,脚本直接把 test_evaluator(...) 的返回值当作标量存入绘图数据,而不同 Sentence Transformers 版本可能返回指标字典;应在所用版本中检查返回契约,明确提取所需的 Spearman 指标,不应把整张字典传入数值计算。其二,os.makedirs(output_path) 未设置 exist_ok=True,重复运行可能因目录已存在而停止。添加该选项前也要决定是否允许覆盖旧指标,最好为每次实验使用独立目录。本文没有将未执行的兼容性建议描述为“修复已通过”。

选择后续示例与部署维度

官方还提供 matryoshka_nli_reduced_dim.py,用于训练最大输出维度为 256 的 NLI 模型;matryoshka_sts.py 用 STSBenchmark 训练集和 CoSENTLoss;2d_matryoshka_nli.py 与 2d_matryoshka_sts.py 则展示联合层数和维度的训练。选择脚本时应先确定数据监督形式和希望优化的成本,不能只看名称相近就替换损失。

部署维度最终应由自己的评估决定:保持数据切分和评估指标一致,分别记录短向量召回质量、候选数量、精排质量、索引大小和检索耗时。原文说明了多维表示的可行性;要证明生产系统的收益,仍须在真实数据分布与检索设置下完成独立实验。

来源与许可:译编基于上述官方文档、训练/评估脚本及其链接资料。Sentence Transformers 项目使用 Apache License 2.0,相关许可证随稿保存在 sources/LICENSE.txt;模型与数据集分别受各自仓库许可约束,项目许可不替代模型或数据许可。全文翻译、转载和配图授权由委托方于 2026-10-05 确认。图由未完纪编辑制作,非原文实验图。本文仅作静态代码审查,未执行训练、下载模型、登录平台或上传模型;未发现问题不代表不存在漏洞。

保留的完整许可声明

以下为原项目适用许可文本,原作者与文档归属及本文改动说明见正文。

Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.
      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.
      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.
      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).
      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. For the purposes
      of this License, Derivative Works shall not include works that remain
      separable from, or merely link (or bind by name) to the interfaces of,
      the Work and Derivative Works thereof.
      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."
      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.
   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.
   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.
   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and
      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and
      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.
      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.
   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.
   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.
   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.
   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.
   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.
   END OF TERMS AND CONDITIONS

   APPENDIX: How to apply the Apache License to your work.
      To apply the Apache License to your work, attach the following
      boilerplate notice, with the fields enclosed by brackets "{}"
      replaced with your own identifying information. (Don't include
      the brackets!)  The text should be enclosed in the appropriate
      comment syntax for the file format. We also recommend that a
      file or class name and description of purpose be included on the
      same "printed page" as the copyright notice for easier
      identification within third-party archives.
   Copyright 2019 Nils Reimers

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
limitations under the License.

Sentence Transformers 上游通知

以下保留所引用上游代码的官方 NOTICE.txt,与原有完整 Apache License 2.0 一起提供;其范围不扩张到第三方模型或数据集。

-------------------------------------------------------------------------------
Sentence Transformers

Copyright 2019-2025
Ubiquitous Knowledge Processing (UKP) Lab
Technische Universität Darmstadt

Copyright 2025-present
Hugging Face, Inc.
-------------------------------------------------------------------------------
© 版权声明
THE END
喜欢就支持一下吧
点赞0 分享
评论 抢沙发

请登录后发表评论

    暂无评论内容