Memory-augmented neural network

Overview

A memory-augmented network stores its working state in a separate store that the model can address, instead of only in one hidden vector or in its weights. The model computes an address from the current input, reads the slots it points to, and, in most designs, writes back to them. The store can be small and rewritten (a few key-value slots), a bank of recent past frames with first-in-first-out eviction, a fixed set of learned prototypes that is trained but not rewritten per input, or a fixed-size set of state tokens that every frame of a stream updates.

The invariant shared by the first three cards is the read: a weighted sum of stored values, where the weights come from inner products between a query and stored keys 1910.13197 §3.2.2 2310.12982 §3.4.1 2302.05160 §3.3. Two of those cards normalise the weights with a softmax across slots. The third applies a sigmoid to each slot on its own, so its weights do not sum to one 2302.05160 eq. 3. The fourth card reads its state with cross-attention inside each decoder block 2501.12387 §3.1. That card does not restate how the attention is normalised, so this page does not say. The designs differ in what is stored, how it is written (if at all), how often it changes, and how far back it reaches.

Four cards cover this concept, in four domains (student modelling, video object segmentation, weakly supervised video anomaly detection and 3D reconstruction), with different metrics. The page therefore describes design choices rather than a settled best design.

Variants

Fixed key-value slot memory with erase and add writes

A key matrix of N slots is fixed after training. A value matrix of N slots holds the state, and each answer rewrites every slot with an erase vector and an add vector, each weighted by that slot's attention 1910.13197 §3.2.3 eq. 5 eq. 7. The model is SKVMN, which builds on the DKVMN key-value memory 1910.13197 §3.1. Card: Knowledge Tracing with Sequential Key-Value Memory Networks. The authors report AUC of 84.0 on Synthetic-5 against 82.7 for DKVMN, and between 2.0 and 2.4 points higher on the four real datasets 1910.13197 tab. 3.

Recurrent state linked along identity chains

SKVMN also runs an LSTM over the exercises, but it links each exercise only to the most recent earlier exercise with the same identity vector, so unrelated steps are skipped 1910.13197 §3.3.1 1910.13197 §3.3.2. The identity rule is not fully specified in the paper (see the card's Pitfalls). The authors say this makes inference faster, but they report no runtime or FLOP numbers 1910.13197 §1 1910.13197 §5. Card: Knowledge Tracing with Sequential Key-Value Memory Networks.

Frame-level memory bank with a bounded readout

Each memory frame contributes keys and values to a pixel memory, which is read by matching the query frame's keys and keeping the top entries. Cutie uses 30 entries and at most five memory frames, the first frame always kept and the rest in first-in-first-out order 2310.12982 §3.4.4. Card: Putting the Object Back into Video Object Segmentation. The base readout is inherited from XMem, which the authors do not claim as new 2310.12982 §3.4.1.

Compact object summary memory

Cutie also keeps N = 16 object summaries. Each summary is a weighted average of object features over the pixels of the memory frames, pooled separately for the object and the background 2310.12982 §3.3 eq. 6 eq. 8. It is updated as a streaming average, so its cost does not grow with video length, which the authors report 2310.12982 §3.3 app. G.2. The object memory is read by object queries inside the object transformer 2310.12982 §3.2.1. Card: Putting the Object Back into Video Object Segmentation.

Dual learned prototype banks with sigmoid matching

Two banks of M = 60 learned prototypes, one for normal and one for abnormal snippets, are read by the same rule. Each snippet is matched to every slot by a sigmoid of a scaled inner product, not a softmax, and the augmented feature is the match-weighted sum of the slots 2302.05160 eq. 3. A snippet's bank match is the mean of its top-K slot matches 2302.05160 eq. 4. A dual memory loss trains each bank to fire on its own videos and stay quiet on the other kind 2302.05160 eq. 5. The banks are trained parameters: the card's method has no write step at test time, so this variant sits at the edge of the definition above. Card: Dual Memory Units with Uncertainty Regulation for Weakly Supervised Video Anomaly Detection. The authors report 86.97% frame AUC on UCF-Crime and 81.66% AP on XD-Violence with I3D RGB features 2302.05160 tab. 1 tab. 2. The repository's shipped records that match these numbers were selected on the test split, so they are not held-out results (see the card's Paper vs code).

Persistent state tokens updated once per image

The state is a fixed set of 768 learned tokens of width 768, shared by all scenes at the start of a stream 2501.12387 §3.1 §3.4. Each new image passes through two decoders that run on the image tokens and the state tokens together. In each block the two sides cross-attend, so the state takes in the new image and the image reads the stored context 2501.12387 §3.1 eq. 2. The state is therefore rewritten on every frame. The card describes this update as attention, not as an erase-and-add rule 2501.12387 §3.1. A virtual camera given as a raymap can read the state without writing to it, to predict an unseen view 2501.12387 §3.2. The authors report 16.58 FPS on KITTI at 512x144 on one A100 2501.12387 tab. 2. Memory use and training time are not reported, and the authors say recurrent training is time-consuming 2501.12387 §5. Card: Continuous 3D Perception Model with Persistent State.

When to use which

For agents that need to recall text across sessions, the memory in agent memory is an external text store, not a learned slot store, and these cards do not cover it.

Implementation sketch

Language-agnostic sketch of a slot read and write. Dimensions follow the slot-memory card. Choices the paper leaves open are marked.

# Read: attention over N fixed key slots  [arxiv:1910.13197 §3.2.1] [eq. 1]
# keys: (N, d_k) fixed after training; values: (N, d_v) per student
def read(keys, values, query_vec):
    # Softmax over inner products; no scaling factor stated  [eq. 1]
    # UR-DMU uses a per-slot sigmoid here instead  [arxiv:2302.05160 eq. 3]
    # CUT3R reads its state by cross-attention in each decoder block  [§3.1]
    w = softmax(keys @ query_vec)            # (N,)
    # Weighted sum of value slots  [arxiv:1910.13197 §3.2.2] [eq. 2]
    r = w @ values                           # (d_v,)
    return w, r

# Write: erase then add, each slot scaled by its weight  [eq. 5] [eq. 7]
def write(values, w, erase, add):            # erase, add: (d_v,) in (0, 1), (-1, 1)
    for i in range(N):
        values[i] = values[i] * (1 - w[i] * erase) + w[i] * add
    return values

# Object summary update: streaming weighted mean  [§3.3] [eq. 6]
# Keep running sums and area; read sums / area. Code adds 1e-4 to the area
# [code: cutie/model/transformer/object_transformer.py#L132]
def update_summary(sums, area, feats, pool_weight):
    sums = sums + pool_weight @ feats        # (N, C)
    area = area + pool_weight.sum(axis=1)    # (N,)
    return sums, area

# Streaming state update, one frame at a time  [arxiv:2501.12387 §3.1]
# state: S = 768 learned tokens of width 768  [§3.4]
# The paper's pose token is left out; see the card's Paper vs code
def stream_step(state, image_tokens):
    # Decoder blocks: state and image tokens cross-attend  [eq. 2]
    state, image_out = decoder(state, image_tokens)
    # Heads read image_out; the updated state goes to the next frame
    return state, image_out

Contradictions

None across cards: the four papers study different tasks and metrics. Within papers:

• SKVMN says it has fewer parameters than DKVMN on ASSISTments2009 and JunyiAcademy 1910.13197 §5.1. At matching (d, N) it has more parameters in 1910.13197 tab. 2 (for example 7.8k against 7k).
• Cutie's abstract describes running time as similar to XMem, while its main table gives 36.4 FPS for Cutie-base against 22.6 for XMem on the same hardware 2310.12982 tab. 1. The card's Pitfalls use the table.
• UR-DMU's paper gives K = floor(M/16) + 1 = 4 for every top-K operation 2302.05160 §4.2. The repository uses that value for the memory read, but 13 for snippet selection and video pooling with 200 training snippets model.py. Which one an implementation follows changes the bank matches and the video score.
• For the normal bank on an abnormal video, UR-DMU's paper sets the target to 1 2302.05160 eq. 5. The repository's target is 1 on the complement of the anomaly mask, which pushes that score towards 0 on abnormal clips train.py. The card records this as a paper-code difference.
• CUT3R's paper describes one learnable pose token placed in front of each image 2501.12387 §3.1 eq. 2. The official code reads the pose token for frames after the first from a learned memory, which the paper does not describe model.py. The card records this as a paper-code difference.

Papers

• Knowledge Tracing with Sequential Key-Value Memory Networks: key-value slot memory with erase and add writes, plus a chain-linked LSTM, for student mastery.
• Dual Memory Units with Uncertainty Regulation for Weakly Supervised Video Anomaly Detection: two learned normal and abnormal prototype banks, read by per-slot sigmoid match and trained with a dual memory loss.
• Putting the Object Back into Video Object Segmentation: frame-level pixel memory with an object-level summary memory and object queries for video segmentation.
• Continuous 3D Perception Model with Persistent State: a fixed set of learned state tokens, updated by decoder cross-attention for each streamed image, giving per-frame pointmaps and camera poses.