μNAS：面向微控制器的约束神经架构搜索

Liberis E, Dudziak Ł, Lane N D. μnas: Constrained neural architecture search for microcontrollers[C]//Proceedings of the 1st Workshop on Machine Learning and Systems. 2021: 70-79.

1. 引言与研究背景

1.1 微控制器的资源约束挑战

物联网（IoT）设备的智能化正在改变我们与周围世界的交互方式。这些设备的核心是微控制器单元（MCU），它们是包含低频处理器、持久程序存储器（Flash）和易失性静态RAM的超小型计算机，全部集成在单一芯片上。2019年，全球MCU出货量估计超过300亿个单元，其广泛应用得益于极低的成本和功耗优势。

然而，这些优势是以计算能力的大幅降低为代价的。本研究聚焦于"中等规格"的物联网级MCU，其SRAM和持久存储空间最多64KB，可用于神经网络推理。这种严苛的资源约束给运行资源密集型计算带来了前所未有的挑战。相比之下，即使是移动设备也拥有比MCU多几个数量级的计算资源。

1.2 现有方法的局限性

传统的模型压缩方法，如剪枝、蒸馏、量化和使用更轻量级的层，虽然能够在保持精度的同时减少模型大小，但大多数方法并不适合生成MCU级别的神经网络。许多技术针对的是移动设备或类似平台，这些平台的计算资源比MCU多出几个数量级。即使是专门设计的资源高效模型，如MobileNet和SqueezeNet，也会超出假定的资源预算10倍以上。

2. 神经网络在MCU上的执行机制

2.1 执行策略详解

μNAS遵循TensorFlow Lite Micro解释器和CMSIS-NN库使用的执行策略，这种策略具有以下特点：

顺序执行原则：算子按预定义顺序逐个执行，完全没有并行性。这与GPU或现代CPU的并行执行形成鲜明对比。

内存管理机制：执行一个算子时遵循三步流程：

• 在SRAM中为输出缓冲区分配/预留内存
• 完全执行算子主体
• 释放不再需要的输入缓冲区

这意味着神经网络的激活矩阵仅存储在SRAM中，而所有静态数据（如网络权重）则从Flash中的可执行二进制文件读取。

量化策略：网络权重和所有计算都量化为8位定点数据类型（int8）。这种量化方法的选择基于四个关键优势：字节对齐（值可以直接加载而无需解码）、不需要芯片上存在浮点运算单元、广泛的行业采用、以及在量化过程中不会造成显著的精度损失。

2.2 分支网络的内存管理挑战

图1：计算图的不同执行路径对峰值内存使用的影响

图1展示了一个具有分支结构的卷积神经网络计算图。网络包含多个Conv2D层，具有不同的配置（1×1卷积、3×3深度可分离卷积等）和一个Concat层。图中显示了两种不同的执行顺序：

• 默认执行路径（左侧）：按照1→2→3→4→5→6→7的顺序执行，产生5216字节的峰值内存使用
• 优化执行路径（右侧）：按照1→2→5→3→4→6→7的顺序执行，峰值内存使用降至4960字节

这256字节的差异虽然看似不大，但在只有64KB总内存的MCU上却至关重要。图中每个节点旁边的尺寸标注（如7×7×32、4×4×16）表示该层输出张量的维度，这直接影响内存占用。

3. 多目标优化的数学框架

3.1 问题形式化

神经架构搜索的核心是一个约束零阶优化问题。给定搜索空间A，我们寻找最优架构α*，使其最大化某个目标函数L，同时满足资源使用约束。验证集精度对于特定架构α由通过梯度下降获得的最优权重θ*决定。

将资源约束视为软约束并重新表述为额外目标，NAS成为多目标优化问题：

$$\alpha^* = \arg\min_{\alpha \in A} L(\alpha) = \arg\min_{\alpha \in A} \begin{cases} 1.0 - \text{ValAccuracy}(\alpha) \\ \text{ModelSize}(\alpha) \\ \text{PeakMemUsage}(\alpha) \\ \text{Latency}(\alpha) \end{cases}$$

3.2 Pareto前沿与标量化

多目标问题的解不是单一值，而是Pareto前沿——一组点的集合，在这些点上，改进一个目标函数必然导致另一个目标函数变差。为了有效探索Pareto前沿，μNAS采用随机标量化方法：

$$L_t(\alpha) = \max_{i \in \{1,2,3,4\}} \left\{ \lambda_i^t \cdot f_i(\alpha) \right\}$$

其中：

• $f_1(\alpha) = 1.0 - \text{ValAccuracy}(\alpha)$
• $f_2(\alpha) = \text{PeakMemUsage}(\alpha)$
• $f_3(\alpha) = \text{ModelSize}(\alpha)$
• $f_4(\alpha) = \text{MACs}(\alpha)$

标量系数$\lambda_i^t$的采样策略为：$1/\lambda_i \sim \text{Uniform}[0; b_i]$，其中$b_i$是用户为第i个目标指定的软上界。

4. 搜索空间的层次化设计

4.1 架构模板

μNAS的搜索空间基于高度参数化的架构模板。一个架构由N个"卷积块"组成（N∈[1,10]），每个块的详细配置如下：

连接模式：每个块可以与前一个块串联或并联连接，提供了灵活的信息流路径。

卷积层配置：每个块包含M个卷积层（M∈[1,3]），每层具有以下可配置参数：

• 可选的2×2最大池化前置操作
• 全卷积或深度可分离卷积的选择
• 核大小K∈{1,3,5,7}
• 输出通道数C∈[1,128]
• 步长S∈{1,2}
• 可选的批归一化和ReLU激活

池化和全连接层：每个卷积块后跟P×P的平均或最大池化（P∈{2,4,6}），然后是F个全连接层（F∈[1,3]），每层有U个单元（U∈[10,256]）。

4.2 搜索空间规模分析

这种细粒度的设计产生了包含$1.15 \times 10^{152}$个可能架构的搜索空间。这个天文数字级的规模既是挑战也是机遇——它提供了找到最优架构的可能性，但也使搜索变得极其困难。

5. 资源约束的精确建模

5.1 峰值内存使用的算法计算

峰值内存使用的准确计算是μNAS的核心创新之一。算法基于以下原理：

工作集定义：在任意时刻t，工作集$W_t$包含所有必须保存在内存中的张量。这包括：

• 当前执行算子的输入张量
• 当前执行算子的输出张量
• 后续算子需要但尚未处理的张量

拓扑排序枚举：对于具有分支的网络，不同的执行顺序产生不同的工作集序列。μNAS枚举所有有效的拓扑排序，计算每种排序的峰值内存使用：

$$\text{PeakMemUsage}(\alpha) = \min_{\pi \in \Pi(\alpha)} \max_{t} \sum_{tensor \in W_t^\pi} \text{size}(tensor)$$

其中$\Pi(\alpha)$是架构α的所有有效拓扑排序的集合。

5.2 模型大小的量化计算

采用8位整数量化后，模型大小计算简化为：

$$\text{ModelSize}(\alpha) = \sum_{l \in \text{layers}(\alpha)} \text{params}(l) \times 1 \text{ byte}$$

5.3 延迟预测模型

图2：MACs作为MCU延迟预测器的有效性验证

图2分为两部分，形成鲜明对比：

上图（GPU延迟）：展示了桌面GPU上FLOPs与实测延迟的关系，数据点分散且呈现非线性模式，延迟范围在0-10毫秒，FLOPs范围在0-400M。这种分散反映了GPU的复杂执行特性，包括并行处理、缓存效应和调度优化。

下图（MCU延迟）：展示了1000个随机采样模型在NUCLEO-H743ZI2 MCU板上的实测延迟与MACs的关系。数据点紧密围绕一条直线分布，R²=0.975的高拟合度证明了线性关系。延迟范围在0-800毫秒，MACs范围在0-80M。

这种线性关系使得延迟预测模型简化为：

$$\text{Latency}(\alpha) \approx k \cdot \text{MACs}(\alpha) + c$$

其中k和c是通过实验确定的常数。

6. 搜索算法的设计与实现

6.1 老化进化算法

老化进化（AE）算法维护一个大小为P的种群，通过以下步骤演化：

1. 选择：从种群中随机采样S个架构
2. 评估：计算每个架构的$L_t(\alpha)$值
3. 变异：对最优架构应用随机变形操作
4. 更新：将新架构加入种群，移除最老的架构

变形操作包括改变层数、修改通道数、调整核大小等，详见表1中的完整列表。

6.2 贝叶斯优化方法

贝叶斯优化使用高斯过程（GP）作为代理模型来近似目标函数。核函数基于最优传输定义两个神经网络之间的距离：

$$k(\alpha_1, \alpha_2) = \exp\left(-\frac{d_{OT}(\alpha_1, \alpha_2)}{2\sigma^2}\right)$$

其中$d_{OT}$是基于最优传输的距离度量。

7. 模型压缩：结构化剪枝

7.1 DPF剪枝算法

μNAS采用DPF（Dynamic Pruning with Feedback）剪枝，基于L2范数识别和移除不重要的通道/单元：

$$\text{Importance}(c_i) = \|\mathbf{w}_i\|_2 = \sqrt{\sum_j w_{ij}^2}$$

其中$\mathbf{w}_i$是第i个通道或单元的权重向量。

7.2 渐进式剪枝策略

剪枝在训练过程中逐步进行，目标稀疏度$s_{target}$通过以下公式在训练周期内线性增加：

$$s_t = s_{initial} + (s_{target} - s_{initial}) \cdot \frac{t - t_{start}}{t_{end} - t_{start}}$$

8. 实验结果的深入分析

8.1 配置比较实验

图3：不同μNAS配置的Pareto前沿比较

图3展示了在Chars74K（左）和MNIST（右）数据集上三种配置的性能：

Chars74K结果：

• 错误率范围：0.10-0.50
• 模型大小范围：0-30000字节
• AE+剪枝配置（绿色点）明显占据了最左下角的位置，表示在相同模型大小下实现了最低的错误率
• 纯AE（蓝色点）次之，而贝叶斯优化（橙色点）的Pareto前沿最差

MNIST结果：

• 错误率范围：0.00-0.10
• 模型大小范围：0-8000字节
• AE+剪枝在极小模型（<500字节）区域展现出卓越性能，达到了<1%的错误率

8.2 资源目标消融研究

图4：资源约束对搜索行为的影响

图4通过三个子图展示了不同约束配置下MNIST模型的分布：

左图（无模型大小约束）：

• 峰值内存使用集中在100-1000字节的低范围
• 模型大小分散在10²-10⁶字节的广泛范围
• 搜索被峰值内存约束主导，产生深而窄的网络

中图（无内存使用约束）：

• 模型大小集中在较低范围
• 峰值内存使用分散广泛
• 搜索倾向于浅而宽的架构

右图（所有约束存在）：

• 两个维度都展现出良好的分布
• 发现了多样化的权衡点
• 实现了资源使用和精度之间的平衡

9. 与现有方法的综合比较

表2详细比较了μNAS与现有方法在多个数据集上的表现。特别值得注意的结果包括：

MNIST：μNAS找到了仅480字节、精度99.19%的模型，相比SpArSe在相似精度下模型大小减少5.7倍。

Speech Commands：在保持95.58%精度的同时，模型大小仅37KB，峰值内存使用21.1KB，MACs仅1.1M——比RENA减少636倍的计算量。

Fashion MNIST：达到93.22%的精度，同时保持63.6KB的模型大小和12.6KB的峰值内存使用。

附录：数学推导

A. 多目标优化的Pareto最优性条件

A.1 Pareto支配关系

对于两个解$\alpha_1, \alpha_2 \in A$，我们定义Pareto支配关系$\prec$如下：

$$\alpha_1 \prec \alpha_2 \Leftrightarrow \forall i: f_i(\alpha_1) \leq f_i(\alpha_2) \land \exists j: f_j(\alpha_1) < f_j(\alpha_2)$$

A.2 Pareto前沿的数学定义

Pareto前沿$\mathcal{P}$定义为所有非支配解的集合：

$$\mathcal{P} = \{\alpha \in A | \nexists \alpha' \in A: \alpha' \prec \alpha\}$$

A.3 随机标量化的收敛性分析

定理：在适当的条件下，随机标量化方法以概率1收敛到完整的Pareto前沿。

证明概要：

令$\Lambda$为所有可能的权重向量$\lambda$的集合。对于Pareto前沿上的任意点$\alpha^* \in \mathcal{P}$，存在权重向量$\lambda^*$使得：

$$\alpha^* = \arg\min_{\alpha \in A} \max_i \{\lambda_i^* f_i(\alpha)\}$$

由于我们从连续分布中采样$\lambda$，随着采样次数$T \rightarrow \infty$：

$$P(\exists t \leq T: \|\lambda^t - \lambda^*\| < \epsilon) \rightarrow 1$$

因此，随机标量化方法能够以任意精度逼近Pareto前沿上的任意点。

B. 峰值内存使用的图算法

B.1 动态规划形式化

令$G = (V, E)$为神经网络的计算图，其中$V$是算子集合，$E$是数据依赖关系。定义状态$s = (v, M)$，其中$v$是当前算子，$M$是内存中的张量集合。

峰值内存使用的递归关系：

$$PMU(s) = \max\left\{\text{mem}(M \cup \text{output}(v)), \max_{v' \in \text{succ}(v)} PMU((v', M'))\right\}$$

其中$M'$是执行$v$后的更新内存状态：

$$M' = (M \setminus \text{inputs}(v)) \cup \text{output}(v) \cup \text{future\_deps}(v)$$

B.2 分支网络的内存调度优化

对于具有分支的网络，最优执行顺序可通过以下整数线性规划（ILP）求解：

决策变量：

• $x_{ij} \in \{0,1\}$：算子$i$是否在时间步$j$执行
• $m_j$：时间步$j$的内存使用量

目标函数：

$$\min \max_j m_j$$

约束条件：

• 每个算子恰好执行一次：$\sum_j x_{ij} = 1, \forall i$
• 依赖关系：$\sum_{k<j} x_{ik} \geq x_{pj}, \forall (p,i) \in E$
• 内存计算：$m_j = \sum_i \text{size}(i) \cdot y_{ij}$

其中$y_{ij}$表示张量$i$在时间$j$是否在内存中。

C. 量化误差分析

C.1 仿射量化的数学形式

8位整数量化将浮点值$r$映射到整数$q$：

$$q = \text{round}\left(\frac{r}{s}\right) + z$$

其中$s$是缩放因子，$z$是零点。反量化操作：

$$\hat{r} = s(q - z)$$

C.2 量化误差界

量化误差$\epsilon = r - \hat{r}$的界限：

$$|\epsilon| \leq \frac{s}{2}$$

对于神经网络层$f(x) = Wx + b$，量化后的输出误差：

$$\|f(x) - \hat{f}(\hat{x})\| \leq \|W\| \cdot \|\epsilon_x\| + \|\epsilon_W\| \cdot \|x\| + \|\epsilon_b\|$$

C.3 累积误差分析

对于L层网络，总误差界：

$$\|\text{output} - \widehat{\text{output}}\| \leq \sum_{l=1}^L \prod_{k=l+1}^L \|W_k\| \cdot \epsilon_l$$

这解释了为什么深度网络需要仔细的量化策略。

D. 搜索算法的理论分析

D.1 老化进化的收敛速度

假设目标函数$L$是Lipschitz连续的，常数为$L_f$。老化进化算法的期望改进率：

$$\mathbb{E}[\Delta L_t] \geq \frac{1}{P} \sum_{i=1}^S P(\text{select } i) \cdot \mathbb{E}[\Delta L | \text{mutate}(i)]$$

在适当的变异率$\mu$下，算法以速率$O(1/\sqrt{t})$收敛到局部最优。

D.2 贝叶斯优化的遗憾界

使用高斯过程的贝叶斯优化，其累积遗憾$R_T$满足：

$$R_T = \sum_{t=1}^T (L(\alpha^*) - L(\alpha_t)) \leq O(\sqrt{T \gamma_T \log T})$$

其中$\gamma_T$是最大信息增益，对于我们的核函数：

$$\gamma_T \leq O((\log T)^{d+1})$$

其中$d$是搜索空间的有效维度。

E. 剪枝算法的理论保证

E.1 结构化剪枝的最优性

给定目标稀疏度$s$，最优剪枝策略是保留重要性得分最高的$(1-s)$比例的通道：

$$\mathcal{C}^* = \arg\max_{\mathcal{C}: |\mathcal{C}|=(1-s)C} \sum_{c \in \mathcal{C}} \|\mathbf{w}_c\|_2^2$$

E.2 剪枝后的性能界

在适当的假设下，剪枝后的网络性能损失界：

$$|L(\alpha_{pruned}) - L(\alpha_{original})| \leq O(s \cdot \max_c \|\nabla_{\mathbf{w}_c} L\|)$$

这表明渐进式剪枝（逐步增加$s$）能够更好地保持性能。

F. 搜索空间大小的精确计算

搜索空间的总大小通过乘法原理计算：

$$|\mathcal{A}| = \prod_{n=1}^{10} \left[2 \cdot \prod_{m=1}^3 \left(2 \cdot 2 \cdot 4 \cdot 128 \cdot 2 \cdot 2 \cdot 2\right) \cdot 2 \cdot 3 \cdot \prod_{f=1}^3 247\right]$$

考虑所有可能的组合，得到$|\mathcal{A}| = 1.15 \times 10^{152}$。

这个庞大的搜索空间使得穷举搜索完全不可行，即使以每秒评估$10^9$个架构的速度，也需要$10^{143}$秒（远超宇宙年龄）才能遍历整个空间。这突出了高效搜索算法的必要性。