How to efficiently compute a pairwise kl divergence matrix of a batch of Gaussian distributions?

username: guzy0324
cooked: <p>Hi! I have a <code>mu</code> of shape <code>[B, H]</code>, and a <code>std</code> of shape <code>[B, H]</code>.<br>
Each row of the <code>mu</code> and <code>std</code> means parameters of a Gaussian distribution. Is there any efficient way to compute pairwise kl divergence matrix of shape <code>[B, B]</code>?</p>
<p>codes to calculate kl divergence of 2 Gaussian distributions that can broadcast on batch dims:</p>
<pre><code class="lang-auto">        k = mu_q.size(1)
        mu_diff = mu_p - mu_q
        mu_diff_sq = torch.mul(mu_diff, mu_diff)
        logdet_std_q = torch.sum(2 * torch.log(torch.clamp(std_q, min=1e-8)), dim=1)
        logdet_std_p = torch.sum(2 * torch.log(torch.clamp(std_p, min=1e-8)), dim=1)
        fs = torch.sum(torch.div(std_q**2, std_p**2), dim=1) + torch.sum(
            torch.div(mu_diff_sq, std_p**2), dim=1
        )
        kl_divergence = (fs - k + logdet_std_p - logdet_std_q) * 0.5
</code></pre>

username: tom
cooked: <p>If you do the computation on <code>mu_p = mu[:, None, :]</code> and <code>mu_q = mu[None, :, :]</code>, similarly for <code>std</code> and replace the <code>dim=1</code> by <code>dim=-1</code> (i.e. the last dimension), you should be good to go thanks to broadcasting. Obviously, it would be even neater if we could avoid materializing [B, B, H] matrices, but that is much harder (and might be done by torch.compile or similar automatically now or in the future).</p>
<p>Best regards</p>
<p>Thomas</p>

username: guzy0324
cooked: <p>It is surprisingly faster than <code>mu @ mu.T</code>.</p>
<p>benchmark codes:</p>
<pre><code class="lang-auto">import torch
import time

def kl_div(mu_q, std_q, mu_p, std_p):
    """Computes the KL divergence between the two given variational distribution.\
        This computes KL(q||p), which is not symmetric. It quantifies how far is\
        The estimated distribution q from the true distribution of p."""
    k = mu_q.size(-1)
    mu_diff = mu_p - mu_q
    mu_diff_sq = torch.mul(mu_diff, mu_diff)
    logdet_std_q = torch.sum(2 * torch.log(torch.clamp(std_q, min=1e-8)), dim=-1)
    logdet_std_p = torch.sum(2 * torch.log(torch.clamp(std_p, min=1e-8)), dim=-1)
    fs = torch.sum(torch.div(std_q**2, std_p**2), dim=-1) + torch.sum(
        torch.div(mu_diff_sq, std_p**2), dim=-1
    )
    kl_divergence = (fs - k + logdet_std_p - logdet_std_q) * 0.5
    return kl_divergence

if __name__ == "__main__":
    B = 100
    D = 256
    times = 100
    mu = torch.randn(B, D).cuda()
    std = torch.randn(B, D).cuda()
    start = time.time()
    for i in range(times):
        kl = kl_div(mu[:, None, :], std[:, None, :], mu[None, :, :], std[None, :, :])
    end = time.time()
    print(end - start, kl.shape)

    start = time.time()
    for i in range(times):
        kl = mu @ mu.T
    end = time.time()
    print(end - start, kl.shape)

    start = time.time()
    for i in range(times):
        kl = (mu[:, None, :] * mu[None, :, :]).sum(dim=-1)
    end = time.time()
    print(end - start, kl.shape)
</code></pre>
<p>outputs:</p>
<pre><code class="lang-auto">0.008843183517456055 torch.Size([100, 100])
0.4114084243774414 torch.Size([100, 100])
0.0010275840759277344 torch.Size([100, 100])
</code></pre>

username: tom
cooked: <p>I think you want <code>time.perf_counter()</code> and <code>torch.cuda.synchronize()</code> before every time taking.<br>
(I have this in my training materials somewhere, but here is an old blog post discussing that in “When timing CUDA models”: <a href="https://lernapparat.de/transformers-pytorch-tvm/" class="inline-onebox" rel="noopener nofollow ugc">Lernapparat - Machine Learning</a> if you ignore the other bits.)</p>
<p>Best regards</p>
<p>Thomas</p>

username: guzy0324
cooked: <aside class="quote no-group" data-username="tom" data-post="4" data-topic="192519">
<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/tom/48/3162_2.png" class="avatar"> tom:</div>
<blockquote>
<p>torch.cuda.synchronize()</p>
</blockquote>
</aside>
<p>I have put the <code>perf_counter</code> and <code>synchronize</code> in my code, but the <code>kl_div</code> is still faster than <code>mu @ mu.T</code>.</p>
<pre><code class="lang-auto">import torch
import time

def kl_div(mu_q, std_q, mu_p, std_p):
    """Computes the KL divergence between the two given variational distribution.\
        This computes KL(q||p), which is not symmetric. It quantifies how far is\
        The estimated distribution q from the true distribution of p."""
    k = mu_q.size(-1)
    mu_diff = mu_p - mu_q
    mu_diff_sq = torch.mul(mu_diff, mu_diff)
    logdet_std_q = torch.sum(2 * torch.log(torch.clamp(std_q, min=1e-8)), dim=-1)
    logdet_std_p = torch.sum(2 * torch.log(torch.clamp(std_p, min=1e-8)), dim=-1)
    fs = torch.sum(torch.div(std_q**2, std_p**2), dim=-1) + torch.sum(
        torch.div(mu_diff_sq, std_p**2), dim=-1
    )
    kl_divergence = (fs - k + logdet_std_p - logdet_std_q) * 0.5
    return kl_divergence

if __name__ == "__main__":
    B = 100
    D = 256
    times = 1
    mu = torch.randn(B, D).cuda()
    std = torch.randn(B, D).cuda()
    torch.cuda.synchronize()

    start = time.perf_counter()
    for i in range(times):
        kl = kl_div(mu.unsqueeze(1), std.unsqueeze(1), mu.unsqueeze(0), std.unsqueeze(0))
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)

    start = time.perf_counter()
    for i in range(times):
        kl = kl_div(mu.view(B, 1, D), std.view(B, 1, D), mu.view(1, B, D), std.view(1, B, D))
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)

    start = time.perf_counter()
    for i in range(times):
        kl = kl_div(mu[:, None, :], std[:, None, :], mu[None, :, :], std[None, :, :])
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)

    start = time.perf_counter()
    for i in range(times):
        kl1 = []
        for m, s in zip(mu, std):
            kl1.append(kl_div(m.unsqueeze(0).expand_as(mu), s.unsqueeze(0).expand_as(std), mu, std))
        kl1 = torch.stack(kl1)
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl1.shape)

    print(kl)
    print(kl1)
    assert ((kl - kl1).abs() &lt; 1e-10).all()

    start = time.perf_counter()
    for i in range(times):
        kl = mu @ mu.T
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)

    start = time.perf_counter()
    for i in range(times):
        kl1 = (mu[:, None, :] * mu[None, :, :]).sum(dim=-1)
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl1.shape)

    print(kl)
    print(kl1)
    assert ((kl - kl1).abs() &lt; 1e-4).all()

    start = time.perf_counter()
    for i in range(times):
        kl = mu.T
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)

    start = time.perf_counter()
    muT = mu.T
    for i in range(times):
        kl = mu @ muT
    torch.cuda.synchronize()
    end = time.perf_counter()
    print(end - start, kl.shape)
</code></pre>
<p>outputs:</p>
<pre><code class="lang-auto">0.000576397986151278 torch.Size([100, 100])
0.00018305005505681038 torch.Size([100, 100])
0.00018039101269096136 torch.Size([100, 100])
0.0067256499314680696 torch.Size([100, 100])
tensor([[     0.0000,  57080.6406,  66168.9922,  ...,  53879.9414,
         368704.2188,  79685.1484],
        [ 39291.7500,      0.0000, 107659.6328,  ..., 142458.5469,
         399184.1562,  47731.7344],
        [ 56672.7148,  88284.0703,      0.0000,  ...,  73333.1875,
         956130.2500,  51419.0156],
        ...,
        [ 79648.5938,  60400.5156,  77861.4609,  ...,      0.0000,
         243438.4062, 108050.6641],
        [151945.2969,  58202.5508,  33863.8906,  ...,  31096.5391,
              0.0000, 103794.0938],
        [ 50227.0508,  78881.3125,  76599.9531,  ...,  58455.1133,
         310935.7500,      0.0000]], device='cuda:0')
tensor([[     0.0000,  57080.6406,  66168.9922,  ...,  53879.9414,
         368704.2188,  79685.1484],
        [ 39291.7500,      0.0000, 107659.6328,  ..., 142458.5469,
         399184.1562,  47731.7344],
        [ 56672.7148,  88284.0703,      0.0000,  ...,  73333.1875,
         956130.2500,  51419.0156],
        ...,
        [ 79648.5938,  60400.5156,  77861.4609,  ...,      0.0000,
         243438.4062, 108050.6641],
        [151945.2969,  58202.5508,  33863.8906,  ...,  31096.5391,
              0.0000, 103794.0938],
        [ 50227.0508,  78881.3125,  76599.9531,  ...,  58455.1133,
         310935.7500,      0.0000]], device='cuda:0')
0.39876252599060535 torch.Size([100, 100])
0.00013514305464923382 torch.Size([100, 100])
tensor([[286.4470,  12.7784, -24.8256,  ..., -16.4116,  -4.7319,  22.7014],
        [ 12.7784, 247.9196,  -5.5962,  ..., -26.5139,  17.1346,   7.6810],
        [-24.8256,  -5.5962, 294.9920,  ..., -31.1086,   8.4197,  -2.6804],
        ...,
        [-16.4116, -26.5139, -31.1086,  ..., 252.4709,  39.1322, -24.4404],
        [ -4.7319,  17.1346,   8.4197,  ...,  39.1322, 250.0481, -23.8120],
        [ 22.7014,   7.6810,  -2.6804,  ..., -24.4404, -23.8120, 245.9375]],
       device='cuda:0')
tensor([[286.4470,  12.7784, -24.8256,  ..., -16.4116,  -4.7319,  22.7014],
        [ 12.7784, 247.9196,  -5.5962,  ..., -26.5139,  17.1346,   7.6810],
        [-24.8256,  -5.5962, 294.9920,  ..., -31.1086,   8.4197,  -2.6804],
        ...,
        [-16.4116, -26.5139, -31.1086,  ..., 252.4708,  39.1322, -24.4404],
        [ -4.7319,  17.1346,   8.4197,  ...,  39.1322, 250.0481, -23.8120],
        [ 22.7014,   7.6810,  -2.6804,  ..., -24.4404, -23.8120, 245.9375]],
       device='cuda:0')
1.8235063180327415e-05 torch.Size([256, 100])
3.278802614659071e-05 torch.Size([100, 100])
</code></pre>
<p>time consumed:<br>
<code>mu @ muT</code> &lt; <code>(mu[:, None, :] * mu[None, :, :]).sum(dim=-1)</code> &lt; <code>kl_div</code> &lt; <code>mu @ mu.T</code></p>

username: tom
cooked: <p>Yeah, but now you can trust your numbers more. <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>