将CodeBERTa压缩到10KB以下

在我写文章的平台上，我永远只输入代码片段，系统就猜错语言（并使用错误的语法高亮）。其他平台也是如此，我一直对此感到困惑，我的意思是，这能有多难？

前沿LLM可以轻松检测语言，但小模型呢？CodeBERTa-language-id正是为此特定任务构建的，准确率高达99.9%，但它也有330 MB。因此，从应用开发者的角度来看，它注定永远在服务器上运行。

在规模的另一端，Highlight.js在浏览器中运行，并且可以检测编程语言。但检测不是主要焦点，准确率相当低。

我们真的不得不在"好但大"和"小但差"之间做出选择吗？还是我们可以将LLM中发现的一些智能捕获到一个非常小的模型中，该模型可以通过网络发送？

在本文中，我将探索将智能缩小到一个小模型的一般想法，以语言检测为例。

让我们先看看一些结果……

1、结果

• 简单模式

在"简单模式"中，我使用CodeSearchNet数据集，其中包含六种编程语言的代码。CodeBERTa-language-id模型在此数据集上获得 99%的F1，提醒一下，该模型为330 MB。

（如果您不知道F1是什么，只需将其视为准确率，但更智能。）

我的模型也获得99%。大小 —— 包括模型参数和用JavaScript编写的推理引擎，零依赖，发送到浏览器 —— 是2 KB。是的，二。

现在，99%很棒，但这个数据集太容易了。每一项都是六种编程语言之一中的一个整洁的小函数，这是一个经典的学术基准（定义明确但与现实脱节）。

所以让我们漫步进入现实世界……

• 困难模式

为了更真实的数据集，我使用The Stack V2的一个子集。我根据medium.com编辑器中检测到的语言过滤了31种语言。这是常规编程语言与SQL、Markdown、YAML、CSS和Shell脚本等东西的混合。我将数据集中的原始代码文件拆分为10行片段。这更能代表用户可能在博客文章或GitHub评论中实际输入的内容，也是一个更难的任务。

在这个数据集上，我的模型获得86%的F1分数，模型现在为14 KB。

Highlight.js得分36%。

我们真的无法看到CodeBERTa-language-id模型在此任务上的表现如何，因为它只知道六种语言。但如果我们只为这六种语言过滤数据集，它获得89%（而我的获得95%，同时缩小到2.5 KB）。

现在，如果您认为我必须做了一些特别聪明的事情来实现这一点，您完全错了，您很快就会看到。

我应该注意，我在这里提到的分数可能低估了它（相信我兄弟）。当您将代码文件拆分为10行片段时，其中一些看起来像垃圾（数字列表，无法知道是什么语言），而且通常您会在不同的语言文件中嵌套一种语言（例如，YAML文件中的shell脚本）。我已经过滤掉其中一些，但模型犯的许多错误是相当可以原谅的。

2、流程

下面我将描述我是如何创建这个模型的，但我将保持高层次。我的目标部分是展示我认为很酷的东西，但也为了使用非常小的ML模型解决给定问题的一般流程。

如果您想要更多详细信息，您可以找到生成模型的代码在这个非常混乱的repo中。

我最初的计划是评估最新和最伟大的模型缩小技术。我从一个非常简单的模型开始，只是为了获得基准分数，但事实证明，基准模型非常棒。

任何ML工作流程的第一步都是弄清楚如何将您的输入转换为ML模型可以使用的数字表示。LLM风格的tokenization是一种方法，但它只适用于某些模型架构，因此不是一个选项。相反，我要求Codex列出我想要检测的所有31种语言中的所有保留词，以及在这些语言中常见发现的符号。它提出了一个包含746个标记的列表，比如"import"、"function"和"::"。

因此，我可以通过提出746个问题来将任何代码片段转换为746个布尔值的列表，比如*"片段是否包含*import"、"片段是否包含*function"*……

是的，这丢弃了很多信号，但它与我的解决问题格言一致："先做愚蠢的"。

将我的代码片段文本转换为值列表（片段的"特征"）后，我在这些特征上训练了一个LogisticRegression模型。如果您没有跟上最新的机器学习技术，您可能不知道"逻辑回归"模型只是最简单可能的ML模型，并且已经存在了近一个世纪。

这个模型具有三个属性，使其特别适合这种类型的问题：

• 首先，它不会遭受与transformer模型相同的上下文窗口限制。因此，虽然CodeBERTa-language-id模型只能查看代码片段的前512个标记，我的小模型将"看到"整个片段。
• 其次，推理需要极少量的代码。只需要~25行基本JavaScript来使用模型预测给定代码片段的语言。
• 第三，它快如闪电。在10,000个示例上训练需要3秒（在CPU上）。因此迭代和实验很快，一旦部署，进行预测需要不到一毫秒（比Highlight.js快得多）。

值得注意的是，当模型足够小可以通过网络发送时，它消除了处理大模型伴随而来的许多困难。您不必担心版本控制、部署、服务或扩展，因为模型只是代码库中的静态资产，推理发生在用户的硬件上。

如果我们退一步，我们可以看到这个系统中"智能"的两个不同部分：

• 在从代码片段中提取特征时，我们拥有来自前沿LLM的"智能"，它对所有31种语言有深刻理解，编码为特征列表（您可以通过人类 + Google + 半天实现相同的事情）。
• 然后是通过查看几千个示例来学习这些特征与编程语言之间关系的逻辑回归模型的"智能"。

3、窥探模型内部

如果您是机器学习的新手，这个特定模型可能是建立ML模型如何工作的直觉的好方法。

如果您跳转到langid.dgapps.io，您可以看到模型的每个参数以及它们如何用于从文本输入（代码片段）通过模型权重到预测（编程语言）。

langid.dgapps.io（在移动设备上不太好，抱歉）

值列（带有标题import、;\n、//等）代表"特征"。片段特征行表示特征是否在提供的代码片段中。下面的数值是模型的权重（对于不存在的特征淡出）。为了进行预测，每行的黑色数字相加，这给每种语言一个分数。在截图中，CSS获得最高分13，因此预测是CSS。

这基本上就是ChatGPT的工作方式，除了它有大约万亿个更多的参数，完全不同的tokenization策略，以及看起来完全不同的架构。

如果您想知道的话：是的，这个模型绝对显示出意识迹象。

4、缩小模型

好的，回到故事……

此时，基准模型看起来很有希望。它比任何其他模型都小，比几乎所有其他模型都表现得更好。

当保存为JSON时，训练的模型在磁盘上为630 KB，但经过一些努力，它将降至10 KB以下（我善意地要求您对此感到惊讶）。

缩小发生在4.5个阶段中……

第一阶段：压缩

这很简单。630 KB的gzip压缩后只有200 KB。从现在开始的所有大小数据都包括压缩。

第二阶段：模型手术

简单模型的一个好处是很容易对它们进行小手术。在我的情况下，我重新排列了模型的权重矩阵，以便早期的列具有更强的预测能力（记住，每列代表一个特征，即某个关键字或符号）。这个变化在数学上是惰性的，但这意味着您可以选择一个稍微性能较差的较小模型，只需取前n列。（一种穷人的Matryoshka表示学习。）

对于下面的图表，我创建了100个模型变体，每个变体具有不同数量的特征，导致具有各种大小和F1分数的模型。

一个实用的方法是有一个后训练步骤，删除特征（模型权重矩阵右侧的列），直到F1分数下降到最佳分数以下某个小的增量。因此，实际上，模型缩小到仅需要它那么大（因此更容易任务的模型更小）。

仅使用746个特征中的400个将模型大小从200 KB减少到113 KB，而不会显着影响F1。

第三阶段：四舍五入

由于模型保存为JSON，一种非常简单的用大小交换性能的方法是简单地将模型的权重四舍五入到更少的小数位，导致JSON文件中更少的实际数字（一种穷人的量化）。四舍五入到1位小数几乎不会移动F1分数，但将模型大小从113 KB减少到14 KB。

嘿，你好，你喜欢关于边际收益的神秘讨论吗？您会喜欢下一段……

四舍五入与gzip算法很好地结合。当四舍五入到1位小数时，模型的12,400个参数仅包含45个唯一值，因此gzip可以用很少的位引用绝大多数参数（引用它上次看到相同值的位置）。从2位小数减少到1位仅将JSON长度减少了15%，但由于所有那些重复值，gzip大小下降了50%。太棒了！

您可能在想：好吧，这变得有点极端，这家伙已经失去了对现实的看法，正在担心单个千字节。在这种情况下，您会讨厌下一点……

第四阶段：量化、原始字节及更多

我们可以更小。将模型参数作为文本在JSON文件中发送很容易但天真。将参数转换为8位整数并发送原始字节将大小从14 KB减少到10 KB。这个下降不是很大，因为前面提到的gzip魔法使JSON形式非常高效。而且由于量化繁琐，以及将模型打包和解包为原始字节，对于这么小的模型我不会费心，但对于更大的模型，这将是值得的。

哦，还有一个最终的边际收益：使用Brotli而不是Gzip来压缩量化的原始字节将大小减少到9 KB。所以现在我可以博客文章标题中说"10 KB以下"。

5、指导

这一切都很好，但是您应该如何知道微型ML模型可能是工作的正确工具呢？

我认为编程问题存在于一个谱系上，从"绝对需要AI"到"容易用常规代码解决"。想想视觉理解的领域：在规模的一端，您可能想要检测图像是否包含一个好人；为此，您需要严肃的AI。在规模的另一端，您可能想要检测图像是否是黑白的；您可以很容易地在常规代码中做到这一点。但是中间的空间呢？如果您想要检测图像是照片还是线条画。也许（在前沿LLM的帮助下），您可以编写一个"特征提取器"，将输入图像转换为数值特征集，并训练一个非常小的模型将这些特征映射到"照片"或"线条画"的预测。

如果您有一个问题，需要对某些输入进行分类或做出预测，并且可以想象输入中驱动输出的有限特征集，那么我会建议探索小型ML模型作为解决方案。

希望我已经设法调整了您对"AI"可以做多小的期望，同时仍然做有用的工作。如果您使用这种方法构建了很酷的东西，我很乐意听到。

原文链接: Shrinking a language detection model to under 10 KB

汇智网翻译整理，转载请标明出处