LLM + RoboticsUnit 4 · Specialised architectures · Day 20

Relational Recurrent Neural Networks

Split an LSTM's memory into slots and let the slots attend to each other at every step. The network gets much better at comparing things it saw at different times.

AuthorsAdam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Théophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, Timothy Lillicrap (DeepMind, UCL)
Year / venue2018 · NeurIPS
Core ideaSelf-attention between memory slots inside a recurrent core
Study time3–4 h

Original: arXiv:1806.01822 · Code, notebook and exercises: 30u30 Day 20

Start here (no maths needed)

A computer that keeps its memories in separate boxes, and lets those boxes talk to each other, gets much better at comparing things it saw at different times.

Think of it like this

Think of a detective's notebook. Each clue gets its own page. When a new clue arrives, the detective flips back through the other pages and asks, "Does this connect to anything I already know?" A normal memory program stuffs all clues onto one messy page. This program keeps a few separate pages, and at every step each page checks the others and the newest clue.

How it works

  1. The program has a few memory slots, like pages in a notebook.
  2. Something new arrives, for example a point on a map.
  3. Each slot looks at the other slots and the new thing, and pays most attention to the ones that matter.
  4. Each slot decides how much old memory to keep and how much new information to write.
  5. This repeats at every step, so the program can later answer questions that need comparing things from different times.

Why robots care

Objects leave a robot's camera view and come back. A robot tidying a room might see a toy, turn away, then see a box. To decide which toy is furthest from the box, it has to remember things it can no longer see and compare them.

Try it

Words to know

Memory slot
One separate box where the program stores part of what it has seen.
Attention
A way for each slot to choose which other slots to listen to most.
Recurrent
Working step by step and carrying memory forward from one step to the next.
Gate
A dial that decides how much old memory to keep and how much new to add.

Going deeper: everything below is the full version, with some maths. On a first read, skim the headings and the "Why it matters for robotics" box, and come back for the rest later.

1. Summary

TL;DR

Memory networks are good at storing and retrieving. This paper asks whether they can also relate what they've stored. It builds the Relational Memory Core (RMC): a matrix of memory slots that runs multi-head self-attention over itself plus the new input at each time step, then updates through LSTM-style gates. On a toy task built to need comparisons across time, LSTM and DNC baselines stayed under 30% accuracy while the RMC reached 91%. It also improved program evaluation, a partially observed Pac-Man game, and word-level language modelling.

2. The core idea

Let M be the memory matrix (one row per slot) and x the new input. One attention step:

Q = M·W_q
K = [M; x]·W_k
V = [M; x]·W_v
M̃ = softmax(Q·Kᵀ / √d_k) · V        # same shape as M

M̃ is a proposed update. It's passed through a row-wise MLP g_ψ (with residual connections) and then gated into memory:

m_{i,t} = σ(f_{i,t} + b̃_f) ∘ m_{i,t−1}  +  σ(i_{i,t}) ∘ g_ψ(m̃_{i,t})
h_{i,t} = σ(o_{i,t}) ∘ tanh(m_{i,t})

f, i, o are the usual forget, input and output gate pre-activations computed from x_t and the previous output. So each slot behaves like an LSTM cell state from Day 2, but its candidate update comes from attention over all slots. The authors note that in practice they didn't need the output gate.

3. Deep dive

Three ways to hold memory

Attention within a step, not over the whole past

A Transformer attends over every past token, so its memory and cost grow with the sequence. The RMC attends only over its N slots and the current input. Cost per step stays fixed, whatever the sequence length. The paper's discussion points out that growing buffers are a problem for online RL, where the stream never ends. Module 05 covers the same trade-off at modern scale: fixed-size recurrent state vs a full attention cache.

Slots versus slot size

Because weights are shared across rows, you can trade many small slots against a few big ones without changing the parameter count. The paper found the best ratio depends on the task. The N-th farthest task wanted many slots. Language modelling used a single large slot. Program evaluation did well with 4 slots and 4 heads.

The N-th farthest task

The model sees eight random 16-dimensional vectors, one per step, each with a label. Then it must answer: "which vector is the n-th farthest from vector m?" It has to remember all eight, compute every distance to m (which may arrive last), and sort those distances. The demo is the same task in 2D.

Interactive: the N-th farthest question

Eight labelled vectors. Pick the reference m and the rank n. Lines show distances to m; the answer is highlighted.

Answer (label)
–
Distances to compare
–

4. Why it matters for robotics

5. How to study it

  1. Read sections 3 and 4 of the paper, and appendix A.1 for the N-th farthest setup (about 1.5 h). Sketch Figure 1 from memory.
  2. In the 30u30 Day 20 folder, read implementation.py and trace one step: input stacked onto memory, attention, MLP, gates.
  3. Do the exercises in order: dot-product attention, gating, one RMC step, the N-th farthest data generator, then training.
  4. Run train_minimal.py and compare against an LSTM baseline on N-th farthest.
  5. Stretch: sweep the number of slots and heads at a fixed total memory size, and see how accuracy changes.

6. Self-check

Why are queries computed from M only, but keys and values from [M; x]?

So the output has exactly one row per memory slot. Each slot can read from other slots and the new input, but the input doesn't get a row of its own in the next memory.

Why doesn't adding memory slots add parameters?

The projection, MLP and gate weights are shared across rows. More rows means more storage and more compute, not more weights.

How does the RMC differ from a Transformer layer over the sequence?

A Transformer attends over all past positions, so cost grows with length. The RMC attends only over its fixed set of slots and the current input, then carries the slots forward recurrently.

Why is N-th farthest hard for a plain LSTM?

It must keep eight vectors separately, then compare distances to one of them, possibly seen last. Packing all of that into one vector and sorting implicitly is hard. The LSTM stayed under 30%.

What does the RMC borrow from the LSTM?

The gating: a forget gate on the old slot value and an input gate on the new candidate, applied per slot. That keeps training stable over long sequences.

7. Further reading