不靠量化也能把 Transformer 塞进单片机,靠的是删块和 BFloat16
编译自 MDPI JLPEA 2025, 15(1), 8 · 开放获取 CC BY 4.0
提到把 Transformer 放进单片机,行业里几乎条件反射地给出同一个答案:量化到 int8,激进一点就 int4。但量化本身要付出代价,量化感知训练需要校准集、需要重训,还需要在精度和位宽之间反复试,模型还没上板,先花掉一轮工程成本。
博洛尼亚大学先进电子系统研究中心与苏黎世联邦理工集成系统实验室的 Alberto Dequino、Luca Bompani、Luca Benini、Francesco Conti 在 Journal of Low Power Electronics and Applications 上给出了另一条路径。他们没有走低精度路线,而是先把模型删瘦,再用 BFloat16 部署到商用 RISC-V 多核单片机上,整个过程不需要量化感知微调,也不需要后训练量化。
这件事的背景是过去几年端侧 AI 的边界在移动。早期 TinyML 处理的是分类和检测这类固定拓扑的小模型,参数量在几十万以内,量化到 int8 就够用。生成式模型进来之后情况变了,模型规模上了一个数量级,算子也从卷积为主变成矩阵乘法为主,而低精度量化在注意力结构上更容易掉精度,这也是论文刻意绕开量化路线的原因之一。
从靠近输出头的一侧开始删
论文的出发点是 Transformer 编码层存在大量冗余。剪枝搜索的策略很直接:系统性地移除编码块中的多头自注意力模块,从最靠近输出头的那一块开始往前删,每次保留一段连续的、位置靠前的编码层。这个顺序有讲究,越靠近输入的编码层承载的通用特征越多,越靠近输出的层越偏向任务特定的处理,删后面的代价更小。
删完之后要微调。MobileBert 用 5 个 epoch,学习率 0.00005,权重衰减 0.01,优化器用 AdamW;TinyViT 用 30 个 epoch,学习率从 0.000000125 变到 0.0000625,权重衰减 0.00000001。评估分别落在 GLUE 基准和 ImageNet-1k 上。TinyViT 因为包含多个结构不同的编码区域,论文还分别搜索了单区域删除、跨多区域同时删除,以及部分删除类 MobileNet 嵌入器这几种组合。
这套方法针对的是 Transformer 上单片机的另一个现实障碍。注意力机制里的 softmax 每个输出都依赖多个输入,需要多趟遍历注意力矩阵才能算完;多头注意力本身的运算量还随序列长度二次增长。论文选择从结构上砍掉冗余层,而不是在数值精度上继续压缩,某种程度上是在回避这些算子本身的优化难题。
删到什么程度
MobileBert 的结果最直观。把编码块删到只剩 1 个,内存需求下降 94.9%,推理延迟改善约 19 倍,代价是精度退化 10.2%,论文强调这个水平仍未退化到随机猜测。如果把编码层数减半、保留 12 个块,内存下降 49.5%,延迟改善约 2 倍,而精度几乎没有下降。这条曲线说明,冗余主要集中在靠近输出的一侧。
图1 MobileBert 保留不同数量编码块时的内存、延迟与精度变化(依据原文第 5.3 节绘制)
TinyViT 的表现略有不同。最激进的配置下,每个编码区域都减到单个注意力块,内存使用下降 57.4%,延迟改善 1.6 倍,精度退化相对更明显,误差增加约两倍。如果只裁剪第三个编码区域,内存下降 20.2%,推断时间改善 1.04 倍,误差增长被限制在 1.1 倍以内。对视觉模型来说,删块换来的内存收益不如语言模型那样戏剧化,但作为一种不需要量化校准的手段,仍然划算。
不量化,那就把十六位浮点用满
模型瘦身之后,部署环节走的是另一条技术路线。论文没有依赖精度敏感的量化方案,而是用 32 位和 16 位浮点格式,重点优化 BFloat16 下的计算内核。硬件平台是 GreenWaves 的 GAP9 片上系统,基于开源 VEGA 架构,包含 9 个 RISC-V 核心,其中 8 个并行计算核加 1 个集群控制器,指令集为 RV32IMFC 并带 XpulpV2 扩展,支持打包 SIMD,可以在 32 位数据通路上一次处理两个连续的 16 位元素。
存储层级是这道题的关键约束。一级紧耦合数据存储器限制在 128 kB,二级 SRAM 为 1.5 MB,另有 4 MB 的 MRAM,数据靠非阻塞 DMA 在两级之间搬运。之所以把力气都花在矩阵乘法上,是因为在多头自注意力里,矩阵乘法占据了 99.72% 的乘累加运算和 97.62% 的活跃周期,其他算子再优化也改不了大局。
论文在矩阵乘法上做了四件事:用硬件循环消除分支惩罚与计数器更新,最多两级循环走硬件;用 U×V 的循环展开显式计算内部块,把寄存器文件的数据复用做到最大;按序列长度维度把输出区域均匀分给各个核心做并行;用 Google OR-Tools 的约束编程自动搜索各算子的最佳分块尺寸。最后一项尤其重要,因为一级缓存太小,分块尺寸选错会让访存成为瓶颈。
分块这件事需要和 DMA 配合才有意义。一级缓存只有 128 kB,装不下整个注意力矩阵,数据必须在两级存储之间来回搬运。非阻塞 DMA 允许计算与搬运重叠,但如果分块顺序设计得不好,核心会频繁停下来等数据。论文用约束编程把分块尺寸当作待解变量交给求解器,把访存次数和并行度写进约束里,省掉了手工调参的漫长试错。不同算子采用不同的分块维度,投影层和注意力矩阵乘法各有各的最优解。
图2 注意力内核的加速来源分解与算子优化重点(依据原文第 5.2、5.5 节绘制)
加速的账怎么算
相对朴素 C 移植版本的注意力内核,FP32 下实现了超过 116 倍的加速,BFloat16 下最高达到 220 倍。拆开来看,8 核并行贡献了约 7.88 倍,FP16 下为 7.97 倍,接近线性的 8 倍理想值;动态分块在并行基础上再平均提升约 5 倍;矩阵乘法的专项优化在 tinyLLAMA 中带来 18.7 倍加速,而整个模型生成 256 个 token 的周期数从 5.5 亿次降到 5100 万次,约 10 倍改善。指数函数用 Schraudolph 近似替代,计算复杂度降低约一半,引入的误差控制在不高于 2%。
论文还部署了一个带 26 万可学习参数的 tinyLLAMA 小模型,在 FP16 下仅占 688 KB,不需要剪枝就能放进 GAP9。在 370 MHz 主频下,它的吞吐量达到 1219 tokens/s,平均功率只有 57 mW,折算下来单 token 能耗为 47.5 微焦,静态寄生功耗约 7.85 mW。这个功耗水平意味着它完全可以靠电池供电长期运行。
这条路能走多远
论文的局限也很清楚。剪枝搜索依赖对模型结构的逐块尝试,不同模型族需要各自设计搜索空间,不是一套参数通用,这在模型结构快速迭代的当下会持续产生人工成本。部署侧的优化深度绑定 RV32IMFC 与 XpulpV2 扩展,换到别的指令集架构上,打包 SIMD 和硬件循环带来的收益要大打折扣。
最激进剪枝配置下的精度退化仍然可观,在精度敏感的任务上未必承受得起。论文也没有给出剪枝与量化联合使用时的表现,而这两条路线在工程上通常是叠加使用的。把这项工作放进端侧 AI 的语境里,它的价值在于提醒了一件事:把大模型塞进小设备,量化不是唯一入口。先利用结构冗余把模型削到合适大小,再用中等精度的浮点格式配合底层内核优化,可以在完全不引入量化校准流程的前提下拿到可用的结果。
论文也给出了与同类工作的横向对照。在 int8 基准下,它的 FP32 版本比 TinyFormer 慢约 3.9 倍,FP16 版本慢约 2.3 倍;但与 PULP-NN 相比,FP32 性能相近,FP16 则快 1.7 倍。这组数字说明高精度格式并非在所有比较中都处于劣势,当底层内核优化足够充分时,16 位浮点相对 int8 的性能差距可以缩到可以接受的范围,而精度上的收益是白拿的。
|
内容来源:MDPI Journal of Low Power Electronics and Applications 2025, 15(1), 8 原文标题:Optimizing BFloat16 Deployment of Tiny Transformers on Ultra-Low Power Extreme Edge SoCs 作者:Alberto Dequino, Luca Bompani, Luca Benini, Francesco Conti 原文链接:https://doi.org/10.3390/jlpea15010008 许可协议:CC BY 4.0(https://creativecommons.org/licenses/by/4.0/) 配图说明:图1、图2 为本文依据原文献数据自行绘制。 改动声明:本文在其基础上进行了中文编译与删节,未改变技术结论与原始数据。 |





