向量化 softmax 交叉熵梯度

Vectorizing softmax cross-entropy gradient

我正在尝试使用(几乎)完全矢量化的操作来实现我自己的神经网络。那里有很多帖子,但我似乎找不到适合所有这三个的帖子:

  1. 在梯度计算中将交叉熵和 softmax 项分开(这样我就可以互换最后的激活和损失)
  2. multi-class classification(y是one-hot encoded)
  3. 所有操作都完全向量化

我的主要问题是:如何使用 dE/da (N x K) 和 da/dz (N x K x K) 达到 dE/dz (N x K)一个完全矢量化的操作?即 我如何向量化 dE_dz_test2?

我的第二个问题是: 有没有更好的写法softmax_derivative?

我用这个作为一次计算一个样本的梯度的参考: http://saitcelebi.com/tut/output/part2.html 这是为了弄清楚如何做反向传播 https://peterroelants.github.io/posts/neural-network-implementation-part04/

def one_hot_encode(y, n_classes):
    y_onehot = np.zeros((len(y), n_classes))
    for i, y_i in enumerate(y):
        y_onehot[i, y_i] = 1
    return y_onehot

def cross_entropy_derivative(y_true, y_pred):
    # dE / da
    # input: N x K
    # output: N x K array
    N = len(y_true)
    return -(y_true / y_pred) / N

def softmax(x):
    # activation (a)
    # input: N x K array
    # output: N x K array
    # https://eli.thegreenplace.net/2016/the-softmax-function-and-its-derivative/
    exp = np.exp(x - np.max(x))
    return exp / np.sum(exp, axis=1)[:, None]

def softmax_derivative(Z):
    # da/dz
    #input: N x K array
    #output: N x K x K array
    #http://saitcelebi.com/tut/output/part2.html
    N, K = Z.shape
    s = softmax(Z)[:, :, np.newaxis]
    a = np.tensordot(s, np.ones((1, K)), axes=([-1],[0]))
    I = np.repeat(np.eye(K, K)[np.newaxis, :, :], N, axis=0)
    b = I - np.tensordot(np.ones((K, 1)), s.T, axes=([-1],[0])).T
    return a * np.swapaxes(b, 1, 2)

def softmax_derivative_test(Z):
    # da/dz
    # non-vectorized softmax gradient calculation
    #http://saitcelebi.com/tut/output/part2.html
    N, K = Z.shape
    da_dz = np.zeros((N, K, K))
    kron_delta = np.eye(K)
    s = softmax(Z)
    for n in range(N):
        for i in range(K):
            for j in range(K):
                da_dz[n, i, j] = s[n, i] * (kron_delta[i, j] - s[n, j])
    return da_dz


def dE_dz_test2(dE_da, da_dz):
    # array (N x K)
    # array (N x K x K)
    # output: array (N x K)
    N, K = dE_da.shape
    dE_dz = np.zeros((N, K))
    for n in range(N):
        dE_dz[n, :] = np.matmul(da_dz[n], dE_da[n, :, np.newaxis]).T
    return dE_dz

def some_type_of_matrix_multiplication_(dE_da, da_dz):
    # how do i get dE/dz from dE_da and da_dz
    pass

X = np.random.rand(100, 2)
W = np.random.rand(2, 4)
y = np.random.randint(0, 4, size=100)
y = one_hot_encode(y, 4)
Z = X @ W
S = softmax(Z)
N, K = Z.shape

# da / dz for softmax
da_dz = softmax_derivative(Z) # (100, 4, 4)
da_dz_test = softmax_derivative_test(Z) # (100, 4, 4) - non vectorized implementation
print(np.isclose(da_dz, da_dz_test).all()) # equivalence test

dE_da = cross_entropy_derivative(y, S) # (100, 4)
dE_dz = some_type_of_matrix_multiplication_(dE_da, da_dz) # what do I do here? *****
dE_dz_test  = (S - y) / N # (100, 4) If you combine dE/da and da/dz terms
dE_dz_test2 = dE_dz_test2(dE_da, da_dz)
print(np.isclose(dE_dz_test, dE_dz_test2).all()) # equivalence test

True

True

这是一种使用 np.einsum 的方法:

def da_dz_pp(z,sm=None):
    if sm is None:
        sm = softmax(z)
    res = np.einsum('ij,ik->ijk',sm,-sm)
    np.einsum('ijj->ij',res)[...] += sm
    return res

def dE_dz_pp(y,z,sm=None):
    if sm is None:
        sm = softmax(z)
    dE_da = cross_entropy_derivative(y,sm)
    da_dz = da_dz_pp(z,sm)
    return np.einsum('ij,ijk->ik',dE_da,da_dz)

它似乎重现了您的代码输出,而且速度更快。