Forecast energy score: compute pairwise distances between tensor "columns"

mishooax: <p>Hi,</p>
<p>I am trying to implement an <a href="https://cran.r-project.org/web/packages/scoringRules/vignettes/article.pdf" rel="noopener nofollow ugc"><em>energy score</em></a> loss for a forecast ensemble.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/c/acd5ee043121828063af82f281ea0313a9baf08d.png" data-download-href="https://discuss.pytorch.org/uploads/default/acd5ee043121828063af82f281ea0313a9baf08d" title="energy-score"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/c/acd5ee043121828063af82f281ea0313a9baf08d.png" alt="energy-score" data-base62-sha1="oEYsYhHa5yp70KxgWvB7NouhW2h" width="690" height="128" data-dominant-color="F4F4F4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">energy-score</span><span class="informations">944×176 6.34 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I’m having trouble vectorizing the second term, which is essentially a sum over pairwise distances between different ensemble members (i != j) of the forecast tensor <code>X</code>. <code>X.shape = (batch_size, n_vars, lat, lon, n_members)</code> and the <code>(i, j)</code> index the <code>n_members</code> dimension (<code>X.shape[-1]</code>).</p>
<p>Here’s the “loopy” version i coded for a single sample:</p>
<pre data-code-wrap="python"><code class="lang-python">    # works only with a batch_size == 1
    tmp_b = torch.zeros(1)
    for i in range(m):
        for j in range(m):
            if i == j:
                continue
            # i am using a more general formula where each 2-norm is raised to the power beta
            tmp_b += torch.pow(torch.linalg.norm((y_pred[..., i] - y_pred[..., j]).flatten(), ord=2), beta)

    # unbiased estimator: divide by (m-1) instead of m as the i = j terms cancel out
    tmp_b *= - 1.0 / (2 * m * (m - 1))
</code></pre>
<p>So I have two questions:</p>
<p>1/ How to vectorize the code above?<br>
2/ How to vectorize the formula for <code>batch_size &gt; 1</code> (essentially this would entail averaging over the energy losses for all samples in the batch)</p>
<p>Thanks!</p>

mishooax: <p>I was able to vectorize the single-sample batch code with the help of <code>torch.cdist</code>:</p>
<pre><code class="lang-python">    tmp_b = - 1.0 / (2 * m * (m - 1)) * torch.sum(
        torch.pow(
            torch.cdist(
                y_pred.reshape(-1, m).T, y_pred.reshape(-1, m).T, p=2
            ), beta
        )
    )
</code></pre>
<p>Question 2 still stands. Anybody know how to vectorize this op for batch sizes &gt; 1? Thanks.</p>

mishooax: <p>Hmm… I <em>think</em> this does it:</p>
<pre><code class="lang-python">    tmp_b = - 1.0 / (2 * m * (m - 1)) * torch.mean(
        torch.sum(
          torch.pow(
              torch.cdist(
                  y_pred_.reshape(bs, -1, m).permute(0,2,1), y_pred_.reshape(bs, -1, m).permute(0,2,1), p=2
              ), beta
          ), axis=(1,2),
        ), dim=0
    )
</code></pre>