Code debugging: How to implement Generalized Dirichlet distributions KL-Divergence in python?

hormone03: <p>I am working on this paper <a href="https://www.sciencedirect.com/science/article/abs/pii/S003132031830311X" rel="noopener nofollow ugc">link to paper</a>. I want to implement the KL-Divergence between two generalized Dirichlet distributions in equation 11 of this paper, see screenshot below:</p>
<p><a href="https://i.stack.imgur.com/ocWUw.png" rel="noopener nofollow ugc"><img src="https://i.stack.imgur.com/ocWUw.png" alt="enter image description here" width="690" height="468"></a><br>
<a href="https://i.stack.imgur.com/sUOFw.png" rel="noopener nofollow ugc"><img src="https://i.stack.imgur.com/sUOFw.png" alt="enter image description here" width="690" height="96"></a></p>
<p>Alpha_1 and Beta_1 are estimated from a decoder network, for example:</p>
<pre><code>decoder_alpha = torch.tensor([[0.2, 0.3, 0.3, 0.4],[0.2, 0.3, 0.3, 0.4],[0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4]])

decoder_beta = torch.tensor([[0.3, 0.6, 0.4, 0.8],[0.2, 0.3, 0.3, 0.4],[0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4], [0.2, 0.3, 0.3, 0.4]])
</code></pre>
<p>Also, Alpha_2 and Beta_2 are priors:</p>
<pre><code>prior_alpha = torch.tensor([[0.1, 0.1, 0.4, 0.1],[0.8, 0.7, 0.1, 0.4],[0.2, 0.8, 0.9, 0.1], [0.1, 0.5, 0.2, 0.4], [0.1, 0.2, 0.1, 0.4], [0.2, 0.1, 0.3, 0.3]])

prior_beta = torch.tensor([[0.7, 0.6, 0.1, 0.2],[0.5, 0.8, 0.1, 0.2],[0.2, 0.8, 0.5, 0.4], [0.2, 0.6, 0.1, 0.4], [0.6, 0.8, 0.3, 0.2], [0.2, 0.6, 0.3, 0.9]])
</code></pre>
<p>Here is my implementation:</p>
<pre><code>        decoderParamSum = decoder_alpha + decoder_beta
        priorParamSum = prior_alpha + prior_beta
        alphaParamsDiff = decoder_alpha - prior_alpha
        numerator = torch.lgamma(decoderParamSum) + torch.lgamma(prior_alpha) + torch.lgamma(prior_beta)
        denomirator = torch.lgamma(decoder_alpha) + torch.lgamma(decoder_beta) + torch.lgamma(priorParamSum)
        firstTerm = torch.sum((numerator - denomirator),dim=1)
        secondTerm = torch.sum((torch.digamma(decoderParamSum)-torch.digamma(decoder_beta)), dim=1)
        secondTerm = torch.reshape(secondTerm, (input.shape[0], 1))
        secondTerm = torch.digamma(decoder_alpha) - torch.digamma(decoder_beta) - secondTerm
        secondTerm = torch.sum(torch.multiply(alphaParamsDiff, secondTerm), dim=1)
        thirdTerm = torch.sum((torch.digamma(decoderParamSum)-torch.digamma(decoder_beta)), dim=1)
        thirdTerm = torch.reshape(thirdTerm,(input.shape[0], 1))
        v1 = torch.cat([decoder_beta[:,:-1] -decoder_alpha[:, 1:] - decoder_beta[:, 1:], decoder_beta[:, -1:] - 1], dim=-1)
        v2 = torch.cat([prior_beta[:,:-1] -prior_alpha[:, 1:] - prior_beta[:, 1:], prior_beta[:, -1:] - 1], dim=-1)
        
        thirdTerm = torch.sum((torch.multiply((v1-v2), thirdTerm)), dim=1)
       
        KLD = firstTerm - secondTerm + thirdTerm
</code></pre>
<p>However, the loss values I’m getting are far from my expectation, with large negative values, and not stable. So, I’m guessing that there is a problem with my implementation.</p>
<p>Can anyone please check out my implementation of the KL-divergence or if there is a python implementation already in existence (I have surfed the internet, but nothing coming up). Thank you in advance.</p>