self_organizing_multiparadigm_networks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dfa6ac80-623d-4e66-9cd5-478cfd20b562",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6c171ab5-9aed-4fcd-abec-cd8b6f8ee15b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T21:44:24.720667Z",
     "iopub.status.busy": "2021-10-22T21:44:24.719511Z",
     "iopub.status.idle": "2021-10-22T21:44:24.761363Z",
     "shell.execute_reply": "2021-10-22T21:44:24.760470Z",
     "shell.execute_reply.started": "2021-10-22T21:44:24.720355Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%config Completer.use_jedi=False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13d444c5-4389-43c1-a633-e573ba2a6b49",
   "metadata": {},
   "source": [
    "Latex headers\n",
    "$\\newcommand{\\E}{\\mathbb{E}}$\n",
    "$\\newcommand{\\Var}{\\mathrm{Var}}$\n",
    "$\\newcommand{\\Cov}{\\mathrm{Cov}}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6bb3918a-49f2-4331-8f4e-359ae8a8370c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T21:44:25.873617Z",
     "iopub.status.busy": "2021-10-22T21:44:25.873452Z",
     "iopub.status.idle": "2021-10-22T21:44:41.086755Z",
     "shell.execute_reply": "2021-10-22T21:44:41.084191Z",
     "shell.execute_reply.started": "2021-10-22T21:44:25.873600Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import tqdm\n",
    "import math\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import salina\n",
    "import wandb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d9803906-a1d2-4e7f-b805-21fb9668d8a7",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T21:44:43.123224Z",
     "iopub.status.busy": "2021-10-22T21:44:43.123024Z",
     "iopub.status.idle": "2021-10-22T21:44:51.726616Z",
     "shell.execute_reply": "2021-10-22T21:44:51.724803Z",
     "shell.execute_reply.started": "2021-10-22T21:44:43.123202Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: \"WANDB_NOTEBOOK_NAME\"=\"critical_multiparadigm_networks\"\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mjacobfv\u001b[0m (use `wandb login --relogin` to force relogin)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                    Syncing run <strong><a href=\"https://wandb.ai/jacobfv/critical_multiparadigm_networks/runs/3r7n6iqc\" target=\"_blank\">playful-sound-3</a></strong> to <a href=\"https://wandb.ai/jacobfv/critical_multiparadigm_networks\" target=\"_blank\">Weights & Biases</a> (<a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">docs</a>).<br/>\n",
       "\n",
       "                "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<button onClick=\"this.nextSibling.style.display='block';this.style.display='none';\">Display W&B run</button><iframe src=\"https://wandb.ai/jacobfv/critical_multiparadigm_networks/runs/3r7n6iqc?jupyter=true\" style=\"border:none;width:100%;height:420px;display:none;\"></iframe>"
      ],
      "text/plain": [
       "<wandb.sdk.wandb_run.Run at 0x7fb6a58fbe20>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%env \"WANDB_NOTEBOOK_NAME\" \"self_organizing_multiparadigm_networks\"\n",
    "wandb.init(project=\"self_organizing_multiparadigm_networks\", entity=\"jacobfv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a0d1cb9-b803-4ee0-b055-3432f35efa65",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T16:33:47.575009Z",
     "iopub.status.busy": "2021-10-22T16:33:47.574782Z",
     "iopub.status.idle": "2021-10-22T16:33:47.580682Z",
     "shell.execute_reply": "2021-10-22T16:33:47.579840Z",
     "shell.execute_reply.started": "2021-10-22T16:33:47.574983Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "wandb.config = {\n",
    "    \"epochs\": 100,\n",
    "    \"batch_size\": 128,\n",
    "    \"sgd_learning_rate\": 0.001,\n",
    "    \"ul_learning_rate\": 0.001,\n",
    "    \"IP_stddev\": 0.05,\n",
    "    \"IP_mean\": 0.1,\n",
    "    \"cm_alpha\": 2.0,\n",
    "    \"lambda_reconstruct_x\": 1.0,\n",
    "    \"lambda_reconstruct_h1\": 0.9,\n",
    "    \"lambda_reconstruct_h2\": 0.8,\n",
    "    \"lambda_reconstruct_h3\": 0.7,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2f5fa61-33b3-47ed-a63e-dd282c8384e6",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Self-Organizing Multi-Paradigm Networks (SOMPNets)\n",
    "\n",
    "Self-organizing multi-paradigm networks (SOMPNets) are a family of deep learning architectures designed for iterative update processing (such as progressive representation refinement). SOMPNets are characterized by \n",
    "- activations displaying critical phenomena including very long correlation distance and ergodicity\n",
    "- layer-local iteratively-applied unsupervised parameter update rules\n",
    "- training by a joint combination of supervised, self-supervised, unsupervised, and reinforcement learnin\n",
    "- pseudo-linear activation transformations allowing linear combination of separate hidden states and parameters from asynchronously trained networks\n",
    "\n",
    "**Ideally keep the bucket close to the edge at all times**"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a4fdc1-b3ff-4489-bb19-82da52ed3ba4",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Design\n",
    "\n",
    "Each SOMPCell in a self-organizing multi-paradigm network updates by a system of fast and slow temporal dynamics\n",
    "\n",
    "### Fast Dyanmics\n",
    "\n",
    "**Activations** may be dense or sparse and are given by\n",
    "$$Z_t = X^1_t W^1 + X^2_t W^2 + \\dots + b + \\zeta$$\n",
    "$$B_t = a_{b,decay} B_{t-1} + Z_t - a_{b,discharge} |Y^{sparse}_{t-1}|$$\n",
    "$$Y^{sparse}_t = \\alpha_{sparse} \\Theta_{\\pm} (\\alpha_{sparse}^{-1} Z)$$\n",
    "$$Y^{dense}_t = \\alpha_{dense} \\tanh (\\alpha_{dense}^{-1}Z_t)$$\n",
    "where\n",
    "- $X^1 \\in \\mathbb{R}^{\\dots \\times t \\times n_1}, X^2 \\in \\mathbb{R}^{\\dots \\times t \\times n_2}, \\dots$ are the inputs,\n",
    "- $W^1 \\in \\mathbb{R}^{n_1 \\times n_{out}}, W^2 \\in \\mathbb{R}^{n_2 \\times n_{out}}, \\dots$ are weights corresponding to each input,\n",
    "- $b \\in \\mathbb{R}^{n_{out}}$ is the bias vector,\n",
    "- $\\zeta \\in \\mathbb{R}^{n_{out}} \\sim \\mathcal{N}(\\mu=\\mu_{\\zeta}, \\Sigma=\\Sigma_{\\zeta})$ is a noise factor,\n",
    "- $Z_t \\in \\mathbb{R}^{\\dots \\times t \\times n_{out}}$ is the input summand\n",
    "- $B_t \\in \\mathbb{R}^{\\dots \\times t \\times n_{out}}$ is the sparse activation bucket\n",
    "- $Y^{\\{sparse,dense\\}} \\in \\mathbb{R}^{\\dots \\times t \\times n_{out}}$ is the output activation\n",
    "\n",
    "and a positive-negative threshold function\n",
    "$$\\Theta_{\\pm}(x) = \\begin{cases}\n",
    "    +1 & \\text{if  } x \\ge 1 \\\\\n",
    "    0  & \\text{if  } -1 < x < 1 \\\\\n",
    "    -1 & \\text{if  } x \\le -1 \n",
    "\\end{cases}$$\n",
    "with gradients flowing through $\\Theta_{\\pm}(x)$ as if it were linear with respect to $x$.\n",
    "\n",
    "**Dense activations:** Temporality rescaling dense activation inputs extends the $\\tanh$ nonsaturating regime to $\\approx( -\\alpha_{dense}/2, \\alpha_{dense}/2)$ which facilitates long propagation distances and helps keep the gradients alive on backpropagation. Within the psuedo-linear regime, activations and weights may be linearly combined facilitating decentralized and asynchronous training. \n",
    "\n",
    "**Sparse activations** utilize a presynaptic bucket $B$ to store signals which gradually decay by a factor of $a_{b,decay}$ at every time step and abruptly decrement by $a_{b,discharge}$ when discharging. \n",
    "\n",
    "In both the dense and sparse case, activations are expressed within a common continuous/discrete $(-\\alpha, +\\alpha)$ open/closed range. Weights $W$ are intiailized with the expectation that inputs share an equivalent domain, and bias $b$ is initialized at $0$. Both trainable variables $W$, $b$ are $L_2$-regularized to minimize exploding gradients and alleviate soft or hard signal saturation. (Conversely, noise covariance $\\Sigma_{\\zeta}$ can be increased with the aim of encouraging weight growth (hence, saturating activations) and in later epochs might be increased (even cyclicly) to help sparsify weights.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6535c3b7-0b28-4713-beec-cf6708e5a572",
   "metadata": {},
   "source": [
    "**Multiparadigm training** blends unsupervised processes, supervisory and reinforcement signals, and self-supervised objectives to optimize the trainable parameters of CMPNets. Being end-to-end differentiable, individual critical-multiparadigm cells (CMPCell's) can be combined with other deep learning layers such as convolutional layers and attention mechanisms. Unsupervised weight update rules compute gradients *on each foreward pass* with an unsupervised update rate that may optionally be modulated by a reward signal. Gradients may be applied immediantly, or they may be accumulated over a time sequence and then applied along with backpropagated gradients by the optimizer. Note that combining unsupervised update rules with supervisory ones requires flipping the gradient in code since the optimizer updates variables in the reverse direction of gradients (however I keep the sign positive in math). I apply the following 7 update rules with the intention that weights and biases should adapt to represent the statistical dynamics underlying their associated activations:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf128e66-3692-4b6b-afd9-65a259cb0827",
   "metadata": {},
   "source": [
    "**1. Spike Timing Dependant Plasticity (STDP)** makes synapse $w^{ji}$ decrease when $y_i$ activates before $x_j$ (since the connection must've not been important) and increase when $x_i$ activates after $y_j$ (since the connection made a significant contribution). It has no effect when one of the values is $0$ and an inverse effect when one value is positive and the other is negative (or vice versa):\n",
    "$$\\Delta w_{STDP,ji} = \\hat{x}_{t-1,j} \\hat{y}_{t,i} - \\hat{x}_{t,j} \\hat{y}_{t-1,i}$$\n",
    "with $\\hat{\\cdot{}}$ representing $\\cdot{} \\div \\alpha$ where $\\alpha$ is corresponding to the layer that produced the activation in heterogeneous architectures. In the extended temporal case, STDP is expressed using Python-style slicing as:\n",
    "$$\\Delta w_{STDP,ji} = \\hat{x}_{:-1,j} \\hat{y}_{1:,i} - \\hat{x}_{1:,j} \\hat{y}_{:-1,i}$$\n",
    "Finally, full matrix deltas are computed as:\n",
    "```python\n",
    "dW_stdp = - (einsum('...j,...i->ji', xhat[..., :-1, :], yhat[..., 1:, :])\n",
    "            - einsum('...j,...i->ji', xhat[..., 1:, :], yhat[..., :-1, :]))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68e42e7e-cb88-417b-995f-31aa97767301",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T21:53:04.723956Z",
     "iopub.status.busy": "2021-10-22T21:53:04.723682Z",
     "iopub.status.idle": "2021-10-22T21:53:04.732861Z",
     "shell.execute_reply": "2021-10-22T21:53:04.731921Z",
     "shell.execute_reply.started": "2021-10-22T21:53:04.723928Z"
    },
    "tags": []
   },
   "source": [
    "**2. Covariance Decay (CD)** makes synapse $w^{ji}$ decay or grow in nonlinear proportion to the absolute covariance between $x_j$ and $y_i$ computed via $\\beta_{cd}$ rolling means:\n",
    "$$\\sigma_{ji} = \\Cov(\\hat{x}_{:-1,j} \\hat{y}_{1:,i}) = \\E_{\\beta_{cd}}{[ \\hat{x}_{t-1,j}\\hat{y}_{t,i} ]} - \\E_{\\beta_{cd}}{[ \\hat{x}_{t,j} ]} \\E_{\\beta_{cd}}{[ \\hat{y}_{t,i} ]}$$\n",
    "The absolute of this covariance factor $|\\sigma|$ is next scaled around $1$ by a learned parameter $a_{cd}$ giving the covariance decay coefficient: $$c_{ji} = a_{cd}(|\\sigma_{ji}|-1)+1$$ Ideally, the covariance decay coefficient $c_{ji}$ could be applied to directly onto its corresponding weight as $w_{ji} \\leftarrow c_{ji}w_{ji}$. However to remain compatible with gradient-based update paradigms, this coefficient is finally expressed as a weight increment to be compounded with other gradients:\n",
    "$$\\Delta w_{CD,ji} = (1-c_{ji}) w_{ji}$$\n",
    "\n",
    "In tensor computations, this is expressed:\n",
    "```python\n",
    "cov = Bmean(xhat[..., :-1, :, None] * yhat[..., 1:, None, :], beta=beta_cd, axis=-3) \\\n",
    "    - Bmean(xhat, axis=-2)[..., :, :, None] * Bmean(yhat, axis=-2)[..., :, None, :]\n",
    "coef = a_cd * cov - a_cd + 1\n",
    "dW_cd = - (1 - coef) * W\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "279b812a-ae78-4178-9b84-dab546bc942a",
   "metadata": {},
   "source": [
    "**3. Structural Plasticity (SP)** randomly adds small values synapses between unconnected neurons by Bernoili probability factor $p_{SP}=\\frac{a_{sp}}{N_x N_y}$ which scales inversely quatratically with respect to the number of input $N_x$ and output $N_y$ dimensions. \n",
    "$$\\Delta w_{SP,ji} = d, \\; \\; \\;d \\sim \\mathcal{B}(\\ \\cdot\\ ; p=\\small{\\frac{a_{sp}}{N_x N_y}})$$\n",
    "$a_{sp}$ is made differentiable by the reparametrization trick:\n",
    "```python\n",
    "dW_SP = -sigmoid(TODO)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "178e9659-0e30-4e63-a272-1913e77face6",
   "metadata": {},
   "source": [
    "**4. Intrinsic plasticity (IP)** homeostatically shift the inputs maintain a mean firing rate $H_{IP} \\sim \\mathcal{N}(\\mu_{IP}=0.1, \\Sigma_{IP}^2 = 0)$ \n",
    "$$\\Delta \\mu_{\\zeta} = \\eta_{IP}[x(t) - H_{IP}]$$\n",
    "TODO\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b99b0af-9861-4f3b-9116-bb24317d3285",
   "metadata": {},
   "source": [
    "**5. Variance-Invariance-Covariance Sequence Regularization (VIC)** aims to build representations that are progressively insensitive to time by applying the following regularization penalties:\n",
    "- local temporal element invariance: Assuming that sequence elements in a local neighborhood represent the same information, maximize the similarity $s$ between $y_{\\dots,t,:}$ and $y_{\\dots,t-1,:},y_{\\dots,t-2,:},\\dots$ assigning exponential weight to nearer elements\n",
    "$$\\mathcal{L}_{VIC} = $$ \n",
    "https://arxiv.org/pdf/2109.00783.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04658a86-1dd6-4d0c-bc60-31414036c601",
   "metadata": {},
   "source": [
    "**6. L2 regularization** prevents weights from growing excessively large unless they serve a meaningful statistical purpose. Formally,\n",
    "$$ \\mathcal{L}_{L2,W} = \\frac{1}{2} \\sum_{ji} w_{ji}^2$$\n",
    "$$ \\mathcal{L}_{L2,b} = \\frac{1}{2} \\sum_{i} b_{i}^2$$\n",
    "and in code,\n",
    "```python\n",
    "add_loss(eta_L2W * (W**2).sum())\n",
    "add_loss(eta_L2W * (b**2).sum())\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6949f530-c06a-485d-9f9c-8d47e338c99c",
   "metadata": {},
   "source": [
    "**7. Mean Regularization (MR)** aim to preserve a mean activation across between SOMPCells by \n",
    "$$\\mathcal{L}_{MR} = (\\sum \\hat{y}-\\mu_{MR})^2$$\n",
    "with gradients backpropagated through all trainable parameters by\n",
    "```python\n",
    "loss_MR = (yhat.sum() - mu_MR)**2\n",
    "add_loss(eta_MR * loss_MR)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b547ba3e-0ebd-4088-8d42-9dd76d71d7f1",
   "metadata": {},
   "source": [
    "**7. Sparsity Regularization (SR)** (which applies only to sparse activation `SOMPCell`'s) aims to tune sparsity by penalizing activation KL-divergence from a Bernoilii distribution\n",
    "$$\\mathcal{L}_{SR} = \\sum \\biggl[ |\\hat{y}| \\log { \\frac{ |\\hat{y}| }{ a_{SR} } } - (1 - \\hat{y})| \\log { \\frac{ 1 - |\\hat{y}| }{ 1 - a_{SR} } } \\biggr]$$\n",
    "with gradients backpropagated through all trainable parameters by\n",
    "```python\n",
    "kl_SR = yhat * log(yhat / a_SR) - (1-yhat) * log((1-yhat)/(1-a_SR))\n",
    "loss_SR = kl_SR.sum()\n",
    "add_loss(eta_SR * loss_SR)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6da7b0b2-32a0-4f19-ad8d-19fb3a9e91db",
   "metadata": {},
   "source": [
    "**8. Supervised learning (SL)** trains weights with respect to a supervised or self-supervised objective as normal. \n",
    "$$ \\Delta W_{SL} = \\frac{\\delta \\mathcal{L}}{\\delta W} \\biggr{|}_{x}$$\n",
    "And in code,\n",
    "```python\n",
    "sl_loss = eta_SL * sl_loss\n",
    "sl_loss.backward()\n",
    "for var in trainable_vars:\n",
    "    var.grad *= eta_SL\n",
    "```\n",
    "\n",
    "Many of the problem domains SOMPNet are applied to lend self-supervised next-sequence-element prediction as a powerfu pretraining objective:\n",
    "```python\n",
    "loss = (ypred - ytrue)**2\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f80cb03c-770a-42fe-9200-0dca7f74d06d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T17:14:39.954773Z",
     "iopub.status.busy": "2021-10-22T17:14:39.954500Z",
     "iopub.status.idle": "2021-10-22T17:14:39.962797Z",
     "shell.execute_reply": "2021-10-22T17:14:39.961895Z",
     "shell.execute_reply.started": "2021-10-22T17:14:39.954746Z"
    },
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "Unsupervised weight updates are applied as gradients which accumulate over steps on a sequence along with backpropagation gradients before being applied to the trainable variables by an optimizer. Putting it all togethor, a foreward pass looks *like* this:\n",
    "```python\n",
    "def foreward(self, inputs, previous, ...):\n",
    "    \n",
    "    # possibly build new weights if new inputs have been added\n",
    "    self.build({k: inputs[k].shape for k in inputs.keys()})\n",
    "    \n",
    "    # foreward computations\n",
    "    zeta = ...\n",
    "    Z = sum(inputs[k] @ self.weights[k] for k in inputs.keys()) + self.bias + zeta\n",
    "    Y = alpha * tanh(Z / alpha)  # the output\n",
    "    # shaped like `previous` but shifted. i.e.: `previous[...,-1, :] == Y[..., -2, :]`\n",
    "    \n",
    "    # slow updates\n",
    "    # NOTE: Gradients are reversed from delta's given in above \n",
    "    # formulae because optimization is a minimization process\n",
    "    # TODO: double-check the above statement with torch.\n",
    "    # 1. STDP\n",
    "    for k in inputs.keys():\n",
    "        dW_stdp = ...\n",
    "        self.weights[k].grad -= lr_stdp * dW_stdp\n",
    "        \n",
    "    # 2. CD\n",
    "    for k in inputs.keys():\n",
    "        dW_cd = ...\n",
    "        self.weights[k].grad -= lr_cd * dW_cd\n",
    "        \n",
    "    # 3. SP\n",
    "    for k in inputs.keys():\n",
    "        dW_sp = ...\n",
    "        self.weights[k].grad -= lr_sp * dW_sp\n",
    "    \n",
    "    # 4. IP\n",
    "    self.bias.grad -= ...\n",
    "    \n",
    "    # 5. VIC\n",
    "    sim = ...\n",
    "    covar = ...\n",
    "    invar = ...\n",
    "    self.add_loss(lambda_vic_sim*sim + lambda_vic_covar*covar + lambda_vic_invar*invar)\n",
    "    \n",
    "    # ... other update rules\n",
    "    \n",
    "    # return output signal with values for next time step\n",
    "    return Y\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97665533-7c2f-45f3-ae0e-9b7c4576e2a9",
   "metadata": {},
   "source": [
    "### Training\n",
    "\n",
    "The objective of weight training is to move optimization into the inner loop. This almost inherently demands cyclic multi-interaction step network toplogies and update processes. I blend gradient descent and unsupervised updates as:\n",
    "```python\n",
    "# `env` can be a dataset providing sequential examples (videos, text, robot task)\n",
    "while not env.done:\n",
    "    # foreward pass and local unsupervised updates\n",
    "    y = model(x)\n",
    "    # supervised loss, critic function\n",
    "    loss = eta_SL * loss(y)\n",
    "    loss.backward()  # accumulate SL gradients\n",
    "    # all gradients have accumulated, now apply them (faster optimization)\n",
    "    opt.apply_gradients()\n",
    "    \n",
    "# or wait until end  of interaction sequence to apply gradients (slower optimization)\n",
    "opt.apply_gradients()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54cf9a9e-4a49-4cd5-860f-f5c196208210",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "### CMPCell\n",
    "TODO: \n",
    "- specify time must be >2\n",
    "- change it to state that grads are added over the interaction steps\n",
    "\n",
    "\n",
    "### CMPNet\n",
    "\n",
    "Since `CMPCell`'s update iteratively, I use `salina` to wrap potentially cyclic layer connectivity graphs into a single seq2seq `CMPNet` which can connect and differentiate with other networks. Updates "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d722e31f-ef0a-4997-a6d0-4e598a3daa1a",
   "metadata": {},
   "source": [
    "## Experiments\n",
    "\n",
    "I test CMPNets on the following experiments:\n",
    "- progressive representation refinement\n",
    "- autoregression\n",
    "- decision making\n",
    "\n",
    "I also observe how well CMPNets perform acting as:\n",
    "- feedforeward dense layers\n",
    "- standard convolution layers {1D, 2D, 3D}\n",
    "- replacing single/multi-head attention key/query weight matrices\n",
    "- performing "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4edbba02-9569-4eb5-ae50-de795a0eae00",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true,
    "tags": []
   },
   "source": [
    "## Slow Dynamics\n",
    "\n",
    "Weights are modified with the intention that they should describe the statistical dynamics underlying actual activation patterns. I interpret signal values as probabilities from $-\\alpha$ (certainly false), $0$ (unknown), to $+\\alpha$ (certainly true) update weights by the following unsupervised rules:\n",
    "\n",
    "1. **Spike Timing Dependant Plasticity (STDP)** makes synapse $w^{ji}$ decrease when $y_i$ activates before $x_j$ (since the connection must've not been important) and increase when $x_i$ activates after $y_j$ (since the connection made a significant contribution). It has no effect when one of the values is $0$:\n",
    "$$\\Delta w_{ji} = \\eta_{STDP} ( x_{t-1,j} y_{t,i} - x_{t,j} y_{t-1,i} )$$\n",
    "In the extended temporal case,\n",
    "```python\n",
    "# input is shaped [.., T, Nin]\n",
    "# output is shaped [..., T, Nout]\n",
    "# W is shaped [Nin, Nout]\n",
    "\n",
    "# dw_ji = eta_STDP * (input_j[:-1]output_i[1:] - input_j[1:]output_i[:-1])\n",
    "dw = eta * (\n",
    "\n",
    "\n",
    ")\n",
    "\n",
    "```\n",
    "\n",
    "- general coorelation updates: $cw_{ji} = Coor(x_\n",
    "\n",
    "```python\n",
    "# spike timing dependant plasticity\n",
    "dW = input[:-1]output[1:] - input[1:]output[:-1]\n",
    "if reward:\n",
    "    dW *= (1+reward)  # or some other function to make rewarding events more memorable\n",
    "```\n",
    "\n",
    "The noise mean component homeostatically shifts to center average activation values around a desired setpoint\n",
    "```python\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1daf94-4f1c-4b55-8ca3-78d1dc859630",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-10-22T17:57:25.587638Z",
     "iopub.status.busy": "2021-10-22T17:57:25.587231Z",
     "iopub.status.idle": "2021-10-22T17:57:25.597771Z",
     "shell.execute_reply": "2021-10-22T17:57:25.596605Z",
     "shell.execute_reply.started": "2021-10-22T17:57:25.587589Z"
    },
    "tags": []
   },
   "source": [
    "## Regularization\n",
    "\n",
    "Weights $W_{.}$ and biases $b$ are $L_2$-regularized to avoid saturating the $\\tanh$ input:\n",
    "$$\\mathcal{L}_{L_2} = \\lambda_{L_2} \\sum_{i,j} w_{ij}^2 + \\lambda_{L_2} \\sum_{i} b_i^2$$ \n",
    "```python\n",
    "add_loss(lambda_L2_reg*( (W**2).sum() + (b**2).sum() ))\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbb1a9a4-bcd7-419e-bfbe-7b8d65f18670",
   "metadata": {},
   "source": [
    "\n",
    "modify W with the assumption that it should describe the statistical dynamics underlying actual activation patterns. Update by\n",
    "\n",
    "w = beta_cW*w - eta_sgd*dw_sgd + eta_stdp*(1+arousal(r))*dW_stdp\n",
    "\n",
    "- gradients: sgd\n",
    "- normalized spike timing dependant plasticity: dW_stdp = j1[:-1]i1[1:] - j1[1:]i1[:-1]\n",
    "- cW from absolute normalized coorelation on shifted time: cW = | Coor(i1[1:], j1[:-1]) |\n",
    "- beta_cW is a rolling mean of cW\n",
    "- i1 is the activation normalized by its scaling parameter to be in (-1,+1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52c1fd12-cddc-4090-9ada-e509835b1ac3",
   "metadata": {},
   "source": [
    "## Theory\n",
    "\n",
    "Critical multi-paradigm networks utilize a rescaled activation $Y = \\alpha \\tanh(\\alpha^{-1} XW + b + k + \\zeta)$. Scaling the activation space around $\\tanh$ extends the nonsaturating regime to $\\approx( -\\alpha/2, \\alpha/2)$ facilitating psuedo-linear activation foreward propagation and gradient backpropagation across deep layer traversals. Weights $W$ and biases $b$ are $L_2$-regularized to further alleviate saturation. Conversely, noise $\\zeta \\sim \\mathcal{N}(\\mu=0, \\sigma^2=\\sigma_{\\zeta}^2)$ encourages larger weights (hence, saturating activations) and in later epochs might be increased (even cyclicly) to help sparsify weights. Recieving and expressing activations in the $(-\\alpha, +\\alpha)$ range allows CMN's to interface and train end-to-end with other deep learning layers including attention mechanisms. CMN's internally utilize unsupervised weight update rules with a nonhomeogeneous reward-modulated update-rate."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e578a6c-5366-485d-9197-372f987c5bae",
   "metadata": {},
   "source": [
    "## Implementation\n",
    "\n",
    "Unsupervised rules should still apply their changes by addition on each timestep. (by default, gradients don't flow through unsupervised weight update rules). The previous activation outputs must be supplied if any other inputs are supplied (this is programmatically a hidden state)\n",
    "\n",
    "```python\n",
    "wandb.config = { ... }  # lowest precedence hyperparameters (optional)\n",
    "\n",
    "...\n",
    "\n",
    "# salina could turn this into a simple rnn\n",
    "# retrieve variables from reuccrent workspace\n",
    "y_next = CMCell({'x': X, 'y': y}, prev=y)  # generate y value for t+1 given values <=t\n",
    "# write variable to reuccrent workspace\n",
    "\n",
    "# open-ended growth\n",
    "y_next = CMCell({'x': X, 'y': y, 'v': V}, prev=y)  # new inputs\n",
    "y_next = CMCell({'y': y, 'v': V}, prev=y)  # optional inputs\n",
    "y_next = CMCell({'x': X, 'y': y, 'v': V}, prev=y, update=False, hparams=overrides)\n",
    "\n",
    "...\n",
    "\n",
    "# build recurrent model from cell\n",
    "model.fit(X,Ytrue)\n",
    "```\n",
    "\n",
    "Hyperparameters are overriden from `wandb.config` (optional, lowest precedence), `CMCell.HPARAM_DEFAULTS`, `self.hparams`, `hparams` (fn arg), and finally specific args and assignments (highest precedence).\n",
    "\n",
    "Input values are passed as a dictionary `{'x1': x1, 'x2': x2, 'y': y}`. All inputs for a single call must have matching batch axes and sequence lengths but can have individually varying input dimensions (i.e.: input must be shaped `[..., T, N1], [..., T, N2], [..., T, N3], ...`). You can change the number of inputs at any time by calling `build` with a corresponding dictionary of input shapes or just passing a new dictionary key in. Keys cannot be deleted. All inputs are optional."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c471fdf3-e4d8-40b6-a449-4d428abee475",
   "metadata": {},
   "source": [
    "## Unsupervised Learning\n",
    "\n",
    "I interpret signal values as probabilities from $-\\alpha$ (certainly false), $0$ (unknown), to $+\\alpha$ (certainly true) and assume neurons follow an exprapolated rule of implication. For each weight $w_{ji} \\in W$ connecting input $j$ to output $i$, elementwise updates proceed as:\n",
    "\n",
    "```python\n",
    "def update_weight(w_ji, x, y, x_prev, y_prev, t):\n",
    "    #  Updates the weight between input j and output i\n",
    "    #  w_ji: current value of weight\n",
    "    #  x: activation value of input j at time t\n",
    "    #  y: activation value of output i at time t\n",
    "    #  x_prev: activation value of input j at time t-1\n",
    "    #  y_prev: activation value of output i at time t-1\n",
    "    #  t: threshold to establish nonzero activation value = small number like 0.2\n",
    "    \n",
    "    # TODO: make these cases exclusive\n",
    "    # maybe identify a few simple expressions that capture everything\n",
    "    \n",
    "    if w_ji == 0:  # weight is dead\n",
    "        dw = 0  # do nothing\n",
    "    if x_prev > t and y > t:  # parent was + and the child was +\n",
    "        dw = 0  # nothing needs to change\n",
    "    if x_prev > t and -t < y < t:  # parent was + and the child was neutral\n",
    "        dw = +1  # the weight mattered so it should be strengthened\n",
    "    if x_prev > t and y < -t:  # parent was + and the child was -\n",
    "        dw = -1  # the weight didn't matter so it should be recycled\n",
    "       \n",
    "       \n",
    "```\n",
    "\n",
    "Bias $b$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "178d8b11-1144-4dc2-b18d-2a5a6608f833",
   "metadata": {},
   "source": [
    "## Supervised Learning\n",
    "\n",
    "Use standard gradient descent to optimize:\n",
    "\n",
    "$$\\arg_{\\theta_{model}} \\min L(x)$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b2e4eb8-42d4-423e-a2a2-70dc021cd452",
   "metadata": {},
   "outputs": [],
   "source": [
    "class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c506d0f-a406-4785-b6f5-8afb777d16cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# in the training loop\n",
    "wandb.log({\"loss\": loss})\n",
    "\n",
    "# Optional\n",
    "wandb.watch(model)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "ai"
  },
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}