Memory Bank

メモリバンクとは、モデルが以前に計算した特徴ベクトル、つまり埋め込みを保存する領域です。新しい入力を、1バッチや1回のフォワードパスに収まるよりも多くの例と比較できるように保持します。手法によって、保存ベクトルは以前の学習バッチ、以前の動画フレーム、または正常画像の参照セットから得られます。類似度検索やアテンションで読み出し、新しいベクトルの到着に応じて更新します。メモリバンクはコンピュータービジョンの3つの分野で重要です。対照学習では1ステップあたり何千もの負例を提供し、動画セグメンテーションとトラッキングでは以前のフレームにおける物体の見え方を記憶し、異常検知では正常データの特徴を保持します。

対照学習におけるメモリバンク#

対照的な自己教師あり学習では、同じ画像の2つのビューを近づけ、別の画像を遠ざけるようエンコーダーを学習します。一般に、負例が多いほど学習信号が向上します。比較対象を増やすには通常バッチサイズを大きくする必要がありますが、GPUメモリが制約となります。Zhirong Wuらは、CVPR 2018でUnsupervised Feature Learning via Non-Parametric Instance Discriminationを発表し、この目的のためのメモリバンクを導入しました。この手法では、ImageNetの学習画像128万枚それぞれについて128次元の特徴を保持し、画像が再び現れるたびに対応する項目を更新して、各ステップで何千もの負例を取り出しました。項目は古いバージョンのエンコーダーによって書き込まれるため、保存された特徴は古くなり、現在のモデルとの一貫性が失われました。

Facebook AI ResearchのKaiming Heらによるモメンタムコントラスト（MoCo）は、キューと2つ目のエンコーダーを使ってこの問題を解決しました。デフォルトで最新のバッチのキー65,536個を保持し、新しいバッチを追加するたびに最も古いものを削除します。キーはモメンタムエンコーダーから生成されます。その重みはメインエンコーダーにゆっくり追従し、古い値の0.999倍にメインエンコーダーの重みの0.001倍を加えて更新されるため、異なるステップで書き込まれたキー間の一貫性が保たれます。

動画セグメンテーションとトラッキングにおけるメモリバンク#

動画では、メモリバンクを使って以前のフレームを参照することで、モデルがオクルージョンや外観の変化に対応できます。時空間メモリネットワーク（STM）は、過去のフレームのキー特徴とバリュー特徴をマスクとともに保存し、それらにアテンションを適用して各新規フレームをセグメント化しました。XMemは、短期メモリと長期メモリを分けることで長時間の動画に対応しました。MetaのSAM 2は、ユーザーがプロンプトを与えたフレームと直近のフレームをメモリバンクに保持します。メモリアテンションは新規フレームの特徴をこれらのフレームに基づいて条件付けるため、物体が視界から出て再び入ったときにもマスクを追従させられます。Ultralyticsパッケージでは、SAM 2はデフォルトでプロンプトを与えたフレームと直近6フレームにアテンションを適用します。また、SAM2DynamicInteractivePredictorをupdate_memory=Trueとともに呼び出すと、プロンプトを与えた画像の物体状態がメモリバンクに追加され、その後の画像を保存済みの情報に基づいてセグメント化します。

マルチオブジェクトトラッカーも、再識別に同じ考え方を使います。DeepSORTは各トラックの直近100個の外観記述子のギャラリーを保持し、新しい検出を、保存された記述子が最も近いトラックに照合します。

異常検知におけるメモリバンク#

Karsten RothらがCVPR 2022で発表したPatchCoreは、欠陥のない製品画像からImageNetで事前学習したCNNが抽出したパッチ特徴のメモリバンクを構築します。テスト時のパッチは、保存済みの最も近い特徴との距離でスコア付けされます。画像スコアには最も悪いパッチスコアが使われるため、正常なパッチのどれにも似ていないものが際立ちます。貪欲なコアセットサブサンプリングによって特徴の広がりを維持しながら一部だけを残し、この手法はMVTec AD産業用欠陥検出ベンチマークで、画像レベルのAUROC最大99.6%を達成しました。

このPyTorchの例では、同じ手法を8×8の手書き数字に適用し、生ピクセルを特徴として使います。数字の3の画像100枚をメモリバンクにし、それ以外の数字はすべて異常として扱います。

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枚では、メモリバンク全体を使った場合、未学習の3の画像83枚を他の数字から分離するROC曲線下面積（AUC）は0.984になります。エントリー数10のコアセットのスコアは0.982で、メモリと検索回数を10分の1に抑えながら、検出性能のほとんどを維持しています。

メモリバンクと関連概念の比較#

• ベクトルデータベース： 埋め込みを永続的に保存し、多数のアプリケーションに類似検索を提供します。メモリバンクは1つのモデルの学習または推論ループ内にあり、通常はGPUメモリ上で、モデルの実行に伴って上書きされます。
• コンテキストウィンドウ： トランスフォーマーが1回の処理でアテンションを適用するトークンです。メモリバンクは、そのウィンドウ外にある別のバッチ、フレーム、画像のベクトルを保持します。
• Transformer-XL： 前のテキストセグメントの隠れ状態を、次のセグメントがアテンションを適用するメモリとしてキャッシュします。これは活性化のメモリバンクであり、言語モデルが1つのウィンドウを超えて扱える範囲を広げます。
• ハードネガティブマイニング： 学習時に最も紛らわしい負例を選びます。メモリバンクは、選択に使える大量の負例を提供します。
• エピソード記憶： AIエージェントや強化学習において、意思決定の指針として検索される過去の経験の保存領域です。メモリバンクは、負例、過去のフレーム、正常なパッチなど、モデル自身の計算に使うベクトルを保存します。