从零实现二分类交叉熵时JAX训练神经网络结果与Keras不一致问题

1. 输入未归一化
   你使用的MNIST像素值范围为0~255，未做归一化直接输入网络会导致Dense层输出值过大，经过ReLU后容易让后续Sigmoid层进入饱和区，梯度变得极小，拖慢训练速度。

2. 自定义交叉熵损失存在不必要的数值误差
   你实现的稳定版损失，是先从Sigmoid输出的概率y_hat逆推回logits，这一步额外引入了clip、log、除法三类非线性操作，不仅会引入计算误差，还会在y_hat接近0/1时，因为clip操作截断梯度，导致训练速度远慢于Keras官方实现的损失函数。
   正确的稳定二分类交叉熵实现应该直接对接全连接层的原始输出（去掉网络末尾的Sigmoid层），在logits层面计算损失，避免多轮转换的误差：

def binary_cross_entropy_from_logits(logits, y):
    return jnp.mean(jnp.maximum(logits, 0) - logits * y + jnp.log(1 + jnp.exp(-jnp.abs(logits))))

3. 可验证的修复步骤：

• 在加载训练数据后添加归一化操作：train_xs = train_xs.astype(jnp.float32) / 255.0
• 替换自定义损失为上述logits层面的实现，同时移除JAX网络最后的stax.Sigmoid层，训练完成后做推理时再手动加sigmoid计算概率和精度
• 以上两步修改后，JAX的训练精度会和Keras完全对齐

内容的提问来源于stack exchange，提问作者Jack M