SNAP-UQ: Self-supervised Next-Activation Prediction for Single-Pass Uncertainty¶

会议: ICLR 2026
arXiv: 2508.12907
代码: 无
领域: audio_speech
关键词: 不确定性估计, TinyML, 单次推理, 自监督, 微控制器部署, OOD 检测

一句话总结¶

SNAP-UQ 提出一种面向 TinyML 场景的单次前向传播不确定性估计方法：在骨干网络的选定层附加微型 int8 预测头，用自监督方式预测下一层的激活统计量，将实际激活与预测之间的偏差（"surprisal"）聚合为不确定性分数，无需额外前向传播、时序缓冲或集成，仅增加几十 KB 闪存即可在微控制器上实现可靠的分布偏移检测和故障检测。

研究背景与动机¶

TinyML 模型越来越多地部署在电池供电的微控制器（MCU）上，用于视觉和音频的私密低延迟感知。但部署后的输入分布不断变化——传感器漂移、光照和声学环境变化、分布内腐蚀（CID）和分布外（OOD）样本交替出现。现代神经网络在这些偏移下经常过度自信，即使在保留集上校准良好也是如此。

在 MCU 上解决不确定性估计面临严峻约束： - MC Dropout 和 Deep Ensembles 需要多次前向传播，延迟和闪存成本成倍增加 - Early-exit ensembles 在推理时仍需额外分类头和内存带宽，且依赖 softmax 信号（在 CID 下脆弱） - 事后校准（Temperature Scaling）在分布偏移下通常失效 - 经典 OOD 检测器（ODIN/G-ODIN）在超紧凑模型上迁移性差

关键洞察：层间动态比 softmax 置信度更早反映分布偏移——特征在类别后验扁平化之前就会相对于网络自身的变换变得非典型。

方法详解¶

整体框架¶

SNAP-UQ 在骨干网络的 2-3 个选定"接驳层"（taps）\(\mathcal{S}\) 处附加微型预测头。每个头从上一层激活的低秩投影预测下一层激活的条件高斯参数（均值 \(\mu_\ell\) 和对角方差 \(\sigma_\ell^2\)），计算标准化误差作为 surprisal 分数，加权聚合后经轻量映射得到最终不确定性估计。

关键设计¶

深度方向的下一层激活预测模型：
- 在每个接驳层 \(\ell \in \mathcal{S}\)，先通过投影器 \(P_\ell\) 压缩上一层激活：\(z_\ell = P_\ell a_{\ell-1} \in \mathbb{R}^{r_\ell}\)（\(r_\ell \ll d_{\ell-1}\)）
- 预测头 \(g_\ell\) 输出对角高斯参数：\((\mu_\ell, \log \sigma_\ell^2) = g_\ell(z_\ell)\)
- 可选低秩+对角协方差：\(\Sigma_\ell = \text{diag}(\sigma_\ell^2) + B_\ell B_\ell^\top\)，通过 Woodbury 恒等式高效计算
- 设计动机：建模层间条件关系 \(a_{\ell-1} \mapsto a_\ell\)，而非传统的无条件类别统计量
自监督训练目标：
- 辅助损失基于对角高斯 NLL：\(\mathcal{L}_{SS} = \frac{1}{|\mathcal{B}|}\sum_{x \in \mathcal{B}} \sum_{\ell \in \mathcal{S}} \frac{1}{2}[\|(a_\ell - \mu_\ell) \odot \sigma_\ell^{-1}\|^2 + \mathbf{1}^\top \log \sigma_\ell^2]\)
- 总损失：\(\mathcal{L} = \mathcal{L}_{clf} + \lambda_{SS}\mathcal{L}_{SS} + \lambda_{reg}\mathcal{R}\)，\(\lambda_{SS}\) 很小（\(10^{-3} \sim 10^{-2}\)）
- 正则化：方差下界（softplus + \(\epsilon^2\)）防止崩塌、尺度控制 \(\mathcal{R}_{var} = \sum_\ell \|\log \sigma_\ell^2\|_1\) 防止过度分散
- 可选 detach 模式：对小骨干网络 stop-grad \(a_\ell\) 避免梯度拔河
单次 surprisal 聚合与映射：
- 标准化误差：\(\bar{e}_\ell(x) = \frac{1}{d_\ell}\|(a_\ell - \mu_\ell) \odot \sigma_\ell^{-1}\|^2\)
- SNAP 分数：\(S(x) = \sum_{\ell \in \mathcal{S}} w_\ell \bar{e}_\ell(x)\)
- 经 logistic 映射融合可选的置信度代理得到最终不确定性：\(U(x) = \sigma(\beta_0 + \beta_1 S(x) + \beta_2 m(x))\)
- 映射参数一次性离线拟合，无需在线标签
MCU 友好的整数实现：
- \(P_\ell\)、\(W_\mu\)、\(W_\sigma\) 量化为 int8
- 用 256 项查找表（LUT）替代 \(\exp(-\frac{1}{2}\log\sigma^2)\) 的指数运算
- 2 个接驳点、\(r_\ell \in [32,128]\) 时，额外计算量 < 骨干的 2%，闪存仅增加几十 KB

损失函数 / 训练策略¶

使用维度归一化避免大维度层主导损失
层权重可设为均匀或反方差加权（\(w_\ell \propto 1/\hat{\text{Var}}[\bar{e}_\ell]\)）
QAT（量化感知训练）在最后 20% epoch 插入伪量化
Student-\(t\) 和 Huberized 变体作为鲁棒性替代方案

实验关键数据¶

主实验——MCU 部署性能¶

平台/任务	方法	Flash (KB)	Peak RAM (KB)	延迟 (ms)	能耗 (mJ)
Big-MCU/SpeechCmd	BASE	220	84	60	2.1
	EE-ens	360	132	85	3.0
	DEEP	290	108	70	2.5
	SNAP-UQ	182	70	52	1.7
Small-MCU/CIFAR-10	BASE	180	92	260	9.5
	EE-ens	OOM	—	—	—
	DEEP	OOM	—	—	—
	SNAP-UQ	158	85	178	6.4

故障检测¶

方法	MNIST ID✓-ID×	SpeechCmd ID✓-ID×	CIFAR-10 ID✓-OOD	SpeechCmd ID✓-OOD
BASE	0.75	0.90	0.90	0.88
EE-ens	0.85	0.90	0.90	0.90
DEEP	0.85	0.91	0.92	0.92
SNAP-UQ	0.90	0.94	0.92	0.94

CID 流监控¶

方法	MNIST-C AUPRC	延迟(帧)	SpeechCmd-C AUPRC	延迟(帧)
BASE	0.54	42	0.52	67
EE-ens	0.63	31	0.59	55
SNAP-UQ	0.66	24	0.65	41

消融实验¶

配置	AUPRC (CIFAR-10-C)	延迟 (ms)
P only, r=32	0.62	88
M+P, r=64	0.70	83
M+P, r=128	0.72	86
M+P+early, r=64	0.71	90

关键发现¶

SNAP-UQ 在 Small-MCU 上是唯一可用的 UQ 方案：EE-ens 和 DEEP 在 CIFAR-10/Small-MCU 都 OOM，SNAP-UQ 正常部署
深度 surprisal 比 softmax 信号更早反应偏移：在流监测实验中，SNAP-UQ 的检测延迟比 EE-ens 短约 25-30%
INT8 量化几乎无损：INT8 头部的 AUPRC 仅比 FP32 下降 0.01，但闪存减少 1.6-2.1 倍
两个接驳点最优：mid+penultimate 组合一致提供最佳精度-延迟权衡，加入 early tap 反而因噪声降低效果

亮点与洞察¶

核心创新在于将"层间动态偏离度"作为不确定性信号——不同于传统基于 softmax/能量/Mahalanobis 的方法，SNAP-UQ 捕捉的是条件性、深度方向的信号
理论分析清晰优雅：Proposition 2.1 证明 SNAP 分数等价于深度方向负对数似然的仿射变换；Proposition 2.2 证明其等价于到条件均值的 Mahalanobis 距离；Proposition 2.3 证明对 BN 尺度变换的不变性
整个设计高度面向工程实际：int8 量化、LUT 避免指数运算、CMSIS-NN 兼容、无时序缓冲
与 ASH、ReAct 等激活整形方法的 head-to-head 对比（Appendix O）显示 SNAP-UQ 在 risk-at-coverage 和 AURC 上全面领先

局限与展望¶

部分固件会融合或省略中间激活，暴露接驳点可能需要运行时修改
对角协方差无法完整捕捉跨通道结构，极端扭曲下可能低估/高估 surprisal
性能对接驳层位置和投影器秩的选择敏感
可选的置信度混合和映射仍需小量标注开发集
仅在四个 benchmark 和两个 MCU 层级上评估，未覆盖更多模态和 tiny transformer 架构

评分¶

新颖性: ⭐⭐⭐⭐⭐
实验充分度: ⭐⭐⭐⭐⭐
写作质量: ⭐⭐⭐⭐⭐
价值: ⭐⭐⭐⭐⭐