EmbBERT:面向极致边缘计算的小型语言模型架构优化 中文标题

【字体: 时间:2026年03月09日 来源:Neural Networks 6.3

编辑推荐:

  本文针对在内存资源极为有限(低于2MB)的微型设备(如可穿戴、IoT设备)上部署NLP模型的难题,提出了一种新型Tiny Language Model (TLM)——EmbBERT。该模型通过集成紧凑嵌入层、流线型前馈块和高效注意力机制,在严格内存预算下实现了接近更大模型(如BERT-Tiny)的性能,并展现了对8比特量化的鲁棒性。这项工作为在资源受限的边缘设备上实现本地智能语言处理提供了高效解决方案。

  
在智能家居、工业自动化和可穿戴设备日益普及的今天,物联网(IoT)系统与各种小型设备正在扮演越来越重要的角色。这些微型设备虽然计算和存储能力有所提升,但通常仍运行在极为苛刻的资源约束下,可用内存往往不足2MB,处理能力也远低于100MHz。这使得在这些设备上设计和部署复杂的机器学习(ML)和深度学习(DL)模型变得异常困难。为了解决这一问题,Tiny Machine Learning (TinyML) 领域应运而生,它致力于在内存、算力和能量预算都极为有限的环境中高效地执行模型推理,从而实现数据的本地处理,降低延迟,提升实时响应能力,并增强隐私保护。
尽管TinyML已在关键词检测、图像分类和目标检测等领域取得了进展,但在微型设备上部署自然语言处理(NLP)模型仍然是一个重大挑战。目前,基于注意力机制的Transformer架构模型(如BERT、RoBERTa、DistilBERT等)在NLP任务上取得了巨大成功,但它们通常依赖数百万甚至数十亿参数,需要巨大的内存资源。即使是经过压缩的变体,如MobileBERT(2530万参数),其规模也远超仅有几MB内存的微控制器所能承载的范围。为了填补这一空白,研究团队引入了一种名为EmbBERT的Tiny Language Model (TLM)。它专为在严格内存预算下(2MB)运行而设计,并通过一系列架构优化,在保持竞争力的准确率的同时,极大地减少了内存占用。
研究人员开展了一系列实验来验证EmbBERT的有效性。他们设计了一个名为TinyNLP的定制化基准测试,包含IMDb、ag_news等多个适用于资源受限环境的现实场景数据集。此外,他们还使用了标准的GLUE基准测试套件进行评估。在实验中,他们将EmbBERT与多个基线模型进行了对比,包括按2MB内存预算缩放的BERT、MAMBA、NanoBERT变体,以及一个仅包含嵌入器的简单模型和一个嵌入器加卷积层的模型。作为参考,他们还对比了比这些模型大十倍、占用约20MB内存的SotA模型BERT-Tiny。
为了开展研究,作者首先重新设计了模型的核心架构。他们采用了一个基于Nano Embedder的高效嵌入块,通过引入降维投影显著减少了嵌入层的参数量。模型的核心是N个高效编码器,每个编码器包含一个归一化层和两条并行路径:一条是高效的注意力机制,另一条是卷积跳跃连接。与标准注意力相比,高效注意力取消了单独的键(Key)和值(Value)投影矩阵,仅通过查询(Query)和未投影的归一化输入计算注意力,从而减少了约66%的投影参数和激活内存。卷积跳跃连接则用深度可分离一维卷积替代了传统的全连接前馈网络,以更低的计算代价实现局部token交互。最终,两条路径的输出通过可学习的加权差分进行聚合。研究还精确分析了各模块的权重和激活内存需求,并给出了计算模型总内存占用的公式。在模型训练方面,支持BERT式预训练的模型在BookCorpus数据集上进行了掩码语言建模(MLM)和下一句预测(NSP)任务预训练。随后,所有模型均在特定下游数据集上进行微调。此外,研究还进行了8位块级后训练量化分析,以评估量化对模型性能和内存占用的影响。
3.1. The EmbBERT Architecture
这部分详细描述了EmbBERT的整体架构。模型主要由两大模块构成:Nano嵌入块和N个高效编码器序列。嵌入块负责将输入token映射为低维、紧凑的表示δ。每个高效编码器则并行处理经过归一化的输入,通过一个单头高效注意力机制捕捉全局依赖,并通过一个卷积跳跃连接块捕捉局部特征。编码器最后将两个路径的输出通过可学习参数λEA和λCS进行加权差分合并,形成最终输出。这种设计旨在最小化参数和激活内存,使其能在1-2 MB的TinyML预算下运行。
3.2. Computing memory requirements of EmbBERT
此部分精确计算了EmbBERT在受限环境中的总内存需求。总内存Mtot是权重内存和激活内存之和,其计算考虑了嵌入器和编码器中各组件(如嵌入层、归一化层、注意力块、卷积块)的权重数量和激活大小。文章推导出了每个组件的内存占用公式,并列成表格以便于计算不同配置下的内存开销。这是确保模型能够在目标设备上部署的关键步骤。
4.1. Baseline Models and Comparisons
为公平评估,研究设定了多个对比基线,包括按2MB内存约束缩放的BERT、MAMBA、NanoBERT,以及两个更简单的模型(仅嵌入器,以及嵌入器加卷积层)。所有模型的超参数和内存占用情况被详细列出,以说明EmbBERT在相同预算下的设计优势,例如其拥有更多的权重但更低的激活占用。
4.2. Pre-training and Fine-tuning
研究详细介绍了模型的训练协议。对于BERT类模型,采用了标准的掩码语言建模(MLM)和下一句预测(NSP)任务在BookCorpus数据集上进行预训练。之后,所有模型都在TinyNLP和GLUE的各任务数据集上进行了微调。对于结构过于简单的Embedder和Embedder+Conv模型,由于无法有效吸收预训练,采用了更长时间的直接任务训练,以确保与其他模型消耗的总计算量相当。
5. Experimental Results
实验结果表明,在TinyNLP基准测试中,EmbBERT的平均准确率达到87.19%,表现最佳,甚至略优于内存占用大10倍的BERT-Tiny(86.99%)。简单的Embedder和Embedder+Conv模型也表现不俗,而缩放的BERT和MAMBA模型则表现较差。在更具挑战性的GLUE基准测试中,EmbBERT再次以63.50的平均分位居榜首,超过了BERT-Tiny(63.16),而其他2MB模型的得分则显著落后。这证明了EmbBERT架构在严格内存限制下的高效性。
5.1. Discussion of the Results
对结果的分析显示,EmbBERT在情感分析、意图识别等语义和情感线索明显的任务上表现出色,在需要处理句法结构关系的任务上也有良好表现。在GLUE的句法理解和语义相似性任务上优势明显,但在需要复杂逻辑推理的任务上(如RTE)仍有不足,这是此类小模型的预期局限。总体而言,EmbBERT是当前TinyML硬件上用于NLP应用的先进语言模型。
6.1. Quantizing EmbBERT
该部分研究了EmbBERT对8位后训练量化的鲁棒性。量化后,模型总内存占用从约2MB降至781KB,实现了2.4倍的内存节省。量化版本的EmbBERT在TinyNLP上平均准确率为88.17%,在GLUE上得分为62.81,性能几乎没有下降。这表明EmbBERT能有效适应量化,进一步增强了其在资源受限硬件上的部署潜力。
6.2. Scaling the EmbBERT architecture
研究还探索了EmbBERT架构的可扩展性,评估了从0.5MB (Nano) 到 40MB (Big) 的多种配置。结果显示,模型性能(以GLUE得分为衡量标准)随模型规模增加而单调提升。EmbBERT-Big (40MB) 的得分超过了更大参数量的BERT-Tiny (20MB)。这证明了EmbBERT架构在从亚兆字节到数十兆字节的内存范围内具有良好的可扩展性,允许开发者根据具体设备的资源限制和性能需求灵活选择模型规模。
7. Ablation Study
消融研究旨在评估EmbBERT各个核心组件(高效注意力、卷积跳跃连接、可学习的加权差分聚合)以及预训练过程对最终性能的贡献。通过依次移除或替换这些组件进行实验,结果显示每个组件都对模型性能有积极贡献,其中预训练过程带来的提升最为显著。这证明了EmbBERT的设计是合理且必要的,每个优化都对在2MB预算下实现高性能至关重要。
本研究成功设计并验证了EmbBERT,这是一种专为内存预算仅为2MB的微型设备设计的创新语言模型。通过整合高效的嵌入块、简化的注意力机制、卷积跳跃连接和可学习的聚合方式,EmbBERT在极端资源约束下实现了与规模大十倍(20MB)的SotA模型(如BERT-Tiny)相媲美的性能。在TinyNLP和GLUE基准测试上的广泛实验证实了其有效性。此外,该模型对8比特量化展现出卓越的鲁棒性,可进一步将内存需求压缩至781kB,同时其架构在从亚兆字节到数十兆字节的范围内具有良好的可扩展性。消融研究证实了所有核心组件和预训练过程对性能的积极贡献。该工作为在资源极度匮乏的边缘设备上实现本地、高效、智能的NLP处理提供了切实可行的解决方案,是迈向“无处不在的智能”的重要一步。相关代码、脚本和检查点均已公开,以确保研究的可复现性。
相关新闻
生物通微信公众号
微信
新浪微博

热点排行

    今日动态 | 人才市场 | 新技术专栏 | 中国科学人 | 云展台 | BioHot | 云讲堂直播 | 会展中心 | 特价专栏 | 技术快讯 | 免费试用

    版权所有 生物通

    Copyright© eBiotrade.com, All Rights Reserved

    联系信箱:

    粤ICP备09063491号