Jason Liu 演示用几百条数据微调开源嵌入模型超越 OpenAI text-embedding-3-small

只需少量示例，经过微调的开源嵌入模型就能以更低的价格提供比 OpenAI 的 text-embedding-3-small 等专有模型更高的准确率。在本文中，我们将解释如何用 Modal 创建一个这样的模型。 首先，我们会介绍微调的基础知识。然后，我们会讨论我们进行的一项实验，以确定在一个简单的问答用例中我们需要多少微调数据。

为什么要微调

开放模型让你起步

定制模型很重要。这正是 Netflix 能够持续推荐更好的电影、Spotify 能够为你的 daylist 找到新歌的原因。通过追踪你是否看完了你选择的电影，或者你是否跳过了一首歌，这些公司积累了数据。他们利用这些数据来改进其内部嵌入模型和推荐系统，从而为你带来更好的建议和更好的体验。这甚至可以带来更多用户的更多参与，进而带来更多数据，进而带来更好的模型，形成一个被称为数据飞轮的良性循环。

数据飞轮：更多用户意味着更多数据，更多数据意味着更好的模型，更好的模型意味着更多用户。

像 Netflix 和 Spotify 这样大型的、以机器学习为先导的组织已经利用数据飞轮从零开始创建了自己的模型，如今他们拥有大量数据。但当你刚刚创办一家新公司或启动一个新项目时，你并不总是拥有所需的数据。在 2010 年代，启动一个数据飞轮需要大量的创造力或资源投入。

但到了 2020 年代，功能强大、采用宽松许可证的通用预训练模型的出现，极大地简化了这一启动步骤。你可以从这些模型之一开始——它们经过训练，能够识别大型、多样化数据集中的模式——并期望它们在你的任务上表现相当不错。

在之前的一篇博客文章中，我们展示了这一点：利用 Modal 的自动扩缩容基础设施，在数百块 GPU 上部署一个现成模型，在不到 15 分钟内嵌入了整个英文维基百科。

微调启动数据飞轮

这些模型以及轻松运行它们的基础设施的可用性，对于刚刚起步、还没有任何用户数据的组织来说是个好消息。但至关重要的是，要尽快转向一个比现成模型性能更好的定制模型。幸运的是，数据积累得很快：只需几十个用户每天与服务交互 3-4 次，就能在几天内产生数百个数据点。

而这正是我们训练一个模型所需的全部数据，该模型在样本数据集上识别文本相似度方面可以击败 OpenAI 的 text-embedding-3-small。

我们在 Modal 上用来通过现成模型创建嵌入的同一套可扩展、无服务器的基础设施，也可以用来定制它，这一过程称为微调。最终结果是一个性能卓越、运营成本显著降低的机器学习应用：这是启动你自己的数据飞轮的第一步。

如何微调：数据集、模型和基础设施

在微调模型时，必须做出许多设计决策。我们在此回顾其中几项。

寻找或创建数据集

尽管机器学习中的许多讨论和研究都围绕模型展开，但任何称职的机器学习工程师都会告诉你，数据集才是最关键的组成部分。

嵌入模型通常是在由成对项目构成的数据集上训练的，其中一些对被标记为“相似”（如同一段落中的句子），而另一些被标记为“不同”（如随机选取的两个句子）。同样的原理可以应用于比句子更长的文本——段落、页面、文档——也可以应用于文本以外的事物——图像、歌曲、用户点击流——或者可以同时应用于多种模态——图像及其说明文字、歌曲及其歌词、用户点击流及已购产品。

我们将使用 Quora 数据集，其中包含来自 Quora 帖子的成对问题，其中一些对是被标记为重复的问题。

Hugging Face 数据集查看器中 Quora 重复数据集的部分示例行。

您可以在此处的交互式查看器中查看该数据集。有些问题对，如“Can I hack my Charter Motorolla DCX3400?”和“How do I hack Motorola DCX3400 for free internet?”，非常相似，但并非重复，即它们是“难负样本”。

综合来看，这使得我们在此训练的模型可能对检索增强生成（RAG）聊天机器人有用。在基于嵌入的 RAG 聊天机器人中，必须在一个大型文本语料库中搜索少量与用户查询“匹配”的段落，即可能包含答案的段落。该数据集将训练模型对问题主题中非常细微的差异保持敏感。近似重复项也可以在检索之前或在训练其他模型之前被移除，这种技术被称为“语义去重”。

选择基础模型

我们在此主要关注许可宽松、权重可获取的模型。这些模型的权重可以像下载和修改开源代码一样下载和修改。因此，我们在此将它们称为“开源”模型，尽管并没有适用于模型的开源促进会认可的“开源”定义。模型通常通过 Hugging Face 基于 git LFS 的模型仓库中心发布，我们将从这里获取模型。

或者，我们也可以使用 API 来微调专有模型，正如一些嵌入 API 服务所提供的那样。除了成本方面的考虑外，我们发现微调模型足够复杂且特定于用例，因此控制训练过程是必要的。

如何在可用模型之间进行选择？每个模型的训练方式不同，并且针对特定的用例。最关键的是，模型是在特定的一种或多种模态（文本、图像、音频、视频等）以及特定数据集上训练的。一旦您缩小到处理您用例中模态的模型，就比较它们在公开基准（如 MTEB）上的表现。除了任务性能外，还要通过公开基准数据审查模型在资源需求和吞吐量/延迟方面的表现（像 Lambda Labs 这样的硬件提供商在这里是不错的资源）。

例如，嵌入维度，即模型输出嵌入中的条目数量，是一个重要的考虑因素。更大的向量可以存储更多信息，从而带来更好的任务性能，但随着我们随时间嵌入更多数据，可能会导致显著更高的成本（成本扩展更像 RAM 而非磁盘）。在微调时，我们可以调整这个维度。

获取训练基础设施

微调模型需要大量的计算资源。即使是那些之后可以在 CPU 上令人满意地运行、甚至是客户端或边缘 CPU 上运行的模型，也常常是在 GPU 上训练的，因为 GPU 在像训练这样易于并行化的工作负载上可以实现高吞吐量。

对于典型的微调任务，我们需要一到八块服务器级 GPU。超过八块 GPU 通常需要将训练分布到多个节点上，这是由于连接性的限制，这会显著增加硬件成本和工程复杂度。

但如今服务器级 GPU 十分稀缺，这意味着购买或租用它们都很昂贵，而且云服务提供商常常要求最低规模和时长的预留。但微调任务不太像生产工作流（始终在线、流量可合理预测），而更像开发工作流（间歇性、不可预测）。综合起来，这些现象导致了大规模的过度分配和过度支出，据 ClearML 和 AI Infrastructure Alliance 的这项调查，各组织报告的峰值利用率平均约为 60%——非峰值时甚至更低。

来源：The State of AI Infrastructure at Scale 2024

Modal 解决了这个问题：它提供自动扩缩容的基础设施，包括 GPU，因此你只需为使用量付费（也就是所谓的“无服务器”）。Modal 还提供 Python 风格、基础设施即代码的接口，使数据科学家和机器学习研究人员能够拥有并控制自己的基础设施。

有了这些资源，我们需要确定如何界定模型训练过程的范围。我们在训练、迭代超参数和数据调整上花费的时间和金钱越多，任务性能就会越好，但收益递减。一般来说，我们建议要么在某个指标上训练到满意为止（例如至少 90% 的准确率），要么选择若干需要满足的指标和一个要最大化的指标（例如在召回率 ≥ 50% 的前提下我们能获得的最高准确率），然后对资源和时间设定硬性上限。

对微调超参数进行网格搜索

确定如何训练一个全新的模型架构或在一个全新的任务上训练是一个研究项目，应当相应地界定范围。但微调更简单——我们可以使用已有的训练配方，比如用于训练该模型的配方（如果它是开源的）。但仍然有实验的空间，包括上面列出的许多考量因素。

我们选择了我们认为最重要的三个作为实验参数：我们应该训练哪个预训练模型、用多少数据、以及用多少输出维度？因为这些实验参数决定了模型中参数（权重和偏置）的取值，所以它们被称为超参数。

探索超参数最简单的方法是为每个超参数定义一组可能的值，然后检查所有组合——即网格搜索。这是一种暴力方法，但正如我们将在下面看到的，它有效且易于并行化。

我们在 Wikipedia 嵌入示例中使用的原始 bge-base-en-v1.5 模型之外又增加了两个模型，并尝试了两种不同的嵌入维度。对于每一种配置，我们测试了不同的数据集大小，从一百个到超过十万个样本不等：

所有其他超参数均保持不变。

接下来，我们使用标准库模块 itertools 提供的 product 函数生成了 model、dataset_size 和 dense_out_features 的所有可能组合，该函数会创建一个迭代器，返回每个输入列表中元素的所有可能组合（即笛卡尔 product）。然后，我们使用这些组合为模型微调过程生成配置对象：

无论我们的微调过程是什么，它接收 config 并生成一个 results 的字典。我们将其包装在一个函数中，并用 @app.function() 装饰该函数，使其能在 Modal 的自动扩缩容基础设施上运行，如下面的伪代码所示。

在此基础上，扩展就像调用 objective.map 来 并行运行实验 一样简单。我们将其包装在一个函数中，并用 @app.local_entrypoint() 装饰该函数，这样我们就可以通过 modal run 从命令行启动实验。

我们的训练过程可以放在单个 GPU 上运行，而且每个实验都是并行运行的，因此我们可以将其扩展到 Modal 允许的最大同时 GPU 工作进程数——在撰写本文时，这个数量可达数千。对于大型训练任务来说，这可能意味着下周才能得到结果，还是午餐后就能得到结果的区别。

用几百个示例击败专有模型

下图总结了我们实验的结果，展示了我们在 Quora 数据集上训练的模型的错误率（预测不正确的比例）与微调期间使用的数据集示例数量之间的关系，三个模型各有一张图。为了便于比较，还展示了 OpenAI text-embedding-3-small 模型的性能。为完整起见，我们展示了测试的两种不同嵌入维度大小，尽管对于其他超参数的任何设置，我们都没有观察到它们之间的性能差异。

在 Quora 数据集上微调的三个模型的错误率与数据集大小的关系。

在这三种情况中，我们看到了一些在微调中常见的模式：

• 对于 jina-embeddings-v2-small-en 模型，错误率高于基线，并且从未下降。一如既往，有可能我们未调整的其他超参数的不同设置会带来该模型性能的提升。这类结果正是你不希望从超参数搜索中得到的，因为不清楚下一步该做什么。
• 对于 all-mpnet-base-v2 模型，仅用 100 个示例后错误率就低于基线，但即使示例数量增加到三个数量级，我们也没有观察到太多改进。
• 对于 bge-base-en-v1.5 模型，错误率一开始高于基线模型，但随着数据增多迅速改善，在 200 个示例时就有说服力地超过了基线，并且在 100,000 个示例时仍在改善。

回顾这些结果，我们会选择继续使用微调后的 bge-base-en-v1.5 模型，尤其是如果我们预计未来能够通过数据飞轮收集更多数据的话。我们很可能会选择 256 维嵌入，因为与 512 维嵌入相比，它们生成和存储的成本更低，而且我们没有观察到使用更大嵌入带来准确率方面的好处。

你可能会反驳说，相对于基线模型的改进在绝对数值上很小——错误率从17%降至13%。但相对而言，这是一个巨大的差异：微调后的模型避免了基线模型所犯错误的整整四分之一。随着错误率的降低，这一现象会变得更加明显：一个可靠性为99%的系统可以用于那些可靠性为95%的系统无法胜任的场景，即使差异的幅度看起来很小。

下一步

在本文中，我们展示了如何微调一个开源嵌入模型，以在简单的问答任务上超越专有模型。我们还讨论了微调模型时需要考虑的因素，以及如何使用Modal进行超参数的网格搜索。我们已经证明，即使只有几百个示例，我们也能实现比专有模型更好的性能。

展望未来，微调的下一步是将这一过程投入实际运行，以便我们可以收集更多数据并对模型进行迭代。通过完全自动化，我们甚至可以将模型转变为一个持续改进的系统，由我们收集的额外数据驱动。

这些将数据转化为对用户有用特征的流程和管道，远比模型本身更能体现机器学习团队的产出。借助开源模型和无服务器基础设施，构建它们比以往任何时候都更容易。