Memory Bank

Банк памяти — это хранилище векторов признаков, или эмбеддингов, которые модель вычислила ранее. Он нужен, чтобы сравнивать новые входные данные с гораздо большим числом примеров, чем можно разместить в одном пакете или обработать за один прямой проход. В зависимости от метода сохранённые векторы берутся из предыдущих обучающих пакетов, предыдущих кадров видео или эталонного набора обычных изображений. Банк считывается с помощью поиска по сходству или механизма внимания и обновляется по мере поступления новых векторов. Банки памяти широко используются в трёх областях компьютерного зрения: контрастивном обучении, где они предоставляют тысячи отрицательных примеров на каждом шаге; сегментации видео и отслеживании объектов, где сохраняют информацию о том, как объект выглядел в предыдущих кадрах; и обнаружении аномалий, где хранят сведения о нормальных данных.

Банки памяти в контрастивном обучении#

Контрастивное самообучение обучает кодировщик сближать два представления одного изображения и отдалять от них представления других изображений; чем больше отрицательных примеров, тем полезнее обучающий сигнал. Обычно для сравнения с большим числом изображений нужен больший размер пакета, ограниченный объёмом памяти GPU. Жиронг Ву и его коллеги предложили использовать для этой цели банк памяти в работе Unsupervised Feature Learning via Non-Parametric Instance Discrimination, представленной на CVPR 2018: в нём хранился 128-мерный признак для каждого из 1,28 миллиона обучающих изображений ImageNet. Каждая запись обновлялась при повторной обработке соответствующего изображения, а на каждом шаге из банка выбирались тысячи отрицательных примеров. Поскольку записи создавались более ранними версиями кодировщика, сохранённые признаки постепенно устаревали и переставали соответствовать текущей модели.

Метод Momentum Contrast (MoCo), разработанный Каймином Хэ и его коллегами из Facebook AI Research, решил эту проблему с помощью очереди и второго кодировщика. В банке хранятся ключи из последних пакетов — по умолчанию 65 536; каждый новый пакет добавляется в очередь, а самый старый удаляется. Ключи поступают от кодировщика с моментумом, веса которого медленно следуют за весами основного кодировщика: новое значение вычисляется как 0,999 от старого значения плюс 0,001 от веса основного кодировщика. Благодаря этому ключи, записанные на разных шагах, остаются согласованными.

Банки памяти для сегментации видео и отслеживания объектов#

В видео банк памяти помогает модели справляться с перекрытиями и изменением внешнего вида, используя информацию из предыдущих кадров. Сети пространственно-временной памяти (STM) хранили ключевые и значимые признаки прошлых кадров вместе с масками и сегментировали каждый новый кадр с помощью внимания к ним; XMem расширил этот подход на длинные видео, добавив отдельные краткосрочную и долгосрочную памяти. SAM 2 от Meta хранит в банке памяти кадры, для которых пользователь задавал запрос, а также самые недавние кадры. Механизм внимания к памяти связывает с ними признаки каждого нового кадра — так маска следует за объектом, который покидает поле зрения, а затем появляется снова. В пакете Ultralytics SAM 2 по умолчанию учитывает кадры, для которых задан запрос, и 6 последних кадров. При вызове SAM2DynamicInteractivePredictor с параметром update_memory=True состояние объекта на заданном изображении добавляется в банк памяти, после чего последующие изображения сегментируются с опорой на сохранённую информацию.

Трекеры нескольких объектов используют ту же идею для повторной идентификации: DeepSORT хранит галерею дескрипторов внешнего вида для каждого трека — до 100 последних — и сопоставляет новые детекции с треком, дескрипторы которого ближе всего.

Банки памяти для обнаружения аномалий#

PatchCore, представленный Карстеном Ротом и его коллегами на CVPR 2022, создаёт банк признаков фрагментов, которые CNN, предобученная на ImageNet, извлекает из изображений изделий без дефектов. Оценка тестового фрагмента определяется расстоянием до ближайшего сохранённого признака, а оценка изображения равна наибольшей оценке его фрагментов. Поэтому выделяется всё, что не похоже ни на один нормальный фрагмент. Жадная выборка ядра позволяет сохранить лишь небольшую долю признаков, охватывающую при этом весь диапазон их значений; на промышленном бенчмарке обнаружения дефектов MVTec AD метод достигал AUROC до 99,6% на уровне изображения.

В этом примере на PyTorch тот же метод применяется к рукописным цифрам размером 8×8, а в качестве признаков используются исходные пиксели: банк памяти формируют 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 тестовых изображениях полный банк отделяет 83 ранее не встречавшиеся цифры 3 от остальных цифр с площадью под ROC-кривой (AUC) 0.984, а оценка для ядра из 10 записей составляет 0.982. Таким образом сохраняется почти всё качество обнаружения при сокращении объёма памяти и числа поисковых запросов в десять раз.

Банк памяти и смежные понятия#

• Векторная база данных: Постоянно хранит эмбеддинги и предоставляет поиск по сходству множеству приложений. Банк памяти находится внутри цикла обучения или вывода одной модели, обычно в памяти GPU, и перезаписывается по мере работы модели.
• Контекстное окно: Токены, к которым трансформер обращается за один проход. Банк памяти хранит векторы за пределами этого окна — из других пакетов, кадров или изображений.
• Transformer-XL: Кэширует скрытые состояния предыдущего текстового сегмента, чтобы следующий сегмент мог к ним обратиться. Это банк памяти активаций, который позволяет языковой модели работать за пределами одного окна.
• Поиск сложных отрицательных примеров: Отбирает наиболее запутанные отрицательные примеры для обучения. Банк памяти предоставляет большой пул отрицательных примеров, из которого можно выбирать.
• Эпизодическая память: В ИИ-агентах и обучении с подкреплением — хранилище прошлых событий, к которым обращаются для принятия решений. Банки памяти хранят векторы, необходимые для вычислений самой модели, например отрицательные примеры, прошлые кадры или нормальные фрагменты.