Neural Turing Machine: Complete Guide — Differentiable External Memory

Summary

The Neural Turing Machine (NTM) is a revolutionary neural architecture proposed by Alex Graves, Greg Wayne, and Ivo Danihelka at DeepMind in 2014. It extends the capabilities of recurrent neural networks by adding an addressable external memory, making it possible to learn algorithmic tasks that were beyond the reach of classical deep learning: sequence copying, association, sorting, and even the execution of simple programs.

The fundamental principle is elegant: a neural controller (usually an LSTM or MLP) interacts with an externalized memory matrix via fully differentiable read and write heads. This differentiability allows training via classical backpropagation — the gradient flows through the attention operations, which radically differentiates NTM from non-differentiable classical external memories.

The NTM subsequently inspired DeepMind’s Differentiable Neural Computer (DNC) (2016), which adds temporal linking mechanisms and dynamic memory allocation. Together, these architectures constitute a fundamental lineage of research in deep learning, foreshadowing the attention mechanisms that would dominate the field a few years later.

Mathematical Principle

A Neural Turing Machine consists of two fundamental elements:

Global Architecture

NTM = neural controller + memory matrix M ∈ ℝ^{N×W}

where N is the number of memory locations and W is the width of each memory vector. The controller can be a recurrent network (LSTM, GRU) or a multilayer perceptron (MLP). At each time step t, the controller:

1. Receives an external input xₜ
2. Reads vectors from memory via the read heads
3. Produces an output yₜ and read/write parameters
4. Modifies memory via the write heads

Read Mechanism

Reading is performed via weighted attention over the memory locations:

rₜ = Σᵢ wₜ(i) · Mₜ(i)

where wₜ is an attention weight vector of dimension N, normalized by softmax (Σᵢ wₜ(i) = 1), and Mₜ(i) denotes the i-th memory vector of width W. This soft reading mechanism is what makes the operation differentiable: instead of reading a single discrete location, a weighted combination of all locations is read, allowing the gradient to flow to each memory position.

Write Mechanism

Writing modifies the memory in two phases: erasure followed by addition.

Erasure phase:

M̃ₜ(i) = Mₜ₋₁(i) · (1 − wₜ(i) · eₜ)

where eₜ ∈ [0,1]^W is the erasure vector generated by the controller. The term (1 − wₜ(i) · eₜ) allows selectively erasing certain dimensions of location i, proportionally to the attention weight wₜ(i).

Addition phase:

Mₜ(i) = M̃ₜ(i) + wₜ(i) · aₜ

where aₜ ∈ ℝ^W is the add vector. The complete write operation combines both:

Mₜ(i) = Mₜ₋₁(i) · (1 − wₜ(i) · eₜ) + wₜ(i) · aₜ

This formulation ensures that only the locations being attended to are modified, and that the modification is proportional to that attention.

Content-Based Addressing

Content-based addressing retrieves a memory location based on its similarity to a key vector emitted by the controller:

wᶜ(i) = softmax( cos(kₜ, Mₜ(i)) / β )

where kₜ ∈ ℝ^W is the key vector produced by the controller, β > 0 is a sharpness parameter that controls the focus of attention, and cos(·,·) denotes cosine similarity:

cos(kₜ, Mₜ(i)) = (kₜ · Mₜ(i)) / (‖kₜ‖ · ‖Mₜ(i)‖)

A high β concentrates attention on the most similar location; a low β produces more diffuse attention.

Location-Based Addressing

Location-based addressing allows sequential traversal of memory, essential for ordered tasks like sequence copying. It operates in three sub-steps on the previous weight vector wₜ₋₁:

1. Interpolation (gating):

w̃ₜ = gₜ · wᶜₜ + (1 − gₜ) · wₜ₋₁

where gₜ ∈ [0,1] determines the blend between content-based addressing and the previous position.

2. Circular shift (convolutional shift):

ŵₜ(j) = Σₖ w̃ₜ(j − k) · sₜ(k)

where sₜ is a shift kernel of odd dimension (typically 3), and the index is taken modulo N (circular rotation). This shift allows focusing on the previous or next adjacent location.

3. Sharpening:

wₜ(i) = ŵₜ(i)^γ / Σⱼ ŵₜ(j)^γ

where γ ≥ 1 is a sharpening parameter that prevents the weight vector from becoming too diffuse through successive convolutions.

Complete Flow at Each Time Step

The complete execution at each iteration t is as follows:

Input xₜ → Controller → Read/Write Parameters
                        ↓
            wₜ(content) via cosine similarity
                        ↓
            wₜ(location) via circular shift
                        ↓
            Read: rₜ = Σᵢ wₜ(i) · Mₜ(i)
                        ↓
            Write: Mₜ(i) = Mₜ₋₁(i)·(1−wₜ(i)·eₜ) + wₜ(i)·aₜ
                        ↓
            Output yₜ

Intuition: The Difference Between Memorizing and Writing

To understand why the NTM represents a conceptual leap, consider this analogy: a classical neural network has its memory embedded in its weights — like a musician playing entirely from memory. They can reproduce what they learned during training, but they cannot memorize new information during execution. If you present them with a sequence they’ve never seen, they must process it in real time without being able to retain it.

An NTM, on the other hand, has an external memory — like a pianist reading a score. The neural controller is the pianist: they know how to play, they know the technique. But the external score is the memory M. The pianist can:

• Write new notes: when they learn something new during the sequence, they write it on the score (write head)
• Reread later: when they need that information, they go back and find it at the appropriate location (read head)
• Modify their notes: they can erase part of it and write over it (erasure + addition)

This is the fundamental difference between memorizing a phone number in your head and writing it on a Post-it note. In the first case, capacity is limited and fragile. In the second, you can store much more, retrieve information precisely, and modify it at will.

This separation between processing capacity (the controller) and storage capacity (the memory) is exactly what allows NTMs to generalize to sequence lengths never seen during training — a well-known weakness of classical LSTMs and GRUs.

Python Implementation with PyTorch

Basic Architecture of a Simplified NTM

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class SoftAttention(nn.Module):
    """
    Soft attention mechanism for read/write heads.
    Combines content-based and location-based addressing.
    """
    def __init__(self, memory_width, shift_range=1):
        super().__init__()
        self.memory_width = memory_width
        self.shift_range = shift_range

    def content_addressing(self, key, memory, beta):
        """
        Content-based addressing: normalized cosine similarity.

        Args:
            key: tensor of shape (batch_size, memory_width)
            memory: tensor of shape (batch_size, memory_size, memory_width)
            beta: sharpness parameter (batch_size, 1)
        Returns:
            weights: tensor of shape (batch_size, memory_size)
        """
        # Normalization for cosine similarity
        key_norm = F.normalize(key, p=2, dim=-1).unsqueeze(1)
        memory_norm = F.normalize(memory, p=2, dim=-1)

        # Cosine similarity
        cos_sim = torch.sum(key_norm * memory_norm, dim=-1)

        # Apply sharpness and softmax
        weights = F.softmax(cos_sim * beta, dim=-1)
        return weights

    def location_addressing(self, prev_weights, shift, gamma):
        """
        Location-based addressing: circular shift + sharpening.

        Args:
            prev_weights: (batch_size, memory_size)
            shift: shift weights (batch_size, 2*shift_range+1)
            gamma: sharpening parameter (batch_size, 1)
        Returns:
            weights: (batch_size, memory_size)
        """
        batch_size, memory_size = prev_weights.shape

        # Circular convolution (shift)
        shifts = torch.arange(-self.shift_range, self.shift_range + 1, device=prev_weights.device)
        shifted = torch.zeros_like(prev_weights)

        for s_idx, s in enumerate(shifts):
            rolled = torch.roll(prev_weights, shifts=int(s), dims=-1)
            shifted += shift[:, s_idx].unsqueeze(1) * rolled

        # Renormalization
        shifted = shifted / (shifted.sum(dim=-1, keepdim=True) + 1e-8)

        # Sharpening
        sharpened = shifted ** gamma
        weights = sharpened / (sharpened.sum(dim=-1, keepdim=True) + 1e-8)
        return weights

    def forward(self, key, memory, beta, gate, shift, gamma, prev_weights):
        """
        Complete attention flow: content -> interpolation -> location -> sharpening.
        """
        # Content-based addressing
        w_content = self.content_addressing(key, memory, beta)

        # Interpolation between content and previous position
        w_gated = gate * w_content + (1 - gate) * prev_weights

        # Location-based addressing
        w_shifted = self.location_addressing(w_gated, shift, gamma)

        return w_shifted


class NTM(nn.Module):
    """
    Simplified Neural Turing Machine with LSTM controller.

    Architecture:
        - Controller: LSTM
        - Memory: fully differentiable N × W matrix
        - 1 read head + 1 write head
        - Hybrid addressing (content + location)
    """
    def __init__(self, input_size, output_size, memory_size=128, memory_width=20,
                 controller_size=100, shift_range=1):
        super().__init__()

        self.input_size = input_size
        self.output_size = output_size
        self.memory_size = memory_size
        self.memory_width = memory_width
        self.controller_size = controller_size
        self.shift_range = shift_range

        # Memory initialized to zero
        self.register_buffer('init_memory', torch.zeros(memory_size, memory_width))

        # LSTM Controller
        self.controller = nn.LSTMCell(
            input_size=input_size + memory_width,
            hidden_size=controller_size
        )

        self._build_read_write_heads()

        self.attention = SoftAttention(memory_width, shift_range)

        self.output_layer = nn.Linear(controller_size, output_size)

        self._init_parameters()

    def _build_read_write_heads(self):
        """Builds the read and write layers."""
        self.read_key_layer = nn.Linear(self.controller_size, self.memory_width)
        self.read_beta_layer = nn.Linear(self.controller_size, 1)
        self.write_key_layer = nn.Linear(self.controller_size, self.memory_width)
        self.write_beta_layer = nn.Linear(self.controller_size, 1)
        self.write_gate_layer = nn.Linear(self.controller_size, 1)
        self.write_shift_layer = nn.Linear(self.controller_size, 2 * self.shift_range + 1)
        self.write_gamma_layer = nn.Linear(self.controller_size, 1)
        self.write_erase_layer = nn.Linear(self.controller_size, self.memory_width)
        self.write_add_layer = nn.Linear(self.controller_size, self.memory_width)
        self.read_gate_layer = nn.Linear(self.controller_size, 1)
        self.read_shift_layer = nn.Linear(self.controller_size, 2 * self.shift_range + 1)
        self.read_gamma_layer = nn.Linear(self.controller_size, 1)

    def _init_parameters(self):
        """Weight initialization — crucial for NTM convergence."""
        for name, param in self.named_parameters():
            if 'weight' in name and param.dim() > 1:
                nn.init.xavier_uniform_(param)
            elif 'bias' in name:
                nn.init.zeros_(param)

        with torch.no_grad():
            self.write_shift_layer.bias.data.fill_(-10)
            self.read_shift_layer.bias.data.fill_(-10)
            self.write_erase_layer.bias.data.fill_(1.0)
            self.read_key_layer.bias.data.fill_(0.0)
            self.write_key_layer.bias.data.fill_(0.0)

    def init_state(self, batch_size, device):
        """Initializes memory and attention weights."""
        memory = self.init_memory.unsqueeze(0).repeat(batch_size, 1, 1).to(device)
        prev_read_weights = torch.ones(batch_size, self.memory_size, device=device) / self.memory_size
        prev_write_weights = torch.ones(batch_size, self.memory_size, device=device) / self.memory_size
        h = torch.zeros(batch_size, self.controller_size, device=device)
        c = torch.zeros(batch_size, self.controller_size, device=device)
        return memory, prev_read_weights, prev_write_weights, (h, c)

    def forward(self, inputs):
        """
        Forward pass of a sequence through the NTM.

        Args:
            inputs: (batch_size, seq_len, input_size)
        Returns:
            outputs: (batch_size, seq_len, output_size)
            memory_history: history of memory states for debugging
        """
        batch_size, seq_len, _ = inputs.shape
        device = inputs.device

        memory, prev_read_w, prev_write_w, hc = self.init_state(batch_size, device)

        outputs = []
        memory_history = []

        for t in range(seq_len):
            x_t = inputs[:, t, :]

            # --- Read ---
            read_key = self.read_key_layer(hc[0])
            read_beta = F.softplus(self.read_beta_layer(hc[0])) + 1e-8
            read_gate = torch.sigmoid(self.read_gate_layer(hc[0]))
            read_shift = F.softmax(self.read_shift_layer(hc[0]), dim=-1)
            read_gamma = F.softplus(self.read_gamma_layer(hc[0])) + 1

            read_weights = self.attention(
                key=read_key, memory=memory, beta=read_beta,
                gate=read_gate.squeeze(-1), shift=read_shift,
                gamma=read_gamma, prev_weights=prev_read_w
            )

            read_vector = torch.sum(read_weights.unsqueeze(-1) * memory, dim=1)

            # --- Controller ---
            controller_input = torch.cat([x_t, read_vector], dim=-1)
            h_new, c_new = self.controller(controller_input, hc)
            hc = (h_new, c_new)

            # --- Write ---
            write_key = self.write_key_layer(h_new)
            write_beta = F.softplus(self.write_beta_layer(h_new)) + 1e-8
            write_gate = torch.sigmoid(self.write_gate_layer(h_new))
            write_shift = F.softmax(self.write_shift_layer(h_new), dim=-1)
            write_gamma = F.softplus(self.write_gamma_layer(h_new)) + 1
            erase_vector = torch.sigmoid(self.write_erase_layer(h_new))
            add_vector = torch.tanh(self.write_add_layer(h_new))

            write_weights = self.attention(
                key=write_key, memory=memory, beta=write_beta,
                gate=write_gate.squeeze(-1), shift=write_shift,
                gamma=write_gamma, prev_weights=prev_write_w
            )

            # Apply write operation
            erase = 1 - write_weights.unsqueeze(-1) * erase_vector.unsqueeze(1)
            memory = memory * erase + write_weights.unsqueeze(-1) * add_vector.unsqueeze(1)

            # --- Output ---
            output = self.output_layer(h_new)
            outputs.append(output)

            prev_read_w = read_weights
            prev_write_w = write_weights
            memory_history.append(memory.detach().clone())

        outputs = torch.stack(outputs, dim=1)
        return outputs, memory_history

Copy Task

The copy task is the canonical test for NTMs: learning to memorize an input sequence and then reproduce it exactly.

class CopyTask:
    """
    Copy task: the NTM must learn to reproduce a sequence.

    Format:
        [Random sequence] + [End delimiter] + [Recall start delimiter]
    The NTM must produce the original sequence after the delimiter.
    """
    def __init__(self, seq_length, input_dim):
        self.seq_length = seq_length
        self.input_dim = input_dim

    def generate_batch(self, batch_size=32, device='cpu'):
        """Generates a batch of training data."""
        sequence = torch.rand(batch_size, self.seq_length, self.input_dim, device=device)

        end_marker = torch.zeros(batch_size, 1, self.input_dim, device=device)
        end_marker[:, :, -1] = 1.0

        start_marker = torch.zeros(batch_size, 1, self.input_dim, device=device)
        start_marker[:, :, -2] = 1.0

        inputs = torch.cat([sequence, start_marker], dim=1)
        targets = torch.cat([sequence, end_marker], dim=1)

        return inputs, targets


def train_copy_task(ntm, num_iterations=10000, batch_size=32, seq_length=8, lr=1e-4):
    """Trains the NTM on the copy task."""
    task = CopyTask(seq_length, ntm.input_size)
    optimizer = torch.optim.Adam(ntm.parameters(), lr=lr)
    criterion = nn.BCEWithLogitsLoss()

    for iteration in range(num_iterations):
        inputs, targets = task.generate_batch(batch_size)
        outputs, _ = ntm(inputs)

        loss = criterion(outputs, targets)
        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(ntm.parameters(), max_norm=10.0)
        optimizer.step()

        if iteration % 500 == 0:
            print(f"Iteration {iteration} | Loss: {loss.item():.6f}")

    print(f"Training complete. Final loss: {loss.item():.6f}")
    return ntm


def evaluate_copy(ntm, seq_length=12):
    """
    Evaluates copy ability on a sequence of unseen length.
    This is the key generalization test for NTMs.
    """
    ntm.eval()
    task = CopyTask(seq_length, ntm.input_size)
    inputs, targets = task.generate_batch(batch_size=1, device='cpu')

    with torch.no_grad():
        outputs, _ = ntm(inputs)
        predicted = (torch.sigmoid(outputs) > 0.5).float()

    accuracy = (predicted == targets).float().mean()
    print(f"Accuracy on sequence of length {seq_length}: {accuracy.item():.4f}")
    return accuracy

Association Task

class AssociationTask:
    """
    Association task: the NTM must learn to associate input -> output
    pairs and then generalize to new associations.

    Format:
        Encoding phase: (key, value) presented sequentially
        Recall phase: a key is given, the NTM must produce the value
    """
    def __init__(self, num_pairs, key_dim, value_dim):
        self.num_pairs = num_pairs
        self.key_dim = key_dim
        self.value_dim = value_dim

    def generate_batch(self, batch_size=32, device='cpu'):
        keys = torch.zeros(batch_size, self.num_pairs, self.key_dim, device=device)
        for b in range(batch_size):
            perm = torch.randperm(self.key_dim)[:self.num_pairs]
            for i, idx in enumerate(perm):
                keys[b, i, idx] = 1.0

        values = torch.rand(batch_size, self.num_pairs, self.value_dim, device=device)

        inputs = torch.cat([keys, keys], dim=1)
        targets = torch.cat([values, values], dim=1)

        return inputs, targets

Hyperparameters

The hyperparameters of an NTM are numerous and their tuning is crucial for convergence:

Practical Recommendations

• For the copy task: N=128, W=20, LSTM controller of size 100, 1 reader, 1 writer are sufficient for sequences of length 10-20.
• For association: N=256, W=30, controller of size 200 are recommended since the task requires more distinct storage capacity.
• Initialization: Initializing shift weights to -10 (favoring zero displacement) and erase vectors to 1.0 (no erasure by default) are proven practices that considerably accelerate convergence.
• Gradient clipping: Essential. Without clipping at a norm of 10, NTMs diverge almost systematically due to the exponentials in the attention softmax.

Advantages and Limitations

Advantages

1. Explicit memory capacity: Unlike LSTMs where memory is implicit in recurrent connections, the NTM has a dedicated, addressable storage space. This allows retaining exactly the necessary information, not just compressed summaries.
2. Generalization to unseen lengths: This is the most remarkable advantage. An NTM trained on sequences of length 8 can often copy sequences of length 20 or more — something a classical LSTM fails at. The computation/memory separation enables this extrapolation.
3. Algorithmic learning: NTMs can learn pseudo-algorithms: sorting, searching, repeated copying. They develop emergent behaviors resembling pointers and loops, without being programmed with them.
4. Fully differentiable: All operations — reading, writing, attention — are smooth and derivable. Training via backpropagation works end-to-end, without the need for reinforcement learning or stochastic policies.

Limitations

1. High computational complexity: Each time step requires computing the cosine similarity between the key vector and all N memory locations. The complexity is O(N × W) per head, per time step. With N=256, that’s 200 times more expensive than an equivalently sized LSTM.
2. Training difficulty: NTMs converge slowly and require meticulous hyperparameter tuning. The loss landscape is very irregular, with many local minima. Initialization is critical.
3. Fixed-size memory: Unlike a computer’s stack which can grow, the NTM’s memory has a predefined size N. If the task requires more locations, the NTM fails. This is the problem that the DNC partially solves with dynamic allocation.
4. Limited scale: NTMs work well on short synthetic tasks but don’t scale to real-world tasks (natural language processing, computer vision). This is why the field evolved toward Transformers and large-scale external attention mechanisms.
5. Difficult interpretability: Although the memory is explicit, understanding what is stored at each location and why remains difficult. Attention weights are dense and their interpretation is not trivial.

4 Concrete Use Cases

1. Copying and Recalling Long Sequences

In dialogue or text generation systems, an NTM can serve as a long-term context memory. When a conversation exceeds the immediate context window, the NTM stores important elements in its external memory and retrieves them when necessary — for example, a user’s name mentioned 200 turns ago.

2. Reasoning on Graphs

NTMs can learn to traverse graphs by using their memory as a workspace. Each visited node can be stored in a memory location, and the read/write heads are used to track the path taken. Research has shown that NTMs can learn graph traversal algorithms like BFS and DFS in a purely learned manner.

3. Robotic Control with Memory

A mobile robot can use an NTM as a cognitive map: the external memory stores explored locations and their properties, while the controller decides the next action. Unlike a classical occupancy grid, the map learned by the NTM is semantic — it encodes not just “there is a wall here” but “this area has been explored and contains an interesting object.”

4. Neural Programming

The NTM is the foundation of neural programming: learning programs from examples. By combining external memory with a sufficiently expressive controller, one can learn to execute simple programs (addition, sorting, binary search) directly from input/output pairs. This approach paves the way for systems that learn not a fixed function, but a generalizable algorithm.

See Also

• Master the Hypocycloid and Grid Points with Python: A Complete Guide
• Python and the Hilbert Hotel Paradigm: Modeling and Innovative Applications