Memory Bank

Memory bank là kho lưu trữ các vector feature, hay embedding, mà model đã tính toán trước đó, được giữ lại để so sánh đầu vào mới với nhiều ví dụ hơn số lượng có thể chứa trong một batch hoặc một lượt forward. Tùy phương pháp, các vector được lưu trữ có thể đến từ batch huấn luyện trước đó, frame video trước đó hoặc một tập tham chiếu gồm ảnh bình thường; memory bank được truy xuất bằng tra cứu độ tương đồng hoặc attention và được cập nhật khi có vector mới. Memory bank giữ vai trò trọng tâm trong ba lĩnh vực thị giác máy tính: học tương phản, nơi chúng cung cấp hàng nghìn ví dụ âm tính mỗi bước; phân đoạn video và tracking, nơi chúng ghi nhớ diện mạo của vật thể trong các frame trước; và phát hiện bất thường, nơi chúng lưu trữ đặc trưng của dữ liệu bình thường.

Memory bank trong học tương phản#

Học tự giám sát tương phản huấn luyện encoder để kéo hai góc nhìn của cùng một hình ảnh lại gần nhau và đẩy các hình ảnh khác ra xa; nhìn chung, càng có nhiều mẫu âm tính thì tín hiệu huấn luyện càng tốt. Việc so sánh với nhiều hình ảnh hơn thường đòi hỏi batch size lớn hơn, nhưng bộ nhớ GPU lại có giới hạn. Zhirong Wu và các cộng sự giới thiệu memory bank cho mục đích này trong bài báo Học feature không giám sát thông qua phân biệt instance phi tham số tại CVPR 2018: memory bank lưu feature 128 chiều cho từng ảnh trong số 1,28 triệu ảnh huấn luyện ImageNet, làm mới từng mục mỗi khi ảnh tương ứng xuất hiện lại và lấy hàng nghìn mẫu âm tính từ đó ở mỗi bước. Do các mục được ghi bởi những phiên bản encoder cũ hơn, feature lưu trữ dần lỗi thời và không còn nhất quán với model hiện tại.

Momentum Contrast (MoCo) của Kaiming He và các cộng sự tại Facebook AI Research khắc phục vấn đề này bằng một hàng đợi và encoder thứ hai. Memory bank lưu các key của những batch gần nhất, mặc định là 65.536 key; mỗi batch mới được đưa vào hàng đợi và batch cũ nhất bị loại bỏ. Key đến từ encoder momentum, có trọng số thay đổi chậm theo encoder chính: cập nhật bằng 0.999 lần giá trị cũ cộng với 0.001 lần trọng số của encoder chính, nhờ đó các key được ghi ở những bước khác nhau vẫn nhất quán với nhau.

Memory bank trong phân đoạn video và tracking#

Trong video, memory bank giúp model xử lý hiện tượng che khuất và thay đổi diện mạo bằng cách tham chiếu các frame trước đó. Mạng Space-Time Memory (STM) lưu feature key và value của các frame trước cùng mask của chúng, rồi phân đoạn mỗi frame mới bằng cách attention đến các thông tin đó; XMem mở rộng phương pháp này cho video dài với bộ nhớ ngắn hạn và dài hạn riêng biệt. SAM 2 của Meta lưu memory bank gồm các frame được người dùng gợi ý và các frame gần nhất; memory attention điều kiện hóa feature của mỗi frame mới dựa trên những frame đó, nhờ vậy mask có thể bám theo vật thể khi vật thể rời khỏi rồi xuất hiện lại trong khung hình. Trong package Ultralytics, SAM 2 mặc định attention đến các frame được gợi ý cộng với 6 frame gần nhất; SAM2DynamicInteractivePredictor thêm trạng thái đối tượng của ảnh được gợi ý vào memory bank khi được gọi với update_memory=True, sau đó phân đoạn các ảnh tiếp theo dựa trên thông tin đã lưu.

Tracker đa đối tượng dùng cùng ý tưởng để tái nhận diện: DeepSORT lưu thư viện gồm 100 descriptor diện mạo gần nhất của mỗi track và ghép phát hiện mới với track có descriptor lưu trữ gần nhất.

Memory bank cho phát hiện bất thường#

PatchCore, do Karsten Roth và các cộng sự trình bày tại CVPR 2022, xây dựng memory bank gồm feature patch được một CNN tiền huấn luyện trên ImageNet trích xuất từ ảnh sản phẩm không có lỗi. Điểm của patch kiểm tra được xác định bằng khoảng cách đến feature lưu trữ gần nhất; điểm ảnh là điểm patch cao nhất, nhờ đó mọi thứ khác biệt với tất cả patch bình thường đều nổi bật. Lấy mẫu coreset tham lam chỉ giữ lại một phần nhỏ feature nhưng vẫn bao phủ phạm vi phân bố của chúng; phương pháp đạt AUROC cấp ảnh lên tới 99.6% trên benchmark phát hiện lỗi công nghiệp MVTec AD.

Ví dụ PyTorch này áp dụng cách làm tương tự cho các chữ số viết tay 8x8, dùng pixel thô làm feature: 100 ảnh chữ số 3 tạo thành memory bank, còn mọi chữ số khác được xem là bất thường:

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}")

Trên 1.697 ảnh kiểm tra, memory bank đầy đủ phân biệt 83 chữ số 3 chưa từng thấy với các chữ số khác, đạt diện tích dưới đường cong ROC (AUC) là 0.984; coreset 10 mục đạt 0.982, giữ được gần như toàn bộ chất lượng phát hiện với chỉ một phần mười bộ nhớ và số lượt tra cứu.

Memory bank và các khái niệm liên quan#

• Cơ sở dữ liệu vector: Lưu embedding lâu dài và cung cấp chức năng tìm kiếm tương đồng cho nhiều ứng dụng. Memory bank nằm bên trong vòng lặp huấn luyện hoặc suy luận của một model, thường ở bộ nhớ GPU, và bị ghi đè khi model chạy.
• Cửa sổ ngữ cảnh: Các token mà transformer attention trong một lượt xử lý. Memory bank lưu vector bên ngoài cửa sổ đó, lấy từ batch, frame hoặc hình ảnh khác.
• Transformer-XL: Lưu trạng thái ẩn của đoạn văn bản trước đó làm bộ nhớ để đoạn tiếp theo attention đến; đây là memory bank các activation giúp mở rộng phạm vi của model ngôn ngữ vượt quá một cửa sổ.
• Khai thác mẫu âm tính khó: Chọn các mẫu âm tính dễ gây nhầm lẫn nhất để huấn luyện. Memory bank cung cấp một tập lớn mẫu âm tính để phương pháp lựa chọn này khai thác.
• Bộ nhớ theo tập: Trong agent AI và học tăng cường, đây là nơi lưu trữ các trải nghiệm trước đó được truy xuất để định hướng quyết định. Memory bank lưu vector phục vụ tính toán của chính model, chẳng hạn như mẫu âm tính, frame trước đó hoặc patch bình thường.