试图理解计算 Torch 中 LogSoftMax 输入梯度的代码

Question

代码来自：https://github.com/torch/nn/blob/master/lib/THNN/generic/LogSoftMax.c

我不明白此代码如何计算梯度 w.r.t 到模块 LogSoftMax 的输入。我感到困惑的是这两个 for 循环在做什么。

for (t = 0; t < nframe; t++)
{
sum = 0;
gradInput_data = gradInput_data0 + dim*t;
output_data = output_data0 + dim*t;
gradOutput_data = gradOutput_data0 + dim*t;

for (d = 0; d < dim; d++)
  sum += gradOutput_data[d];

for (d = 0; d < dim; d++)
  gradInput_data[d] = gradOutput_data[d] - exp(output_data[d])*sum;
 }
}

Answer 1

在前进时我们有（x = 输入向量，y = 输出向量，f = logsoftmax，i = 第 i 个分量）：

yi = f(xi)
   = log( exp(xi) / sum_j(exp(xj)) )
   = xi - log( sum_j(exp(xj)) )

计算 f 的雅可比 Jf 时，您有（第 i 行）：

dyi/dxi = 1 - exp(xi) / sum_j(exp(xj))

对于与 i 不同的 k：

dyi/dxk = - exp(xk) / sum_j(exp(xj))

这为 Jf 提供：

1-E(x1)     -E(x2)     -E(x3)    ...
 -E(x1)    1-E(x2)     -E(x3)    ...
 -E(x1)     -E(x2)    1-E(x3)    ...
...

与E(xi) = exp(xi) / sum_j(exp(xj))

如果我们将 gradInput 命名为梯度 w.r.t 输入，将 gradOutput 命名为梯度 w.r.t 输出，反向传播给出（链式规则）：

gradInputi = sum_j( gradOutputj . dyj/dxi )

这相当于：

gradInput = transpose(Jf) . gradOutput

最后给出第 i 个分量：

gradInputi = gradOutputi - E(xi) . sum_j( gradOutputj )

所以第一个循环预先计算 sum_j( gradOutputj )，最后一个循环预先计算上述术语，即 grad 的第 i 个分量。输入 - 除了 Torch 实现中的指数项缺少 1 / sum_j(exp(xj))（ 上面的微积分可能应该仔细检查，即使它听起来正确并解释了当前的实现 ） .

UPDATE：缺失 1 / sum_j(exp(xj))项没有问题。由于 jacobian 是根据 output 值计算的，并且由于此先前计算的输出恰好是 log-softmax 分布，因此该分布的 sum-exp 为 1:

sum_j(exp(outputj)) = sum_j(exp( log(exp(inputj) / sum_k(exp(inputk) ))
                    = sum_j(         exp(inputj) / sum_k(exp(inputk)  )
                    = 1

所以没有必要在实现中明确这个术语，它给出（对于 x = 输出）：

gradInputi = gradOutputi - exp(outputi) . sum_j( gradOutputj )

试图理解计算 Torch 中 LogSoftMax 输入梯度的代码

Trying to understand code that computes the gradient wrt to the input for LogSoftMax in Torch

mathematical-optimization

torch

gradient-descent

softmax