4、稀疏化训练原理：权重剪枝、掩码机制、梯度更新与稀疏正则化方法

TinyML 模型部署在微控制器（MCU）上时，面临 Flash 容量（通常 256KB~1MB）和 RAM（通常 64KB~512KB）的硬性约束。稀疏化训练的核心目标，是在训练阶段就让模型权重产生大量零值，从而利用 CSR（Compressed Sparse Row）或 Bitmask 等格式压缩存储，并跳过零值乘加运算以降低推理延迟。本章从权重剪枝的数学本质出发，深入剖析掩码机制如何与反向传播协同工作，并对比 L1 正则化、硬阈值正则化（HTR）与 ADMM 三种主流稀疏正则化方法的适用场景。

训练阶段 剪枝决策 推理部署 权重 W 掩码 M 前向: W⊙M 阈值 τ 或稀疏度 S |W| < τ → M=0 CSR/Bitmask 存储 跳过零值 MAC 反向传播: 仅更新 M=1 的权重 稀疏正则化方法对比 L1 正则化 λ|W| 软收缩 渐进稀疏 硬阈值正则化 Top-K 剪枝 立即置零 ADMM 交替方向乘子 精确稀疏约束

上图展示了稀疏化训练的整体框架。训练阶段维护权重矩阵 W 和二进制掩码矩阵 M（0/1），前向传播计算 W⊙M（逐元素相乘）；剪枝决策阶段根据阈值或目标稀疏度更新 M；推理阶段则利用稀疏格式存储和计算。反向传播时，梯度仅通过 M=1 的位置回传，M=0 的位置梯度被屏蔽，这是掩码机制的核心。

4.1 权重剪枝的数学本质与粒度选择

权重剪枝的本质是寻找一个稀疏权重矩阵 Ws，使其与原始稠密矩阵 W 在某个范数意义下误差最小：min ||W - Ws||F，同时满足 ||Ws||0 ≤ k（k 为目标非零元素个数）。实际工程中，我们通常不直接求解这个 NP-hard 问题，而是采用贪心策略：按绝对值大小排序，保留 Top-k 个权重，其余置零。

剪枝粒度直接影响压缩率和硬件效率：

  • 非结构化剪枝（细粒度）：每个权重独立判断是否剪掉。压缩率最高，但产生的稀疏模式不规则，MCU 上需要查表或位图索引，可能引入额外寻址开销。
  • 通道剪枝（粗粒度）：整个卷积核或全连接层的一整列/行被移除。硬件实现简单，但精度损失较大，通常需要配合重训练。
  • 块状剪枝（N:M 模式）：每 M 个连续权重中保留 N 个（如 2:4）。NVIDIA 的 Ampere 架构和部分 MCU DSP 扩展支持这种模式，兼顾压缩率和规则性。

对于 TinyML 场景，推荐 2:4 块状剪枝或 4:8 模式，因为 Cortex-M 系列内核的 SIMD 指令可以高效处理这类规则稀疏模式，避免逐元素分支判断。

4.2 掩码机制：前向传播与反向传播的协同

掩码机制是稀疏化训练的核心工程实现。训练过程中，我们维护一个与权重同形状的二进制张量 M，其更新规则如下：

# 伪代码：每 T 个 epoch 更新一次掩码
def update_mask(weights, target_sparsity, T):
    if epoch % T == 0:
        threshold = torch.quantile(torch.abs(weights), target_sparsity)
        mask = (torch.abs(weights) > threshold).float()
    return mask

# 前向传播
output = torch.mul(weights, mask)  # 或 weights * mask

# 反向传播（关键：梯度直通）
grad_weights = grad_output * input  # 正常计算梯度
grad_weights = grad_weights * mask  # 屏蔽被剪枝位置的梯度

这里有一个关键设计决策：被剪枝的权重是否允许在后续训练中“复活”？ 有两种策略：

  1. 硬掩码（冻结）：一旦 M=0，该权重永远不再参与前向和反向。优点是训练稳定，但可能错过最优解。
  2. 软掩码（允许复活）：每次掩码更新时重新计算所有权重的绝对值，之前被剪掉的权重如果重新变大，可以恢复为 1。这需要保留完整的权重值（而非直接置零），仅在推理时压缩。

工程实践中，软掩码策略在 TinyML 中更常用，因为训练阶段 Flash/RAM 充足，可以保留完整精度权重；推理阶段再根据最终掩码生成稀疏模型。但要注意，软掩码需要额外的内存存储完整权重，训练时的内存峰值会略高。

4.3 稀疏正则化方法：从 L1 到 ADMM

掩码机制解决了“如何剪”的问题，而稀疏正则化解决“如何让权重自然趋向于零”的问题。三种主流方法各有优劣：

方法 数学形式 稀疏度控制 收敛速度 适用场景
L1 正则化 Loss + λ∑|w| 间接（通过 λ 调节） 慢（需要长时间训练） 小模型、精度要求不高的场景
硬阈值正则化（HTR） 每 T 步 Top-K 剪枝 直接（精确控制稀疏度） 快（立即生效） 大多数 TinyML 场景，推荐
ADMM 增广拉格朗日 + 投影 精确（硬约束） 中（需要交替优化） 对稀疏度有严格要求的场景

L1 正则化的梯度更新公式为：w ← w - η(∂L/∂w + λ·sign(w))。它会让所有权重以恒定速率向零收缩，但小权重收缩快、大权重收缩慢，最终产生自然的稀疏分布。缺点是 λ 需要仔细调参，且训练时间较长。

硬阈值正则化（HTR）是最直接的方法：每 T 个 epoch，将绝对值最小的 k 个权重置零。其变体包括：

  • 渐进剪枝：从稀疏度 0 开始，每 T 步增加 Δs，直到目标稀疏度。避免一次性大幅剪枝导致精度崩溃。
  • 自动学习率调整：剪枝后模型容量下降，需要降低学习率（通常乘以 0.1~0.5）以稳定训练。

ADMM（交替方向乘子法）将稀疏约束转化为：min f(W) + g(Z)，s.t. W=Z，其中 g(Z) 是指示函数（Z 满足稀疏度约束时为 0，否则为 ∞）。通过增广拉格朗日乘子迭代求解。ADMM 的优点是能精确满足稀疏度约束，但需要额外维护对偶变量，内存开销增加约 30%，在 MCU 训练（如果有）时需谨慎。

实战建议：对于大多数 TinyML 项目，推荐使用 渐进硬阈值正则化 + 软掩码 的组合。具体做法：每 5 个 epoch 更新一次掩码，每次增加 5% 稀疏度，同时将学习率衰减 20%。这样既保证了训练稳定性，又能精确控制最终稀疏度。

4.4 案例拆解：关键词唤醒模型（KWS）的稀疏化训练

我们为一个 10 类关键词唤醒任务训练 CNN 模型，原始模型参数量 180K，部署在 Cortex-M4 上（Flash 256KB，RAM 64KB）。初始模型大小 720KB（FP32），远超 Flash 限制。

初始目标：将模型压缩到 100KB 以内，同时保持准确率下降不超过 2%。

训练配置：

  • 数据集：Google Speech Commands v2（35K 条 1 秒音频）
  • 基线准确率：94.2%
  • 稀疏化策略：渐进 HTR，从第 10 个 epoch 开始，每 5 epoch 增加 10% 稀疏度，目标稀疏度 85%
  • 掩码策略：软掩码，允许权重复活
  • 学习率：初始 0.001，每 10 epoch 衰减 0.5

训练过程关键指标：

Epoch 稀疏度 验证准确率 损失值 备注
10 0% 93.8% 0.21 开始剪枝
15 10% 93.5% 0.23 轻微下降
25 50% 92.9% 0.27 下降 0.9%
35 80% 92.1% 0.32 下降 1.7%
40 85% 91.8% 0.35 最终结果

分析判断：在稀疏度 50% 之前，准确率下降平缓（<0.5%）；超过 70% 后，每增加 5% 稀疏度，准确率下降约 0.3%。最终 85% 稀疏度时准确率 91.8%，下降 2.4%，略超目标。我们采取两个补救措施：

  1. 对最后两层全连接层降低剪枝力度（目标稀疏度 70%），因为靠近输出的层对分类边界更敏感。
  2. 增加 5 个 epoch 的微调，学习率降至 0.0001，仅更新未被剪枝的权重。

结果对比：调整后准确率回升至 92.5%，模型大小从 720KB 压缩至 108KB（CSR 格式），Flash 占用满足要求。推理时间从 12.3ms 降至 4.1ms（跳过零值 MAC），RAM 占用从 48KB 降至 22KB。

风险与代价：训练时间增加了约 40%（因为需要多次掩码更新和微调）；CSR 格式在 MCU 上需要额外的索引数组，增加了 5% 的 Flash 开销；稀疏度 85% 时，部分卷积核完全被剪掉，导致特征提取能力下降，需要确保每层至少保留 30% 的通道。

常见误区：不要在整个训练过程中使用固定掩码。如果掩码从不更新，被剪掉的权重永远无法恢复，一旦初始剪枝决策错误（例如剪掉了重要的低频特征），模型精度将永久受损。建议至少每 5 个 epoch 重新计算一次掩码。

4.5 稀疏化训练的调参策略与总结

稀疏化训练的核心超参数包括：目标稀疏度、剪枝步长（Δs）、掩码更新周期（T）、学习率衰减策略。以下是推荐的调参顺序：

  1. 先确定目标稀疏度：根据 Flash 容量反推。例如 Flash 256KB，模型原始 720KB，需要压缩 2.8 倍，对应稀疏度约 65%（考虑 CSR 索引开销）。
  2. 设置渐进剪枝步长：Δs 建议在 5%~10% 之间。步长太大（>20%）会导致精度骤降；太小则训练时间过长。
  3. 掩码更新周期 T：通常 3~10 个 epoch。T 太小，掩码频繁变化导致训练不稳定；T 太大，剪枝决策滞后。
  4. 学习率联动：每次增加稀疏度时，将学习率乘以 0.7~0.9。稀疏度越高，模型容量越小，需要更小的学习率避免震荡。

最后需要强调的是，稀疏化训练不是孤立的技术，它需要与量化（第 5 章）、蒸馏（第 6 章）结合使用。通常先做稀疏化（减少参数量），再做量化（减少位宽），最后进行联合微调。在 TinyML 工程中，稀疏化带来的收益不仅体现在存储上，更重要的是推理速度的提升——Cortex-M 系列内核的 DSP 指令可以跳过零值，实现 2~4 倍的加速。

经验总结：稀疏化训练的本质是“在训练过程中动态地决定哪些权重重要”。掩码机制是骨架，正则化方法是策略，而学习率调度是润滑剂。三者缺一不可。建议在项目初期就建立完整的稀疏化训练流水线，而不是在训练完成后才做剪枝——后者通常会导致精度大幅下降。
← 上一章 📖 返回目录 下一章 →