lettuce-emb-512d-v1 — 512维蒸馏嵌入模型 (ONNX)

Zeolit/lettuce-emb-512d-v1 是一个定制的512维句子嵌入模型，通过从 BAAI/bge-m3 蒸馏训练而成，专门针对对话/角色扮演记忆用例进行了优化。

它基于 DistilBERT（512词元上下文）构建，配有512维投影头，并导出为 ONNX 格式，以实现快速、可移植的推理（桌面端 + 移动端）。

该模型旨在为 LettuceAI（一个专注于角色扮演的LLM客户端）的记忆系统提供支持，但也可用作通用语义嵌入模型。

🏆 性能亮点

检索性能 (ArguAna)

总体：在小于1亿参数模型中排名第1

这意味着：

• 击败所有"微型/小型"嵌入模型（MiniLM、gte-small、e5-small、Snowflake-xs）
• 仅被1亿参数以上的模型超越（gte-base: 0.46, bge-base: 0.58, bge-small: 0.51, bge-m3: 0.68）
• 对于6600万参数来说性能卓越 — 远超同级水平

语义相似度性能 (STS基准测试)

优势：

• STS13 (0.649)：匹配/击败2-3倍大小的模型
• STS15 (0.676)：与顶级小模型有竞争力（e5-small: 0.78, gte-small: 0.77）
• 平衡的Pearson分数（0.616-0.644）显示一致的相关性质量

背景：

• 在1000万-1亿参数模型的语义相似度任务中处于中游水平
• 权衡：针对检索而非完美的STS分数进行了优化
• 这是一个有意的设计选择，适用于实际应用中找到正确文档比完美相似度评分更重要的情况

综合评估

✅ 顶级检索 — 6600万参数中最佳
✅ 有竞争力的语义相似度 — 适合聚类、去重
✅ 高效 — 卓越的尺寸性能比
✅ 实用 — 针对现实世界的RAG和搜索用例进行了优化

模型摘要

• 基础编码器： distilbert-base-uncased
• 参数量： 6600万
• 教师模型： BAAI/bge-m3
• 嵌入维度： 512
• 最大序列长度： 512 词元
• 格式： ONNX (FP32)
• 池化方式： 最后一层隐藏状态的均值池化 + L2归一化
• 领域： 通用英文文本，额外关注对话/角色扮演数据

该模型的训练目标：

• 在嵌入空间中匹配BGE-M3的相似度结构（批次级知识蒸馏）
• 将语义相关的文本拉近（对比损失）
• 捕捉对话、角色驱动和角色扮演风格文本中的细微差别

训练与蒸馏

高层训练设置：

• 学生模型： DistilBERT，在池化输出之上有一个512维投影头
• 教师模型： BGE-M3（1024维多语言嵌入模型）
• 损失函数：
  • 几何蒸馏： 学生和教师批次相似度矩阵之间的MSE
  • 对比损失： 正样本对的 1 - cosine(student_emb_1, student_emb_2)
• 批次大小： 16
• 最大序列长度： 512
• 优化器： AdamW，lr = 2e-5
• 轮次： 2（在组合数据集上）

数据集

学生模型在以下混合数据上训练：

1. 自定义角色扮演语料库 (JSONL)
   RP风格对话日志中的连续行被转换为正样本对：

   • ("utterance_i", "utterance_{i+1}")

2. stsb_multi_mt (en)
   标准语义相似度数据集，用作额外的通用监督。

这些共同教会模型：

• 通用句子级语义相似度
• 对话轮次与对话流程
• 角色扮演特定语义（情感细微差别、关系、"场景连续性"）

使用方法 (Python + ONNX Runtime)

1. 安装依赖

pip install onnxruntime transformers

2. 基础嵌入示例

import onnxruntime as ort
import numpy as np
from transformers import AutoTokenizer

model_id = "Zeolit/lettuce-emb-512d-v1"

tokenizer = AutoTokenizer.from_pretrained(model_id)
session = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

MAX_SEQ_LEN = 512

def embed(texts):
    if isinstance(texts, str):
        texts = [texts]

    enc = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=MAX_SEQ_LEN,
        return_tensors="np",
    )

    outputs = session.run(
        ["sentence_embedding"],
        {
            "input_ids": enc["input_ids"],
            "attention_mask": enc["attention_mask"],
        },
    )[0]
    # outputs: (batch_size, 512)
    return outputs

def cosine(a, b):
    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))

sents = [
    "The hero shielded his companion from the arrow.",
    "He risked everything to save her life.",
    "The market is bustling with traders today.",
]

emb = embed(sents)
print(emb.shape)  # (3, 512)
print("cos(0,1) =", cosine(emb[0], emb[1]))
print("cos(0,2) =", cosine(emb[0], emb[2]))

🧩 预期用途

适用于：

• 对话/角色扮演记忆检索（主要用例）
• 语义搜索和文档检索
• RAG（检索增强生成）系统
• 聚类和重复检测
• 轻量级、本地嵌入，支持：
  • 桌面客户端（Tauri、Electron等）
  • 移动客户端（通过ONNX Runtime）

优势：

• 出色的检索性能 — 高精度找到相关文档
• 高效尺寸 — 6600万参数，推理速度快
• 长上下文 — 512词元比典型小模型捕捉更多信息

权衡：

• 更多针对检索优化，而非完美的语义相似度评分
• STS性能扎实但与更大模型相比不是顶级
• 这是一个有意的设计选择，适用于实际应用中找到正确文档比完美相似度分数更重要的情况

非目标：

• 高风险安全关键应用
• 与BGE-M3相当的多语言性能（该模型主要专注于英语）
• 跨数百万向量的超大规模检索（考虑更大的模型或专门系统）

限制与偏见

• 该模型在英文文本上训练，可能无法很好地泛化到其他语言。
• 角色扮演风格语料库可能包含有偏见、非正式或情绪化的内容。嵌入可能反映这些模式。
• 该模型不执行任何安全过滤；它只是将文本编码为向量。
• 在某些STS任务（STS12、STS14）上较弱，这些任务更偏向重度蒸馏或对比模型。

在生产环境中使用前，请始终在您的具体应用场景中评估该模型。

基准测试

详细检索指标 (ArguAna)

核心指标：

召回性能：

排名质量：

与相似模型的对比

差距分析：

• 对比小于1亿参数模型： nDCG@10提升10-20%
• 对比1亿参数模型： 尽管尺寸小34%，仅落后gte-base 14%
• 对比教师模型 (bge-m3)： 落后42%，但尺寸仅为11.6%

语义相似度详情 (STS)

STS12 (2012基准测试)：

STS13 (2013基准测试)：

STS14 (2014基准测试)：

STS15 (2015基准测试)：

与小模型对比 (STS)

性能说明：

• STS13优势： 基本匹配MiniLM-L12（两者均约0.65）
• STS15优势： 接近顶级表现者（0.676 vs 0.75-0.78范围）
• STS12/14权衡： 落后领先者10-15% — 有改进空间
• 总体： 扎实的中游性能，针对检索而非峰值STS分数进行了优化

许可证

该模型以 Apache-2.0 许可证发布。

它派生自：

• distilbert-base-uncased (Apache-2.0) 作为学生骨干
• BAAI/bge-m3 (MIT) 作为蒸馏教师

致谢

• HuggingFace Transformers
• ONNX Runtime
• BAAI / BGE
• LettuceAI项目，激励了专注于角色扮演的嵌入模型的开发