ixaxaar/pytorch-dnc

Differentiable Neural Computers, Sparse Access Memory and Sparse Differentiable Neural Computers, for Pytorch

📚 pytorch‑dnc とは何か?

pytorch‑dnc は、PyTorch用の純粋なPythonライブラリで、いくつかのメモリ拡張ニューラルネットワークアーキテクチャを実装しています:

アーキテクチャ 論文(元論文) 追加機能
DNC(Differentiable Neural Computer) Graves et al., Nature 2016 再帰的なコントローラーと、微分可能な読み書きが可能な行列値の外部メモリ。
SDNC(Sparse DNC) Rae et al., NeurIPS 2018 DNCと同様だが、スパースな読み書きを使用して数千のメモリスロットまでスケーリング可能。
SAM(Sparse Access Memory) Rae et al., NeurIPS 2018 任意のRNNコントローラーに接続可能な独立したスパースメモリモジュール。

このライブラリを使えば、これらのモジュールのいずれかをPyTorchモデルに組み込み、アルゴリズム的推論(シーケンスのコピー、加算、arg-maxなど)を必要とするタスク上でエンドツーエンドで学習できます。


🚀 クイックスタート

# PyPIからパッケージをインストール
pip install dnc

# (オプション)GPU加速のスパース演算にはFAISSが必要
conda install -c pytorch faiss-gpu

またはソースからインストール:

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

🛠️ モジュールの使い方

3つのクラスはすべて同じコンストラクタシグネチャを持ちます(ほとんどの引数には意味のあるデフォルト値があります)。以下は古典的なDNCの最小限の例です:

import torch
from dnc import DNC

# 64次元入力と128次元隠れ状態を持つDNCを作成
model = DNC(
    input_size=64,
    hidden_size=128,
    nr_cells=100,          # 100個のメモリスロット
    cell_size=32,          # 各スロットは32次元ベクトルを格納
    read_heads=4,
    batch_first=True,
    device=torch.device('cuda:0')
)

# 初期隠れ状態(コントローラー、メモリ、読み取りベクトル)— モデルが遅延して作成するようにする
h = (None, None, None)

# ランダムなバッチをフォワード(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   # 新しいエピソードの開始時にメモリをクリア
)

SDNCSAM でも同じパターンが使えます。追加の引数は sparse_readstemporal_reads のみです。

デバッグモード

コンストラクタに debug=True を渡してください。この場合、フォワード呼び出しは3番目の値としてNumPy配列の辞書を返します。これは memorylink_matrixread_weights などの生のメモリ行列を含みます。これらのデータは Visdom などのツールで可視化でき、ネットワークが外部メモリをどのように使っているかを確認できます。


📊 リポジトリに同梱された例題タスク

リポジトリには、古典的なDNC実験を再現できる実行可能なスクリプトが含まれています:

タスク テスト内容 実行方法
コピー課題 任意長の入力シーケンスを格納・再現する能力。 python ./tasks/copy_task.py -cuda 0 -optim adam -sequence_max_length 8
加算課題 長いシーケンス内に配置された2つの数値を合計する学習(元の「加算を学ぶ」ベンチマーク)。 python ./tasks/add_task.py …(スクリプトヘルプを参照)
argmax課題 シーケンス内の最大要素のインデックスを見つける。 python ./tasks/argmax_task.py …

すべてのタスクは豊富なコマンドラインオプション(学習率、最適化手法、メモリサイズ、カリキュラム学習など)を受け入れます。コピー課題では、Visdom サーバーを起動(pip install visdom && python -m visdom.server)して、学習中のメモリ行列のヒートマップを監視することもできます。


🏗️ コード構成

├─ dnc/                # コア実装(DNC、SDNC、SAM)
├─ tasks/              # コピー/加算/argmax用のトレーニングスクリプト
├─ docs/               # アーキテクチャ図とスクリーンショット
├─ requirements.txt    # Python依存関係
└─ tests/              # pytestスイート

コアモジュールは純粋なPyTorchで構成されており、FAISS はSDNCとSAMで使用されるGPU加速スパース読み書きのみに依存しています。


📦 誰がこのライブラリを使うべきか?

  • ニューラルチューリングマシンやその他の微分可能なデータ構造を実験している研究者。
  • シーケンスからシーケンスへのモデルに即時接続可能な外部メモリが必要な実務家(例:プログラム合成、長文の推論など)。
  • DNC論文の概念を具体的かつ実行可能な形で示したい教育者。

📚 補足リーディング

  • オリジナルDNC論文Hybrid computing using a neural network with dynamic external memory(Graves et al., Nature 2016)。
  • スパースメモリ論文Scaling Memory‑Augmented Neural Networks with Sparse Reads and Writes(Rae et al., NeurIPS 2018)。

✅ TL;DR

pytorch‑dnc は、DNC、SDNC、SAM用の準備完了済みで、ドキュメントが充実したPyTorchモジュールを提供し、例題トレーニングスクリプトと外部メモリの中身を覗けるデバッグモードを備えています。pip install dnc でインストールし、モデルに DNC/SDNC/SAM を組み込み、長距離メモリを必要とするアルゴリズムタスクの学習を開始できます。

関連

  • プロジェクト
  • プロジェクト
  • プロジェクト
  • プロジェクト
  • Dispatch