hw2 LogSumExp梯度公式推导

[LittleHeroZZZX]

您好，请问下面这段hw2中计算LogSumExp梯度的代码不需要计算max(z)？ https://github.com/kcxain/dlsys/blob/0ce32e1322a268ab0d6b64156dd1e62b0bb834c3/hw2/python/needle/ops.py#L395-L411

在这段代码中，计算梯度时似乎并没有使用 Z_max。我推导出的梯度公式是：

exp(z - max(z)) / sum(exp(z - max(z)))

即使用节点的输出 f 来简化计算，得到的表达式依然是 f - max(z)。我对这个表达式的含义有些困惑，能否请您解释一下其中原因？

[yofufufufu]

为了避免SoftmaxLoss计算过程中的exp计算出现溢出，提供数值稳定的实现，进行如下处理：

\begin{aligned}
\ell_\text{softmax}(z,y) &= \log \sum_{i=1}^k \exp z_i - z_y \\ 
&= \log \sum_{i=1}^k \exp (z_i-\max z +\max{z}) - z_y \\ 
&= \log \sum_{i=1}^k {(\exp (z_i-\max z)}.\exp {\max z)} - z_y \\
&= \log {(\sum_{i=1}^k \exp {(z_i-\max z)}.\exp {\max z})} - z_y \\
&= \log \sum_{i=1}^k \exp {(z_i-\max z)} + \log \exp \max z - z_y \\
&= \log \sum_{i=1}^k \exp {(z_i-\max z)} + \max z - z_y
\end{aligned}

其中 $zi-\max{z}$ 一定小于等于0，因此是数值稳定的。这也就是为什么homework2中要求先实现一个LogSumExp算子的原因。 显然，该算子的反向传播过程，与 $$\log \sum{i=1}^k \exp z_i$$ 是一致的。我们就按照后者的来求，因此：

\begin{aligned}
gradient_{z_i} &= (\sum_{i=1}^k \exp z_i)^{-1} . \exp{z_i} \\
&= \exp(\log \sum_{i=1}^k \exp z_i)^{-1} . \exp{z_i} \\ 
&= \exp(z_i - \log \sum_{i=1}^k \exp z_i)
\end{aligned}

所以，该算子的反向传播过程，adjoints的计算只要用到正向传播的输入和输出即可（莫名的简洁😱）

[LittleHeroZZZX]

非常感谢您一针见血的解释，我理解这部分内容啦🙏