Key Value Memory Networks

This repo contains the implementation of Key Value Memory Networks for Directly Reading Documents in Tensorflow. The model is tested on bAbI.

Tutorial

• There is a must-read tutorial on Memory Networks for NLP from Jason Weston @ ICML 2016.

• [Video] [Slides] Sumit Chopra, from Facebook AI, gave a lecture about Reasoning, Attention and Memory at Deep Learning Summer School 2016.

Get Started

git clone https://github.com/siyuanzhao/key-value-memory-networks.git

mkdir ./key-value-memory-networks/logs
mkdir ./key-value-memory-networks/data/
cd ./key-value-memory-networks/data
wget http://www.thespermwhale.com/jaseweston/babi/tasks_1-20_v1-2.tar.gz
tar xzvf ./tasks_1-20_v1-2.tar.gz

cd ../
python single.py

Usage

Running a single bAbI task

# Train the model on a single task <task_id>
python single.py --task_id <task_id>

There are serval flags within single.py. Below is an example of training the model on task 20 with specific learning rate, feature_size and epochs.

python single.py --task_id 20 --learning_rate 0.005 --feature_size 40 --epochs 200

Check all avaiable flags with the following command.

python single.py -h

Running a joint model on all bAbI tasks

python joint.py

There are also serval flags within joint.py. Below is an example of training the joint model with specific learning rate, feature_size and epochs.

python joint.py --learning_rate 0.005 --feature_size 40 --epochs 200

Check all avaiable flags with the following command.

python joint.py -h

Results

Bag of Words

Joint Model

The model is jointly trained on 20 tasks (1k training examples / weakly supervised) with following hyperparameters.

• BATCH_SIZE=50
• EMBEDDING_SIZE=40
• EPOCHS=200
• EPSILON=0.1
• FEATURE_SIZE=50
• HOPS=3
• L2_LAMBDA=0.1
• LEARNING_RATE=0.001
• MAX_GRAD_NORM=20.0
• MEMORY_SIZE=50
• READER=bow

python joint.py

• results on 10k training examples are here

RNN(GRU)

Requirements

• tensorflow 1.10
• scikit-learn 0.19
• six 1.10.0