can NALU learn to do gaussian elimination?

.gitignore

@@ -1,4 +1,3 @@
-*.ipynb
 ckpt/
 logs/
 plots/

Gaussian Elimination.ipynb

@@ -0,0 +1,391 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%load_ext autoreload\n",
+    "%autoreload 2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import math\n",
+    "import random\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "import torch.nn.functional as F\n",
+    "\n",
+    "from models import *"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def train(model, optimizer, data, target, num_iters):\n",
+    "    for i in range(num_iters):\n",
+    "        out = model(data)\n",
+    "        loss = F.mse_loss(out, target)\n",
+    "        mea = torch.mean(torch.abs(target - out))\n",
+    "        optimizer.zero_grad()\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        if i % 5000 == 0:\n",
+    "            print(\"\\t{}/{}: loss: {:.3f} - mea: {:.3f}\".format(\n",
+    "                i+1, num_iters, loss.item(), mea.item())\n",
+    "            )"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Permutation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# permute the first column with the third\n",
+    "\n",
+    "A = torch.from_numpy(np.array([\n",
+    "    [0, 1, -1],\n",
+    "    [3, -1, 1],\n",
+    "    [1, 1, -2],\n",
+    "])).float()\n",
+    "\n",
+    "B = torch.from_numpy(np.array([\n",
+    "    [-1, 1, -0],\n",
+    "    [1, -1, 3],\n",
+    "    [-2, 1, 1],\n",
+    "])).float()\n",
+    "\n",
+    "P = torch.from_numpy(np.array([\n",
+    "    [0, 0, 1],\n",
+    "    [0, 1, 0],\n",
+    "    [1, 0, 0],\n",
+    "])).float()\n",
+    "\n",
+    "assert torch.allclose(torch.matmul(A, P), B)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\t1/40000: loss: 1.444 - mea: 0.927\n",
+      "\t5001/40000: loss: 0.077 - mea: 0.177\n",
+      "\t10001/40000: loss: 0.000 - mea: 0.007\n",
+      "\t15001/40000: loss: 0.000 - mea: 0.000\n",
+      "\t20001/40000: loss: 0.000 - mea: 0.000\n",
+      "\t25001/40000: loss: 0.000 - mea: 0.000\n",
+      "\t30001/40000: loss: 0.000 - mea: 0.000\n",
+      "\t35001/40000: loss: 0.000 - mea: 0.000\n",
+      "tensor([[ 0.0000, -0.0000,  1.0000],\n",
+      "        [-0.0001,  1.0000, -0.0001],\n",
+      "        [ 1.0000,  0.0000, -0.0000]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "net = NeuralAccumulatorCell(3, 3)\n",
+    "optim = torch.optim.RMSprop(net.parameters(), lr=1e-4)\n",
+    "\n",
+    "train(net, optim, A, B, int(4e4))\n",
+    "\n",
+    "print(net.W.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "actual: \n",
+      "tensor([[ 0.0000, -0.0001,  1.0000],\n",
+      "        [-0.0000,  1.0000,  0.0000],\n",
+      "        [ 1.0000, -0.0001, -0.0000]])\n",
+      "\n",
+      "expected: \n",
+      "tensor([[ 0.,  0.,  1.],\n",
+      "        [ 0.,  1.,  0.],\n",
+      "        [ 1.,  0.,  0.]])\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"actual: \\n{}\\n\".format(net.W.transpose(0, 1)))\n",
+    "print(\"expected: \\n{}\\n\".format(P))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Column Scaling"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# scale the first column by 5\n",
+    "\n",
+    "A = torch.from_numpy(np.array([\n",
+    "    [0, 1, -1],\n",
+    "    [3, -1, 1],\n",
+    "    [1, 1, -2],\n",
+    "])).float()\n",
+    "\n",
+    "B = torch.from_numpy(np.array([\n",
+    "    [0, 1, -1],\n",
+    "    [15, -1, 1],\n",
+    "    [5, 1, -2],\n",
+    "])).float()\n",
+    "\n",
+    "P = torch.from_numpy(np.array([\n",
+    "    [5, 0, 0],\n",
+    "    [0, 1, 0],\n",
+    "    [0, 0, 1],\n",
+    "])).float()\n",
+    "\n",
+    "assert torch.allclose(torch.matmul(A, P), B)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\t1/90000: loss: 28.833 - mea: 3.000\n",
+      "\t5001/90000: loss: 21.383 - mea: 2.246\n",
+      "\t10001/90000: loss: 16.247 - mea: 1.900\n",
+      "\t15001/90000: loss: 11.851 - mea: 1.626\n",
+      "\t20001/90000: loss: 8.166 - mea: 1.367\n",
+      "\t25001/90000: loss: 5.191 - mea: 1.109\n",
+      "\t30001/90000: loss: 2.924 - mea: 0.852\n",
+      "\t35001/90000: loss: 1.366 - mea: 0.595\n",
+      "\t40001/90000: loss: 0.511 - mea: 0.340\n",
+      "\t45001/90000: loss: 0.207 - mea: 0.227\n",
+      "\t50001/90000: loss: 0.148 - mea: 0.184\n",
+      "\t55001/90000: loss: 0.103 - mea: 0.153\n",
+      "\t60001/90000: loss: 0.065 - mea: 0.122\n",
+      "\t65001/90000: loss: 0.036 - mea: 0.091\n",
+      "\t70001/90000: loss: 0.016 - mea: 0.060\n",
+      "\t75001/90000: loss: 0.004 - mea: 0.030\n",
+      "\t80001/90000: loss: 0.000 - mea: 0.001\n",
+      "\t85001/90000: loss: 0.000 - mea: 0.000\n",
+      "tensor([[ 5.0001e+00, -5.1951e-05,  5.1859e-05],\n",
+      "        [-5.0003e-05,  1.0000e+00, -5.0024e-05],\n",
+      "        [-5.0004e-05,  4.9928e-05,  9.9995e-01]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "net = NeuralAccumulatorCell(3, 3)\n",
+    "optim = torch.optim.RMSprop(net.parameters(), lr=1e-4)\n",
+    "\n",
+    "train(net, optim, A, B, int(9e4))\n",
+    "\n",
+    "print(net.W.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "actual: \n",
+      "tensor([[ 5.0001e+00, -5.0003e-05, -5.0004e-05],\n",
+      "        [-5.1951e-05,  1.0000e+00,  4.9928e-05],\n",
+      "        [ 5.1859e-05, -5.0024e-05,  9.9995e-01]])\n",
+      "\n",
+      "expected: \n",
+      "tensor([[ 5.,  0.,  0.],\n",
+      "        [ 0.,  1.,  0.],\n",
+      "        [ 0.,  0.,  1.]])\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"actual: \\n{}\\n\".format(net.W.transpose(0, 1)))\n",
+    "print(\"expected: \\n{}\\n\".format(P))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Column Elimination"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def basis_vec(k, n):\n",
+    "    \"\"\"Creates the k'th standard basis vector in R^n.\"\"\"\n",
+    "    error_msg = \"[!] k cannot exceed {}.\".format(n)\n",
+    "    assert (k < n), error_msg\n",
+    "    b = np.zeros([n, 1])\n",
+    "    b[k] = 1\n",
+    "    return b"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# add -3x the second column to the first => P = (I - (c)(e_k)(e_l.T))\n",
+    "\n",
+    "A = torch.from_numpy(np.array([\n",
+    "    [3, 1, -1],\n",
+    "    [3, -1, 1],\n",
+    "    [1, 1, -2],\n",
+    "])).float()\n",
+    "\n",
+    "B = torch.from_numpy(np.array([\n",
+    "    [0, 1, -1],\n",
+    "    [6, -1, 1],\n",
+    "    [-2, 1, -2],\n",
+    "])).float()\n",
+    "\n",
+    "P = torch.from_numpy(\n",
+    "    np.eye(3) + (-3)*basis_vec(1, 3).dot(basis_vec(0, 3).T)\n",
+    ").float()\n",
+    "\n",
+    "assert torch.allclose(torch.matmul(A, P), B)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\t1/90000: loss: 5.444 - mea: 1.667\n",
+      "\t5001/90000: loss: 1.509 - mea: 0.679\n",
+      "\t10001/90000: loss: 0.226 - mea: 0.247\n",
+      "\t15001/90000: loss: 0.084 - mea: 0.165\n",
+      "\t20001/90000: loss: 0.038 - mea: 0.111\n",
+      "\t25001/90000: loss: 0.010 - mea: 0.057\n",
+      "\t30001/90000: loss: 0.000 - mea: 0.005\n",
+      "\t35001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t40001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t45001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t50001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t55001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t60001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t65001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t70001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t75001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t80001/90000: loss: 0.000 - mea: 0.000\n",
+      "\t85001/90000: loss: 0.000 - mea: 0.000\n",
+      "tensor([[ 9.9995e-01, -3.0001e+00,  5.0084e-05],\n",
+      "        [ 4.9894e-05,  1.0000e+00, -4.9870e-05],\n",
+      "        [ 4.9999e-05,  4.9925e-05,  9.9995e-01]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "net = NeuralAccumulatorCell(3, 3)\n",
+    "optim = torch.optim.RMSprop(net.parameters(), lr=1e-4)\n",
+    "\n",
+    "train(net, optim, A, B, int(9e4))\n",
+    "\n",
+    "print(net.W.data)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "actual: \n",
+      "tensor([[ 9.9995e-01,  4.9894e-05,  4.9999e-05],\n",
+      "        [-3.0001e+00,  1.0000e+00,  4.9925e-05],\n",
+      "        [ 5.0084e-05, -4.9870e-05,  9.9995e-01]])\n",
+      "\n",
+      "expected: \n",
+      "tensor([[ 1.,  0.,  0.],\n",
+      "        [-3.,  1.,  0.],\n",
+      "        [ 0.,  0.,  1.]])\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"actual: \\n{}\\n\".format(net.W.transpose(0, 1)))\n",
+    "print(\"expected: \\n{}\\n\".format(P))"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.6.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

failures.py

@@ -58,7 +58,7 @@ def main():
         mses = []
         for i in range(100):
             net = MLP(4, 1, 8, 1, non_lin)
-            optim = torch.optim.Adam(net.parameters(), lr=LEARNING_RATE)
+            optim = torch.optim.RMSprop(net.parameters(), lr=LEARNING_RATE)
             train(net, optim, train_data, NUM_ITERS)
             mses.append(test(net, test_data))
         all_mses.append(torch.cat(mses, dim=1).mean(dim=1))