为什么视觉变形金刚键和查询linar层不合并为一个矩阵

Question

我研究一些视觉变压器代码（例如vit-pytorch） 并在注意模块中发现：

#x is input
key=nn.Linear(...,bias=False)(x)
query=nn.Linear(...,bias=False)(x)
similar_matrix=torch.matmul(query,key.transpose(...))

由于可以将线性视为矩阵，所以我

key=K^T @ x
query=Q^T @ x
similar_matrix = query^T @ key = x^T @ (Q @ K^T) @ x
(K,Q means learnable matrix, @ means matmul, ^T means transpose)

在这里得到q @ k^t，我认为它们可以合并到矩阵中，以减少参数的量和计算，

为什么不这样做呢？是因为训练效果不好吗？

Answer 1

让我们清除一些事情。

由于bias = false这确实是正确的想法。这意味着只有权重，但偏差将是none。快速让我们检查一下

import torch
import torch.nn as nn

m = nn.Linear(3, 4,bias=False)
n = nn.Linear(3,4,bias=False)

if hasattr(m, 'weight'):
    print(m.weight) # prints the weight
else:
    print("doesn't exist")
if hasattr(m, 'bias'):
    print(m.bias) # None
else:
    print("doesn't exist")

，让我们运行此片段，它遵循您的想法，

x = torch.randint(1,5,(2,3)).float()
m = nn.Linear(3, 4,bias=False)
n = nn.Linear(3,4,bias=False)

K = m.weight
Q = n.weight

k = m(x) # [email protected]
q = n(x) # [email protected]
print([email protected])
print( [email protected] @  K @x.T) 


#output (it would be same in your case as per the initialized input)
tensor([[-4.6655, -5.2234],
        [-6.8665, -7.6535]], grad_fn=<MmBackward0>)
tensor([[-4.6655, -5.2234],
        [-6.8665, -7.6535]], grad_fn=<MmBackward0>)

因此这两个术语都是等效的。 （确认的）。
但是，参数的数量将与其他方法完全相同。由于我们使用相同的重量矩阵和东西。

哪个更快？

timeit("m(x) @  n(x).T", globals())
timeit("[email protected] @  K @x.T", globals())

 #Output
1000 loops, best of 3: 50.1 usec per loop
1000 loops, best of 3: 31.1 usec per loop

因此，第一个花费的时间比第二个时间多。如果我们增加了大小，那么

x = torch.randint(1,5,(2000,3000)).float()
m = nn.Linear(3000, 4000,bias=False)
n = nn.Linear(3000,4000,bias=False)

K = m.weight
Q = n.weight
timeit("m(x) @  n(x).T", globals())
timeit("[email protected] @  K @x.T", globals())
# output
1 loops, best of 3: 2.93 sec per loop
1 loops, best of 3: 2.73 sec per loop

在这里也可以看到第二种情况下的性能显着提高。因此，理想情况下，人们可以像您猜到的那样使用第二个绩效提高。

但是主要的事情是，为什么他们不遵循的原因可以归因于在使用bias的未来更改的尝试中，并且在这种情况下，第一个代码将更加可读性。并符合现有标准。在这种小情况下，它可能是合适的，但对于较大的型号，这可能很丑陋。为了更好的可读性和未来的范围，我相信他们选择了第一个。