TinyML:面向无处不在的边缘AI - 技术深度解析

Soro S. TinyML for ubiquitous edge AI[J]. arXiv preprint arXiv:2102.01255, 2021.

引言:从云计算到边缘智能的演进

在过去的二十年间,我们见证了计算范式的重大转变。数据生产以指数级速度增长,从2010年的2 ZB增长到2020年的64 ZB,预计到2025年将达到175 ZB。这种爆炸式增长催生了云计算时代,其特点是拥有几乎无限的计算资源和存储能力。深度学习革命正是在这种背景下发生的,ResNet-152等模型包含超过6000万个参数,需要11.3 GFLOPs的计算能力才能完成一次前向传播。

然而,随着物联网设备的普及,一个新的挑战出现了:如何在功耗仅为毫瓦级别、内存仅有几百KB的微控制器上运行智能算法?这就是TinyML要解决的核心问题。目前全球已部署约2500亿个微控制器,每年新增400-500亿个,这些设备构成了一个庞大的、尚未充分利用的计算网络。

图1详解:计算技术的三次浪潮

图1展示了从云ML到移动ML再到TinyML的演进过程,这个演进图清晰地展现了三个不同计算时代的特征:

云机器学习时代(2006年起)的特点是使用大规模服务器集群,图中左侧显示了一个典型的数据中心机房,充满了蓝色LED灯光的服务器机架。这个时代的代表性神经网络包括DNN(深度神经网络)、大型模型(参数量达16-32GB)、数百万到数十亿的参数规模。硬件主要依赖GPU、TPU和FPGA等高性能计算设备,代表性网络架构包括AlexNet、Inception、ResNet和VGGnet。这些模型主要关注准确性,使用TFLOP级别的计算能力,数据主要存储在云端,共享率仅为1%。

移动机器学习时代(2016年起)的标志是智能手机和平板电脑的普及,图中中间展示了一部典型的智能手机。这个时代开始关注优化算法和轻量级CNN,面临内存限制(通常8GB RAM)和应用大小限制。计算能力降至GFLOP级别,开始在准确性和效率之间寻求平衡。硬件采用SoC(系统级芯片)和NPU(神经网络处理单元),代表性模型包括MobileNet v1/v2、ShuffleNet和SqueezeNet。数据类型扩展到图片、音频、点击和GPS信息,本地处理率提升至5%。

TinyML时代(2019年起)则进入了超低功耗嵌入式计算领域,图中右侧展示了一个典型的微控制器开发板(看起来像Arduino Nano或类似设备)。这个时代面临极其严苛的资源约束:仅有约100KB的RAM、使用MCU配合硬件加速器、传感器种类丰富(CMOS相机、IR相机、音频、IMU、温度、化学和加速度计等)。数据处理主要在物理世界的边缘完成,本地处理率高达95%。

图2深度分析:神经网络模型的准确率-复杂度权衡

图2是一个散点图,横轴表示操作数(G-FLOPs),纵轴表示Top-1准确率(%),每个圆圈代表一个不同的神经网络模型,圆圈的大小表示模型参数量的多少。

这张图揭示了几个重要的技术洞察。首先,在图的左下角,我们看到SqueezeNet v1.1和SqueezeNet v1.0,它们的准确率约为55-58%,但计算复杂度极低,仅需要不到5 GFLOPs。这些模型通过使用Fire模块(包含squeeze和expand层)实现了极致的压缩。

图的中部区域密集分布着多个模型。MobileNet系列(包括MobileNet v1和v2)展现了移动端优化的成果,它们在约70%的准确率水平上,仅需要不到5 GFLOPs的计算量。这些模型通过深度可分离卷积(Depthwise Separable Convolutions)实现了计算效率的突破,将标准卷积分解为深度卷积和逐点卷积,计算复杂度从DK×DK×M×N×DF×DFD_K \times D_K \times M \times N \times D_F \times D_FDK×DK×M×N×DF×DF降低到DK×DK×M×DF×DF+M×N×DF×DFD_K \times D_K \times M \times D_F \times D_F + M \times N \times D_F \times D_FDK×DK×M×DF×DF+M×N×DF×DF

VGG系列模型(VGG-11、VGG-13、VGG-16、VGG-19)形成了一条几乎垂直的线,准确率从约69%提升到72%,但计算复杂度从7 GFLOPs急剧增加到20 GFLOPs。这些模型的圆圈特别大,表明它们的参数量巨大(VGG-19有1.43亿参数),展现了早期深度学习"暴力美学"的特点。

图的右上角是高精度模型区域。NASNet-A-Large达到了最高的准确率(超过82%),但需要超过23 GFLOPs的计算量。SENet-154和ResNet系列也在这个区域,它们通过残差连接和注意力机制实现了高准确率。特别值得注意的是,ResNet-152虽然层数很深,但通过残差连接避免了梯度消失问题,使得如此深的网络成为可能。

图3详解:模型压缩流水线架构

图3展示了一个完整的模型压缩流程,从原始网络到最终的压缩模型,包含三个主要阶段:剪枝、量化和编码。

第一阶段是剪枝(Pruning),图中左侧显示了原始网络经过几个步骤:首先进行连接性训练(Train Connectivity),然后剪除连接(Prune Connections),最后重新训练权重(Train Weights)。这个过程将原始网络的大小减少了9-13倍,但保持了相同的准确率。剪枝过程中,小于阈值的权重被设为零,网络从密集连接变为稀疏连接。

第二阶段是量化(Quantization),图中中间部分展示了"less bits per weight"的过程。这包括聚类权重(Cluster the Weights)、生成码本(Generate Code Book)、使用码本量化权重(Quantize the Weights with Code Book)和重新训练码本(Retrain Code Book)。这个阶段实现了27-31倍的压缩,同时保持准确率。量化将32位浮点权重转换为更少位数的表示,例如8位或更少。

第三阶段是Huffman编码(Huffman Encoding),图中右侧显示了编码权重(Encode Weights)和编码索引(Encode Index)的过程。这个最终阶段又实现了35-49倍的压缩。Huffman编码利用权重分布的不均匀性,为频繁出现的值分配较短的编码,为罕见的值分配较长的编码。

整个流水线的关键在于每个阶段都保持了模型的准确率,同时显著减少了存储需求。从原始网络到最终压缩模型,总压缩比可以达到35-49倍,这使得原本需要几百MB存储的模型可以压缩到几MB甚至更小。

图4解析:TensorFlow Lite唤醒词检测系统架构

图4展示了一个完整的唤醒词检测应用的系统架构,分为几个关键组件:

系统的输入端是设备麦克风(Device microphone),它持续捕获环境中的音频信号。音频数据首先经过音频提供器(Audio provider),该模块负责从麦克风捕获音频样本。这个过程通常以16kHz的采样率进行,每个样本16位,形成连续的音频流。

接下来是特征提供器(Feature provider),它将原始音频数据转换为频谱图。这个转换通常使用短时傅里叶变换(STFT),将时域信号转换为时频表示。对于一个长度为N的音频帧,STFT计算为:
X(m,k)=∑n=0N−1x[n+mH]w[n]e−j2πkn/NX(m,k) = \sum_{n=0}^{N-1} x[n+mH]w[n]e^{-j2\pi kn/N}X(m,k)=n=0N1x[n+mH]w[n]ej2πkn/N
其中w[n]是窗函数,H是跳跃大小,k是频率索引。

TF Lite解释器(TF Lite interpreter)运行训练好的神经网络模型。模型通常是一个小型CNN,专门训练用于识别特定的唤醒词(如"yes"、“no”、“silence"和"unknown”)。模型的输出是每个类别的概率分布。

识别命令模块(Recognize commands)使用推理输出来决定是否检测到命令词。它实施了时间平滑和置信度阈值,以减少误报。只有当连续多帧都检测到相同的命令词且置信度超过阈值时,才会触发命令。

最后,命令响应器(Command responder)根据识别到的命令采取行动,例如点亮设备LED或触发其他功能。整个系统形成一个主循环(Main loop),持续处理音频输入并响应检测到的命令。

微控制器机器学习的数学基础

在资源受限的环境中运行深度学习模型需要深入理解计算和内存需求。考虑一个标准的卷积层,其计算复杂度可以精确表示为:

对于输入特征图大小为Hin×Win×CinH_{in} \times W_{in} \times C_{in}Hin×Win×Cin,卷积核大小为K×KK \times KK×K,输出通道数为CoutC_{out}Cout,stride为SSS,padding为PPP的卷积层,输出特征图大小为:
Hout=⌊Hin+2P−KS⌋+1H_{out} = \lfloor \frac{H_{in} + 2P - K}{S} \rfloor + 1Hout=SHin+2PK+1
Wout=⌊Win+2P−KS⌋+1W_{out} = \lfloor \frac{W_{in} + 2P - K}{S} \rfloor + 1Wout=SWin+2PK+1

总的乘加运算(MACs)数量为:
MACs=Hout×Wout×K2×Cin×CoutMACs = H_{out} \times W_{out} \times K^2 \times C_{in} \times C_{out}MACs=Hout×Wout×K2×Cin×Cout

内存需求包括存储权重和激活值。权重内存为:
Mweights=K2×Cin×Cout×bweightM_{weights} = K^2 \times C_{in} \times C_{out} \times b_{weight}Mweights=K2×Cin×Cout×bweight
其中bweightb_{weight}bweight是每个权重的位数。

峰值激活内存(假设使用流式处理)为:
Mactivation=max(Hin×Win×Cin,Hout×Wout×Cout)×bactivationM_{activation} = max(H_{in} \times W_{in} \times C_{in}, H_{out} \times W_{out} \times C_{out}) \times b_{activation}Mactivation=max(Hin×Win×Cin,Hout×Wout×Cout)×bactivation

对于一个典型的MCU(如ARM Cortex-M4,运行频率48MHz),理论峰值性能约为48 MOPS(百万次操作/秒)。如果我们需要10 FPS的实时推理,那么每帧可用的操作数为:
Opsper_frame=48×10610=4.8×106Ops_{per\_frame} = \frac{48 \times 10^6}{10} = 4.8 \times 10^6Opsper_frame=1048×106=4.8×106

这严格限制了可以部署的模型复杂度。例如,MobileNet v1的第一层(224×224×3输入,3×3卷积,32输出通道)需要约1080万次操作,已经超出了单帧的计算预算。

模型压缩的高级数学方法

量化理论的深入分析

量化是将连续值映射到离散值的过程。对于神经网络权重量化,我们通常使用均匀量化或非均匀量化。均匀量化的数学表示为:

给定权重w∈[wmin,wmax]w \in [w_{min}, w_{max}]w[wmin,wmax],使用bbb位量化,量化步长为:
Δ=wmax−wmin2b−1\Delta = \frac{w_{max} - w_{min}}{2^b - 1}Δ=2b1wmaxwmin

量化函数为:
Q(w)=round(w−wminΔ)×Δ+wminQ(w) = \text{round}\left(\frac{w - w_{min}}{\Delta}\right) \times \Delta + w_{min}Q(w)=round(Δwwmin)×Δ+wmin

量化误差的期望值为零(假设误差均匀分布),方差为:
σq2=Δ212\sigma_q^2 = \frac{\Delta^2}{12}σq2=12Δ2

对于神经网络的前向传播,量化引入的误差会逐层累积。考虑一个L层网络,第l层的输出为:
y(l)=f(W(l)x(l−1)+b(l))y^{(l)} = f(W^{(l)}x^{(l-1)} + b^{(l)})y(l)=f(W(l)x(l1)+b(l))

量化后变为:
y~(l)=f(Q(W(l))x(l−1)+Q(b(l)))\tilde{y}^{(l)} = f(Q(W^{(l)})x^{(l-1)} + Q(b^{(l)}))y~(l)=f(Q(W(l))x(l1)+Q(b(l)))

误差传播可以通过泰勒展开分析:
∥y~(L)−y(L)∥≤∑l=1L∥J(l)∥⋅∥Q(W(l))−W(l)∥⋅∥x(l−1)∥\|\tilde{y}^{(L)} - y^{(L)}\| \leq \sum_{l=1}^{L} \|J^{(l)}\| \cdot \|Q(W^{(l)}) - W^{(l)}\| \cdot \|x^{(l-1)}\|y~(L)y(L)l=1LJ(l)Q(W(l))W(l)x(l1)

其中J(l)J^{(l)}J(l)是第l层的雅可比矩阵。

剪枝的理论基础

最优脑损伤(Optimal Brain Damage, OBD)算法使用泰勒展开来估计剪枝的影响。损失函数的变化可以近似为:
δE=∑igiδwi+12∑ihii(δwi)2+12∑i≠jhijδwiδwj+O(∥δw∥3)\delta E = \sum_i g_i \delta w_i + \frac{1}{2}\sum_i h_{ii}(\delta w_i)^2 + \frac{1}{2}\sum_{i \neq j} h_{ij}\delta w_i \delta w_j + O(\|\delta w\|^3)δE=igiδwi+21ihii(δwi)2+21i=jhijδwiδwj+O(δw3)

其中gi=∂E∂wig_i = \frac{\partial E}{\partial w_i}gi=wiE是梯度,hij=∂2E∂wi∂wjh_{ij} = \frac{\partial^2 E}{\partial w_i \partial w_j}hij=wiwj2E是Hessian矩阵元素。

在训练收敛后,gi≈0g_i \approx 0gi0,且假设Hessian是对角的(即hij=0h_{ij} = 0hij=0i≠ji \neq ji=j),则剪除权重wiw_iwi(设δwi=−wi\delta w_i = -w_iδwi=wi)导致的损失增加为:
δEi≈12hiiwi2\delta E_i \approx \frac{1}{2}h_{ii}w_i^2δEi21hiiwi2

定义显著性(saliency)为:
si=hiiwi22s_i = \frac{h_{ii}w_i^2}{2}si=2hiiwi2

剪枝策略是移除显著性最小的权重。

知识蒸馏的信息论视角

知识蒸馏可以从信息论角度理解。教师网络的软目标包含了关于类间相似性的信息。使用温度T的softmax函数:
qi=exp⁡(zi/T)∑jexp⁡(zj/T)q_i = \frac{\exp(z_i/T)}{\sum_j \exp(z_j/T)}qi=jexp(zj/T)exp(zi/T)

当T增大时,概率分布变得更加平滑,保留了更多的类间关系信息。

蒸馏损失函数结合了硬目标和软目标:
L=αH(y,σ(zs))+(1−α)T2H(σ(zt/T),σ(zs/T))L = \alpha H(y, \sigma(z_s)) + (1-\alpha)T^2 H(\sigma(z_t/T), \sigma(z_s/T))L=αH(y,σ(zs))+(1α)T2H(σ(zt/T),σ(zs/T))

其中H是交叉熵,第二项乘以T2T^2T2是为了保持梯度的量级。

信息瓶颈理论表明,好的表示应该压缩输入信息同时保留预测相关信息:
min⁡I(X;Z)−βI(Z;Y)\min I(X; Z) - \beta I(Z; Y)minI(X;Z)βI(Z;Y)

其中Z是中间表示,I是互信息,β控制压缩和预测之间的权衡。

能耗模型与优化

MCU的能耗可以建模为:
Etotal=Ecompute+Ememory+EidleE_{total} = E_{compute} + E_{memory} + E_{idle}Etotal=Ecompute+Ememory+Eidle

计算能耗与操作数成正比:
Ecompute=α⋅Ops⋅Vdd2E_{compute} = \alpha \cdot Ops \cdot V_{dd}^2Ecompute=αOpsVdd2

其中α是与技术节点相关的常数,VddV_{dd}Vdd是供电电压。

内存访问能耗取决于内存层次:
Ememory=∑ini⋅eiE_{memory} = \sum_{i} n_i \cdot e_iEmemory=iniei

其中nin_ini是对第i级内存的访问次数,eie_iei是每次访问的能耗。典型值为:

  • SRAM读取:5 pJ/字节
  • Flash读取:50 pJ/字节
  • DRAM读取:640 pJ/字节

对于电池供电的设备,如果使用容量为C的电池,电压为V,效率为η,则运行时间为:
Tbattery=C⋅V⋅ηPavgT_{battery} = \frac{C \cdot V \cdot \eta}{P_{avg}}Tbattery=PavgCVη

其中平均功耗Pavg=Etotal⋅finferenceP_{avg} = E_{total} \cdot f_{inference}Pavg=Etotalfinferencefinferencef_{inference}finference是推理频率。

TinyML框架的软件优化技术

算子融合

算子融合减少了中间结果的内存访问。例如,Conv-BatchNorm-ReLU可以融合为一个算子:
y=max⁡(0,γConv(x)−μσ+β)y = \max(0, \gamma \frac{Conv(x) - \mu}{\sigma} + \beta)y=max(0,γσConv(x)μ+β)

可以预计算融合参数:
Wfused=γσWconvW_{fused} = \frac{\gamma}{\sigma} W_{conv}Wfused=σγWconv
bfused=γbconv−μσ+βb_{fused} = \gamma \frac{b_{conv} - \mu}{\sigma} + \betabfused=γσbconvμ+β

这样融合后的操作变为:
y=max⁡(0,Wfused∗x+bfused)y = \max(0, W_{fused} * x + b_{fused})y=max(0,Wfusedx+bfused)

内存布局优化

不同的内存布局影响缓存命中率。常见布局包括:

  • NCHW:批次-通道-高度-宽度
  • NHWC:批次-高度-宽度-通道
  • CHWN:通道-高度-宽度-批次

对于深度可分离卷积,NHWC布局通常更高效,因为深度卷积在通道维度上是独立的。

循环优化

循环展开可以减少分支预测开销:

// 原始循环
for(int i = 0; i < N; i++) {
    y[i] = w[i] * x[i];
}

// 4倍展开
for(int i = 0; i < N; i += 4) {
    y[i] = w[i] * x[i];
    y[i+1] = w[i+1] * x[i+1];
    y[i+2] = w[i+2] * x[i+2];
    y[i+3] = w[i+3] * x[i+3];
}

循环分块优化数据局部性:

// 分块矩阵乘法
for(int i = 0; i < N; i += B) {
    for(int j = 0; j < N; j += B) {
        for(int k = 0; k < N; k += B) {
            // 处理B×B的块
            matmul_block(C[i:i+B][j:j+B], 
                        A[i:i+B][k:k+B], 
                        B[k:k+B][j:j+B]);
        }
    }
}

实际应用案例分析

语音关键词检测

一个典型的关键词检测系统使用如表2所示的网络架构。对于检测"Hey Siri"这样的唤醒词,系统需要:

  1. 音频预处理:40ms窗口,25ms重叠,产生40维MFCC特征
  2. 网络架构:1个卷积层(8个3×3滤波器)+ 2个全连接层(32和4个神经元)
  3. 内存需求
    • 权重:(3×3×1×8 + 8) + (8×49×32 + 32) + (32×4 + 4) ≈ 13KB
    • 激活:max(40×49, 8×49, 32, 4) × 1字节 ≈ 2KB
  4. 计算需求:约500K操作/秒,在48MHz MCU上占用约1%的计算资源

视觉唤醒词

视觉唤醒词数据集包含115,000张96×96灰度图像,分为"person"和"no person"两类。典型的MobileNet v1模型配置:

  • 深度乘数:0.25
  • 输入分辨率:96×96
  • 参数量:230KB
  • 准确率:85%
  • 推理时间:约200ms(在ARM Cortex-M4上)

未来发展趋势与挑战

TinyML技术的发展面临着多方面的机遇和挑战。硬件方面,新型存算一体架构有望突破冯·诺依曼瓶颈,将计算能效提升10-100倍。模拟计算和脉冲神经网络等新范式也在探索中。

软件方面,自动化神经架构搜索(NAS)正在专门为MCU设计优化的网络结构。联邦学习使得边缘设备能够协同训练模型,而不需要共享原始数据。增量学习和持续学习技术使模型能够适应新的数据分布。

标准化工作也在进行中,TinyMLPerf基准测试套件提供了统一的性能评估标准,包括延迟、准确率、能耗和内存使用等指标。

安全和隐私是另一个重要考虑。边缘推理虽然避免了数据上传到云端,但模型本身可能泄露训练数据的信息。差分隐私、安全多方计算等技术正在被研究以保护模型和数据的隐私。

结论

TinyML代表了机器学习发展的新前沿,它将智能带到了计算的最边缘。通过巧妙的算法设计、硬件优化和软件工程,我们能够在仅有几百KB内存和毫瓦级功耗的设备上运行复杂的神经网络。这不仅仅是技术上的成就,更重要的是它开启了全新的应用可能性。

从技术角度看,TinyML涉及多个学科的交叉融合:机器学习提供算法基础,计算机体系结构优化硬件设计,嵌入式系统工程确保实际可部署性,信号处理提供高效的特征提取方法。这种跨学科的特性既是挑战也是机遇,需要研究者具备更广泛的知识背景。

展望未来,随着5G和6G网络的部署,边缘计算将变得更加重要。TinyML设备可以形成分布式智能网络,在本地处理数据的同时,通过高速网络协同工作。这种计算范式将使我们能够构建真正的智能环境,其中每个传感器都具有理解和响应其环境的能力。


附录A:数学推导

A.1 深度可分离卷积的计算复杂度分析

标准卷积的计算复杂度推导:

设输入特征图为X∈RH×W×CinX \in \mathbb{R}^{H \times W \times C_{in}}XRH×W×Cin,卷积核为K∈Rk×k×Cin×CoutK \in \mathbb{R}^{k \times k \times C_{in} \times C_{out}}KRk×k×Cin×Cout,输出特征图为Y∈RH′×W′×CoutY \in \mathbb{R}^{H' \times W' \times C_{out}}YRH×W×Cout

标准卷积操作可以表示为:
Yi,j,m=∑a=0k−1∑b=0k−1∑c=0Cin−1Ka,b,c,m⋅Xi+a,j+b,c+bmY_{i,j,m} = \sum_{a=0}^{k-1} \sum_{b=0}^{k-1} \sum_{c=0}^{C_{in}-1} K_{a,b,c,m} \cdot X_{i+a,j+b,c} + b_mYi,j,m=a=0k1b=0k1c=0Cin1Ka,b,c,mXi+a,j+b,c+bm

计算复杂度为:
FLOPsstd=H′×W′×k2×Cin×Cout\text{FLOPs}_{std} = H' \times W' \times k^2 \times C_{in} \times C_{out}FLOPsstd=H×W×k2×Cin×Cout

深度可分离卷积分为两步:

步骤1:深度卷积(Depthwise Convolution)
对每个输入通道独立应用一个k×kk \times kk×k的卷积核:
Zi,j,c=∑a=0k−1∑b=0k−1Ka,b,cdw⋅Xi+a,j+b,cZ_{i,j,c} = \sum_{a=0}^{k-1} \sum_{b=0}^{k-1} K^{dw}_{a,b,c} \cdot X_{i+a,j+b,c}Zi,j,c=a=0k1b=0k1Ka,b,cdwXi+a,j+b,c

计算复杂度:
FLOPsdw=H′×W′×k2×Cin\text{FLOPs}_{dw} = H' \times W' \times k^2 \times C_{in}FLOPsdw=H×W×k2×Cin

步骤2:逐点卷积(Pointwise Convolution)
使用1×11 \times 11×1卷积将深度卷积的输出投影到CoutC_{out}Cout维:
Yi,j,m=∑c=0Cin−1Kc,mpw⋅Zi,j,cY_{i,j,m} = \sum_{c=0}^{C_{in}-1} K^{pw}_{c,m} \cdot Z_{i,j,c}Yi,j,m=c=0Cin1Kc,mpwZi,j,c

计算复杂度:
FLOPspw=H′×W′×Cin×Cout\text{FLOPs}_{pw} = H' \times W' \times C_{in} \times C_{out}FLOPspw=H×W×Cin×Cout

总计算复杂度:
FLOPsds=FLOPsdw+FLOPspw=H′×W′×(k2×Cin+Cin×Cout)\text{FLOPs}_{ds} = \text{FLOPs}_{dw} + \text{FLOPs}_{pw} = H' \times W' \times (k^2 \times C_{in} + C_{in} \times C_{out})FLOPsds=FLOPsdw+FLOPspw=H×W×(k2×Cin+Cin×Cout)

压缩比:
ρ=FLOPsdsFLOPsstd=k2×Cin+Cin×Coutk2×Cin×Cout=1Cout+1k2\rho = \frac{\text{FLOPs}_{ds}}{\text{FLOPs}_{std}} = \frac{k^2 \times C_{in} + C_{in} \times C_{out}}{k^2 \times C_{in} \times C_{out}} = \frac{1}{C_{out}} + \frac{1}{k^2}ρ=FLOPsstdFLOPsds=k2×Cin×Coutk2×Cin+Cin×Cout=Cout1+k21

对于典型值k=3k=3k=3Cout=256C_{out}=256Cout=256,压缩比约为ρ≈1/256+1/9≈0.115\rho \approx 1/256 + 1/9 \approx 0.115ρ1/256+1/90.115,即计算量减少约8.7倍。

A.2 量化感知训练的梯度计算

在量化感知训练中,前向传播使用量化权重,但反向传播需要特殊处理。使用直通估计器(Straight-Through Estimator, STE):

前向传播:
w^=Q(w)=s⋅clip(round(ws),−2b−1,2b−1−1)\hat{w} = Q(w) = s \cdot \text{clip}\left(\text{round}\left(\frac{w}{s}\right), -2^{b-1}, 2^{b-1}-1\right)w^=Q(w)=sclip(round(sw),2b1,2b11)

其中sss是量化尺度,bbb是位宽。

反向传播使用STE:
∂L∂w=∂L∂w^⋅1∣w∣≤α\frac{\partial L}{\partial w} = \frac{\partial L}{\partial \hat{w}} \cdot \mathbb{1}_{|w| \leq \alpha}wL=w^L1wα

其中1\mathbb{1}1是指示函数,α\alphaα是量化范围。

量化尺度的优化可以通过学习得到。定义可学习的量化参数α\alphaα
s=α2b−1−1s = \frac{\alpha}{2^{b-1}-1}s=2b11α

使用重参数化技巧:
α=σ(β)\alpha = \sigma(\beta)α=σ(β)

其中σ\sigmaσ是sigmoid函数,β\betaβ是可学习参数。

梯度计算:
∂L∂β=∂L∂w^⋅∂w^∂s⋅∂s∂α⋅∂α∂β\frac{\partial L}{\partial \beta} = \frac{\partial L}{\partial \hat{w}} \cdot \frac{\partial \hat{w}}{\partial s} \cdot \frac{\partial s}{\partial \alpha} \cdot \frac{\partial \alpha}{\partial \beta}βL=w^Lsw^αsβα

A.3 批归一化折叠的详细推导

批归一化操作:
y=γx−μσ2+ϵ+βy = \gamma \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \betay=γσ2+ϵ xμ+β

对于卷积层输出z=Wx+bz = Wx + bz=Wx+b,应用批归一化:
y=γWx+b−μσ2+ϵ+βy = \gamma \frac{Wx + b - \mu}{\sqrt{\sigma^2 + \epsilon}} + \betay=γσ2+ϵ Wx+bμ+β

定义折叠后的权重和偏置:
Wfold=γσ2+ϵWW_{fold} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} WWfold=σ2+ϵ γW
bfold=γσ2+ϵ(b−μ)+βb_{fold} = \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}(b - \mu) + \betabfold=σ2+ϵ γ(bμ)+β

验证:
Wfoldx+bfold=γσ2+ϵWx+γσ2+ϵ(b−μ)+β=γWx+b−μσ2+ϵ+β=y\begin{align} W_{fold}x + b_{fold} &= \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}} Wx + \frac{\gamma}{\sqrt{\sigma^2 + \epsilon}}(b - \mu) + \beta \\ &= \gamma \frac{Wx + b - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta \\ &= y \end{align}Wfoldx+bfold=σ2+ϵ γWx+σ2+ϵ γ(bμ)+β=γσ2+ϵ Wx+bμ+β=y

这种折叠在推理时节省了归一化操作的计算和内存开销。

A.4 知识蒸馏的KL散度分析

教师模型和学生模型的软标签分别为pTp^TpTpSp^SpS,使用温度TTT
piT=exp⁡(ziT/T)∑jexp⁡(zjT/T)p_i^T = \frac{\exp(z_i^T/T)}{\sum_j \exp(z_j^T/T)}piT=jexp(zjT/T)exp(ziT/T)
piS=exp⁡(ziS/T)∑jexp⁡(zjS/T)p_i^S = \frac{\exp(z_i^S/T)}{\sum_j \exp(z_j^S/T)}piS=jexp(zjS/T)exp(ziS/T)

KL散度:
DKL(pT∣∣pS)=∑ipiTlog⁡piTpiSD_{KL}(p^T || p^S) = \sum_i p_i^T \log \frac{p_i^T}{p_i^S}DKL(pT∣∣pS)=ipiTlogpiSpiT

展开:
DKL=∑ipiTlog⁡piT−∑ipiTlog⁡piSD_{KL} = \sum_i p_i^T \log p_i^T - \sum_i p_i^T \log p_i^SDKL=ipiTlogpiTipiTlogpiS

第一项是教师分布的熵(常数),第二项是交叉熵。

T→∞T \to \inftyT时,使用泰勒展开:
piT≈1N+ziT−zˉTNTp_i^T \approx \frac{1}{N} + \frac{z_i^T - \bar{z}^T}{NT}piTN1+NTziTzˉT

其中NNN是类别数,zˉT\bar{z}^TzˉT是logits的均值。

这表明高温度下的蒸馏主要匹配logits的相对大小,而不是绝对值。

A.5 剪枝的最优化理论

将剪枝问题表述为一个约束优化问题:
min⁡w,mL(w⊙m)s.t.∥m∥0≤k\min_{w,m} L(w \odot m) \quad \text{s.t.} \quad \|m\|_0 \leq kw,mminL(wm)s.t.m0k

其中m∈{0,1}nm \in \{0,1\}^nm{0,1}n是掩码向量,⊙\odot是逐元素乘积,kkk是保留的权重数。

使用拉格朗日松弛:
L(w,m,λ)=L(w⊙m)+λ(∥m∥0−k)\mathcal{L}(w,m,\lambda) = L(w \odot m) + \lambda(\|m\|_0 - k)L(w,m,λ)=L(wm)+λ(m0k)

这是一个NP难问题。常用的近似方法包括:

  1. 幅度剪枝:选择∣wi∣|w_i|wi最大的kkk个权重
  2. 梯度剪枝:选择∣wi⋅gi∣|w_i \cdot g_i|wigi最大的kkk个权重
  3. Hessian剪枝:选择∣wi2⋅hii∣|w_i^2 \cdot h_{ii}|wi2hii最大的kkk个权重

Fisher信息可以作为Hessian的近似:
Fii=Ex∼pdata[(∂log⁡p(y∣x,w)∂wi)2]F_{ii} = \mathbb{E}_{x \sim p_{data}}\left[\left(\frac{\partial \log p(y|x,w)}{\partial w_i}\right)^2\right]Fii=Expdata[(wilogp(yx,w))2]

A.6 能耗建模的详细分析

动态能耗模型:
Edynamic=αCLVdd2fE_{dynamic} = \alpha C_L V_{dd}^2 fEdynamic=αCLVdd2f

其中α\alphaα是活动因子,CLC_LCL是负载电容,VddV_{dd}Vdd是供电电压,fff是频率。

静态能耗(漏电流):
Estatic=IleakVddtE_{static} = I_{leak} V_{dd} tEstatic=IleakVddt

总能耗:
Etotal=Edynamic+EstaticE_{total} = E_{dynamic} + E_{static}Etotal=Edynamic+Estatic

对于神经网络推理,能耗可以分解为:
Einference=∑l=1L(Ecompute(l)+Ememory(l))E_{inference} = \sum_{l=1}^{L} (E_{compute}^{(l)} + E_{memory}^{(l)})Einference=l=1L(Ecompute(l)+Ememory(l))

每层的计算能耗:
Ecompute(l)=nMAC(l)⋅eMACE_{compute}^{(l)} = n_{MAC}^{(l)} \cdot e_{MAC}Ecompute(l)=nMAC(l)eMAC

内存访问能耗:
Ememory(l)=nread(l)⋅eread+nwrite(l)⋅ewriteE_{memory}^{(l)} = n_{read}^{(l)} \cdot e_{read} + n_{write}^{(l)} \cdot e_{write}Ememory(l)=nread(l)eread+nwrite(l)ewrite

数据重用可以显著降低能耗。重用因子定义为:
R=计算次数内存访问次数R = \frac{\text{计算次数}}{\text{内存访问次数}}R=内存访问次数计算次数

优化目标是最大化RRR,这导致了各种优化的数据流架构。

A.7 联邦学习在TinyML中的应用

联邦平均算法(FedAvg):

设有KKK个边缘设备,每个设备kkk有本地数据集DkD_kDk

本地更新:
wkt+1=wkt−η∇Lk(wkt)w_k^{t+1} = w_k^t - \eta \nabla L_k(w_k^t)wkt+1=wktηLk(wkt)

全局聚合:
wt+1=∑k=1K∣Dk∣∣D∣wkt+1w^{t+1} = \sum_{k=1}^{K} \frac{|D_k|}{|D|} w_k^{t+1}wt+1=k=1KDDkwkt+1

通信成本:
Ccomm=2⋅rounds⋅K⋅size(w)C_{comm} = 2 \cdot \text{rounds} \cdot K \cdot \text{size}(w)Ccomm=2roundsKsize(w)

对于TinyML设备,关键挑战是:

  1. 有限的本地计算能力
  2. 不可靠的网络连接
  3. 非独立同分布(non-IID)数据

压缩技术可以减少通信开销:

  • 梯度量化:g~=Q(g)\tilde{g} = Q(g)g~=Q(g)
  • 稀疏化:只传输top-k梯度
  • 误差累积:et+1=g−g~+ete^{t+1} = g - \tilde{g} + e^tet+1=gg~+et

A.8 神经架构搜索(NAS)for TinyML

搜索空间定义:
A={a∣latency(a)≤L,memory(a)≤M,energy(a)≤E}\mathcal{A} = \{a | \text{latency}(a) \leq L, \text{memory}(a) \leq M, \text{energy}(a) \leq E\}A={alatency(a)L,memory(a)M,energy(a)E}

多目标优化:
min⁡a∈A{−accuracy(a),latency(a),energy(a)}\min_{a \in \mathcal{A}} \{-\text{accuracy}(a), \text{latency}(a), \text{energy}(a)\}aAmin{accuracy(a),latency(a),energy(a)}

使用强化学习的NAS:

  • 状态:当前部分架构
  • 动作:添加一个层/操作
  • 奖励:r=accuracy−λ1⋅latency−λ2⋅energyr = \text{accuracy} - \lambda_1 \cdot \text{latency} - \lambda_2 \cdot \text{energy}r=accuracyλ1latencyλ2energy

可微分NAS使用连续松弛:
oˉ(i,j)=∑o∈Oexp⁡(αo(i,j))∑o′∈Oexp⁡(αo′(i,j))⋅o(x)\bar{o}^{(i,j)} = \sum_{o \in \mathcal{O}} \frac{\exp(\alpha_o^{(i,j)})}{\sum_{o' \in \mathcal{O}} \exp(\alpha_{o'}^{(i,j)})} \cdot o(x)oˉ(i,j)=oOoOexp(αo(i,j))exp(αo(i,j))o(x)

其中αo(i,j)\alpha_o^{(i,j)}αo(i,j)是架构参数,通过梯度下降优化。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐