Stochastic matrix parameterization

username: leo2k
cooked: <p>Hi all, I’m currently implementing a neural net that output a tensor of shape (batch_size, N, N) where NxN matrices have their column sum and row sum =1. Can anyone give me a help, please</p>

username: KFrank
cooked: <p>Hi Leonard!</p>
<aside class="quote no-group" data-username="leo2k" data-post="1" data-topic="145085" data-full="true">
<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/leo2k/48/31216_2.png" class="avatar"> leo2k:</div>
<blockquote>
<p>NxN matrices have their column sum and row sum =1.</p>
</blockquote>
</aside>
<p>If you want to <em>train</em> your network to produce such matrices, you can<br>
add to your loss function a term that penalizes violations of your desired<br>
constraints.  The matrices usually won’t satisfy the constraints exactly, but<br>
if the penalty term is weighted heavily enough (and you train enough),<br>
they should come close.</p>
<p>You could use:</p>
<pre data-code-wrap="python"><code class="lang-python">penalty = ((x.sum (dim = 2) - 1)**2).sum() + ((x.sum (dim = 1) - 1)**2).sum()
</code></pre>
<p>If you want to coerce the output of your network to satisfy the constraints,<br>
you can simply modify the matrices after they have been produced.  Simply<br>
forcing the rows and columns to sum to one is straightforward.</p>
<p>I’m not sure of your use case.  If you want the rows and columns of your<br>
matrices to be proper row-wise and column-wise (discrete) probability<br>
distributions – that is, not only sum to one, but have the individual matrix<br>
elements lie between zero and one – that is more involved, and I don’t<br>
know of a clean, symmetrical approach (other than employing iteration).</p>
<p>The obvious straightforward method for coercing the row and column<br>
sums, and a clunky, ad hoc approach to coercing row- and column-wise<br>
probability distributions are illustrated in the following script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2022)

def fix_sums (x):    # treats rows and columns symmetrically
    N = x.shape[1]   # assumes that x is of shape [batch_size, N, N]
    x = x + (1 - x.sum (dim = 2).unsqueeze (2)) / N   # row sums
    x = x + (1 - x.sum (dim = 1).unsqueeze (1)) / N   # column sums
    return x

t = torch.randn (3, 4, 4)
u = fix_sums (t)
print ('fix_sums:')
print ('check rows:    ', torch.allclose (u.sum (dim = 2), torch.ones (1)))
print ('check columns: ', torch.allclose (u.sum (dim = 1), torch.ones (1)))

def fix_probs (x):   # does not treat rows and columns symmetrically
    N = x.shape[1]   # assumes that x is of shape [batch_size, N, N]
    b1 = torch.tensor (1.0 / N)   # ad hoc buffer to maintain  0 &lt;= x &lt;= 1
    b2 = torch.tensor (1.5 / N)   #    when column-sums-correction is applied
    rmin = x.min (dim = 2)[0]
    rlo = torch.max (torch.min (rmin, b2), b1)
    rmax = x.max (dim = 2)[0]
    rhi = torch.min (torch.max (rmax, 1 - b2), 1 - b1)
    m = (rhi - rlo) / (rmax - rmin)
    d = (rmax * rlo - rmin * rhi) / (rhi - rlo)
    x = m.unsqueeze (2) * (x + d.unsqueeze (2))   # 0 &lt;= x &lt;= 1
    x = x / x.sum (dim = 2).unsqueeze (2)             # row sums
    x = x + (1 - x.sum (dim = 1).unsqueeze (1)) / N   # column sums
    return x

v = fix_probs (t)
print ('fix_probs:')
print ('check &gt;= 0:    ', torch.all (v &gt;= 0))
print ('check &lt;= 1:    ', torch.all (v &lt;= 1))
print ('check rows:    ', torch.allclose (v.sum (dim = 2), torch.ones (1)))
print ('check columns: ', torch.allclose (v.sum (dim = 1), torch.ones (1)))
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">1.10.2
fix_sums:
check rows:     True
check columns:  True
fix_probs:
check &gt;= 0:     tensor(True)
check &lt;= 1:     tensor(True)
check rows:     True
check columns:  True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: leo2k
cooked: <p>Oh, thank you so much <a class="mention" href="/u/kfrank">@KFrank</a>, this helps me a lot. Btw, is this algorithm to derive a doubly stochastic matrix well known, or you just come up with this method right away brilliantly?</p>

username: leo2k
cooked: <p>I’ve tried the Sinkhorn-Knopp algorithm yet it requires an extra memory footprint <a href="https://github.com/lucidrains/sinkhorn-transformer" class="inline-onebox" rel="noopener nofollow ugc">GitHub - lucidrains/sinkhorn-transformer: Sinkhorn Transformer - Practical implementation of Sparse Sinkhorn Attention</a></p>

username: leo2k
cooked: <p>However, I got this assertion error in training when using your method <a class="mention" href="/u/kfrank">@KFrank</a></p>
<p><code>assert torch.all (x &gt;= 0), f"{x}"</code><br>
It seems x&gt;=0 is not alway satisfied</p>

username: KFrank
cooked: <p>Hi Leonard!</p>
<aside class="quote no-group" data-username="leo2k" data-post="5" data-topic="145085" data-full="true">
<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/leo2k/48/31216_2.png" class="avatar"> leo2k:</div>
<blockquote>
<p><code>assert torch.all (x &gt;= 0), f"{x}"</code><br>
It seems x&gt;=0 is not alway satisfied</p>
</blockquote>
</aside>
<p>Yes, the algorithm I posted is ad hoc.</p>
<pre data-code-wrap="python"><code class="lang-python">    b1 = torch.tensor (1.0 / N)   # ad hoc buffer to maintain  0 &lt;= x &lt;= 1
    b2 = torch.tensor (1.5 / N)   #    when column-sums-correction is applied
</code></pre>
<p><code>b1</code> and <code>b2</code> are “fix-up” parameters that define how much to shrink the<br>
rows inside of <code>[0.0, 1.0]</code> so that they don’t go back outside of that<br>
range when the columns are subsequently modified to sum to one.</p>
<p>They may well need to depend on <code>N</code> (the size of <code>x</code>) and possibly on<br>
the range of the elements in <code>x</code>.</p>
<p>I expect that if the <code>assert</code> fails, you could, as a fix-up, simply apply the<br>
<code>fix_probs()</code> function a second (or third, or fourth …) time.  But that<br>
certainly would be clunky.</p>
<p>Best.</p>
<p>K. Frank</p>

username: KFrank
cooked: <p>Hi Leonard!</p>
<aside class="quote no-group" data-username="leo2k" data-post="3" data-topic="145085" data-full="true">
<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/leo2k/48/31216_2.png" class="avatar"> leo2k:</div>
<blockquote>
<p>is this algorithm to derive a doubly stochastic matrix well known</p>
</blockquote>
</aside>
<p>I don’t know of such an algorithm – what I posted is ad hoc.</p>
<p>Part of the problem is that the goal of the algorithm is not well defined.</p>
<p>Consider the one-dimensional case of producing a “singly-stochastic<br>
vector:”</p>
<p>Let’s say you start with a vector of arbitrary values and want to map it<br>
to a singly-stochastic vector.  You could use the:</p>
<pre data-code-wrap="python"><code class="lang-python">    x = m.unsqueeze (2) * (x + d.unsqueeze (2))   # 0 &lt;= x &lt;= 1
    x = x / x.sum (dim = 2).unsqueeze (2)             # row sums
</code></pre>
<p>piece of the <code>fix_probs()</code> function (without the <code>b1</code> / <code>b2</code> “buffer”).  This has<br>
the advantage that if your vector already does represent a valid (discrete)<br>
probability distribution, it won’t be changed.</p>
<p>But, instead, you could apply <code>softmax()</code> to your vector – the standard<br>
way to convert a vector of <em>raw-score logits</em> to probabilities.  This has<br>
the disadvantage that applying <code>softmax()</code> to a valid set of probabilities<br>
does <em>not</em> leave them unchanged.  Which is right?  That really depends<br>
on what the meaning of the original input vector is.</p>
<p>To illustrate this issue with a silly exaggeration, you can convert your<br>
input matrix to a doubly-stochastic matrix by replacing all of its elements<br>
with <code>1 / N</code>.  Of course, the result is now independent of the input matrix.</p>
<p>But if you reject this <code>1 / N</code> proposal, you have to provide some concrete<br>
criterion for how the output matrix is supposed to depend on the input.<br>
Without such a criterion, the <code>1 / N</code> proposal would be a legitimate solution.</p>
<p>Best.</p>
<p>K. Frank</p>