ixaxaar/pytorch-dnc

📚 What is pytorch‑dnc?

pytorch‑dnc is a pure‑Python library that implements several memory‑augmented neural network architectures for PyTorch:

Architecture	Paper (original)	What it adds
DNC (Differentiable Neural Computer)	Graves et al., Nature 2016	A recurrent controller plus an external matrix‑valued memory that can be read/written differentiably.
SDNC (Sparse DNC)	Rae et al., NeurIPS 2018	Same as DNC but uses sparse reads/writes to scale to thousands of memory slots.
SAM (Sparse Access Memory)	Rae et al., NeurIPS 2018	A stand‑alone sparse memory module that can be plugged into any RNN controller.

The library lets you drop one of these modules into a PyTorch model and train it end‑to‑end on tasks that require algorithmic reasoning (copying sequences, addition, arg‑max, etc.).

---

🚀 Quick start

# Install the package from PyPI
pip install dnc

# (optional) GPU‑accelerated sparse ops need FAISS
conda install -c pytorch faiss-gpu

Or install from source:

git clone https://github.com/ixaxaar/pytorch-dnc
cd pytorch-dnc
pip install -r requirements.txt
pip install -e .

---

🛠️ How to use the modules

All three classes share the same constructor signature (most arguments have sensible defaults). Below is a minimal example for the classic DNC:

import torch
from dnc import DNC

# Create a DNC with a 64‑dim input and 128‑dim hidden state
model = DNC(
    input_size=64,
    hidden_size=128,
    nr_cells=100,          # 100 memory slots
    cell_size=32,          # each slot stores a 32‑dim vector
    read_heads=4,
    batch_first=True,
    device=torch.device('cuda:0')
)

# Initial hidden state (controller, memory, read vectors) – let the model create them lazily
h = (None, None, None)

# Forward a random batch (seq_len=10, batch=4)
output, (h_ctrl, h_mem, h_read) = model(
    torch.randn(10, 4, 64),  # (seq_len, batch, input_dim)
    h,
    reset_experience=True   # clear memory at the start of a new episode
)

The same pattern works for SDNC and SAM; the only extra arguments are the sparsity‑related ones (sparse_reads, temporal_reads).

Debug mode

Pass debug=True to the constructor. The forward call then returns a third value – a dictionary of NumPy arrays that contain the raw memory matrices (memory, link_matrix, read_weights, …). These can be visualised with tools like Visdom to inspect how the network is using its external memory.

---

📊 Example tasks shipped with the repo

The repository includes ready‑to‑run scripts that reproduce the classic DNC experiments:

Task	What it tests	How to run
Copy task	Ability to store and reproduce an input sequence of arbitrary length.	python ./tasks/copy_task.py -cuda 0 -optim adam -sequence_max_length 8
Addition task	Learning to sum two numbers placed in a long sequence (the original “learning to add” benchmark).	python ./tasks/add_task.py … (see script help)
Argmax task	Finding the index of the maximum element in a sequence.	python ./tasks/argmax_task.py …

All tasks accept a rich set of command‑line options (learning rate, optimizer, memory size, curriculum learning, etc.). For the copy task you can also launch a Visdom server (pip install visdom && python -m visdom.server) to watch a heat‑map of the memory matrix during training.

---

🏗️ Code organization

├─ dnc/                # core implementations (DNC, SDNC, SAM)
├─ tasks/              # training scripts for copy / add / argmax
├─ docs/               # architecture diagram and screenshots
├─ requirements.txt    # Python dependencies
└─ tests/              # pytest suite

The core modules are pure PyTorch and rely on FAISS only for the GPU‑accelerated sparse reads/writes used by SDNC and SAM.

---

📦 Who might use this?

• Researchers experimenting with neural Turing machines or other differentiable data structures.
• Practitioners who need a plug‑and‑play external memory for sequence‑to‑sequence models, e.g., program synthesis or reasoning over long contexts.
• Educators looking for a concrete, runnable implementation to demonstrate the concepts from the DNC papers.

---

📚 Further reading

• Original DNC paper – Hybrid computing using a neural network with dynamic external memory (Graves et al., Nature 2016).
• Sparse memory paper – Scaling Memory‑Augmented Neural Networks with Sparse Reads and Writes (Rae et al., NeurIPS 2018).

---

✅ TL;DR

pytorch‑dnc gives you ready‑made, well‑documented PyTorch modules for DNC, SDNC, and SAM, complete with example training scripts and a debug mode that lets you peek inside the external memory. Install via pip install dnc, drop a DNC/SDNC/SAM into your model, and start training on algorithmic tasks that need long‑range memory.