Jensen Shannon Divergence

username: kko
cooked: <p>Hi,</p>
<p>I am trying to implement Jensen Shannon Divergence (JSD) in Pytorch:<br>
class JSD(torch.nn.Module)</p>
<pre><code>def forward(self, P, Q):
    kld = KLDivLoss().cuda()
    M = 0.5 * (P + Q)
    return 0.5 * (kld(P, M) + kld(Q, M))
</code></pre>
<p>When I run the above code, I am getting the below error:<br>
AssertionError: nn criterions don’t compute the gradient w.r.t. targets - please mark these variables as volatile or not requiring gradients.</p>
<p>I guess KL expect second term to be not requiring gradients. But, in JSD M term also contains gradient requiring variable. Is there an easy way of dealing with this? Or, should I write my KLD function?</p>
<p>Thanks.</p>

username: jekbradbury
cooked: <p>I think the easiest way would be to write your own KLDiv function just like this one.</p>

username: abhidipbhattacharyya
cooked: <p>Hello,</p>
<p>Are you able to implement Jensen Shannon  Divergence? If so kindly share the code. It will be really helpful.</p>
<p>Thanks</p>

username: Aryan_Asadian
cooked: <p>Hope this helps.</p>
<pre><code class="lang-auto">class JSD(nn.Module):
    
    def __init__(self):
        super(JSD, self).__init__()
    
    def forward(self, net_1_logits, net_2_logits):
        net_1_probs =  F.softmax(net_1_logits, dim=1)
        net_2_probs=  F.softmax(net_2_logits, dim=1)

        m = 0.5 * (net_1_probs + net_1_probs)
        loss = 0.0
        loss += F.kl_div(F.log_softmax(net_1_logits, dim=1), total_m, reduction="batchmean") 
        loss += F.kl_div(F.log_softmax(net_2_logits, dim=1), total_m, reduction="batchmean") 
     
        return (0.5 * loss)
</code></pre>

username: san_zhang
cooked: <p>So these two variables， net_1_probs ， m are not used？<br>
And where is the total_m</p>

username: jeff-hykin
cooked: <p>I had to modify the example to this:</p>
<p>Note the function is not designed to handle batches of inputs (matrix arguments), although it might.</p>
<pre><code class="lang-python">def jenson_shannon_divergence(net_1_logits, net_2_logits):
    from torch.functional import F
    net_1_probs = F.softmax(net_1_logits, dim=0)
    net_2_probs = F.softmax(net_2_logits, dim=0)
    
    total_m = 0.5 * (net_1_probs + net_1_probs)
    
    loss = 0.0
    loss += F.kl_div(F.log_softmax(net_1_logits, dim=0), total_m, reduction="batchmean") 
    loss += F.kl_div(F.log_softmax(net_2_logits, dim=0), total_m, reduction="batchmean") 
    return (0.5 * loss)
</code></pre>

username: Amin_Jun
cooked: <p>Hey, Here’s a couple of things I thought worthy of mentioning here. First, both codes are only using:</p>
<pre><code class="lang-auto">total_m = 0.5 * (net_1_probs + net_1_probs)
</code></pre>
<p>The correct formulation is:</p>
<pre><code class="lang-auto">total_m = 0.5 * (net_1_probs + net_2_probs)
</code></pre>
<p>.</p>
<p>Also, based on <a class="mention" href="/u/jeff-hykin">@jeff-hykin</a> and <a class="mention" href="/u/aryan_asadian">@Aryan_Asadian</a> implementations, here’s mine. It is easier for me to use modules that are instances of <code>nn.Module</code> similar to <a class="mention" href="/u/aryan_asadian">@Aryan_Asadian</a>’s implementation, because I can have forward/backward hooks.</p>
<pre><code class="lang-auto">class JSD(nn.Module):
    def __init__(self):
        super(JSD, self).__init__()
        self.kl = nn.KLDivLoss(reduction='batchmean', log_target=True)

    def forward(self, p: torch.tensor, q: torch.tensor):
        p, q = p.view(-1, p.size(-1)), q.view(-1, q.size(-1))
        m = (0.5 * (p + q)).log()
        return 0.5 * (self.kl(p.log(), m) + self.kl(q.log(), m))
</code></pre>
<p>Note that I am taking <code>softmax</code> before passing <code>p</code> and <code>q</code> to my <code>JSD</code> instance. Also, note that this implementation works with matrices as well, since in the beginning I’m flattening both tensors.<br>
Also, note that I’m passing <code>log_target=True</code>, which means the <code>m</code> should be in <code>log-space</code>. This makes the implementation slightly faster, because we’re computing the <code>m.log()</code> only once. Hope that this helps.</p>

username: Renly_Hou
cooked: <p><a class="mention" href="/u/amin_jun">@Amin_Jun</a> you did a wonderful job based on the previous answer. However, your implementation is still slightly problematic, which doesn’t guarantee the range of JS-divergence between 0 to 1. The KL-divergence function in pytorch is counterintuitive. KL(a,b) needs to be written in torch.nn.KLDivLoss()(b,a). So the correct one should be:</p>
<pre><code class="lang-auto">class JSD(nn.Module):
    def __init__(self):
        super(JSD, self).__init__()
        self.kl = nn.KLDivLoss(reduction='batchmean', log_target=True)

    def forward(self, p: torch.tensor, q: torch.tensor):
        p, q = p.view(-1, p.size(-1)), q.view(-1, q.size(-1))
        m = (0.5 * (p + q)).log()
        return 0.5 * (self.kl(m, p.log()) + self.kl(m, q.log()))
</code></pre>

username: Amin_Jun
cooked: <p><a class="mention" href="/u/renly_hou">@Renly_Hou</a> You’re absolutely correct! Thanks for catching this. You saved me from many experiments.</p>

username: alexneakameni
cooked: <p>Is this implementation correct if we assume that my inputs aren’t <code>softmax</code> ?<br>
According to this link : <a href="https://pytorch.org/docs/stable/generated/torch.nn.functional.log_softmax.html#torch.nn.functional.log_softmax" rel="noopener nofollow ugc">torch.log_softmax</a>, it’s recommended to directly use log_softmax instead of log(softmax…)</p>
<pre><code class="lang-auto">class JSD(nn.Module):
    def __init__(self):
        super(JSD, self).__init__()
        self.kl = nn.KLDivLoss(reduction='batchmean', log_target=True)

    def forward(self, p: torch.tensor, q: torch.tensor):
        p, q = p.view(-1, p.size(-1)).log_softmax(-1), q.view(-1, q.size(-1)).log_softmax(-1)
        m = (0.5 * (p + q))
        return 0.5 * (self.kl(m, p) + self.kl(m, q))
</code></pre>

username: NikhilMank
cooked: <p>What are the inputs for the variables p and q? Do I have to send a data point from two datasets as a tensors or a batch of data from two datasets as tensors?</p>