Pytorch different result when using `torch.matmul` and `for-loop` to pass input through linear layers

username: schow
cooked: <p>Hi All,</p>
<p>This is my first post here. I have been at this for about two days now. I am working on a model that takes on an input <code>x</code> and passes it through several linear layers, and concatenates the results across the dimension of the number of linear layers as follows:</p>
<pre><code class="lang-python">num_heads = 4
n, c, h, w = 2, 32, 64, 64
x = torch.rand((n*c*h*w)).reshape(n,c,h,w)

x = x.view(n, c, h*w).transpose(1, 2)   # n, h*w, c
heads_q = []
for i in range(num_heads):
    l = nn.Linear(c, c // num_heads, bias=False)
    l.weight.requires_grad_(False)
    heads_q.append(l)

heads_q_res = torch.cat([
        h(x).unsqueeze(0) for h in heads_q
    ], dim=0).detach()  # num_heads, n, h*w, c//num_heads
heads_q_res = heads_q_res.transpose(2,3)    # num_heads, n, c//num_heads, h*w
</code></pre>
<p>I want to get the same results as <code>heads_q_res</code> without having to use a for-loop to pass the input <code>x</code> through each of the linear layers. Thus, I have modified the above code as follows:</p>
<pre><code class="lang-python">num_heads = 4
n, c, h, w = 2, 32, 64, 64
x = torch.rand((n*c*h*w)).reshape(n,c,h,w)

x = x.view(n, c, h*w).transpose(1, 2)   # n, h*w, c

heads_q = []
Wqs = []
for i in range(num_heads):
    l = nn.Linear(c, c // num_heads, bias=False)
    l.weight.requires_grad_(False)
    heads_q.append(l)
    Wqs.append(l.weight.unsqueeze(0).unsqueeze(0))
Wqs = torch.cat(Wqs, dim=0)  # num_heads, 1, c//num_heads, c
print(f'Wqs.shape: {Wqs.shape}')


for i in range(num_heads):
    assert torch.all(Wqs[i] == heads_q[i].weight.unsqueeze(0))
print('All weights check out.')

heads_q_res = torch.cat([
        h(x).unsqueeze(0) for h in heads_q
    ], dim=0).detach()  # num_heads, n, h*w, c//num_heads
heads_q_res = heads_q_res.transpose(2,3)    # num_heads, n, c//num_heads, h*w

_stacked = torch.matmul(Wqs, x.transpose(1,2))    # num_heads, n, c//num_heads, h*w
assert (torch.allclose(heads_q_res, _stacked)), f'\n{torch.abs(heads_q_res-_stacked)}'
</code></pre>
<p>here, <code>_stacked</code> represents the operation of passing the input <code>x</code> through all the linear layers at the same time. But unfortunately, the assertion statement in the last line of the above code block throws an assertion error:</p>
<pre><code class="lang-auto">_stacked.shape:    torch.Size([4, 2, 8, 4096])
heads_q_res.shape: torch.Size([4, 2, 8, 4096])
Traceback (most recent call last):
  File "f.py", line 54, in &lt;module&gt;
    assert (torch.allclose(heads_q_res, _stacked)), f'\n{torch.abs(heads_q_res-_stacked)}'
AssertionError:
tensor([[[[0.0000e+00, 7.4506e-09, 5.9605e-08,  ..., 0.0000e+00,
           2.9802e-08, 0.0000e+00],
          [0.0000e+00, 2.9802e-08, 8.9407e-08,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.1176e-08, 1.4901e-08, 1.1921e-07,  ..., 3.7253e-09,
           0.0000e+00, 0.0000e+00],
          ...,
          [2.9802e-08, 3.7253e-09, 1.4901e-08,  ..., 7.4506e-09,
           0.0000e+00, 1.4901e-08],
          [0.0000e+00, 5.9605e-08, 0.0000e+00,  ..., 5.9605e-08,
           0.0000e+00, 0.0000e+00],
          [2.9802e-08, 0.0000e+00, 5.9605e-08,  ..., 1.4901e-08,
           0.0000e+00, 0.0000e+00]],

         [[1.4901e-08, 2.2352e-08, 5.9605e-08,  ..., 4.4703e-08,
           2.9802e-08, 0.0000e+00],
          [2.6077e-08, 1.4901e-08, 2.9802e-08,  ..., 2.9802e-08,
           2.9802e-08, 1.4901e-08],
          [0.0000e+00, 2.9802e-08, 0.0000e+00,  ..., 1.4901e-08,
           2.9802e-08, 5.9605e-08],
          ...,
          [2.9802e-08, 4.4703e-08, 2.9802e-08,  ..., 1.4901e-08,
           4.4703e-08, 3.3528e-08],
          [0.0000e+00, 0.0000e+00, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [1.4901e-08, 1.4901e-08, 7.4506e-09,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00]]],


        [[[1.3039e-08, 2.9802e-08, 2.9802e-08,  ..., 3.7253e-08,
           1.4901e-08, 2.2352e-08],
          [0.0000e+00, 8.9407e-08, 0.0000e+00,  ..., 3.1665e-08,
           6.7055e-08, 2.9802e-08],
          [0.0000e+00, 2.9802e-08, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 2.9802e-08],
          ...,
          [2.9802e-08, 0.0000e+00, 1.1921e-07,  ..., 5.9605e-08,
           2.9802e-08, 5.9605e-08],
          [0.0000e+00, 4.4703e-08, 0.0000e+00,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [5.9605e-08, 2.9802e-08, 1.4901e-08,  ..., 5.9605e-08,
           0.0000e+00, 0.0000e+00]],

         [[3.7253e-09, 7.4506e-09, 5.9605e-08,  ..., 3.6554e-08,
           5.2154e-08, 0.0000e+00],
          [4.4703e-08, 1.4901e-08, 0.0000e+00,  ..., 2.9802e-08,
           1.4901e-08, 2.9802e-08],
          [2.2352e-08, 7.4506e-09, 0.0000e+00,  ..., 2.9802e-08,
           2.9802e-08, 5.2154e-08],
          ...,
          [0.0000e+00, 2.9802e-08, 2.9802e-08,  ..., 5.9605e-08,
           5.9605e-08, 0.0000e+00],
          [0.0000e+00, 0.0000e+00, 5.9605e-08,  ..., 5.9605e-08,
           0.0000e+00, 0.0000e+00],
          [5.9605e-08, 0.0000e+00, 2.9802e-08,  ..., 5.9605e-08,
           0.0000e+00, 0.0000e+00]]],


        [[[0.0000e+00, 1.4901e-08, 0.0000e+00,  ..., 0.0000e+00,
           8.9407e-08, 2.9802e-08],
          [0.0000e+00, 5.9605e-08, 0.0000e+00,  ..., 0.0000e+00,
           2.9802e-08, 5.9605e-08],
          [2.9802e-08, 2.9802e-08, 1.1921e-07,  ..., 5.9605e-08,
           0.0000e+00, 5.9605e-08],
          ...,
          [1.4901e-08, 1.4901e-08, 0.0000e+00,  ..., 2.9802e-08,
           7.4506e-09, 1.4901e-08],
          [0.0000e+00, 2.2352e-08, 0.0000e+00,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00],
          [2.9802e-08, 0.0000e+00, 1.4901e-08,  ..., 1.4901e-08,
           0.0000e+00, 0.0000e+00]],

         [[0.0000e+00, 5.2154e-08, 0.0000e+00,  ..., 8.9407e-08,
           1.1921e-07, 5.9605e-08],
          [2.2352e-08, 0.0000e+00, 1.4901e-08,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00],
          [0.0000e+00, 5.9605e-08, 5.9605e-08,  ..., 0.0000e+00,
           5.9605e-08, 0.0000e+00],
          ...,
          [1.4901e-08, 4.4703e-08, 1.1176e-08,  ..., 1.4901e-08,
           1.4901e-08, 0.0000e+00],
          [2.9802e-08, 0.0000e+00, 0.0000e+00,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00],
          [2.9802e-08, 1.4901e-08, 1.4901e-08,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00]]],


        [[[7.4506e-09, 5.9605e-08, 2.2352e-08,  ..., 1.1176e-08,
           4.4703e-08, 1.4901e-08],
          [0.0000e+00, 2.9802e-08, 0.0000e+00,  ..., 0.0000e+00,
           5.2154e-08, 4.4703e-08],
          [2.9802e-08, 4.4703e-08, 5.9605e-08,  ..., 0.0000e+00,
           2.9802e-08, 5.9605e-08],
          ...,
          [2.9802e-08, 0.0000e+00, 4.0978e-08,  ..., 2.9802e-08,
           1.4901e-08, 2.9802e-08],
          [2.9802e-08, 0.0000e+00, 0.0000e+00,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00],
          [2.9802e-08, 0.0000e+00, 5.9605e-08,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00]],

         [[2.9802e-08, 7.4506e-09, 7.4506e-08,  ..., 4.4703e-08,
           5.9605e-08, 4.4703e-08],
          [2.9802e-08, 0.0000e+00, 1.4901e-08,  ..., 0.0000e+00,
           1.4901e-08, 5.9605e-08],
          [5.9605e-08, 2.9802e-08, 2.9802e-08,  ..., 3.6322e-08,
           0.0000e+00, 2.9802e-08],
          ...,
          [4.4703e-08, 7.4506e-09, 1.8626e-08,  ..., 1.4901e-08,
           2.6077e-08, 2.2352e-08],
          [0.0000e+00, 1.4901e-08, 5.9605e-08,  ..., 0.0000e+00,
           0.0000e+00, 0.0000e+00],
          [2.9802e-08, 2.9802e-08, 2.9802e-08,  ..., 2.9802e-08,
           0.0000e+00, 0.0000e+00]]]])
</code></pre>
<p>I have also tried replacing</p>
<pre><code class="lang-auto">_stacked = torch.matmul(Wqs, x.transpose(1,2))    # num_heads, n, c//num_heads, h*w
</code></pre>
<p>with</p>
<pre><code class="lang-auto">_stacked = torch.matmul(x.unsqueeze(1), Wqs.squeeze(1).transpose(1,2)).permute(1, 0, 3, 2)  # num_heads, n, c//num_heads, h*w
</code></pre>
<p>but to no avail.</p>
<p>However, I do not get this error when I initialize <code>x</code> as <code>x = torch.ones((n*c*h*w)).reshape(n,c,h,w)</code>. I am completely lost here, and have no clue on how to fix this error. I would like to be able to pass all the assertions without relying on the initializer for <code>x</code>, any help would be appreciated.</p>
<p>Thank you.</p>

username: yuri
cooked: <p>you should do<br>
<code>_stacked = torch.matmul(x, Wqs.transpose(3, 2)).transpose(3, 2)</code><br>
instead of what you have done<br>
please refer to <a href="https://pytorch.org/docs/stable/generated/torch.nn.Linear.html" rel="noopener nofollow ugc">torch.linear</a> to perform the same operation for _stacked</p>

username: schow
cooked: <aside class="quote no-group" data-username="yuri" data-post="2" data-topic="137949">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/71c47a/48.png" class="avatar"> yuri:</div>
<blockquote>
<p>ame oper</p>
</blockquote>
</aside>
<p>hi <a class="mention" href="/u/yuri">@yuri</a> using <code>_stacked = torch.matmul(x, Wqs.transpose(3, 2)).transpose(3, 2)</code> also throws an assertion error. I am beginning to wonder if it might be an issue with numerical precision as the absolute value of the errors are in the range of 10^-7, some are even an order of magnitude lower. Either that or <code>matmul</code> might be doing something undocumented.</p>

username: yuri
cooked: <aside class="quote no-group" data-username="yuri" data-post="2" data-topic="137949">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/71c47a/48.png" class="avatar"> yuri:</div>
<blockquote>
<p>_stacked = torch.matmul(x, Wqs.transpose(3, 2)).transpose(3, 2)</p>
</blockquote>
</aside>
<p>please be aware that the two operations are not the same. the second one is equivalent to passing through a linear layer. the one you posted is not<br>
check (heads_q_res -_stacked) 10-8 precision should be fine</p>

username: yuri
cooked: <p>you can start by<code> x = 10000*x</code> to better compare both operations</p>

username: InnovArul
cooked: <aside class="quote no-group" data-username="yuri" data-post="4" data-topic="137949">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/y/71c47a/48.png" class="avatar"> yuri:</div>
<blockquote>
<p>please be aware that the two operations are not the same. the second one is equivalent to passing through a linear layer. the one you posted is not</p>
</blockquote>
</aside>
<p>In my understanding, both operations achieve the same result. No?</p>
<p>Assuming <code>A = Wqs, B = x, C = _stacked</code></p>
<aside class="quote no-group" data-username="schow" data-post="1" data-topic="137949">
<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/schow/48/32206_2.png" class="avatar"> schow:</div>
<blockquote>
<p><code>_stacked = torch.matmul(Wqs, x.transpose(1,2))</code></p>
</blockquote>
</aside>
<p>This says <code>C = A.(B^T)</code>.</p>
<aside class="quote no-group" data-username="schow" data-post="3" data-topic="137949">
<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/schow/48/32206_2.png" class="avatar"> schow:</div>
<blockquote>
<p>_stacked = torch.matmul(x, Wqs.transpose(3, 2)).transpose(3, 2)</p>
</blockquote>
</aside>
<p>The change that <a class="mention" href="/u/yuri">@yuri</a> suggested does <code>C= (B.(A^T))^T = A.(B^T)</code>.</p>

username: InnovArul
cooked: <p>Just to add to the answer, by default, pytorch creates <code>float</code> tensors (which have ~7 digit precision during operations).<br>
If you would like to have more precision, convert the data to <code>double</code> before performing the operation.</p>

username: yuri
cooked: <p>I don’t think so since this is a ‘‘partial’’ transpose on the 4th and 3rd dimension of the tensor<br>
check<br>
<code>torch.matmul(Wqs, x.transpose(1,2)) -  torch.matmul(x, Wqs.transpose(3, 2)).transpose(3, 2)</code><br>
is not a zero matrix</p>

username: InnovArul
cooked: <p>Probably, check with a tolerance (<code>atol</code> in <code>torch.allclose()</code>) considering it is a floating-point operation, (or) print it and see?</p>