stella 模型
stella是一个通用的中文文本编码模型,目前有两个版本:base 和 large,两个版本的模型均支持1024的输入长度。
完整的训练思路和训练过程已记录在博客,欢迎阅读讨论。
训练数据:
- 开源数据(wudao_base_200GB[1]、m3e[2]和simclue[3]),重点选择了长度大于512的文本
- 在通用语料库上使用大型语言模型构造了一批(问题, 段落)和(句子, 段落)数据
训练方法:
- 对比学习损失函数
- 带有难负例的对比学习损失函数(分别基于bm25和向量构造了难负例)
- EWC(弹性权重整合)[4]
- cosent损失[5]
- 每种类型的数据使用一个迭代器,分别计算损失进行更新
初始权重: stella-base-zh和stella-large-zh分别以piccolo-base-zh[6]和piccolo-large-zh作为基础模型,512-1024的位置嵌入使用层次分解位置编码[7]进行初始化。 感谢商汤科技研究院开源的piccolo系列模型。
stella是一个通用中文文本编码模型,目前有两个版本:base和large,两个版本均支持1024的输入长度。
训练数据主要包括:
- 开源训练数据(wudao_base_200GB、m3e和simclue),重点选择了长度大于512的文本。
- 在通用语料库上使用大型语言模型构造的一批(问题,段落)和(句子,段落)数据。
损失函数主要包括:
- 对比学习损失函数
- 带有难负例的对比学习损失函数(基于bm25和向量构造难负例)
- EWC(弹性权重整合)
- cosent损失
模型权重初始化: stella-base-zh和stella-large-zh分别使用piccolo-base-zh和piccolo-large-zh作为基础模型,512-1024位置嵌入使用层次分解位置编码的初始化策略。
训练策略: 每种类型的数据使用一个迭代器,分别计算损失。
项目文件说明
./run_train.py # 训练脚本
./src/add_new_pos_embed.py # 扩展现有模型长度的脚本
./src/run_eval_stella.py # 评估cmteb效果的脚本
指标
C-MTEB 排行榜
stella模型在C-MTEB[8]的结果,评测脚本请参见博客。
模型名称 | 模型大小 (GB) | 维度 | 序列长度 | 平均 (35) | 分类 (9) | 聚类 (4) | 对分类 (2) | 重排序 (4) | 检索 (8) | 语义相似度 (8) |
---|---|---|---|---|---|---|---|---|---|---|
stella-large-zh | 0.65 | 1024 | 1024 | 64.54 | 67.62 | 48.65 | 78.72 | 65.98 | 71.02 | 58.3 |
stella-base-zh | 0.2 | 768 | 1024 | 64.16 | 67.77 | 48.7 | 76.09 | 66.95 | 71.07 | 56.54 |
piccolo-large-zh | 0.65 | 1024 | 512 | 64.11 | 67.03 | 47.04 | 78.38 | 65.98 | 70.93 | 58.02 |
bge-large-zh | 1.3 | 1024 | 512 | 63.96 | 68.32 | 48.39 | 78.94 | 65.11 | 71.52 | 54.98 |
piccolo-base-zh | 0.2 | 768 | 512 | 63.66 | 66.98 | 47.12 | 76.61 | 66.68 | 71.2 | 55.9 |
bge-large-zh-no-instruct | 1.3 | 1024 | 512 | 63.4 | 68.58 | 50.01 | 76.77 | 64.9 | 70.54 | 53 |
bge-base-zh | 0.41 | 768 | 512 | 62.8 | 67.07 | 47.64 | 77.5 | 64.91 | 69.53 | 54.12 |
长文本评估
通过实际观察发现,C-MTEB的评测数据长度基本都小于512, 更关键的是那些长度大于512的文本,其重点都在前半部分。 这里以CMRC2018的数据为例说明这个问题:
问题:《无双大蛇z》是谁旗下ω-force开发的动作游戏?
段落:《无双大蛇z》是光荣旗下ω-force开发的动作游戏,于2009年3月12日登陆索尼playstation3,并于2009年11月27日推......
段落长度为800多,大于512,但对于这个问题而言只需要前面40个字就足以检索,多余的内容对模型来说是一种噪声,反而降低了效果。 简言之,现有数据集的两个问题: 1)长度大于512的数据过少 2)即便大于512,对于检索而言也只需要前512的文本内容 导致无法准确评估模型的长文本编码能力。
为解决这个问题,我们收集了相关开源数据并使用规则进行过滤,最终整理了6份长文本测试集,分别是:
- CMRC2018,通用百科
- CAIL,法律阅读理解
- DRCD,繁体百科,已转简体
- Military,军工问答
- Squad,英文阅读理解,已转中文
- Multifieldqa_zh,清华的大模型长文本理解能力评测数据[9]
处理规则是选取答案在512长度之后的文本,短的测试数据会进行欠采样,长短文本占比约为1:2,因此模型既需要理解短文本也需要理解长文本。 除了Military数据集,我们提供了其他5个测试数据的下载地址:https://drive.google.com/file/d/1WC6EWaCbVgz-vPMDFH4TwAMkLyh5WNcN/view?usp=sharing
评测指标为Recall@5,结果如下:
数据集 | piccolo-base-zh | piccolo-large-zh | bge-base-zh | bge-large-zh | stella-base-zh | stella-large-zh |
---|---|---|---|---|---|---|
CMRC2018 | 94.34 | 93.82 | 91.56 | 93.12 | 96.08 | 95.56 |
CAIL | 28.04 | 33.64 | 31.22 | 33.94 | 34.62 | 37.18 |
DRCD | 78.25 | 77.9 | 78.34 | 80.26 | 86.14 | 84.58 |
Military | 76.61 | 73.06 | 75.65 | 75.81 | 83.71 | 80.48 |
Squad | 91.21 | 86.61 | 87.87 | 90.38 | 93.31 | 91.21 |
Multifieldqa_zh | 81.41 | 83.92 | 83.92 | 83.42 | 79.9 | 80.4 |
平均 | 74.98 | 74.83 | 74.76 | 76.15 | 78.96 | 78.24 |
注意: 由于长文本评估数据数量有限,在构建时也使用了训练部分数据。如果您自行评估,请注意模型的训练数据以避免数据泄露。 |
使用方法
本模型是在piccolo基础上训练的,因此使用方法与piccolo完全相同。
注意:在stella中,指令中的冒号是英文冒号,即查询:
和结果:
。
在sentence-transformer库中的使用方法:
# 对于短对短数据集,以下是通用的使用方式
from sentence_transformers import SentenceTransformer
sentences = ["数据1", "数据2"]
model = SentenceTransformer('infgrad/stella-base-zh')
print(model.max_seq_length)
embeddings_1 = model.encode(sentences, normalize_embeddings=True)
embeddings_2 = model.encode(sentences, normalize_embeddings=True)
similarity = embeddings_1 @ embeddings_2.T
print(similarity)
# 如果是短对长数据集,建议添加指令,以帮助模型更好地进行检索。
# 注意指令中使用的是英文冒号
直接使用transformers库:
from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize
model = AutoModel.from_pretrained('infgrad/stella-base-zh')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh')
sentences = ["数据1", "数据ABCDEFGH"]
batch_data = tokenizer(
batch_text_or_text_pairs=sentences,
padding="longest",
return_tensors="pt",
max_length=1024,
truncation=True,
)
attention_mask = batch_data["attention_mask"]
model_output = model(**batch_data)
last_hidden = model_output.last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
vectors = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
vectors = normalize(vectors, norm="l2", axis=1, )
print(vectors.shape) # 2,768
训练细节
硬件: 单卡A100-80GB
环境: torch1.13.*; transformers-trainer + deepspeed + gradient-checkpointing
学习率: 1e-6
批量大小: base模型为1024,额外增加20%的难负例;large模型为768,额外增加20%的难负例
数据量: 约100万,其中用LLM构造的数据约有200K。LLM模型大小为13b
待办事项
评测的稳定性:
在评测过程中发现聚类任务的结果与官方结果不完全一致,存在约±0.0x的小差异,基本可以忽略不计,不影响评测结论。
但是不完全一致仍然难以理解,本人尝试了bge和piccolo系列的模型,都存在这个问题。个人猜测可能与使用的库、批量大小等环境有关。
更高质量的长文本训练和测试数据: 训练数据大多是用13b模型构造的,肯定存在噪声。 测试数据基本都是从机器阅读理解数据整理而来,因此问题都是事实型的,不符合真实分布。
域外性能: 尽管最近出现了许多向量编码模型,但对于不太通用的领域,这些模型(包括stella、openai和cohere)的效果都不如BM25。
参考文献
- https://www.scidb.cn/en/detail?dataSetId=c6a3fe684227415a9db8e21bac4a15ab
- https://github.com/wangyuxinwhy/uniem
- https://github.com/CLUEbenchmark/SimCLUE
- https://arxiv.org/abs/1612.00796
- https://kexue.fm/archives/8847
- https://huggingface.co/sensenova/piccolo-base-zh
- https://kexue.fm/archives/7947
- https://github.com/FlagOpen/FlagEmbedding
- https://github.com/THUDM/LongBench