Why detach probs in symmetric_kl function (deberta shift)

Sohaib_Ahmed: <p>In the official implementation of sift by Microsoft/Deberta <a href="https://github.com/microsoft/DeBERTa/blob/master/DeBERTa/sift/sift.py" rel="noopener nofollow ugc">here</a>, they implemented symmetric_kl function as:</p>
<pre><code class="lang-auto">def symmetric_kl(logits, target):
  logit_stu = logits.view(-1, logits.size(-1)).float()
  logit_tea = target.view(-1, target.size(-1)).float()
  logprob_stu = F.log_softmax(logit_stu, -1)
  logprob_tea = F.log_softmax(logit_tea, -1)
  prob_tea = logprob_tea.exp().detach()
  prob_stu = logprob_stu.exp().detach()
  floss = ((prob_tea*(-logprob_stu)).sum(-1))    # Cross Entropy
  bloss = ((prob_stu*(-logprob_tea)).sum(-1))    # Cross Entropy
  loss = floss + bloss
  return loss
</code></pre>
<p>My question is why do you need to detach() when computing prob_tea and prob_stu. I checked and found that its presence affects the gradients computed.</p>

KFrank: <p>Hi Sohaib!</p>
<aside class="quote no-group" data-username="Sohaib_Ahmed" data-post="1" data-topic="201944" data-full="true">
<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/sohaib_ahmed/48/62378_2.png" class="avatar"> Sohaib_Ahmed:</div>
<blockquote>
<p>they implemented symmetric_kl function as:</p>
</blockquote>
</aside>
<p>The Kullback-Leibler divergence is a particular measure of the dissimilarity<br>
of two probability distributions.  It is closely related to cross entropy.  It takes<br>
on its minimum of zero when the two probability distributions are the same.</p>
<p>The <code>symmetric_kl()</code> you are asking about is basically a symmetrized version,<br>
<code>kl (p1, p2) + kl (p2, p1)</code>.</p>
<aside class="quote no-group">
<blockquote>
<pre><code class="lang-auto">  floss = ((prob_tea*(-logprob_stu)).sum(-1))    # Cross Entropy
  bloss = ((prob_stu*(-logprob_tea)).sum(-1))    # Cross Entropy
</code></pre>
</blockquote>
</aside>
<p><code>floss</code> and <code>bloss</code> are indeed the two cross entropies and they are added<br>
together to get the symmetrized result.</p>
<aside class="quote no-group">
<blockquote>
<p>My question is why do you need to detach() when computing prob_tea and prob_stu.</p>
</blockquote>
</aside>
<p>Leaving aside the details of working in log space (that is, working with<br>
log-probabilities rather than probabilities), let’s consider the cross entropy<br>
of two probability distributions, <code>p1</code> and <code>p2</code>.</p>
<p>For fixed <code>p2</code>, <code>cross_entropy (p1, p2)</code> takes on its minimum when <code>p1</code><br>
equals <code>p2</code>.  But for fixed <code>p1</code>, it <em>does not</em> take on its minimum when <code>p2</code> is<br>
equal to <code>p1</code>.  (Rather, for fixed <code>p1</code>, it takes on its minimum when <code>p2</code> has<br>
100% probability for its entry that corresponds to the largest probability<br>
in <code>p1</code> and 0% for all of its other entries.)</p>
<p>Consider <code>floss</code>:  The part of its gradient due to <code>logprob_stu</code> drives the<br>
probability distribution given (in log space) by <code>logits</code> to match <code>target</code>.<br>
This is what we want.  But, if we didn’t detach <code>prob_tea</code>, the part of<br>
<code>floss</code>’s gradient due to <code>prob_tea</code> would drive target to concentrate all<br>
of its probability in the single element that corresponded to the element<br>
of <code>logits</code> with the largest probability.  This is not what we want and would<br>
work at cross-purposes to the <code>logprob_stu</code> part of the gradient.  This is<br>
why we detach <code>prob_tea</code>.  This causes <code>prob_tea</code>’s gradient – which we<br>
don’t want – to be thrown away.</p>
<p>The same story applies to <code>bloss</code>, but with <code>logits</code> and <code>target</code> switched.<br>
For <code>bloss</code>, the part of its gradient due to <code>prob_stu</code> is unhelpful, so we<br>
detach <code>prob_stu</code>.</p>
<p>Best.</p>
<p>K. Frank</p>