Optimized MultivariateNormal with diagonal covariance matrix

juniorrojas: <p>The current <code>MultivariateNormal</code> implementation is adding a significant overhead to my code when I use large batch sizes and after looking into the source, it seems the main cause is in one helper function used in the <code>log_prob</code> method, which runs an explicit for loop for a batch operation:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/multivariate_normal.py#L55" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/multivariate_normal.py#L55" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/multivariate_normal.py#L55</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="45" style="counter-reset: li-counter 44 ;">
<li>return flat_bmat_inv.reshape(bmat.shape)</li>
<li>
</li>
<li>
</li>
<li>def _batch_trtrs_lower(bb, bA):</li>
<li>"""</li>
<li>Applies `torch.trtrs` for batches of matrices. `bb` and `bA` should have</li>
<li>the same batch shape.</li>
<li>"""</li>
<li>flat_b = bb.reshape((-1,) + bb.shape[-2:])</li>
<li>flat_A = bA.reshape((-1,) + bA.shape[-2:])</li>
<li class="selected">flat_X = torch.stack([torch.trtrs(b, A, upper=False)[0] for b, A in zip(flat_b, flat_A)])</li>
<li>return flat_X.reshape(bb.shape)</li>
<li>
</li>
<li>
</li>
<li>def _batch_mahalanobis(bL, bx):</li>
<li>r"""</li>
<li>Computes the squared Mahalanobis distance :math:`\mathbf{x}^\top\mathbf{M}^{-1}\mathbf{x}`</li>
<li>for a factored :math:`\mathbf{M} = \mathbf{L}\mathbf{L}^\top`.</li>
<li>
</li>
<li>Accepts batches for both bL and bx. They are not necessarily assumed to have the same batch</li>
<li>shape, but `bL` one should be able to broadcasted to `bx` one.</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>The other issue is that I just need a diagonal covariance for my current model, but the <code>MultivariateNormal</code> object is very general and runs some unnecessary computations that could be optimized for a diagonal covariance, like <code>torch.trtrs</code>. Would it make sense to have a <code>MultivariateNormal</code> implementation with some optimizations for strictly diagonal covariances? I’ve noticed there’s a new <code>LowRankMultivariateNormal</code> in the master branch that hasn’t made it into the stable release yet. I believe that implementation might be more suitable, the constructor takes a <code>cov_diag</code> explicitly, but it also takes a <code>cov_factor</code>, which might run some unnecessary computations for a strictly diagonal covariance as well:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/lowrank_multivariate_normal.py#L91" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/lowrank_multivariate_normal.py#L91" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/eef083e4774c52161148883c53ce5a1f081e64ca/torch/distributions/lowrank_multivariate_normal.py#L91</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="81" style="counter-reset: li-counter 80 ;">
<li>    Thanks to these formulas, we just need to compute the determinant and inverse of</li>
<li>    the small size "capacitance" matrix::</li>
<li>        capacitance = I + cov_factor.T @ inv(cov_diag) @ cov_factor</li>
<li>"""</li>
<li>arg_constraints = {"loc": constraints.real,</li>
<li>                   "cov_factor": constraints.real,</li>
<li>                   "cov_diag": constraints.positive}</li>
<li>support = constraints.real</li>
<li>has_rsample = True</li>
<li>
</li>
<li class="selected">def __init__(self, loc, cov_factor, cov_diag, validate_args=None):</li>
<li>    if loc.dim() &lt; 1:</li>
<li>        raise ValueError("loc must be at least one-dimensional.")</li>
<li>    event_shape = loc.shape[-1:]</li>
<li>    if cov_factor.dim() &lt; 2:</li>
<li>        raise ValueError("cov_factor must be at least two-dimensional, "</li>
<li>                         "with optional leading batch dimensions")</li>
<li>    if cov_factor.shape[-2:-1] != event_shape:</li>
<li>        raise ValueError("cov_factor must be a batch of matrices with shape {} x m"</li>
<li>                         .format(event_shape[0]))</li>
<li>    if cov_diag.shape[-1:] != event_shape:</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>What’s the recommended approach to create an efficient multivariate normal distribution with a strictly diagonal covariance matrix?</p>

talesa: <p>You can just use <code>torch.distributions.Normal</code> in that case</p>
<pre><code class="lang-auto">import torch
n = 2
d = 5
diagonal = torch.rand(d) + 1.
mu = torch.rand(n, d)
p1 = torch.distributions.Normal(mu, diagonal.reshape(1, d))
p2 = torch.distributions.MultivariateNormal(mu, scale_tril=torch.diag(diagonal).reshape(1, d, d))
x = torch.rand((n,d))
print(p1.log_prob(x).sum(dim=1) - p2.log_prob(x))
</code></pre>

JakobHavtorn: <p>The above answer is the way to go, but it can result in confusing issues if you’re new to <code>torch.distributions</code>.</p>
<p>In order to get the correct Kullback-Leibler divergence (and the correct shape of <code>.log_prob</code>), we need to wrap the <code>Normal</code> in the <code>Independent</code> class which reinterprets some number of batch dimensions as event dimensions. This is not done by the default <code>Normal</code> which simply assumes all dimensions to be batch-dimensions which is generally not the behaviour you want when you’re effectively defining a multivariate normal with diagonal covariance. See below and <a href="https://pytorch.org/docs/stable/distributions.html#independent" rel="nofollow noopener">https://pytorch.org/docs/stable/distributions.html#independent</a></p>
<pre><code class="lang-auto">&gt;&gt;&gt; loc = torch.zeros(3)
&gt;&gt;&gt; scale = torch.ones(3)
&gt;&gt;&gt; mvn = MultivariateNormal(loc, scale_tril=torch.diag(scale))
&gt;&gt;&gt; [mvn.batch_shape, mvn.event_shape]
[torch.Size(()), torch.Size((3,))]
&gt;&gt;&gt; normal = Normal(loc, scale)
&gt;&gt;&gt; [normal.batch_shape, normal.event_shape]
[torch.Size((3,)), torch.Size(())]
&gt;&gt;&gt; diagn = Independent(normal, 1)
&gt;&gt;&gt; [diagn.batch_shape, diagn.event_shape]
[torch.Size(()), torch.Size((3,))]
</code></pre>

talesa: <p>Please use JakobHavtorn’s approach, not mine, that’s the most appropriate way to go as of now.</p>

Dane_Mitrev: <p>It is clear that axis-aligned gaussian can be created using <code>torch.distributions.Normal</code> and if we need to compute KL we can wrap it into <code>Independent</code> with reinterpretation.</p>
<p>My questions is: can we achieve the same result for a given <code>mean</code> and <code>sigma</code> using the <code>torch.distributions.MultivariateNormal</code> ?</p>
<p>example case:</p>
<pre><code class="lang-auto">batch_size = 5
event dims = 2

&gt;&gt;&gt; mu = torch.ones([5,2])
&gt;&gt;&gt; log_sig = torch.ones([5,2])
&gt;&gt;&gt; indep = Independent(Normal(loc=mu, scale=torch.exp(log_sig)), 1)
&gt;&gt;&gt; indep.event_shape, indep.batch_shape
(torch.Size([2]), torch.Size([5]))

&gt;&gt;&gt; mvn = Multivariate(mu, scale_tril=torch.diag(log_sig))  # THIS FAILS obviously
</code></pre>
<p>How can we create a <code>MultivariateNormal</code> instance out of the given <code>mu</code> and <code>sigma</code> with batch_size=5 ?</p>
<p>How can a <code>MultivariateNormal</code> be instantiated so that it infers a <code>batch_shape=torch.Size([5])</code> in this case?</p>

Dane_Mitrev: <p>EDIT: I simplified the question here:</p>
<aside class="quote quote-modified" data-post="1" data-topic="102070">
  <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/dane_mitrev/48/31366_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-instantiate-axis-aligned-multivariate-gaussian/102070">How to instantiate axis-aligned multivariate Gaussian?</a> 
  </div>
  <blockquote>
    It is straightforward to do the following and wrap a Gaussian into an Independent. So if we are given a 
loc = [batch_size, event_shape]
scale = [batch_size]


loc = torch.zeros(5, 3, 2)
scale = torch.ones(2)
normal = Normal(loc, scale)
normal.batch_shape, normal.event_shape
(torch.Size([5, 3, 2]), torch.Size([]))

ind = Independent(normal, 1)
ind.batch_shape, ind.event_shape
(torch.Size([5, 3]), torch.Size([2]))

This Independent distribution is identical to a MultivariateNormal defined as: 
mv…
  </blockquote>
</aside>