1
$\begingroup$

I am trying to find the derivative of the log softmax function :

$LS(z) = log(\frac {e^{z-c}} {\sum_{i=0}^n e^{z_i-c}}) = z - c - log({\sum_{i=0}^n e^{z_i-c}})$

(c = max(z) )

with respect to the input vector z. However it seems I have made a mistake somewhere. Here is what I have attempted out so far:

The partial derivative of logsoftmax with respect to an element of z

${\partial {LS}\over\partial{z_i}}= 1 - {\partial{ log({\sum_{i=0}^n e^{z_i-c}})} \over \partial {z_i} }$

$ = 1 - {1\over{{\sum_{i=0}^n e^{z_i-c}}}} * {\partial{ {\sum_{i=0}^n e^{z_i-c}}}\over \partial z_i} $

$ = 1 - {e^{z-c}\over {\sum_{i=0}^n e^{z_i-c}}} $

$ = 1 - Softmax(z)_i$

Now, Since I am trying to use this for autodifferentiation, I can then substitute the Softmax with logsoftmax for efficiency since I already compute it in the forward pass:

$ = 1 - e^{LS(z)_i}$

So in vectorized form, we then get:

$ {\partial {LS}\over\partial{z}} = \vec 1 - \vec {e^{LS(z)}}$

In numpy code:

def logsoftmax(z):
    c = np.max(z, axis=1)
    lexpsum = c + np.log(np.exp(z - c.reshape((-1,1))).sum(axis=1))
    out = z - lexpsum.reshape((-1,1))

def logsoftmax_backward(out, out.grad):
    # out is the logsoftmax we computed and out.grad is the incoming gradient which will be 
    # multiplied to this gradient as per the chain rule
    return out.grad * (1 - np.exp(out))

# z = np.array([[1., 2., 3.]])
# My output:
# logsoftmax: [[-2.4076061 , -1.4076061 , -0.40760612]]
# logsoftmax_backward: [[0.90996945, 0.75527155, 0.33475912]]
# Gradient using pytorch: [[ 0.7299,  0.2658, -0.9957]]

It seems I am missing a normalizing constant somewhere, or I have done something completely wrong, but I can't figure it out.

$\endgroup$

1 Answer 1

1
$\begingroup$

$ \def\K{{\large\ell}} \def\o{{\tt1}}\def\p{\partial} \def\L{\left}\def\R{\right} \def\LR#1{\L(#1\R)} \def\BR#1{\Big(#1\Big)} \def\vec#1{\operatorname{vec}\LR{#1}} \def\diag#1{\operatorname{diag}\LR{#1}} \def\Diag#1{\operatorname{Diag}\LR{#1}} \def\trace#1{\operatorname{Tr}\LR{#1}} \def\qiq{\quad\implies\quad} \def\grad#1#2{\frac{\p #1}{\p #2}} \def\c#1{\color{red}{#1}} $Denote the all-ones vector by $\o$ and define an offset vector $$y\;=\;z-c \quad\;\qiq dy = dz$$ Then calculate the gradient of the element-wise exponential function $$\eqalign{ p &= \exp(y) &\qiq P = \Diag{p} \\ dp &= P\,dy = P\,dz \\ \grad{p}{z} &= P \\ }$$ Next, calculate the gradient of the element-wise softmax function $$\eqalign{ s &= \frac{p}{\o^Tp} &\qiq S = \Diag{s} = \frac {P}{\o^Tp} \\ ds &= \frac{dp}{{\o^Tp}} - \frac{p\LR{\o^Tdp}}{\LR{\o^Tp}^2} \\ &= \frac{P\,dz}{{\o^Tp}} - \frac{s\LR{\o^TP\,dz}}{{\o^Tp}} \\ &= \BR{S-ss^T}\,dz \\ \grad{s}{z} &= \BR{S-ss^T} \\ }$$ Finally, tackle the element-wise logarithm of the softmax function $$\eqalign{ \K &= \log(s) \\ d\K &= S^{-1}\,ds \\ &= S^{-1}\BR{S-ss^T}\,dz \\ &= \BR{I-\o s^T}\,dz \\ \grad{\K}{z} &= \BR{I-\o s^T} \\ }$$

$\endgroup$

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .