谷歌新作：面向灵活感知的视觉记忆

Towards flexible perception with visual memory

https://arxiv.org/pdf/2408.08172

问题：Training a neural network is a monolithic endeavor, akin to carving knowledge into stone: once the process is completed, editing the knowledge in a network is nearly impossible, since all information is distributed across the network’s weights.

核心：We here explore a simple, compelling alternative by marrying the representational power of deep neural networks with the flexibility of a database.

方法：Decomposing the task of image classification into image similarity (from a pre-trained embedding) and search (via fast nearest neighbor retrieval from a knowledge database), we build a simple and flexible visual memory that has the following key capabilities:

(1.) The ability to flexibly add data across scales: from individual samples all the way to entire classes and billion-scale data;

(2.) The ability to remove data through unlearning and memory pruning;

(3.) An interpretable decision-mechanism on which we can intervene to control its behavior.

Taken together, these capabilities comprehensively demonstrate the benefits of an explicit visual memory.

We hope that it might contribute to a conversation on how knowledge should be represented in deep vision models—beyond carving it in “stone” weights.

训练神经网络是一项整体性的工作，类似于将知识刻入石头：一旦这个过程完成，由于所有信息都分布在网络的权重中，在网络中编辑知识几乎是不可能的。

本文探索了一种简单而有力的替代方案，即将深度神经网络的表示能力与数据库的灵活性结合起来。

通过将图像分类任务分解为图像相似性（基于预训练嵌入）和搜索（通过从知识数据库快速检索最近邻），构建了一个简单且灵活的视觉记忆系统，该系统具有以下关键功能：

（1）能够灵活地跨尺度添加数据：从单个样本到整个类别乃至十亿级的数据；

（2）通过遗忘和内存修剪来删除数据的能力；

（3）可解释的决策机制，可以对其进行干预以控制其行为。

综合这些功能，全面展示了显式视觉记忆的优势。

希望这有助于推动关于如何在深度视觉模型中表示知识的讨论——而不仅仅是将其“刻入”石头般的权重中。

构建基于检索的视觉记忆进行分类

2.1 构建视觉记忆

模型与特征提取：视觉记忆通过从预训练图像编码器中提取特征来构建，这些编码器对图像数据集进行编码，生成特征映射。这些特征映射及其对应的标签对存储在数据库中，形成视觉记忆。文中使用的编码器包括DinoV2和CLIP，它们分别提供了不同层级的特征表示。

特征空间选择：对于DinoV2和CLIP，选择它们的最后一层图像嵌入层作为特征空间。这是因为这些层能够捕捉到丰富的图像表示，适用于后续的分类任务。

数据集与特征存储：给定一个包含（图像，标签）对的训练数据集，使用预训练的图像编码器提取特征，并将这些特征及其标签存储在数据库中，形成视觉记忆。例如，使用ImageNet-1K训练集提取特征并存储。

2.2 使用视觉记忆进行检索式分类

查询图像特征提取：对于给定的查询图像，使用相同的预训练图像编码器提取其特征映射。

最近邻检索：在视觉记忆中，使用余弦距离作为相似性度量，检索与查询图像特征最接近的k个邻居（即（特征，标签）对）。

邻居聚合与分类：为每个邻居分配一个权重，并根据标签聚合分数。最终将最高分数的标签分配给查询图像作为预测结果。文中提出了多种聚合方法，包括等权重投票（Plurality Voting）、距离投票（Distance Voting）、Softmax投票（Softmax Voting）和Rank投票（Rank Voting）。

Rank投票方法：提出了一种新的、简单的聚合方法Rank投票，其中邻居的权重基于其在排序集中的排名。实验结果表明，该方法在性能上优于其他基线方法，特别是在邻居数量较多时。

高效检索实现：对于小数据集如ImageNet，使用矩阵乘法在GPU/TPU上进行快速推理。对于大规模数据集，使用ScaNN进行快速且可扩展的最近邻搜索，以支持十亿级别的视觉记忆。

视觉记忆的能力

3.1. 灵活的终身学习：添加新的OOD类别（包括数据和标签）

无论是端到端训练（监督模型）还是使用线性分类器（自监督模型）的标准分类器，在不重新训练的情况下处理新信息都存在困难。在现有模型中添加新类别或更改标签通常需要重新训练或微调模型的部分。然而，基于检索的视觉记忆能够以自然和灵活的方式处理这些变化，符合终身学习的要求。

作者通过在预训练的DinoV2ViT-L14模型的视觉记忆中添加来自NINCO数据集的64个新类别的数据及其标签（除了ImageNet训练集之外），来展示这一点。NINCO是一个专用的分布外（OOD）数据集，设计用于与现有的ImageNet标签和样本没有重叠。该实验展示了模型能够在不重新训练或微调的情况下，将其知识转移到新的、未见过的概念上。

3.2. 灵活地在计算和内存之间进行权衡

基于检索的视觉记忆允许在计算和内存资源之间进行灵活的权衡。根据当前任务，可以在推理过程中从视觉记忆中检索的邻居数量进行调整，以平衡准确性和计算时间。例如，使用更多的邻居可以提高分类性能，但会增加计算成本。相反，使用较少的邻居可以减少计算需求，但可能会牺牲一些准确性。

作者认为，这种灵活性对于需要在性能和资源之间进行权衡的实际应用来说是一项重要的能力。他们通过分析不同聚合方法实现的准确性-计算权衡，并展示了他们提出的RankVoting方法实现了有利的权衡，来证明了这一点。

3.3. 无需训练即可灵活添加十亿级数据

基于检索的视觉记忆的一个关键优势是其能够扩展到大型数据集而无需额外训练。随着视觉记忆规模的增长，只需从预训练的编码器中提取特征并将其添加到数据库中即可添加新数据。这消除了使用整个数据集重新训练模型的必要性，对于十亿级数据集来说，这可能会非常昂贵。

作者通过使用ScaNN框架进行大规模快速最近邻搜索来展示这一能力。他们表明，基于检索的视觉记忆可以轻松扩展到十亿级数据，使其能够集成和利用大量信息而无需重新训练模型。这一特性对于需要不断更新知识库的应用特别有用，如产品分类或医学图像分析。

3.4. 灵活地删除数据:机器学习

视觉记忆系统如何灵活地移除数据，实现机器遗忘的功能。这是当前深度学习模型所缺乏的一个重要能力，因为一旦模型被训练，其内部的知识（通过权重分布）就很难被修改或删除。然而，基于检索的视觉记忆系统提供了一种简单而有效的解决方案。

灵活性：检索式视觉记忆系统能够直接从记忆数据库中移除特定数据（图像和标签对），无需重新训练整个模型。这允许模型在实际应用中更加灵活地适应变化，如移除包含偏见、不安全内容或已过时的数据。

完美遗忘：实验表明，这种机制能够实现“完美遗忘”，即模型在移除相关数据后，不会在新任务中受到这些数据的影响。这是通过简单地从视觉记忆数据库中删除对应的条目来实现的，不会影响其他数据的表示。

应用前景：机器遗忘能力对于隐私保护、内容审核等应用场景至关重要。检索式视觉记忆系统为这些需求提供了一种可行的技术路径。

3.5. 灵活的数据选择:内存修剪

内存修剪：内存修剪是一种通过移除对模型性能贡献较小的数据来优化内存使用的方法。在检索式视觉记忆系统中，可以通过分析每个数据点的贡献度（例如，通过其被检索的频率或对新任务性能的影响）来执行内存修剪。

实验验证：文档通过实验验证了内存修剪的有效性。通过移除视觉记忆数据库中贡献较小的数据点，不仅减少了内存使用，还保持了甚至提高了模型的分类性能。这表明内存修剪是一种有效的数据选择策略。

灵活性与控制：内存修剪允许研究人员和开发者更灵活地控制视觉记忆系统的内容，使其更加专注于对模型性能有重要影响的数据。这种能力为模型的优化和调试提供了更大的灵活性和控制力。

3.6 灵活扩展与压缩

灵活扩展能力：视觉记忆能够灵活地在不同规模上增加数据。这包括从单个样本到整个类别，甚至是扩展到数十亿规模的数据集。例如，通过使用高效的最近邻搜索库（如ScaNN），可以轻松地将视觉记忆扩展到包含数亿个样本的数据库，而无需对模型进行重新训练。这种能力使得视觉记忆能够快速适应新数据，满足不断变化的需求。

灵活压缩能力：与扩展能力相对应，视觉记忆同样支持通过删除或合并样本来实现压缩。在需要减小内存占用或处理冗余信息时，可以通过删除不相关的样本或使用聚类等技术来合并相似的样本，从而有效地压缩视觉记忆的大小。这种压缩能力有助于降低存储成本和提高检索效率。

3.7 可解释性与可控性

可解释性提升：视觉记忆提供了一种直观且可解释的分类机制。通过检查视觉记忆中用于分类的最近邻样本，用户可以理解模型为何做出特定决策。这种透明度有助于建立对模型行为的信任，并允许对错误分类进行调试。此外，通过可视化最近邻样本，研究人员可以洞察模型如何捕获数据中的模式和异常。

可控性增强：视觉记忆还赋予了对模型行为的精细控制能力。通过编辑或删除视觉记忆中的特定样本，可以直接影响模型的预测结果。例如，在出现数据泄露或偏见问题时，可以从视觉记忆中删除相关样本，从而消除潜在的负面影响。此外，通过调整最近邻检索算法或聚合策略的参数，可以进一步控制模型的性能和灵敏度。这种可控性使得视觉记忆成为一种灵活且强大的工具，能够应对各种复杂的应用场景。