How to calculate log_softmax for list of tensors without breaking autograd

username: Aeryan
cooked: <p>I’m trying to calculate the <em>log_softmax</em> function of a list of tensors, i.e., a list [t_1, t_2, …, t_n] where each <em>t_i</em> is of type <em>torch.tensor</em> and each <em>t_i</em> can be of a different, arbitrary shape. <strong>I do not want to apply the log_softmax function to each <em>t_i</em> separately, but to all of them as if they were part of the same unique tensor.</strong> The output of this function should be a list of tensors with the same shape as the input. Lastly, as I will apply this function to the end layer of a neural network, <strong>I want to be able to differentiate this function</strong>, i.e., the gradients must flow through it.</p>
<p>Pytorch provides the class <a href="https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html" rel="noopener nofollow ugc">torch.nn.LogSoftmax</a>, but I cannot use it as it expects a single tensor as input, instead of a list of tensors. Additionally, I want to calculate the log_softmax function efficiently and in a stable way. To achieve that, I want to use the <a href="https://blog.feedly.com/tricks-of-the-trade-logsumexp/" rel="noopener nofollow ugc"><em>log-sum-exp</em> trick</a>. Lastly, I want to ignore the last value of the first element of the list (see code snippet below), i.e., not apply log_softmax to it.</p>
<p>This is my current implementation:</p>
<pre><code class="lang-auto">def log_softmax(pred_tensors):
    minus_inf = -1000 # Constant that represents minus infinity

    # Calculate the max value
    c = max([preds.amax() if preds is not None else minus_inf for preds in pred_tensors])

    # Calculate log(sum(e^(x_i-c)))
    log_sum_exp = 0
    for r in range(len(pred_tensors)):
        if pred_tensors[r] is not None:
            
            # Arity 0 -&gt; ignore nullary predicate corresponding to termination condition
            curr_sum =  torch.sum(torch.exp(pred_tensors[r][:-1] - c))   if r == 0 else \
                        torch.sum(torch.exp(pred_tensors[r] - c))
            log_sum_exp += curr_sum
            
    log_sum_exp = torch.log(log_sum_exp)
        
    # Calculate log_softmax (apply log_softmax to the original tensor) (except to the termination condition)
    for r in range(len(pred_tensors)):
        if pred_tensors[r] is not None:
            # Arity 0 -&gt; ignore nullary predicate corresponding to termination condition
            if r == 0:
                pred_tensors[r][:-1] -= log_sum_exp + c 
            else:    
                pred_tensors[r] -= log_sum_exp + c


    return pred_tensors
</code></pre>
<p>I have tested it and it works. However, I think my implementation <strong>may be breaking the autograd</strong> of Pytorch, in lines <code>c = max([preds.amax() if preds is not None else minus_inf for preds in pred_tensors])</code> and <code>log_sum_exp += curr_sum</code>.</p>
<p>So, my questions are: <strong>Is my implementation really breaking autograd? If it is, can you provide an alternative implementation that works with autograd?</strong></p>

username: InnovArul
cooked: <aside class="quote no-group" data-username="Aeryan" data-post="1" data-topic="151247">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>log_sum_exp += curr_sum</p>
</blockquote>
</aside>
<p>This looks ok to me.</p>
<aside class="quote no-group" data-username="Aeryan" data-post="1" data-topic="151247">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>c = max([preds.amax() if preds is not None else minus_inf for preds in pred_tensors])</p>
</blockquote>
</aside>
<p>I think that this code should be under <code>with torch.no_grad()</code>, as <code>c</code> is supposed to be a constant without any effect on backpropagation.</p>
<p>With the below code, you can calculate <code>logsumexp</code> on multiple tensors. But I am not sure if it will help your use case.</p>
<pre><code class="lang-auto">x = torch.randn(5,3)
y = torch.randn(5,6)
z = torch.randn(5,9)
composed = torch.cat([x, y, z], dim=-1)
logsumexp = torch.logsumexp(composed, dim=-1, keepdim=True)
</code></pre>

username: Aeryan
cooked: <p>Thank you very much for your answer!</p>
<p>So, if I have understood you correctly, you say that the operation <code>log_sum_exp += curr_sum</code> would correctly calculate the gradients, i.e., not break autograd. I suppose that’s because, even though I am using the inplace operation <em>+=</em> instead of something like <em>torch.sum</em>, <em>torch.sum</em> is called under the hood and, thus, the gradient can be calculated for the sum by using <em>+=</em>. Is that correct?</p>
<p>About the constant <em>c</em>, you are right. I thought I had to calculate the gradient for it as it uses the tensor values <code>preds in pred_tensors</code> but, since it is treated as a constant (i.e., I could have chosen any other value for <em>c</em> such as 0), I suppose I can ignore the gradients for that part.</p>
<p>Finally, about your <code>logsumexp</code> implementation with <code>torch.cat()</code>. It does not work for me as the tensors can have <em>arbitrary</em> shapes, and <code>torch.cat()</code> needs for all the dimensions to be the same except one. For example, imagine I want to calculate <code>logsumexp</code> for the tensors <code>x = torch.randn(5,3)</code> and <code>y = torch.randn(3, 7)</code>. You can’t do <code>composed = torch.cat([x,y])</code> in this case, that’s why I am stuck with my manual implementation.</p>

username: InnovArul
cooked: <aside class="quote no-group" data-username="Aeryan" data-post="3" data-topic="151247">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>even though I am using the inplace operation <em>+=</em> instead of something like <em>torch.sum</em>, <em>torch.sum</em> is called under the hood</p>
</blockquote>
</aside>
<p>Yes, you are right.</p>
<aside class="quote no-group" data-username="Aeryan" data-post="3" data-topic="151247">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>I want to calculate <code>logsumexp</code> for the tensors <code>x = torch.randn(5,3)</code> and <code>y = torch.randn(3, 7)</code></p>
</blockquote>
</aside>
<p>Can you elaborate a bit more on the context? Why are these tensors of different size?<br>
In your code, it looks like you are considering all the elements of the (possibly) multi-dimensional tensors, as if they are vectors.<br>
So, the below code might work?</p>
<pre><code class="lang-auto">x = torch.randn(5,3)
y = torch.randn(5,6)
z = torch.randn(5,9)
composed = torch.cat([x.flatten(), y.flatten(), z.flatten()], dim=-1)
logsumexp = torch.logsumexp(composed, dim=-1, keepdim=True)
</code></pre>

username: Aeryan
cooked: <p>So, about the context. The reason why I have a list with tensors of different sizes is because they correspond to the outputs of a <a href="https://arxiv.org/abs/1904.11694" rel="noopener nofollow ugc">Neural Logic Machine</a>. I want to sample one of those (grounded) predicates, so all the values must be positive and sum to 1. That’s why I am using logsoftmax.</p>
<p>Regarding the code, before your answer I did not know Pytorch provided the function <code>logsumexp</code>. So, after using such function and adapting my code with yours, this is the result:</p>
<pre><code class="lang-auto"> def _log_softmax(self, pred_tensors):
        # Remove the nullary predicate associated with the termination condition, so that it does not
        # affect the log_softmax computation
        term_cond_value = pred_tensors[0][-1]
        pred_tensors[0] = pred_tensors[0][:-1]
        
        # Calculate log_sum_exp of all the values in the tensors of the list
        # 1) flatten each tensor in the list
        # 2) concatenate them as a unique tensor
        # 3) calculate log_sum_exp
        log_sum_exp = torch.logsumexp(torch.cat([preds.flatten() if preds is not None else torch.empty(0, dtype=torch.float32) for preds in pred_tensors]), dim=-1)
    
        # Use log_sum_exp to calculate the log_softmax of the tensors in the list
        for r in range(len(pred_tensors)):
            if pred_tensors[r] is not None:
                pred_tensors[r] -= log_sum_exp
    
        # Append the nullary predicate corresponding to the termination condition
        pred_tensors[0] = torch.cat([pred_tensors[0], term_cond_value.reshape(1)]) # We need reshape() to transform from tensor of dimension 0 to dimension 1
        
        return pred_tensors
</code></pre>
<p>As you can observe, I use the function <code>logsumexp</code> to easily calculate the logsoftmax of my list of tensors. Since I do not want to change the value <code>pred_tensors[0][-1]</code>, I simply remove it from the first tensor of the list, perform the calculations and append it after the calculations. Additionally, I ignore those elements of the list corresponding to None (instead of a valid tensor), in line <code>[preds.flatten() if preds is not None else torch.empty(0, dtype=torch.float32) for preds in pred_tensors]</code>.</p>
<p>I have tested this code and is able to correctly calculate the logsoftmax function. Do you think my implementation works with autograd or is there any part of it which could potentially break the gradients?</p>

username: InnovArul
cooked: <aside class="quote no-group" data-username="Aeryan" data-post="5" data-topic="151247">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/aeryan/48/46900_2.png" class="avatar"> Aeryan:</div>
<blockquote>
<p>Do you think my implementation works with autograd or is there any part of it which could potentially break the gradients?</p>
</blockquote>
</aside>
<p>It looks fine to me.</p>

username: Aeryan
cooked: <p>Thank you very much for your help! I mark the issue as solved, then.</p>