1. ClusterLookup

方法来源：STEGO: Unsupervised Semantic Segmentation by Distilling Feature Correspondences

class ClusterLookup(nn.Module):

    def __init__(self, dim: int, n_classes: int):
        super(ClusterLookup, self).__init__()
        self.n_classes = n_classes
        self.dim = dim
        self.clusters = torch.nn.Parameter(torch.randn(n_classes, dim))

    def reset_parameters(self):
        with torch.no_grad():
            self.clusters.copy_(torch.randn(self.n_classes, self.dim))

    def forward(self, x, alpha, log_probs=False):
        normed_clusters = F.normalize(self.clusters, dim=1)
        normed_features = F.normalize(x, dim=1)
        inner_products = torch.einsum("bchw,nc->bnhw", normed_features, normed_clusters)

        if alpha is None:
            cluster_probs = F.one_hot(torch.argmax(inner_products, dim=1), self.clusters.shape[0]) \
                .permute(0, 3, 1, 2).to(torch.float32)
        else:
            cluster_probs = nn.functional.softmax(inner_products * alpha, dim=1)

        cluster_loss = -(cluster_probs * inner_products).sum(1).mean()
        if log_probs:
            return nn.functional.log_softmax(inner_products * alpha, dim=1)
        else:
            return cluster_loss, cluster_probs

摘要：无监督语义分割的目的是在没有任何形式的注释的情况下，将具有语义意义的类别发现和定位。为了解决这个任务，算法必须为每个像素产生既具有语义意义又足够紧凑的特征，以形成不同的簇。与以前使用单一的端到端框架来实现这一点的工作不同，我们建议将特征学习从簇集群紧化中分离出来。根据经验，我们表明，当前的无监督特征学习框架已经产生了密集的特征，其相关性是语义一致的。这一观察结果促使我们设计STEGO（基于能量的图优化的自监督变换器，Self-supervised Transformer with Energy-based Graph Optimization），这是一个新的框架，将无监督特征提取为高质量的离散语义标签。STEGO的核心是一种新的对比损失函数，它鼓励特征形成紧凑的集群，同时保持它们在整个语料库中的关系。STEGO在CocoStuff（+14 mIoU）和城市景观（+9 mIoU）语义分割挑战上，都比之前的技术水平有了显著的改进。

引言：与之前的方法相比，我们利用了无监督特征学习框架中的预训练特征，并专注于将它们提取成一个紧凑和离散的结构，同时保持它们在图像语料库中的关系。这是由于观察到无监督特征之间的相关性，如DINO学习到的特征（Caron et al.，2021），在同一图像内和跨图像集合的语义上已经是一致的。

因此，我们引入了STEGO（基于能量的自监督变压器），它能够在没有人工监督的情况下联合发现和分割对象。STEGO利用一种新的对比损失，将预先训练过的无监督视觉特征提取为语义簇。STEGO大大改进了现有技术，是缩小与监督分割系统的差距的相当大的一步。我们包括一个简短的视频，详细介绍了在https://aka.ms/stego-video的工作。具体来说，我们做出了以下贡献：

结果表明，无监督深度网络特征与真实语义标签基本一致。
介绍了一种新的基于转换器的无监督语义分割架构STEGO。
证明STEGO在协同（+14 mIoU）和城市景观（+9 mIoU）分割挑战上都取得了最先进的性能。
通过对CocoStuff数据集的消融研究来证明STEGO的设计。

方法：

基于特征对应关系预测类的共现性

自我监督视觉特征学习的最新进展已经产生了具有强大的和语义相关的特征的方法，从而改进了各种下游任务。虽然大多数研究的目标是为图像生成单个向量，但许多研究表明，中间密集的特征是语义相关的（汉密尔顿等人，2021；Collins等人，2018；Zhou等人，2016）。为了使用这些信息，我们将重点放在密集特征图之间的“相关性体积”（Teed & Deng，2020）上。对于卷积架构或转换器架构，这些密集的特征映射可以是特定层的激活映射。此外，变压器中的Q、K或V矩阵也可以作为候选特征，尽管我们发现这些注意张量在实践中表现得不太好。更正式地说，假设f∈R CHW，g∈R CIJ是两种不同图像的特征张量，其中C表示通道维数，（H，W），（I，J）表示空间维数。我们形成了特征对应张量： \[F_{hwij}:=\sum_{c}\frac{f_{chw}}{|f_{hw}|}\frac{g_{cij}}{|g_{ij}|},\] 其项表示特征张量f的空间位置（h、w）处的特征与特征张量g的位置（i、j）之间的余弦相似度。在特殊情况下，这些对应度量同一图像的两个区域之间的相似性。