Memory Bank
记忆库会存储先前批次、视频帧或参考图像中的嵌入,使模型能够将新输入与远超单个批次所能容纳数量的示例进行比较。
记忆库用于存储模型先前计算出的特征向量或嵌入,以便将新输入与远超单个批次或单次前向传播所能容纳的示例进行比较。根据所用方法,存储的向量可能来自之前的训练批次、之前的视频帧,或一组正常图像;记忆库通过相似度查询或注意力读取,并在新向量到达时更新。记忆库在计算机视觉的三个领域中都很重要:对比学习用它在每一步提供数千个负样本;视频分割和跟踪用它记住物体在较早帧中的外观;异常检测则用它保存正常数据的特征。
对比学习中的记忆库#
对比式自监督学习会训练编码器,使同一图像的两个视图彼此靠近,并与其他图像拉开距离;通常负样本越多,训练信号越好。要与更多图像比较,通常需要增大批量大小,但这受限于 GPU内存。Zhirong Wu 等人在 CVPR 2018 上发表的论文《通过非参数实例判别进行无监督特征学习》提出了为此用途设计的记忆库:它为 128 万张 ImageNet训练图像中的每一张保存一个 128 维特征;每次再次看到图像时都会刷新对应条目,并且每一步都会从记忆库中抽取数千个负样本。由于条目由较早版本的编码器写入,存储的特征会逐渐过时,与当前模型不再一致。
Facebook AI Research 的 Kaiming He 等人提出的动量对比(MoCo)通过队列和第二个编码器解决了这个问题。记忆库保存最近几个批次的键,默认数量为 65,536;每加入一个新批次,就会移除最旧的批次。键来自动量编码器,其权重会缓慢跟随主编码器:新权重为旧权重的 0.999 倍加上主编码器权重的 0.001 倍,因此不同步骤写入的键彼此保持一致。
视频分割与跟踪中的记忆库#
在视频中,记忆库让模型可以参考较早的帧,以处理遮挡和外观变化。时空记忆网络(STM)会存储过去帧的键和值特征及其掩码,并通过关注这些内容来分割每个新帧;XMem 则通过单独的短期和长期记忆,将这一方法扩展到长视频。Meta 的 SAM 2会将用户提示的帧和最近的帧保存在记忆库中;其记忆注意力会以这些帧为条件处理每个新帧的特征,因此即使物体离开视野后又重新出现,掩码也能继续跟踪它。在 Ultralytics 软件包中,SAM 2默认会关注提示帧以及最近 6 帧;使用 update_memory=True调用 SAM2DynamicInteractivePredictor时,系统会将提示图像中物体的状态加入记忆库,然后根据存储的信息分割后续图像。
多目标跟踪器也会用相同的思路进行重识别:DeepSORT 会为每条轨迹保留最近 100 个外观描述符,并将新的检测结果与已存储描述符最接近的轨迹匹配。
用于异常检测的记忆库#
Karsten Roth 等人在 CVPR 2022 上发表的 PatchCore会构建一个记忆库,存储 ImageNet 预训练 CNN 从无缺陷产品图像中提取的图块特征。测试图块的得分取决于它与最近存储特征之间的距离,而图像得分则是所有图块得分中的最大值,因此任何与正常图块都不相似的内容都会显现出来。贪心核心集下采样只保留一小部分特征,同时仍能覆盖其分布范围;该方法在 MVTec AD 工业缺陷检测基准测试中,图像级 AUROC 最高达到 99.6%。
这个 PyTorch示例将相同方法用于 8x8 手写数字,并使用原始像素作为特征:100 张数字 3 的图像构成记忆库,其他所有数字都视为异常:
import torch
from sklearn.datasets import load_digits
from sklearn.metrics import roc_auc_score
X, y = load_digits(return_X_y=True)
X, y = torch.tensor(X / 16, dtype=torch.float32), torch.tensor(y)
normal = X[y == 3] # images of the digit 3 play the role of defect-free parts
memory_bank = normal[:100] # memory bank: stored features of 100 normal images
test = torch.cat([normal[100:], X[y != 3]]) # unseen 3s plus every other digit as anomalies
is_anomaly = torch.cat([torch.zeros(len(normal) - 100), torch.ones(len(X) - len(normal))])
# Greedy coreset: repeatedly keep the stored feature farthest from everything kept so far
coreset = [0]
distance = torch.cdist(memory_bank, memory_bank[:1])[:, 0]
for _ in range(9):
coreset.append(int(distance.argmax()))
distance = torch.minimum(distance, torch.cdist(memory_bank, memory_bank[coreset[-1:]])[:, 0])
for name, bank in [("Full memory bank, 100 entries", memory_bank), ("Coreset, 10 entries", memory_bank[coreset])]:
scores = torch.cdist(test, bank).min(dim=1).values # anomaly score: distance to the nearest stored feature
print(f"{name}: AUROC {roc_auc_score(is_anomaly, scores):.3f}")在 1,697 张测试图像中,完整记忆库以 0.984 的ROC 曲线下面积(AUC)区分出 83 张未见过的数字 3 与其他数字;包含 10 个条目的核心集得分为 0.982,只使用十分之一的内存和查询次数,几乎保留了全部检测能力。
记忆库与相关概念的区别#
- 向量数据库:持久化存储嵌入,并为许多应用提供相似度搜索。记忆库位于单个模型的训练或推理循环内部,通常保存在 GPU 内存中,并会随着模型运行而被覆盖。
- 上下文窗口:Transformer 在一次处理中关注的词元。记忆库存储窗口之外的向量,这些向量来自其他批次、帧或图像。
- Transformer-XL:缓存前一个文本片段的隐藏状态,作为下一个片段可以关注的记忆;这是一种激活记忆库,可将语言模型的处理范围扩展到单个窗口之外。
- 困难负样本挖掘:挑选最容易混淆的负样本用于训练。记忆库提供了一个大型负样本池,供此类筛选方法使用。
- 情景记忆:在 AI 智能体和强化学习中,存储过去的经验并在需要时检索,以指导决策。记忆库存储模型自身计算所用的向量,例如负样本、过去的帧或正常图块。










