如何将softmax层用作神经网络的中间层之一并正确地反向传播?

Question

我目前正在用c++制作我自己的机器学习库，作为练习来帮助我提高编码技能并提高我对机器学习的理解。我目前正在从零开始制作一个视觉变压器，以便更好地理解变压器以及如何将它们用于图像。我试图重新创建反向传递的部分代码如下:

energy = torch.einsum('bhqd, bhkd -> bhqk', queries, keys) # batch, num_heads, query_len, key_len

scaling = self.emb_size ** (1/2)
att = F.softmax(energy, dim=-1) / scaling

，其中查询和键是四维张量，softmax函数应用于每个通道，energy是应用于每个单独矩阵的矩阵乘法。为了重建反向传播，我试图用更小的二维矩阵来做。

我将q定义为3x2矩阵，k定义为2x3矩阵。因此，如果我将它们相乘并得到一个矩阵a，我将得到一个3x3矩阵，一旦应用softmax函数，维度保持不变。

我遇到的麻烦是找到q和k相对于损失的梯度。这就是我如何找到a的渐变，在这个例子中是dLdA:

for i in range(len(a)):
for j in range(len(a)):
if i == j:
dLdA[i,j] = a[i] * (1-a[i])
else: 
dLdA[i,j] = -a[i] * a[j]

为了更容易读，我把它写成python。从那里我将得到dLdA的9x9矩阵。从那里我需要得到dLdQ和dLdK。其中，dLdQ=dLdA*dAdQ,dLdK=dLdA*dAdK。如果我想从那里计算dAdK的雅可比矩阵，我就会得到这个:

dAdK = 
[dA1dK1 dA1dK2 ... dA1dK6]
.
.
.
[dA9dK1 dA9dK2 ... dA9dK6]

因为a中有9个总元素，k中有6个总元素。其中dAdK是一个9x6矩阵，dLdA是一个9x9矩阵，当你做矩阵乘法时，你得到一个9x6矩阵，但k是一个2x3矩阵。为了正确反向传播，我还缺什么?

Answer 1

所以我基于所有的数学来回答这个问题。如果你想深入了解数学原理，我建议你阅读这本书。这就是使用numpy以更简单的方式实现的方式，这就是它的工作方式:

class Linear:
def __init__(self, in_rows, out_rows):
self.Weight = np.random.rand(out_rows, in_rows)
self.Bias = np.random.rand(out_rows, 1)
self.prev = None
self.LR = 0.01
def forward(self, x):
self.prev = x
a = np.dot(self.Weight, x)
a += self.Bias
return a
def backward(self, dx):
db = np.sum(dx, axis=len(dx.shape)-1)
db = db.reshape(self.Bias.shape)
self.Bias -= (db*self.LR)
dw = np.dot(dx, self.prev.T)
dl = np.dot(self.Weight.T, dx)
self.Weight -= (dw * self.LR)
return dl
class Softmax:
def forward(self, x):
mx = np.max(x, axis=1, keepdims=True)
x = x - mx  # log-sum-exp trick
e = np.exp(x)
probs = e / np.sum(np.exp(x), axis=1, keepdims=True)
return probs

def jacobian(self, a):
r = a.shape[0]
c = a.shape[1]
output = np.zeros((r,c,r,c))
for i in range(output.shape[0]):
for j in range(output.shape[1]):
for r in range(output.shape[2]):
for s in range(output.shape[3]):
if(i == r and j == s):
output[i][j][r][s] += a[i][j]*(1-a[i][j])
else:
output[i][j][r][s] -= a[i][j]*a[r][s]
return output
def backward(self, dx):
jac = self.jacobian(dx)
jac = np.sum(jac, axis=(0,1))
return jac
def train(x, wanted, m):
model = [Linear(3,4), Linear(4,4), Softmax(), Linear(4,3)]
for i in range(0, m):
curr = x
for layer in model:
curr = layer.forward(curr)
dL = curr - wanted
if(i % 10 == 0):
print("current error sum: ", np.sum(dL))
for layer in reversed(model):
dL = layer.backward(dL)

if __name__ == '__main__':
x = np.random.rand(3,2)
wanted = np.random.rand(3,2)

train(x, wanted, 100)

我得到的一个输出如下:

current error sum:  9.846307612361286
current error sum:  5.5827143820787395
current error sum:  3.073972734037421
current error sum:  1.622572236217964
current error sum:  0.7876682647260167
current error sum:  0.3118719708795148
current error sum:  0.04493155201662072

表示这个错误是可以被纠正的。显然，可以制作一个更好的网络，并且可以而且应该使用更好的误差函数。这只是一个使用计算机科学来实现上面提到的帖子的演示。这也是我在没有进行优化的情况下写出所有内容的另一个原因。需要注意的一点是我在反向函数中加入了雅可比矩阵，在4D张量的前两个维度上加入了雅可比矩阵把它变成了一个2D张量。如果用dA/dW和dA/dx的四维雅可比矩阵分别求出dL/dW和dL/dx，然后求和，所做的运算得到相同的答案。这样做是为了使程序更有效，更容易显示单独的层，以及如何将其转换为通常的torch.nn.Linear层。