Stella: 一个强大的中文文本编码模型

stella

Stella模型简介

Stella是一个由SmartLi8开发的通用中文文本编码模型，旨在为各种自然语言处理任务提供高质量的文本表示。目前，Stella有两个版本可供选择：base和large。这两个版本的模型都支持1024的输入长度，这使得Stella在处理长文本时具有显著优势。

模型特点

长文本支持：Stella模型支持1024的输入长度，这是其最显著的特点之一。在当前大多数模型仅支持512长度的情况下，Stella的长文本处理能力为许多应用场景带来了新的可能性。
多版本选择：Stella提供base和large两个版本，用户可以根据自己的需求和资源选择合适的模型。
优秀的性能：在多个评测基准上，Stella展现出了优秀的性能，特别是在长文本处理方面。
通用性强：Stella经过精心设计和训练，可以适用于多种中文自然语言处理任务。

训练数据与方法

训练数据

Stella的训练数据主要包括以下几个部分：

开源数据：
- wudao_base_200GB
- m3e
- simclue 特别注意的是，在选择这些开源数据时，开发团队着重挑选了长度大于512的文本，以增强模型的长文本处理能力。
构造数据：开发团队在通用语料库上使用大型语言模型(LLM)构造了一批(question, paragraph)和(sentence, paragraph)数据。这种方法可以生成更加多样化和高质量的训练样本。

训练方法

Stella的训练过程采用了多种先进的技术和方法：

对比学习损失函数：这是一种常用的方法，可以帮助模型学习更好的文本表示。
带有难负例的对比学习损失函数：开发团队分别基于bm25和vector构造了难负例，这可以进一步提高模型的性能。
EWC (Elastic Weights Consolidation)：这是一种防止模型在学习新任务时遗忘旧任务的技术。
cosent loss：这是一种专门为文本表示学习设计的损失函数。
多迭代器训练：每一种类型的数据使用一个迭代器，分别计算loss进行更新。这种方法可以更好地平衡不同类型数据的贡献。

初始权重

Stella-base-zh和Stella-large-zh分别以piccolo-base-zh和piccolo-large-zh作为基础模型。对于512-1024的position embedding，开发团队使用了层次分解位置编码进行初始化。这种方法可以有效地扩展模型的输入长度，同时保持良好的性能。

Stella模型架构

评测结果

C-MTEB评测

Stella模型在C-MTEB（Chinese Multi-Task Evaluation Benchmark）上表现出色。以下是Stella-large-zh和Stella-base-zh的评测结果：

Model Name	Model Size (GB)	Dimension	Sequence Length	Average (35)
stella-large-zh	0.65	1024	1024	64.54
stella-base-zh	0.2	768	1024	64.16

这些结果显示，Stella模型在各种任务上都表现出色，包括分类、聚类、对重排序、检索和语义相似度等。特别值得注意的是，Stella模型在支持1024序列长度的同时，仍然保持了较小的模型大小，这使得它在实际应用中具有很强的竞争力。

长文本评测

考虑到C-MTEB评测数据主要针对短文本，开发团队特别整理了6份长文本测试集，以更好地评估模型的长文本处理能力。这些测试集包括：

CMRC2018（通用百科）
CAIL（法律阅读理解）
DRCD（繁体百科，已转简体）
Military（军工问答）
Squad（英文阅读理解，已转中文）
Multifieldqa_zh（清华的大模型长文本理解能力评测数据）

在这些长文本测试集上，Stella模型同样表现出色。以下是部分评测结果：

Dataset	stella-base-zh	stella-large-zh
CMRC2018	96.08	95.56
CAIL	34.62	37.18
DRCD	86.14	84.58

这些结果充分展示了Stella模型在处理长文本时的优秀性能。

长文本评测结果

使用方法

Stella模型的使用方法与piccolo模型完全一致，这使得用户可以轻松地将现有的piccolo代码迁移到Stella上。以下是几种常见的使用方式：

使用sentence-transformers库

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

需要注意的是，在使用Stella模型时，instruction中的冒号应该使用英文冒号，例如"查询: "和"结果: "。

训练细节

为了帮助研究者更好地理解和复现Stella模型，以下是一些关键的训练细节：

硬件：单卡A100-80GB
环境：torch1.13.*; transformers-trainer + deepspeed + gradient-checkpointing
学习率：1e-6
batch_size：base模型为1024，额外增加20%的难负例；large模型为768，额外增加20%的难负例
数据量：约100万，其中用LLM构造的数据约有200K。LLM模型大小为13b

这些详细信息可以帮助其他研究者更好地理解Stella模型的训练过程，并可能在此基础上进行进一步的改进和创新。

未来展望

尽管Stella模型已经展现出了优秀的性能，但开发团队仍然认识到有一些需要改进的地方：

评测的稳定性：在评测过程中发现Clustering任务的结果会有小幅度的波动。虽然这种波动很小（±0.0x），不影响整体评测结论，但团队仍在努力找出原因并提高评测的稳定性。
更高质量的长文本训练和测试数据：目前的训练数据主要是用13b模型构造的，可能存在一些噪声。测试数据主要来自MRC（Machine Reading Comprehension）数据集，问题类型较为单一。团队计划收集和构造更加多样化、高质量的长文本数据。
OOD（Out-of-Domain）性能：对于一些非常专业或特殊的领域，当前的向量编码模型（包括Stella、OpenAI和Cohere）的效果可能不如传统的BM25算法。提高模型在这些领域的表现是未来的一个重要研究方向。

未来研究方向