在PyTorch中实施软最近邻损失

数据集的类邻域可以通过软最近邻损失来学习

在本文中，我们讨论实现软最近邻损失的方法，我们在这里也讨论了这个话题。

表示学习是通过深度神经网络学习给定数据集中最显著特征的任务。它通常是在监督学习范 paradigm下进行的一项隐式任务，并且它是深度学习成功的关键因素（Krizhevsky et al., 2012；He et al., 2016；Simonyan et al., 2014）。换句话说，表示学习自动化了特征提取过程。通过这样做，我们可以将学习到的表示用于分类、回归和合成等下游任务。

我们还可以影响学到的表示如何形成以满足特定的应用场景。在分类的情况下，表示被设定为使同一类别的数据点聚集在一起，而在生成（例如GANs）中，表示被设定为使真实数据和合成数据的点聚集在一起。

同样地，我们一直在使用主成分分析（PCA）为下游任务编码特征。然而，在PCA编码表示中没有任何类别或标签信息，因此在下游任务上的性能可能进一步提高。我们可以通过学习数据集的邻域结构来改进编码表示，即哪些特征被聚类在一起，这些聚类会暗示这些特征属于半监督学习文献中的聚类假设所述同一类别（Chapelle et al., 2009）。

为了将邻域结构整合到表示中，引入了流形学习技术，例如局部线性嵌入或LLE（Roweis & Saul, 2000），邻域成分分析或NCA（Hinton et al., 2004），和t-分布邻域嵌入或t-SNE（Maaten & Hinton, 2008）。

然而，前述的流形学习技术都有各自的缺点。例如，LLE和NCA编码的是线性嵌入而不是非线性嵌入。同时，t-SNE嵌入结果依赖于所使用的超参数而产生不同的结构。

为了避免这些缺点，我们可以使用改进的NCA算法，即软最近邻损失或SNNL（Salakhutdinov & Hinton, 2007；Frosst et al., 2019）。SNNL通过引入非线性性能改进NCA算法，并且它是为神经网络的每个隐藏层计算的，而不仅仅是在最后一个编码层上。该损失函数用于优化数据集中的数据点的纠缠。

在这个背景下，纠缠被定义为类似类别的数据点彼此之间的接近程度与不同类别的数据点相比。低纠缠意味着类似类别的数据点之间相互距离要比不同类别的数据点更近（见图1）。拥有这样一组数据点将使下游任务更容易完成，并且性能更好。Frosst等人（2019）通过引入温度因子T扩展了SNNL目标。因此，最终的损失函数如下所示：

其中d是神经网络的原始输入特征或隐藏层表示的距离度量，T是与隐藏层中数据点之间的距离成正比的温度因子。在这个实现中，我们使用余弦距离作为距离度量以更稳定的计算。

本文的目的是帮助读者理解并实现软最近邻损失，因此我们将分析损失函数以更好地理解它。

距离度量

我们首先需要计算数据点之间的距离，这些数据点可以是原始输入特征或网络的隐藏层表示。

在我们的实现中，我们使用余弦距离度量（图3）进行更稳定的计算。暂时忽略上图中的ij和ik表示的子集，我们只关注计算输入数据点之间的余弦距离。我们通过以下PyTorch代码实现：

normalized_a = torch.nn.functional.normalize(features, dim=1, p=2)normalized_b = torch.nn.functional.normalize(features, dim=1, p=2)normalized_b = torch.conj(normalized_b).Tproduct = torch.matmul(normalized_a, normalized_b)distance_matrix = torch.sub(torch.tensor(1.0), product)

在上面的代码片段中，我们首先在第1和第2行通过欧几里得范数对输入特征进行归一化。然后在第3行，我们得到第二组归一化输入特征的共轭转置。我们计算共轭转置以考虑复向量。在第4和第5行，我们计算输入特征的余弦相似度和距离。

具体来说，考虑以下特征集合：

tensor([[ 1.0999, -0.9438,  0.7996, -0.4247],        [ 1.2150, -0.2953,  0.0417, -1.2913],        [ 1.3218,  0.4214, -0.1541,  0.0961],        [-0.7253,  1.1685, -0.1070,  1.3683]])

使用我们上面定义的距离度量，我们得到以下距离矩阵：

tensor([[ 0.0000e+00,  2.8502e-01,  6.2687e-01,  1.7732e+00],        [ 2.8502e-01,  0.0000e+00,  4.6293e-01,  1.8581e+00],        [ 6.2687e-01,  4.6293e-01, -1.1921e-07,  1.1171e+00],        [ 1.7732e+00,  1.8581e+00,  1.1171e+00, -1.1921e-07]])

抽样概率

我们现在可以计算代表选择每个特征的概率的矩阵，给定其与所有其他特征的配对距离。这仅仅是基于i和j或k点之间的距离选择i点的概率。

我们可以通过以下代码计算：

pairwise_distance_matrix = torch.exp(    -(distance_matrix / temperature)) - torch.eye(features.shape[0]).to(model.device)

此代码首先计算距离矩阵除以温度因子后的负指数，将值缩放为正值。温度因子决定如何控制对点对之间距离的重要性，例如，在较低温度下，损失被小距离主导，而实际距离较远的表示之间的距离变得不太相关。

在减去torch.eye(features.shape[0])（也称为对角矩阵）之前，张量如下：

tensor([[1.0000, 0.7520, 0.5343, 0.1698],        [0.7520, 1.0000, 0.6294, 0.1560],        [0.5343, 0.6294, 1.0000, 0.3272],        [0.1698, 0.1560, 0.3272, 1.0000]])

我们从距离矩阵中减去对角矩阵，以去除所有自相似性项（即每个点到自身的距离或相似度）。

接下来，我们可以通过以下代码计算每对数据点的抽样概率：

pick_probability = pairwise_distance_matrix / (    torch.sum(pairwise_distance_matrix, 1).view(-1, 1)    + stability_epsilon)

掩码抽样概率

到目前为止，我们计算的抽样概率不包含任何标签信息。我们通过使用数据集标签对抽样概率进行掩码来将标签信息整合到抽样概率中。

首先，我们必须从标签向量中导出一个配对矩阵：

masking_matrix = torch.squeeze(    torch.eq(labels, labels.unsqueeze(1)).float())

我们应用掩码矩阵来使用标签信息，将属于同一类别的点的概率分离出来：

masked_pick_probability = pick_probability * masking_matrix

接下来，我们通过计算每行的掩码抽样概率之和来计算特定特征的总概率：

summed_masked_pick_probability = torch.sum(masked_pick_probability, dim=1)

最后，我们可以计算抽样概率的对数，为了方便计算，额外引入一个计算稳定性的变量，然后取平均值作为网络的最近邻损失：

snnl = torch.mean(    -torch.log(summed_masked_pick_probability + stability_epsilon)

现在，我们可以将这些组件串联在一起，形成一个前向传递函数，计算深度神经网络的所有层上的软最近邻损失：

def forward(    self,    model: torch.nn.Module,    features: torch.Tensor,    labels: torch.Tensor,    outputs: torch.Tensor,    epoch: int,) -> Tuple:    if self.use_annealing:        self.temperature = 1.0 / ((1.0 + epoch) ** 0.55)    primary_loss = self.primary_criterion(        outputs, features if self.unsupervised else labels    )    activations = self.compute_activations(model=model, features=features)    layers_snnl = []    for key, value in activations.items():        value = value[:, : self.code_units]        distance_matrix = self.pairwise_cosine_distance(features=value)        pairwise_distance_matrix = self.normalize_distance_matrix(            features=value, distance_matrix=distance_matrix        )        pick_probability = self.compute_sampling_probability(            pairwise_distance_matrix        )        summed_masked_pick_probability = self.mask_sampling_probability(            labels, pick_probability        )        snnl = torch.mean(            -torch.log(self.stability_epsilon + summed_masked_pick_probability)        )        layers_snnl.append(snnl)        snn_loss = torch.stack(layers_snnl).sum()    train_loss = torch.add(primary_loss, torch.mul(self.factor, snn_loss))    return train_loss, primary_loss, snn_loss

可视化分离表示

我们使用软最近邻损失训练了一个自动编码器，并可视化了其学习到的分离表示。这个自动编码器具有 (x-500–500–2000-d-2000–500–500-x) 单元，并在MNIST、Fashion-MNIST和EMNIST-Balanced数据集的一个小标记子集上进行了训练。这是为了模拟标记示例的稀缺性，因为自动编码器应该是非监督模型。

为了更容易、更清晰地可视化EMNIST-Balanced数据集，我们只可视化了任意选择的10个簇。从上面的图中我们可以看到，通过聚类色散和正确的簇分配，潜在编码表示变得更加有利于聚类。

结束语

在本文中，我们剖析了软最近邻损失函数，以及如何在PyTorch中实现它。

软最近邻损失最初由Salakhutdinov和Hinton（2007）引入，它用于计算自动编码器的潜在编码（瓶颈）表示的损失，然后将该表示用于下游k最近邻分类任务。

Frosst、Papernot和Hinton（2019）通过引入温度因子以及在神经网络的所有层之间计算损失，扩展了软最近邻损失的应用。

最后，我们使用一个退火温度因子对软最近邻损失进行优化，以进一步改善网络学习到的分离表示，并加快分离过程（Agarap和Azcarraga，2020）。

完整的代码实现可在GitLab上找到。

参考文献

• Agarap, Abien Fred, 和 Arnulfo P. Azcarraga. “Improving k-means clustering performance with disentangled internal representations.” 2020 International Joint Conference on Neural Networks (IJCNN). IEEE, 2020.
• Chapelle, Olivier, Bernhard Scholkopf, 和 Alexander Zien. “Semi-supervised learning (chapelle, o. et al., eds.; 2006)[book reviews].” IEEE Transactions on Neural Networks 20.3 (2009): 542–542.
• Frosst, Nicholas, Nicolas Papernot, 和 Geoffrey Hinton. “Analyzing and improving representations with the soft nearest neighbor loss.” International conference on machine learning. PMLR, 2019.
• Goldberger, Jacob, et al. “Neighbourhood components analysis.” Advances in neural information processing systems. 2005.
• He, Kaiming, et al. “Deep residual learning for image recognition.” Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.
• Hinton, G., et al. “Neighborhood components analysis.” Proc. NIPS. 2004.
• Krizhevsky, Alex, Ilya Sutskever, 和 Geoffrey E. Hinton. “Imagenet classification with deep convolutional neural networks.” Advances in neural information processing systems 25 (2012).
• Roweis, Sam T., 和 Lawrence K. Saul. “Nonlinear dimensionality reduction by locally linear embedding.” science 290.5500 (2000): 2323–2326.
• Salakhutdinov, Ruslan, 和 Geoff Hinton. “Learning a nonlinear embedding by preserving class neighbourhood structure.” Artificial Intelligence and Statistics. 2007.
• Simonyan, Karen, 和 Andrew Zisserman. “Very deep convolutional networks for large-scale image recognition.” arXiv preprint arXiv:1409.1556 (2014).
• Van der Maaten, Laurens, 和 Geoffrey Hinton. “Visualizing data using t-SNE.” Journal of machine learning research 9.11 (2008).