Optax 的 scale_by_lbfgs 在计算权重倒数时，因浮点精度问题对极小非零值取倒数易导致 NaN；可通过双重 jnp.where 掩码策略规避无效除法，确保梯度计算数值稳定。
optax 的 `scale_by_lbfgs` 在计算权重倒数时，因浮点精度问题对极小非零值取倒数易导致 nan；可通过双重 `jnp.where` 掩码策略规避无效除法，确保梯度计算数值稳定。
Optax 中的 scale_by_lbfgs 是实现 L-BFGS 优化的关键变换器，其核心逻辑之一是根据参数更新与梯度差分的内积 vdot_diff_params_updates 计算缩放权重：weight = 1.0 / vdot_diff_params_updates。然而，当该内积值极小（如 1e-38）但非零时，直接取倒数会溢出为 inf，后续参与梯度传播（尤其在 JAX 的 jit 或高阶微分中）极易触发 NaN——这并非数学错误，而是 JAX 自动微分对未定义操作（如 0/0 或 inf * 0）的敏感性所致（参见 JAX FAQ）。
标准修复方案是采用双重掩码保护（double-where pattern）：首先用布尔掩码识别安全可除区域，再确保分母在所有分支下均有定义且非零。以下是推荐的鲁棒实现方式：
def scale_by_lbfgs(
# ... 其他参数保持不变
):
def update_fn(
updates, state, params=None
):
# ... 前置计算得到 vdot_diff_params_updates
# 假设 vdot_diff_params_updates 已定义
# ✅ 安全倒数计算：双重 where 避免任何路径出现 0 除
is_zero = vdot_diff_params_updates == 0.0
# 分母：若原值为 0，则替换为 1.0（不影响 weight=0 的结果）；否则保留原值
safe_denom = jnp.where(is_zero, 1.0, vdot_diff_params_updates)
# 权重：仅在非零处计算倒数，其余置 0
weight = jnp.where(is_zero, 0.0, 1.0 / safe_denom)
# 后续使用 weight 进行更新...
return updates, state
return update_fn
⚠️ 关键注意事项：
- 不要简单使用
jnp.clip(vdot_diff_params_updates, a_min=1e-12, a_max=None)，因为裁剪会破坏 L-BFGS 的曲率估计一致性，影响收敛性； - 也不要依赖
jnp.nan_to_num(..., nan=0.0)，它无法阻止反向传播中NaN的生成（仅处理前向输出）； - 此修复完全兼容 JAX 的
jit、grad和value_and_grad，无需修改状态结构或额外梯度裁剪； - 若问题仍存在，建议检查上游梯度是否已含
NaN（可用jnp.isfinite(updates).all()断言），因为 LBFGS 本身不引入新奇点，通常根源在模型或损失函数的数值不稳定。
综上，通过双重 jnp.where 实现条件化安全除法，是在不牺牲算法理论性质的前提下，解决 Optax LBFGS 数值脆弱性的简洁、高效且符合 JAX 编程范式的实践方案。
