设备上机器学习的高效序列建模

1726316768839.jpg

随着计算密集型应用的兴起、出于隐私和安全原因需要在设备上保留某些数据,以及在网络连接不可用时提供服务的需求,对移动设备、平板电脑等设备上机器学习 (ML) 模型推理的需求 日益增长。然而,设备上推理带来了从建模到平台支持要求等一系列挑战。这些挑战与如何设计不同的架构来优化内存和计算,同时仍试图保持模型的质量有关。从平台的角度来看,问题在于识别操作并在其基础上进行构建,以便能够很好地适用于不同的产品用例。

在之前的研究中,我们将一种用于生成嵌入的新技术(称为基于投影的嵌入)与QRNN(pQRNN)等高效架构相结合,并证明它们能够胜任许多分类问题。使用蒸馏技术增强这些技术可以进一步提高端到端质量。虽然这是一种有效的方法,但它无法扩展到更大、更广泛的词汇表(即可以输入到模型的所有可能的 Unicode 或单词标记)。此外,投影操作本身的输出不包含可训练的权重,无法利用预训练模型的优势。

ByT5 中提出的无标记模型是设备上建模的一个很好的起点,它可以解决预训练和可扩展性问题,而无需增加模型的大小。这是可能的,因为这些方法将文本输入视为字节流(每个字节的值范围从 0 到 255),这可以将嵌入表的词汇量从约 30,000 个减少到 256 个。虽然 ByT5 为设备上建模提供了一种引人注目的替代方案,但从字级表示到字节流表示会线性增加序列长度;平均字长为四个字符,单个字符最多有四个字节,字节序列长度与字长成比例增加。这可能导致推理延迟和计算成本显著增加。

我们通过为SeqFlowLite库开发和发布三个新的字节流序列模型(ByteQRNN、ByteTransformer 和 ByteFunnelTransformer)来解决此问题,所有这些模型都可以在无监督数据上进行预训练,并且可以针对特定任务进行微调。这些模型利用了Charformer引入的最新创新,包括基于快速字符 Transformer 的模型,该模型使用基于梯度的子字标记化 (GBST) 方法直接在字节级别运行,以及“软”标记化方法,它使我们能够学习标记边界并缩短序列长度。在这篇文章中,我们重点介绍 ByteQRNN,并证明预训练的 ByteQRNN 模型的性能与BERT相当,尽管它小 300 倍。

序列模型架构

我们利用 pQRNN、ByT5 和 Charformer 以及平台优化,例如训练量化(跟踪模型激活的最小和最大浮点值以及用于量化推理模型的权重),将模型大小缩小四分之一,以开发一个名为 ByteQRNN 的端到端模型(如下所示)。首先,我们使用ByteSplitter操作将输入字符串拆分为字节流,并将其提供给词汇量为 259(256 + 3 个额外元标记)的较小嵌入表。

嵌入层的输出被馈送到 GBST 层,该层配备了训练内量化功能,将字节级表示与子词标记化的效率相结合,同时实现潜在子词的端到端学习。我们通过枚举每个子词块长度并将其与每个步幅标记位置(即以固定间隔选择的标记位置)的分数(使用量化密集层计算)相结合来“软”标记字节流序列。接下来,我们将字节流下采样为可管理的序列长度并将其馈送到编码器层。

GBST 层的输出可以下采样到较低的序列长度,以实现高效的编码器计算,也可以由Funnel Transformer等编码器使用,该编码器汇集查询长度并减少自注意力计算以创建 ByteFunnelTransformer 模型。端到端模型中的编码器可以替换为任何其他编码器层,例如SeqFlowLite 库中的Transformer,以创建 ByteTransformer 模型。

使用字节流输入的通用端到端序列模型图。ByteQRNN 模型使用来自 SeqFlowLite 库的 QRNN 编码器。

除了输入嵌入(即上述嵌入层的输出)之外,我们更进一步构建了一个有效的序列到序列 ( seq2seq ) 模型。我们采用 ByteQRNN 并添加基于 Transformer 的解码器模型以及与之配合使用的量化波束搜索(或树探索)来实现这一点。量化波束搜索模块通过使用前一个和当前概率的对数和计算最可能的波束(即可能的输出序列)并返回得到的顶部波束,从而减少了生成解码器输出时的推理延迟。与典型的单精度浮点格式 ( float32 ) 模型相比,此处系统使用更高效的 8 位整数 ( uint8 ) 格式。

解码器 Transformer 模型使用合并注意力子层 (MAtt) 将解码器自注意力的复杂度从二次方降低到线性,从而降低端到端延迟。对于每个解码步骤,与传统 Transformer 解码器的不断增加的缓存大小相比,MAtt 使用固定大小的缓存进行解码器自注意力。下图说明了波束搜索模块如何与解码器层交互,以使用边缘设备(例如手机、平板电脑等)在设备上生成输出 token。

云服务器解码和设备上(边缘设备)实现的比较。左:云服务器束搜索采用基于 Transformer 的解码器模型,在 float32 中具有二次时间自注意力,每个解码步骤的缓存大小都会增加。右:边缘设备实现采用量化束搜索模块以及固定大小的缓存和线性时间自注意力计算。

评估

在开发 ByteQRNN 之后,我们使用曲线下面积(AUC) 指标评估其在civil_comments数据集上的性能,并将其与预训练的 ByteQRNN 和 BERT 进行比较(如下所示)。我们证明,经过微调的 ByteQRNN 提高了整体质量,并使其性能更接近 BERT 模型,尽管其体积小了 300 倍。由于 SeqFlowLite 模型支持训练量化,可将模型大小缩小四分之一,因此生成的模型可以很好地扩展到低计算设备。我们选择了与任务相关的多语言数据源来预训练 BERT 和字节流模型,以实现最佳性能。

在 civil_comments 数据集上对 ByteQRNN 与微调后的 ByteQRNN 和 BERT 进行比较。

结论

继我们之前对 pQRNN 的研究之后,我们评估了字节流模型在设备上的使用情况,以实现预训练,从而提高模型在设备上部署的性能。我们对有预训练和无预训练的 ByteQRNN 进行了评估,并证明预训练的 ByteQRNN 的性能与 BERT 相当,尽管体积小 300 倍。除了 ByteQRNN,我们还发布了 ByteTransformer 和 ByteFunnelTransformer,这两个模型使用不同的编码器,以及合并的注意力解码器模型和波束搜索驱动程序,通过SeqFlowLite库运行推理。我们希望这些模型能为研究人员和产品开发人员提供宝贵的资源,用于未来的设备部署。

致谢

我们要感谢 Khoa Trinh、Jeongwoo Ko、Peter Young 和 Yicheng Fan 帮助开源和评估模型。感谢 Prabhu Kaliamoorthi 的所有头脑风暴和创意。感谢 Vinh Tran、Jai Gupta 和 Yi Tay 对预训练字节流模型的帮助。感谢 Ruoxin Sang、Haoyu Zhang、Ce Zheng、Chuanhao Zhuge 和 Jieying Luo 对 TPU 训练的帮助。非常感谢 Erik Vee、Ravi Kumar 和 Learn2Compress 领导层对项目的赞助以及他们的支持和鼓励。最后,我们要感谢 Tom Small 为本文提供的动画人物。

版权声明

本文仅代表作者观点,不代表本站立场。
本文系作者授权发表,未经许可,不得转载。

评论