Sentence Transformers truncate_embeddings 怎样截取嵌入向量维度

如果向量检索已经能跑起来，但索引存储、内存占用或接口字段又要求更短的向量，可以用 sentence_transformers.util.truncate_embeddings 做维度截取。它接收一个嵌入矩阵和目标维度，返回同样保留样本行数、但列数缩短的新矩阵；传入 None 时则不做截取。

官方地址：https://sbert.net/docs/package_reference/util/tensor.html

最小用法是 truncate_embeddings(embeddings, 128)。先确认原始矩阵是“样本数 × 原始维度”，再选择不超过原始维度的目标值，并让向量索引、相似度函数和持久化字段一起接受新的形状。这个函数只截取已有向量，不会重新训练模型或自动评估业务效果。

先把“截取维度”与“截取样本”分开

一次编码通常会得到二维嵌入矩阵，例如 3 条文本得到一个形状为 (3, 768) 的结果。这里的 3 是样本数，768 才是每个句向量的维度。truncate_dim=128 只会改变第二个数字，结果应当是 (3, 128)；它不会删除第二、第三条文本。

我在排查“向量表变小了但检索结果不对”的故障时，最先确认的就是这两个轴。若把批次切片写成 embeddings[:128]，切掉的是行；使用工具函数的第二个参数，才是在表达“每条向量保留多少维”。

用 truncate_embeddings 截取到目标维度

输入可以是 NumPy 数组，也可以是 PyTorch Tensor。下面的示例同时打印原始和截取后的形状，并用同一模型计算相似度，方便把“形状变化”和“仍能参与后续计算”分开观察。

import numpy as np
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import truncate_embeddings

# 使用支持 Matryoshka 训练的模型，便于演示不同维度的向量表示。
model = SentenceTransformer("tomaarsen/mpnet-base-nli-matryoshka")
texts = ["今天天气很好", "今天是晴天", "我需要去上班"]

# 先得到完整的句向量，再确认输入确实是二维矩阵。
embeddings = model.encode(texts, convert_to_numpy=True)
print("原始形状:", embeddings.shape)

# 第二个参数是每条句向量要保留的目标维度，不是样本条数。
target_dim = 128
short_embeddings = truncate_embeddings(embeddings, target_dim)
print("截取后形状:", short_embeddings.shape)

# 截取后仍是一个可参与相似度运算的嵌入矩阵。
similarity = model.similarity(short_embeddings, short_embeddings)
print("相似度矩阵形状:", tuple(similarity.shape))

对于这段示例，关键输出关系是 (3, 768) 变成 (3, 128)。不要把这段代码理解为模型重新生成了 128 维特征：模型先产出完整向量，工具再按目标维度截取。

让下游索引和存储字段一起换维度

维度截取本身通常很简单，真正容易出错的是只修改了编码代码，却没有修改下游组件。原来按 768 维建立的索引，不能直接拿 128 维向量去查询；数据库字段、批量接口、缓存键对应的序列化格式，也要同步约定。

如果需要逐步迁移，可以给新索引使用独立的版本标识，先用同一批文本同时生成完整向量和截短向量，比较业务指标，再决定是否切换默认查询路径。不要把“接口返回 200”当成召回质量没有变化的证据。

None、输入类型和模型适配的边界

truncate_dim=None 表示不截取

文档把 None 定义为不做截取。它适合在配置项可选时保持统一调用路径，但不会替你推断一个更合适的维度。

NumPy 和 Tensor 都可以作为输入

如果上游使用 convert_to_numpy=True，得到的是 NumPy 数组；如果保留 PyTorch Tensor，则可以直接传入。不要在同一条流水线里无意间混用设备上的 Tensor 和 CPU 上的数组，尤其是在后续还要把结果送回 GPU 的场景中。

模型是否适合短维度，要单独判断

官方示例使用 tomaarsen/mpnet-base-nli-matryoshka，并展示 768 维与 128 维表示的相似度计算。这个示例说明该模型可以用来演示可截取表示，但不等于所有嵌入模型在任意目标维度下都有相同的效果。对于没有针对截取维度训练或评估过的模型，应把目标维度当作需要基准测试的工程参数，而不是默认的质量无损开关。

排查“截完以后结果变差”的顺序

1. 打印输入和输出的形状，确认变化发生在列维度。
2. 确认写入索引的全部向量都使用同一个目标维度。
3. 用固定的小型评测集比较原始维度与目标维度的召回、相似度分布和阈值。
4. 检查查询向量和库中向量是否由相同的模型、相同的截取策略生成。
5. 若指标下降明显，先回退到旧索引，再重新选择目标维度或模型，而不是继续调低阈值掩盖问题。

常见问题

truncate_embeddings 会重新计算嵌入吗？

不会。它处理已经生成的嵌入矩阵，并返回截取后的结果；模型编码和维度截取是两个阶段。

为什么不能把 embeddings[:128] 当成截取 128 维？

二维数组的第一个切片轴通常是样本轴，embeddings[:128] 表示最多保留 128 条样本，不是保留 128 个维度。要截取维度，应使用 truncate_embeddings(embeddings, 128)。

截取后能直接查询原来的向量索引吗？

通常不能直接混用。索引维度、查询向量维度和库中向量维度必须采用同一约定；切换维度时应建立或迁移到匹配的新索引。

目标维度越小越好吗？

不是。维度越小通常越节省存储和计算，但可能损失区分能力。应使用固定评测集比较业务指标，再选择能接受的目标维度。

总结：先把输入确认成“样本数 × 原始维度”，再调用 truncate_embeddings(embeddings, target_dim) 截取列维度；随后同步调整索引、存储和评测配置。它是一个明确的表示变换工具，不是自动压缩质量的保证。