Relational recurrent neural networks

Host: Let's explore how artificial intelligence can learn to connect the dots over time. We are starting with a research paper that introduces relational recurrent neural networks. Guest: That sounds like a highly technical concept, so who exactly is behind this work? Host: It is a major collaboration out of London, driven by a large team at DeepMind and researchers from University College London. The ten-person author list includes notable AI experts like Adam Santoro, Oriol Vinyals, and Timothy Lillicrap. Guest: DeepMind is definitely a big name in the AI space. Just from that phrase relational recurrent neural networks, what can we expect this team to be focusing on? Host: It hints at the combination of two main ideas. Recurrent neural networks are standard AI models used to process sequential data, like sentences or a timeline of events. Guest: And where does the relational piece fit into that? Host: That points to the core focus of their research, which is giving these networks a built-in capacity to understand how different pieces of information relate to one another. Guest: Got it, so this team is essentially trying to help AI reason about relationships within an unfolding sequence of data. Host: Precisely. And the fact that their author list explicitly notes equal contributions shows just how intensive and collaborative this joint effort was.

Host: Let's take a high-level look at how artificial intelligence handles long-term memory and complex connections. Memory-based neural networks are great at remembering information over time, but they actually struggle with a concept called relational reasoning. Guest: What exactly does relational reasoning mean for a neural network? Host: It is the ability to understand how different pieces of remembered information connect to each other. The researchers confirmed that standard AI memory architectures often fail when tasks heavily rely on making those connections. Guest: So how do they fix that blind spot in the memory? Host: They built a new module called a Relational Memory Core, or RMC. It uses a mechanism called multi-head dot product attention, which basically allows different memories to interact directly with one another. Guest: Does allowing those memories to interact actually improve the network's performance? Host: It does, especially on tasks that require reasoning across a sequence of information over time. They saw massive gains in reinforcement learning domains, such as the game Mini PacMan. Guest: Did they test this new core on anything besides gaming environments? Host: Yes, they also applied it to program evaluation and language modeling. The RMC actually achieved state-of-the-art results on several major text databases, including the WikiText-103 and Project Gutenberg datasets.

Host: To kick things off, let's explore how artificial intelligence handles memory compared to the way we humans do. Just like we recall past events to make decisions today, neural networks use systems like LSTMs to successfully store and retrieve sequential data. Guest: It makes sense that AI needs to remember things to process information over time, but are those current memory systems missing something important? Host: They are indeed. While current networks are highly proficient at filing information away and fetching it later, they aren't explicitly designed to look at how different memories interact with each other. Guest: What do you mean by memories interacting? Is that like connecting two separate facts to draw a completely new conclusion? Host: Spot on, that is exactly what we call relational reasoning over time. To improve this, the authors introduce a new architecture called the Relational Memory Core, or RMC, which specifically biases the AI to relate memories to one another. Guest: How does the RMC actually get those isolated memories to talk to each other? Host: It uses a mathematical technique called multi-head dot product attention. Essentially, this mechanism allows the network to simultaneously compare multiple pieces of stored information and explicitly calculate how they relate. Guest: That sounds powerful, but how do the researchers prove this new RMC approach is actually better at connecting the dots? Host: They first built a custom toy task that specifically stresses relational reasoning, proving that older models struggle while the RMC succeeds. After that, they apply the RMC to much harder real-world problems, like complex reinforcement learning, analyzing computer programs, and processing massive language datasets.

Host: We're diving into how artificial intelligence learns to connect the dots, a concept known as relational reasoning. It's the process of understanding how different entities are linked to accomplish a goal, like figuring out which tree in a park is closest to a specific bench. Guest: So instead of just identifying the absolute location of one tree, the system has to actively compare the distances between multiple objects? Host: Exactly, because you couldn't solve that problem by looking at one tree in isolation. In neural networks, we build in assumptions—called inductive biases—to help them find these relationships, like how a basic image model naturally compares neighboring pixels. Guest: But what if the things you need to connect aren't right next to each other in the data? Host: That's where specific architectural designs come in, like Message Passing networks that link data points as connected nodes, or Relation Networks that compute links between distant pairs. And this concept isn't just for physical space; it applies to time as well. Guest: Temporal relational reasoning? Like comparing a piece of information from the start of a sequence to something at the very end? Host: Spot on. Attention mechanisms do exactly this by focusing on past hidden states based on their actual content, allowing the model to relate events across time rather than just relying on what happened most recently. Guest: How does that compare to older sequence models, like LSTMs? Host: Older models definitely have a harder time with this. LSTMs pack all historical information into one single continuous memory vector, which makes it really difficult for the network to keep specific memories distinct and relate them to one another later on.

Host: We are going to explore how we can design a system that stores distinct pieces of information and allows them to interact with one another. The guiding principle here is to use a fixed set of memory slots, kind of like cubbies, but give them a way to communicate at any given moment. Guest: How exactly do these separate memory slots communicate with each other? Host: They use something called multi-head dot product attention, or self-attention, which is a technique borrowed from Transformer models. Unlike older approaches that look back across all past time steps, this system applies attention between the memories at a single, current time step. Guest: So they are only exchanging information about the present moment. How does this self-attention actually move the data around? Host: It uses three learned components for each memory slot: a Query, a Key, and a Value. You can think of a Query as what a memory slot is looking for, the Key as a label for the data it currently holds, and the Value as its actual underlying information. Guest: I see, so one slot's Query tries to find a match with another slot's Key? Host: Exactly, and it does this through a mathematical operation called a scaled dot-product. When a Query and Key match well, the model creates a weighted average of the corresponding Values to propose an update for that memory slot. Guest: That makes sense, so each slot gets continuously updated with relevant information from the others. But why is it called "multi-head" attention? Host: Multi-head means the model runs several of these attention operations independently at the same time, using unique parameters for each "head". This is incredibly useful because it allows a single memory slot to share completely different types of information with different targets simultaneously.

Host: It is time to explore how an AI system actually takes in new information and stores it alongside what it already knows. Imagine we have a matrix of existing memories, and when a new piece of information arrives, we need a way to combine them efficiently. Guest: If the model just keeps adding new information to its memory matrix, wouldn't that matrix eventually grow too large to compute? Host: That is exactly the problem they solve by using a clever trick with the attention mechanism. When calculating the updated memory, the model uses only the existing memory as the queries, but it concatenates the old memory and the new input together to serve as the keys and values. Guest: So it looks at both the old and new information, but forces the output into the existing memory slots? Host: Exactly, which means the updated memory matrix always stays the exact same size as the original one. To handle an ongoing stream of new inputs over time, they embedded this update step into a modified LSTM, which is a classic network used for sequential data. Guest: How does a standard LSTM handle a whole matrix of memories instead of its usual single flat array? Host: They treat the memory matrix as a two-dimensional grid of cell states, where each row operates almost identically to a normal LSTM cell. Crucially, the model shares the exact same learning parameters across all of those memory rows. Guest: Oh, so if the parameters are shared, you could add more memory rows without increasing the number of weights the model has to learn? Host: Spot on, and this lets researchers tune the overall storage capacity depending on the specific task. For instance, they found that language modeling works best with fewer, but much larger, memory slots. Guest: That makes sense, but does it use standard LSTM gates to decide what gets saved or forgotten in those slots? Host: It can, but the authors also suggest a variation called memory gating. Instead of calculating gates for every single tiny element within a memory, the model can use a scalar value to gate an entire memory row at once.

Host: To really see what a new architecture is capable of, we have to put it through a rigorous set of exams. The researchers tested the Relational Memory Core, or RMC, across a variety of supervised and reinforcement learning tasks. Guest: Since the RMC has all those tunable parameters like memory size and attention heads, what kind of tests did they use to measure its performance? Host: One major test is the "Nth Farthest" task, where the model is fed a sequence of vectors and asked a question like, "What is the third farthest vector from vector M?" Guest: That sounds tricky, because it can't just copy and paste an answer it memorized. Host: Right, it forces explicit relational reasoning. The model has to calculate all the pairwise distances to that reference vector, and then sort those distances, rather than just sorting the vectors themselves. Guest: Got it. Did they test it on anything more structured than just random vectors? Host: Yes, they used a "Learning to Execute" dataset made of algorithmic pseudo-code. The model has to read the code and accurately predict the numerical output of executing it. Guest: How does reading pseudo-code test relational memory? Host: Think of mathematical and logical operators as defining relationships between different variables. Successfully running the code proves the model understands those symbolic relationships. Guest: That makes perfect sense. What about the reinforcement learning side you mentioned? Host: For that, they dropped the agent into a game of Mini Pacman, but with a major catch. The agent can only see a tiny five-by-five window around itself. Guest: So if it can't see the whole board, it has to rely entirely on its memory to survive? Host: Exactly. It has to remember where the food was and use that information to predict where the ghosts might be moving in the unseen parts of the maze.

Host: We are turning our attention now to how AI models learn to predict and generate text. Specifically, the authors are looking at word-based language modeling, which is the task of predicting the next word in a sequence based on all the words that came before it. Guest: That sounds a lot like the autocomplete feature on my phone's keyboard. Is that the same basic idea? Host: Exactly, predictive keyboards and search phrase completions are direct applications of this math. But modeling that probability is also a foundational component for bigger systems like machine translation and speech recognition. Guest: So why are the researchers testing their new model, the RMC, on this specific task? Host: Language is essentially a massive, sequential reasoning puzzle. By having the RMC predict the next word, they can evaluate how well it processes information over time across huge amounts of natural data, and then compare it to top-performing baseline models like LSTMs. Guest: What kind of data are they using to train it? Are they just feeding it random sentences from the internet? Host: They are actually very intentional, choosing datasets with whole, connected sentences ranging from a hundred million to four billion training tokens. They use Wikipedia articles, books from Project Gutenberg, and news stories from a dataset called GigaWord. Guest: That is a massive amount of reading, and it covers a pretty wide range of writing styles. Host: It really does, and to handle all those different styles, they set the model's vocabulary size to about two hundred and fifty thousand words. That gives the system enough breadth to recognize rare words and numeric values across everything from an old novel to a modern news report.

Host: It is time to see how these models actually held up when put to the test. In the first experiment involving 16-dimensional vector inputs, the Relational Memory Core, or RMC, completely blew the standard baseline models out of the water. Guest: How big of a performance gap are we talking about between the RMC and older models like the LSTM? Host: It was a stark difference. The baseline models couldn't even break thirty percent accuracy, but the RMC consistently hit ninety-one percent. It even held its ground when they doubled the difficulty to 32-dimensional vectors, though it was a bit more sensitive to its initial setup. Guest: Did the researchers look under the hood to see how the RMC was managing its memory to get those scores? Host: They did, by tracking its attention weights during the task. They noticed that before the model spots a specific reference vector, it neatly shuttles incoming information into just one or two memory slots. Guest: What happens once that target reference vector finally shows up? Host: Its behavior completely changes. Suddenly, all the memory slots shift their focus to those specific locations where the target was just stored. Guest: That sounds like a pretty clear window into how it thinks. Can they map out exactly where every piece of data lives? Host: Actually, they have to be careful with those conclusions. Even after a single round of attention, the memory gets highly distributed across the system, making it really hard to confidently say how the information is compartmentalized. Guest: That makes sense. Besides those vector tests, how did the RMC do on the actual program evaluation, like the Learning to Execute tasks? Host: It matched or beat all the baseline models there as well. Interestingly, the researchers found that it performed best when they trained it without a common technique called "teacher forcing." Guest: Teacher forcing? Does that mean feeding the model the correct ground-truth answers step-by-step during training? Host: Exactly. By taking away that constant hand-holding, the model was forced to generalize better on its own. This worked especially well here because the final answers in these programming tasks are unambiguous, meaning the model didn't need those step-by-step hints to eventually find the right path.

Host: Let's see how this model actually stacks up against the competition when put to the test. In the game Mini Pacman, the Relational Memory Core—or RMC—scored about 100 points higher than a standard LSTM when given a limited view, and nearly doubled the LSTM's score when it could see the whole board. Guest: Doubling the score is a huge jump just from giving it the full picture! Did it show that same kind of dominance in the language tasks? Host: It did perform consistently better across three different language modeling tasks. The researchers measured a drop in "perplexity," which essentially means the model was less confused, improving on the best published results by about 5 to 12 percent. Guest: A 12 percent improvement in language tasks sounds solid, especially compared to established models. Why was it able to outperform them? Host: A big factor was how it handled context. The RMC scored highly even when it was given very few context words to work with, whereas an LSTM really needed a much larger window of preceding words to do well. Guest: So the RMC is just much better at picking up on short-term relationships between words that are close together? Host: Exactly, it captures those short-term relations accurately without needing a massive run-up of text. When breaking the results down, they found the RMC was particularly good at predicting frequent words, which is what ultimately drove that overall performance boost.

Host: As we wrap up our look at this model, it's time to reflect on how it compares to other methods and what these results actually mean. Other approaches have tried using a growing buffer of previous states to remember long sequences, but that gets really tricky when time is unbounded, like in real-world reinforcement learning. Guest: Because a continuously growing buffer would eventually just run out of memory, right? Host: Exactly, which is why the authors suggest a potential compromise for future work: collecting memories perfectly for a short chunk of time, then compressing them into a recurrent core before moving on. Guest: That makes sense. But do they know exactly how their specific design choices caused the model to reason better? Host: They are actually very careful about that, pointing out that while the model clearly performs better on complex tasks, they can't make concrete claims about the exact internal causal mechanics yet. Guest: So it's more about proving that the overall function improved, even if the exact "why" is still a bit of a black box? Host: Spot on, though they did notice some interesting trade-offs regarding how the memory is structured into rows, or slots, and columns, or attention heads. For instance, in language modeling, they had to stick to just one massive memory slot because the task already required so many parameters. Guest: Oh, so there isn't one universally perfect setup, and the ideal ratio of memory slots to memory size just depends on the specific task. Host: Precisely. The ultimate takeaway here is that explicitly allowing memories to interact gives recurrent neural networks a much-needed boost in relational reasoning, whether they are evaluating programs or solving reinforcement learning tasks.

Host: Before diving into the technical details, it is always good to recognize the village it takes to put a major research project together. The authors kick things off by thanking a specific group of people for their feedback, discussions, and support. Guest: Is there anyone in particular they mention? Host: Yes, they list several key researchers, including Caglar Gulcehre, Matt Botvinick, and Charles Blundell. These are all prominent figures on the DeepMind team. Guest: DeepMind is the AI research lab behind major breakthroughs like AlphaGo, right? Host: Exactly. It shows that this work didn't happen in a vacuum, but rather within a highly collaborative environment at one of the world's leading AI organizations. Guest: So they aren't just thanking friends and family here, but rather their professional peers. Host: Right, they specifically highlight the critical feedback provided by these colleagues. In complex AI research, bouncing ideas off experts across different disciplines is absolutely essential for shaping the final work. Guest: That makes sense, since a new model usually requires a lot of rigorous internal review before it ever reaches the public. Host: You hit the nail on the head. Acknowledging that internal support network perfectly sets the stage for the heavy intellectual lifting the authors did.

Host: To really understand how this system operates, we need to look under the hood at how the architecture is put together. The authors provide a helpful glossary for their configurations, centering around something called the memory matrix, or matrix M. Guest: I'm guessing matrix M is where the model stores its information, but how exactly do they measure its size? Host: They measure it using "memory slots," which are essentially the rows in that matrix representing individual memories. If you multiply the number of slots by the size of each memory, you get the "total units," meaning the absolute number of elements in the entire matrix. Guest: Okay, so that covers the storage side of things. I also see terms related to attention, like "num heads"—what does that mean here? Host: "Num heads" refers to the number of attention heads, which dictates how many unique sets of queries, keys, and values the model generates to search through those memories. Guest: And what about the term "num blocks"? Does that also relate to how the model searches its memory? Host: Yes, "num blocks" tells us the number of attention iterations the model performs at each single time-step. Guest: That makes sense, which just leaves the last term, "gate style," which the text says can be per unit or per memory slot. Host: Gating is the mathematical mechanism that controls what information gets updated or kept. The gate style simply defines if that filter is applied at the micro level of individual units, or at the broader level of whole memory slots.

Host: Let's dive into exactly how the researchers set up their experimental test, specifically for a challenge called the Nth Farthest task. To see how well these memory models work, they fed them sequences of eight random, 16-dimensional vectors. Guest: Okay, so the models receive these eight random vectors one by one. Do they get any other information to help them solve the task? Host: Yes, they combined several pieces of information into a single 40-dimensional input for each time step. Along with the 16 numbers of the vector itself, the model receives a unique label, plus the actual instructions for the task—the values of 'n' and 'm'—which tell it exactly what it needs to find. Guest: Wait, you mentioned the labels are unique. Do those labels just tell the model what time-step it's currently on? Host: Actually, no. The labels are sampled randomly, so they explicitly do not match the order the vectors arrive in. This forces the model to actually bind the specific vector to its label in memory, rather than just relying on the sequence order. Guest: That makes sense. And they tested this setup on a few different memory models, right? How did they ensure it was a fair fight? Host: They compared their Relational Memory Core, or RMC, against standard LSTMs and DNCs. To keep it fair, all models used the exact same learning rate optimizer, and their outputs all passed through an identical four-layer neural network before making a final prediction. Guest: Did they have to heavily tweak the internal settings of these models to get the best results? Host: Surprisingly, changing things like the hidden unit sizes for the LSTM or the memory size for the DNC didn't seem to impact their performance much. For their own RMC model, they tested various combinations of memory slots and attention heads. Guest: So by keeping the output processing and tuning process so standardized, they could see which memory structure actually handled the task best. Host: Exactly. They even compared the models based on actual running time, ensuring the focus remained entirely on the efficiency of each core memory mechanism.

Host: We are turning our attention now to how the researchers put this model to the test with complex programming tasks. To see how well the Relational Memory Core handles working memory and symbols, they evaluated it using the Learning to Execute dataset. Guest: What exactly is in that dataset? Are we talking about full software applications? Host: Not quite full software, but rather mini-programs. They are short snippets of code that run in linear time with constant memory, testing fundamental operations like addition, control flow, and for-loops. Guest: How did they train the model to process those code snippets? Host: They used an encoder-decoder architecture. The encoder reads the input sequence and passes its final state to a decoder, which then generates the program's output step by step. They also varied the difficulty during training by changing the length of the terminal values and how deeply nested the operations were. Guest: Did they have to give the decoder the correct answers step by step to keep it on track? Host: That is actually a critical detail in their setup. Normally, models use a technique called teacher forcing, where they are fed the true correct answer from the previous step to help predict the next one. Guest: But they decided to do something different for this evaluation? Host: Exactly. They found the model actually performed best when they removed that dependency on the ground truth. By forcing the model to rely solely on its own past predictions during training, it yielded the strongest results. Guest: So taking off the training wheels made it more robust. Did they set up other models to compare it against? Host: Yes, they established a rigorous comparison against baseline models like standard LSTMs, Dynamic Neural Computers, and Entity Networks. They tested various configurations, attention heads, and memory sizes for everything, making sure to report only the absolute best performing setup for each model to keep the race fair.

Host: Let's look at some concrete examples of how we actually test these AI memory models. The text gives us specific samples of the problems they had to solve, which fall into two main buckets: programmatic tasks and memorization tasks. Guest: When you say programmatic tasks, are we talking about the models writing actual code? Host: Not exactly writing it, but evaluating and predicting the output of it. They were fed things like nested addition, if-else control statements, and even loops, while the memorization tasks tested if they could simply copy, reverse, or double a sequence. Guest: How did the different models handle all that complex logic? Host: The Relational Memory Core, or RMC, really stood out by reaching high accuracy very quickly. It showed excellent data efficiency, especially compared to the standard LSTM baseline model. Guest: Did the other baseline models struggle with specific types of tasks? Host: They did, which is exactly what makes the RMC's performance so impressive. The DNC model stumbled on the control tasks, and the EntNet had trouble with basic memorization like copying and doubling sequences. Guest: So the RMC scored well across the board, but did it achieve that just by being a massive model? Host: Actually, it's quite the opposite. Because of how the RMC manages its internal memory, the best LSTM models ended up needing far more parameters than the top RMC models just to try and keep up.

Host: Let's step into the virtual testing ground where this AI is trained, a custom environment called Viewport Box World. It's essentially a grid-based puzzle game designed to test an AI agent's ability to plan and reason. Guest: What exactly does this puzzle world look like to the AI? Host: It's a simple 14 by 14 grid of pixels where the AI, represented by a dark grey pixel, moves around to pick up colored keys. These keys are then used to open matching colored locks to eventually reach a rewarded white gem. Guest: Are the locks and keys scattered randomly, or is there a pattern? Host: They are paired up in little two-pixel blocks. The right pixel shows the lock color, and the left pixel shows the color of the new key you'll get if you open that lock. Guest: That sounds pretty straightforward, but you called it Viewport Box World. What does the viewport part mean? Host: That's the main twist, because the agent can't see the whole board. It only sees a tiny five-by-five window around itself, plus a little indicator showing which key it's currently holding. Guest: Oh, so it has to explore blindly and actually remember where the different colored locks and keys are? Host: Exactly, and because the locations are completely randomized, the AI has to rely on abstract relationships in its memory rather than just memorizing a spatial layout. This is why standard visual algorithms, like CNNs, usually struggle here. Guest: If it has limited vision, what happens if it grabs the wrong key or takes a bad path? Host: That's where the advanced planning comes in. Keys disappear once used, and the procedurally generated map is full of dead-end distractor paths, so the agent has to mentally link the right sequence of colors together before acting so it doesn't get stuck.

Host: We are diving into the exact mechanics of how the researchers built and trained their AI for the Mini Pacman environment. They generated these levels using a pool of twenty different colors for keys and locks. Guest: What determines how hard a specific level is for the AI? Host: Difficulty comes down to the length of the correct path to the winning gem, plus the number of fake branches, or distractors, meant to confuse it. During training, they kept the correct paths to five steps or fewer so an untrained agent had a tiny chance of winning just by random luck. Guest: How does the AI visually process the maze so it doesn't get tricked by those fake branches? Host: It uses a neural network to scan the screen, plus two special layers that tag the exact X and Y coordinates of every pixel. All that spatial and visual data is passed into a Relational Memory Core, which acts as the agent's working memory. Guest: Does it just play one game at a time to learn the rules? Host: Actually, it uses a distributed architecture called IMPALA, where one hundred different "actor" agents play games simultaneously to gather experience. All one hundred actors constantly feed their gameplay data to a single "learner" that figures out the best overall policy. Guest: That is a massive amount of practice data, so how does the game reward good strategy? Host: The agent earns a big ten-point reward for grabbing the gem, and one point for opening a correct box along the solution path. But the game is unforgiving, because opening a distractor box costs a point and immediately ends the level.

Results. We trained an Importance Weighted Actor-Learner Architectures agent augmented with the RMC on BoxWorld levels that required opening at least 1 and up to 5 boxes. The number of distractor branches was randomly sampled from 0 to 5. This agent achieved high performance in the task, correctly solving 98% of the levels after 1e9 steps. The same agent augmented instead with a ConvLSTM performed significantly worse, reaching only 73%.

Language Modeling. We trained the Recurrent Memory Core with Adam, using a learning rate of 0.001 and gradients were clipped to have a maximum L2 norm of 0.1. Backpropagation-through-time was truncated to a window-length of 100. The model was trained with 6 Nvidia Tesla P100 GPUs synchronously. Each GPU trained with a batch of 64 and so the total batch size was 384. We used 512 (with 0.5 dropout) as the word embedding sizes, and tied the word embedding matrix parameters to the output softmax. We swept over the following model architecture parameters: Total units in memory {1000, 1500, 2000, 2500, 3000}. Attention heads {1, 2, 3, 4, 5}. Number of memories {1,2}. MLP layers {1, 2, 3, 4, 5}. Attention blocks {1, 2, 3, 4}. and chose 2500 total units, 4 heads, 1 memory, a 5-layer MLP, and 1 attention block based upon validation error on WikiText-103. We used these same parameters for GigaWord and Project Gutenberg without additional sweeps, due to the expense of training.

Figure 11: Validation perplexity on WikiText-103. LSTM comparison from [32]. Visual display of data may not match numbers from table 2 because of curve smoothing. Figure 12: Perplexity as a function of test unroll length. Increase in perplexity when models are unrolled for shorter sequence lengths at test time without state transfer between unrolls. Perplexities are compared against the 'best' perplexity where the model is unrolled continuously over the full test set. We see that both models incorporate little information beyond 500 words. Furthermore, the RMC has a smaller gain in perplexity (drop in performance) when unrolled over shorter time steps in comparison to the LSTM, e.g. a regression of 1 perplexity for the RMC vs 5 for the LSTM at 100 time steps. This suggests it is focusing on more recent words in the text. Table 3: Test perplexity split by word frequency on GigaWord v5. Words are bucketed by the number of times they occur in training set, > 10K contains the most frequent words.