新闻 | News
[2024-04-06] 开源puff系列模型,专门针对检索和语义匹配任务,更多的考虑泛化性和私有通用测试集效果,向量维度可变,中英双语。
[2024-02-27] 开源stella-mrl-large-zh-v3.5-1792d模型,支持向量可变维度。
[2024-02-17] 开源stella v3系列、dialogue编码模型和相关训练数据。
[2023-10-19] 开源stella-base-en-v2 使用简单,不需要任何前缀文本。
[2023-10-12] 开源stella-base-zh-v2和stella-large-zh-v2, 效果更好且使用简单,不需要任何前缀文本。
[2023-09-11] 开源stella-base-zh和stella-large-zh
欢迎去本人主页查看最新模型,并提出您的宝贵意见!
stella model
stella是一个通用的文本编码模型,主要有以下模型:
Model Name | Model Size (GB) | Dimension | Sequence Length | Language | Need instruction for retrieval? |
---|---|---|---|---|---|
stella-base-en-v2 | 0.2 | 768 | 512 | English | No |
stella-large-zh-v2 | 0.65 | 1024 | 1024 | Chinese | No |
stella-base-zh-v2 | 0.2 | 768 | 1024 | Chinese | No |
stella-large-zh | 0.65 | 1024 | 1024 | Chinese | Yes |
stella-base-zh | 0.2 | 768 | 1024 | Chinese | Yes |
完整的训练思路和训练过程已记录在博客1和博客2,欢迎阅读讨论。
训练数据:
- 开源数据(wudao_base_200GB[1]、m3e[2]和simclue[3]),着重挑选了长度大于512的文本
- 在通用语料库上使用LLM构造一批(question, paragraph)和(sentence, paragraph)数据
训练方法:
- 对比学习损失函数
- 带有难负例的对比学习损失函数(分别基于bm25和vector构造了难负例)
- EWC(Elastic Weights Consolidation)[4]
- cosent loss[5]
- 每一种类型的数据一个迭代器,分别计算loss进行更新
stella-v2在stella模型的基础上,使用了更多的训练数据,同时知识蒸馏等方法去除了前置的instruction(
比如piccolo的查询:
, 结果:
, e5的query:
和passage:
)。
初始权重:
stella-base-zh和stella-large-zh分别以piccolo-base-zh[6]和piccolo-large-zh作为基础模型,512-1024的position
embedding使用层次分解位置编码[7]进行初始化。
感谢商汤科技研究院开源的piccolo系列模型。
stella is a general-purpose text encoder, which mainly includes the following models:
Model Name | Model Size (GB) | Dimension | Sequence Length | Language | Need instruction for retrieval? |
---|---|---|---|---|---|
stella-base-en-v2 | 0.2 | 768 | 512 | English | No |
stella-large-zh-v2 | 0.65 | 1024 | 1024 | Chinese | No |
stella-base-zh-v2 | 0.2 | 768 | 1024 | Chinese | No |
stella-large-zh | 0.65 | 1024 | 1024 | Chinese | Yes |
stella-base-zh | 0.2 | 768 | 1024 | Chinese | Yes |
The training data mainly includes:
- Open-source training data (wudao_base_200GB, m3e, and simclue), with a focus on selecting texts with lengths greater than 512.
- A batch of (question, paragraph) and (sentence, paragraph) data constructed on a general corpus using LLM.
The loss functions mainly include:
- Contrastive learning loss function
- Contrastive learning loss function with hard negative examples (based on bm25 and vector hard negatives)
- EWC (Elastic Weights Consolidation)
- cosent loss
Model weight initialization:
stella-base-zh and stella-large-zh use piccolo-base-zh and piccolo-large-zh as the base models, respectively, and the
512-1024 position embedding uses the initialization strategy of hierarchical decomposed position encoding.
Training strategy:
One iterator for each type of data, separately calculating the loss.
Based on stella models, stella-v2 use more training data and remove instruction by Knowledge Distillation.
Metric
C-MTEB leaderboard (Chinese)
Model Name | Model Size (GB) | Dimension | Sequence Length | Average (35) | Classification (9) | Clustering (4) | Pair Classification (2) | Reranking (4) | Retrieval (8) | STS (8) |
---|---|---|---|---|---|---|---|---|---|---|
stella-large-zh-v2 | 0.65 | 1024 | 1024 | 65.13 | 69.05 | 49.16 | 82.68 | 66.41 | 70.14 | 58.66 |
stella-base-zh-v2 | 0.2 | 768 | 1024 | 64.36 | 68.29 | 49.4 | 79.95 | 66.1 | 70.08 | 56.92 |
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 |
MTEB leaderboard (English)
Model Name | Model Size (GB) | Dimension | Sequence Length | Average (56) | Classification (12) | Clustering (11) | Pair Classification (3) | Reranking (4) | Retrieval (15) | STS (10) | Summarization (1) |
---|---|---|---|---|---|---|---|---|---|---|---|
stella-base-en-v2 | 0.2 | 768 | 512 | 62.61 | 75.28 | 44.9 | 86.45 | 58.77 | 50.1 | 83.02 | 32.52 |
Reproduce our results
C-MTEB:
import torch
import numpy as np
from typing import List
from mteb import MTEB
from sentence_transformers import SentenceTransformer
class FastTextEncoder():
def __init__(self, model_name):
self.model = SentenceTransformer(model_name).cuda().half().eval()
self.model.max_seq_length = 512
def encode(
self,
input_texts: List[str],
*args,
**kwargs
):
new_sens = list(set(input_texts))
new_sens.sort(key=lambda x: len(x), reverse=True)
vecs = self.model.encode(
new_sens, normalize_embeddings=True, convert_to_numpy=True, batch_size=256
).astype(np.float32)
sen2arrid = {sen: idx for idx, sen in enumerate(new_sens)}
vecs = vecs[[sen2arrid[sen] for sen in input_texts]]
torch.cuda.empty_cache()
return vecs
if __name__ == '__main__':
model_name = "infgrad/stella-base-zh-v2"
output_folder = "zh_mteb_results/stella-base-zh-v2"
task_names = [t.description["name"] for t in MTEB(task_langs=['zh', 'zh-CN']).tasks]
model = FastTextEncoder(model_name)
for task in task_names:
MTEB(tasks=[task], task_langs=['zh', 'zh-CN']).run(model, output_folder=output_folder)
MTEB:
You can use official script to reproduce our result. scripts/run_mteb_english.py
Evaluation for long text
经过实际观察发现,C-MTEB的评测数据长度基本都是小于512的, 更致命的是那些长度大于512的文本,其重点都在前半部分 这里以CMRC2018的数据为例说明这个问题:
question: 《无双大蛇z》是谁旗下ω-force开发的动作游戏?
passage:《无双大蛇z》是光荣旗下ω-force开发的动作游戏,于2009年3月12日登陆索尼playstation3,并于2009年11月27日推......
passage长度为800多,大于512,但是对于这个question而言只需要前面40个字就足以检索,多的内容对于模型而言是一种噪声,反而降低了效果。
简言之,现有数据集的2个问题:
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, 结果如下:
Dataset | 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 |
Average | 74.98 | 74.83 | 74.76 | 76.15 | 78.96 | 78.24 |
注意: 因为长文本评测数据数量稀少,所以构造时也使用了train部分,如果自行评测,请注意模型的训练数据以免数据泄露。
Usage
stella 中文系列模型
stella-base-zh 和 stella-large-zh: 本模型是在piccolo基础上训练的,因此用法和piccolo完全一致
,即在检索重排任务上给query和passage加上查询:
和结果:
。对于短短匹配不需要做任何操作。
stella-base-zh-v2 和 stella-large-zh-v2: 本模型使用简单,任何使用场景中都不需要加前缀文本。
stella中文系列模型均使用mean pooling做为文本向量。
在sentence-transformer库中的使用方法:
from sentence_transformers import SentenceTransformer
sentences = ["数据1", "数据2"]
model = SentenceTransformer('infgrad/stella-base-zh-v2')
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-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-zh-v2')
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
stella models for English
Using Sentence-Transformers:
from sentence_transformers import SentenceTransformer
sentences = ["one car come", "one car go"]
model = SentenceTransformer('infgrad/stella-base-en-v2')
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)
Using HuggingFace Transformers:
from transformers import AutoModel, AutoTokenizer
from sklearn.preprocessing import normalize
model = AutoModel.from_pretrained('infgrad/stella-base-en-v2')
tokenizer = AutoTokenizer.from_pretrained('infgrad/stella-base-en-v2')
sentences = ["one car come", "one car go"]
batch_data = tokenizer(
batch_text_or_text_pairs=sentences,
padding="longest",
return_tensors="pt",
max_length=512,
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
Training Detail
硬件: 单卡A100-80GB
环境: torch1.13.*; transformers-trainer + deepspeed + gradient-checkpointing
学习率: 1e-6
batch_size: base模型为1024,额外增加20%的难负例;large模型为768,额外增加20%的难负例
数据量: 第一版模型约100万,其中用LLM构造的数据约有200K. LLM模型大小为13b。v2系列模型到了2000万训练数据。
ToDoList
评测的稳定性: 评测过程中发现Clustering任务会和官方的结果不一致,大约有±0.0x的小差距,原因是聚类代码没有设置random_seed,差距可以忽略不计,不影响评测结论。
更高质量的长文本训练和测试数据: 训练数据多是用13b模型构造的,肯定会存在噪声。 测试数据基本都是从mrc数据整理来的,所以问题都是factoid类型,不符合真实分布。
OOD的性能: 虽然近期出现了很多向量编码模型,但是对于不是那么通用的domain,这一众模型包括stella、openai和cohere, 它们的效果均比不上BM25。
Reference
- 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
- Downloads last month
- 6,283
Space using infgrad/stella-large-zh-v2 1
Evaluation results
- cos_sim_pearson on MTEB AFQMCvalidation set self-reported47.344
- cos_sim_spearman on MTEB AFQMCvalidation set self-reported49.947
- euclidean_pearson on MTEB AFQMCvalidation set self-reported48.129
- euclidean_spearman on MTEB AFQMCvalidation set self-reported49.947
- manhattan_pearson on MTEB AFQMCvalidation set self-reported48.084
- manhattan_spearman on MTEB AFQMCvalidation set self-reported49.902
- cos_sim_pearson on MTEB ATECtest set self-reported50.980
- cos_sim_spearman on MTEB ATECtest set self-reported53.119
- euclidean_pearson on MTEB ATECtest set self-reported55.126
- euclidean_spearman on MTEB ATECtest set self-reported53.119