陷阱四：Transformer中的Softmax数值溢出——嵌入式下的“温柔一刀”

说实话，Softmax这个函数，在PC上跑Transformer的时候，我从来没正眼看过它。不就是个归一化嘛，算就完了。但当我第一次把BEV模型里的Transformer搬到嵌入式设备上时，这玩意儿给了我一个结结实实的下马威。

模型推理结果突然全乱了。不是精度下降，是直接输出NaN。我当时第一反应是“量化参数设错了？”查了半天，最后定位到Softmax。嗯，这里要注意，Softmax在嵌入式上，尤其是定点或低精度浮点环境下，是个典型的“温柔一刀”——平时没事，关键时刻捅你一刀。

问题根源：指数函数的“爆炸”天性

Softmax的公式大家应该都熟：

softmax(x_i) = exp(x_i) / sum(exp(x_j))

关键就在这个 exp(x_i)。指数函数增长极快。当输入值稍微大一点，比如 x_i = 89，exp(89) 就已经是 1.6e38 了。这在FP32里还能勉强撑住，但到了FP16或者INT8的世界里，直接就溢出了。

为什么会这样？因为嵌入式设备为了省电和省带宽，普遍用FP16甚至INT8推理。FP16的最大值只有65504。你想想看，exp(12) 就已经是16万了，远超FP16的表示范围。结果就是——溢出，变成Inf，然后整个Softmax输出全是NaN。

我在项目中遇到过这种情况：模型在PC上跑FP32精度，AP指标有45%。交叉编译到嵌入式板子上，用FP16推理，AP直接掉到12%。一开始我以为是量化精度损失，后来发现是Softmax输出全乱了，导致后续的注意力权重全部失效。

解法一：数值稳定的Softmax——减最大值

这个技巧在数学课本上就有，但真正在嵌入式里用起来，感受完全不同。标准做法是：

def stable_softmax(x):
    x_max = max(x)
    exp_x = [exp(i - x_max) for i in x]
    sum_exp = sum(exp_x)
    return [e / sum_exp for e in exp_x]

减去最大值后，所有指数函数的输入都 ≤ 0，exp(0) = 1，exp(负数) 快速衰减到0。这样就不会溢出了。

但这里有个坑——我一开始以为这样就万事大吉了。直到我在一个BEV模型的Cross-Attention里，发现输入到Softmax的向量里，最大值和最小值差距极大。减完最大值后，有些值变成了 exp(-128) 这种极小的数。在FP16下，这个数直接下溢成0了。

警告： 减最大值能解决上溢，但解决不了下溢。当输入动态范围过大时，部分值会下溢为0，导致注意力权重“死掉”——某些位置永远得不到关注。

解法二：Log-Softmax与近似计算

对于下溢问题，我个人的习惯是改用Log-Softmax。尤其是在计算注意力权重时，很多时候我们并不需要精确的概率值，而是需要相对大小。

def log_softmax(x):
    x_max = max(x)
    return [i - x_max - log(sum(exp(j - x_max) for j in x)) for i in x]

Log-Softmax的输出范围是 (-inf, 0]，不会溢出。而且后续如果要做交叉熵损失，Log-Softmax天然就是数值稳定的。

但要注意，在嵌入式上实现 log 和 exp 函数本身也有开销。我建议用查表法或者多项式近似。比如用 exp(x) ≈ 1 + x + x^2/2 这种泰勒展开，在输入范围 [-5, 0] 内精度足够了。

技巧： 我曾经在ARM Cortex-M7上做过测试，用查表法实现exp，256个表项，精度误差控制在0.1%以内，速度比标准库函数快了8倍。对于BEV这种实时性要求高的场景，值得一试。

解法三：混合精度——该省省，该花花

还有一种思路，我最近在几个量产项目里用的比较多——混合精度。具体来说：

  • Softmax的输入和输出用FP16存储（省带宽）
  • 但计算 exp 和 sum 时，临时提升到FP32
  • 计算完再转回FP16输出

这样既避免了溢出，又没增加太多存储开销。因为Softmax的计算量在整个Transformer里占比很小，临时用FP32算一下，对整体功耗影响微乎其微。

我建议你在部署框架里这样配置：

// 伪代码示例
tensor input_fp16 = load_from_previous_layer();
tensor input_fp32 = cast_to_fp32(input_fp16);
tensor output_fp32 = softmax(input_fp32);  // 在FP32下计算
tensor output_fp16 = cast_to_fp16(output_fp32);
store_to_next_layer(output_fp16);

这个方案在NVIDIA Jetson和Qualcomm SNPE上都验证过，效果很好。

解法四：裁剪——简单粗暴但有效

如果上面的方法都嫌麻烦，还有一个更直接的——对Softmax的输入做裁剪。说白了，就是把输入值限制在一个安全范围内。

对于FP16，安全范围大约是 [-11, 11]。因为 exp(11) ≈ 59874，离FP16上限65504还有一点余量。对于INT8量化后的Softmax，范围更窄，一般需要裁剪到 [-5, 5] 左右。

精度类型 安全输入范围 推荐裁剪阈值
FP32 [-88, 88] [-80, 80]
FP16 [-11, 11] [-10, 10]
INT8 (对称量化) [-5, 5] [-5, 5]

裁剪会损失一点精度，但换来的是绝对稳定。我在一个ADAS项目中，BEV模型的注意力头有8个，其中2个头的输入值经常跑到15以上。裁剪到10之后，AP指标只掉了0.3%，但再也没出现过NaN。

核心要点： 嵌入式部署Transformer，Softmax的数值稳定性是第一道坎。减最大值是基础操作，混合精度是工程最优解，裁剪是兜底方案。三者结合使用，基本能覆盖所有场景。

避坑指南

我曾经踩过一个坑，分享出来给大家提个醒：

当时我在调试一个BEV模型的时序，发现Softmax计算占了整个Attention层40%的时间。我以为是exp函数太慢，优化了半天查表法。后来才发现，是因为输入值范围太大，减最大值后，部分值下溢为0，导致sum_exp里有很多0，除法的时候出现了除零异常——框架自动加了保护逻辑，反而更慢了。

所以，数值溢出不只是精度问题，还会引发性能问题。你想想看，一个除零异常，框架要捕获、要处理、要回退，这些额外开销在嵌入式上都是实打实的延迟。

我的建议是：在模型部署前，先跑一遍所有中间层的数值分布统计。看看Softmax的输入到底在什么范围。如果最大值超过安全阈值，就提前做裁剪或混合精度。别等到跑起来出问题了再回头查，那太被动了。

好了，这一章就到这里。下一章我们聊聊BEV模型里另一个让人头疼的问题——多模态特征对齐中的“维度灾难”。

← 上一章 📖 返回目录 下一章 →