pytorch 如何通过 argmax 反向传播？

Posted 2023-02-16

【中文标题】pytorch 如何通过 argmax 反向传播？

【英文标题】：How does pytorch backprop through argmax?

【发布时间】：2019-07-24 23:06:42

【问题描述】：

我在 pytorch 中使用质心位置的梯度下降而不是期望最大化来构建 Kmeans。损失是每个点到其最近质心的平方距离之和。为了确定离每个点最近的质心，我使用了 argmin，它在任何地方都不可微。但是，pytorch 仍然能够反向传播和更新权重（质心位置），在数据上提供与 sklearn kmeans 相似的性能。

任何想法这是如何工作的，或者我如何在 pytorch 中解决这个问题？对 pytorch github 的讨论表明 argmax 不可微：https://github.com/pytorch/pytorch/issues/1339。

下面的示例代码（在随机点上）：

import numpy as np
import torch

num_pts, batch_size, n_dims, num_clusters, lr = 1000, 100, 200, 20, 1e-5

# generate random points
vector = torch.from_numpy(np.random.rand(num_pts, n_dims)).float()

# randomly pick starting centroids
idx = np.random.choice(num_pts, size=num_clusters)
kmean_centroids = vector[idx][:,None,:] # [num_clusters,1,n_dims]
kmean_centroids = torch.tensor(kmean_centroids, requires_grad=True)

for t in range(4001):
    # get batch
    idx = np.random.choice(num_pts, size=batch_size)
    vector_batch = vector[idx]

    distances = vector_batch - kmean_centroids # [num_clusters, #pts, #dims]
    distances = torch.sum(distances**2, dim=2) # [num_clusters, #pts]

    # argmin
    membership = torch.min(distances, 0)[1] # [#pts]

    # cluster distances
    cluster_loss = 0
    for i in range(num_clusters):
        subset = torch.transpose(distances,0,1)[membership==i]
        if len(subset)!=0: # to prevent NaN
            cluster_loss += torch.sum(subset[:,i])

    cluster_loss.backward()
    print(cluster_loss.item())

    with torch.no_grad():
        kmean_centroids -= lr * kmean_centroids.grad
        kmean_centroids.grad.zero_()

【问题讨论】：

Argmax 是不可微分的。但是您可以尝试一些技巧，例如homes.cs.washington.edu/~hapeng/paper/peng2018backprop.pdf，该论文还引用了类似思路中的其他一些工作，试图通过某种 argmax/sparsemax 进行反向传播。免责声明：我个人没有处理过此类问题。

【参考方案1】：

正如 alvas 在 cmets 中指出的那样，argmax 是不可微分的。然而，一旦你计算它并将每个数据点分配给一个集群，损失相对于这些集群位置的导数是明确定义的。这就是你的算法所做的。

为什么有效？如果您只有一个集群（因此 argmax 操作无关紧要），您的损失函数将是二次的，最小值为数据点的平均值。现在有多个集群，您可以看到您的损失函数是分段的（在更高维度上考虑体积） - 对于任何一组质心[C1, C2, C3, ...]，每个数据点都分配给某个质心CN，并且损失是本地的 二次方。该局部性的范围由所有替代质心[C1', C2', C3', ...] 给出，来自argmax 的分配保持不变；在这个区域内，argmax 可以被视为一个常数，而不是一个函数，因此 loss 的导数是明确定义的。

现在，实际上，您不太可能将 argmax 视为常数，但您仍然可以将幼稚的“argmax-is-a-constant”导数视为大致指向最小值，因为大多数数据点都是在迭代之间可能确实属于同一个集群。一旦你足够接近一个局部最小值，使得这些点不再改变它们的分配，这个过程就可以收敛到一个最小值。

另一种更具理论性的看待它的方式是，您正在对期望最大化进行近似。通常，您将有“计算分配”步骤，这由argmax 反映，而“最小化”步骤归结为在给定当前分配的情况下找到最小化集群中心。最小值由d(loss)/d([C1, C2, ...]) == 0 给出，对于二次损失，它是通过每个集群内的数据点分析给出的。在您的实现中，您正在求解相同的方程，但使用梯度下降步骤。事实上，如果您使用二阶（牛顿）更新方案而不是一阶梯度下降，您将隐式复制准确的基线 EM 方案。

【讨论】：

感谢您的精彩回答，与 EM 的比较很有帮助。所以因为没有通过 argmin 的反向传播，所以当 pytorch 反向传播损失到质心位置时，集群分配被视为常量？

是的。一般来说，非浮点张量不能携带.grad参数。这意味着membership 不能被反向传播并被视为常量。

【参考方案2】：

想象一下：

t = torch.tensor([-0.0627,  0.1373,  0.0616, -1.7994,  0.8853, 
                  -0.0656,  1.0034,  0.6974,  -0.2919, -0.0456])
torch.argmax(t).item() # outputs 6

我们增加t[0] 一些，δ 接近 0，这会更新 argmax 吗？它不会，所以我们一直在处理 0 梯度。忽略这一层，或者假设它被冻结了。

argmin 或任何其他因变量为离散步长的函数也是如此。