Jax训练过程中随机出现NaN值的调试方法求助

先排查四元数姿态估计场景特有数值风险 #

这类场景是NaN高发区，优先检查以下高频问题：

• 所有四元数归一化操作，必须在分母加极小epsilon避免除零，正确写法参考：q / (jnp.linalg.norm(q) + 1e-12)，训练过程中梯度偏大很容易把四元数参数推为接近零的向量，无epsilon的归一化会直接产生inf，后续运算转NaN。
• 所有反三角函数（acos/asin）的输入必须做截断，约束到[-1 + 1e-7, 1 - 1e-7]区间，四元数求角度差时的点积结果很容易因为浮点误差超出±1范围，直接输入反三角函数会立刻出NaN。
• 损失函数中如果有对数、开方运算，必须对输入做下限截断，避免输入为0或负数产生非法值。

通用梯度NaN调试方法 #

你开启的jax_debug_nans输出栈信息冗余，推荐用以下更高效的调试手段：

• 拆分计算图定位问题节点：用jax.checkpoint把损失函数拆分为前向推理、损失计算两个独立模块，分别打印两个模块的输出和对应梯度，先定位NaN是在前向计算过程产生，还是反向求导阶段产生。
• 临时切换精度验证：将全局运算精度临时切换为float64跑训练，如果NaN出现的步数延后或者完全消失，说明是float32精度不足导致的数值溢出，后续只需在四元数运算、损失计算等关键节点临时升精度到float64，计算完成再转回float32即可。
• 实时打印关键节点数值：用jax.debug.print在你怀疑的高风险节点打印数值（比如四元数的模长、反三角函数的输入值、损失值），训练过程中会实时输出，第一次出现NaN对应的节点就是问题根源，可读性远高于默认的NaN报错栈。

快速验证方案 #

从你给出的训练日志看，前6步损失正常下降，第7步突然出NaN，大概率是梯度过大导致的参数爆炸，可以先做以下快速验证：

• 先加全局梯度裁剪：grads = jax.tree_map(lambda g: jnp.clip(g, -10, 10), grads)，如果加裁剪后不再出NaN，说明是梯度过大问题，后续下调学习率或者更换自适应优化器（比如AdamW）即可。
• 直接把学习率降到当前的1/10再跑训练，很多时候梯度发散都是初始学习率设置过高导致的。

内容的提问来源于stack exchange，提问作者Simon B