incontext_3.4.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b02306",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from matplotlib.patches import Ellipse, Circle\n",
    "import math\n",
    "\n",
    "\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "\n",
    "torch.set_default_tensor_type(torch.DoubleTensor)\n",
    "\n",
    "\n",
    "class BioLinear(nn.Module):\n",
    "\n",
    "    def __init__(self, in_dim, out_dim, in_fold=1, out_fold=1, in_head=1, out_head=1):\n",
    "        super(BioLinear, self).__init__()\n",
    "        self.in_dim = in_dim\n",
    "        self.out_dim = out_dim\n",
    "        self.linear = nn.Linear(in_dim, out_dim)\n",
    "        self.in_fold = in_fold\n",
    "        self.out_fold = out_fold\n",
    "        self.in_head = in_head\n",
    "        self.out_head = out_head\n",
    "        assert in_dim % in_fold == 0\n",
    "        assert out_dim % out_fold == 0\n",
    "        #compute in_cor, shape: (in_dim)\n",
    "        in_dim_fold = int(in_dim/in_fold)\n",
    "        out_dim_fold = int(out_dim/out_fold)\n",
    "        self.in_coordinates = torch.tensor(list(np.linspace(1/(2*in_dim_fold), 1-1/(2*in_dim_fold), num=in_dim_fold))*in_fold, dtype=torch.float)\n",
    "        self.out_coordinates = torch.tensor(list(np.linspace(1/(2*out_dim_fold), 1-1/(2*out_dim_fold), num=out_dim_fold))*out_fold, dtype=torch.float)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.linear(x)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07c4197d",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BioMLP(nn.Module):\n",
    "    def __init__(self, in_dim=2, out_dim=2, w=2, depth=2, shp=None):\n",
    "        super(BioMLP, self).__init__()\n",
    "        if shp == None:\n",
    "            shp = [in_dim] + [w]*(depth-1) + [out_dim]\n",
    "            self.in_dim = in_dim\n",
    "            self.out_dim = out_dim\n",
    "            self.depth = depth\n",
    "                 \n",
    "        else:\n",
    "            self.in_dim = shp[0]\n",
    "            self.out_dim = shp[-1]\n",
    "            self.depth = len(shp) - 1\n",
    "        linear_list = []\n",
    "        for i in range(self.depth):\n",
    "            linear_list.append(BioLinear(shp[i], shp[i+1]))\n",
    "        self.linears = nn.ModuleList(linear_list)\n",
    "        self.shp = shp\n",
    "    \n",
    "\n",
    "    def forward(self, x):\n",
    "        f = lambda x: 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))\n",
    "        #f = torch.nn.SiLU()\n",
    "        for i in range(self.depth-1):\n",
    "            x = f(self.linears[i](x))\n",
    "        x = self.linears[-1](x)\n",
    "        return x\n",
    "    \n",
    "    def get_linear_layers(self):\n",
    "        return self.linears"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7143931",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BioAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, n_head=2, n_embed=6):\n",
    "        super().__init__()\n",
    "        assert n_embed % n_head == 0\n",
    "        # key, query, value projections for all heads, but in a batch\n",
    "        self.l_attn = BioLinear(n_embed, 3*n_embed, out_fold=3, out_head=n_head)\n",
    "        # output projection\n",
    "        self.l_proj = BioLinear(n_embed, n_embed, in_fold=1, in_head=n_head)\n",
    "        #self.l_proj = BioLinear(n_embed, n_embed)\n",
    "        # regularization\n",
    "        self.n_head = n_head\n",
    "        self.n_embed = n_embed\n",
    "\n",
    "    def forward(self, x):\n",
    "        # B: batch size; T: sequence length; C: embedding dimensionality (n_embd)\n",
    "        B, T, C = x.size()\n",
    "\n",
    "        # query, key, value\n",
    "        x = self.l_attn(x)\n",
    "        q, k, v = x[:,:,:C], x[:,:,C:2*C], x[:,:,2*C:3*C]\n",
    "        n_head = self.n_head\n",
    "        assert C % n_head == 0\n",
    "        q = q.reshape(B, T, n_head, int(C/n_head))\n",
    "        k = k.reshape(B, T, n_head, int(C/n_head))\n",
    "        v = v.reshape(B, T, n_head, int(C/n_head))\n",
    "\n",
    "        # (causal) self-attention\n",
    "        attn = torch.einsum('ijhl,ikhl->ijkh', q, k)/np.sqrt(int(C/n_head))\n",
    "        mask = torch.ones(T,T)*float('-inf')\n",
    "        mask = torch.tril(mask, diagonal=-1).permute(1,0).unsqueeze(dim=0).unsqueeze(dim=3)\n",
    "        attn = attn + mask\n",
    "        attn = nn.Softmax(dim=2)(attn)\n",
    "        attn = torch.einsum('ijkl,iklh->ijlh', attn, v)\n",
    "        attn = attn.reshape(B, T, C)\n",
    "\n",
    "        # output projection\n",
    "        y = self.l_proj(attn)\n",
    "        return y\n",
    "    \n",
    "    def get_linear_layers(self):\n",
    "        return [self.l_attn, self.l_proj]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80fd7ef1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BioBlock(nn.Module):\n",
    "    # A transformer block\n",
    "    def __init__(self, n_head=2, n_embed=6):\n",
    "        super().__init__()\n",
    "        self.n_head = n_head\n",
    "        self.n_embed = n_embed\n",
    "        self.ln_1 = nn.LayerNorm(n_embed)\n",
    "        self.attn = BioAttention(n_head=n_head, n_embed=n_embed)\n",
    "        self.ln_2 = nn.LayerNorm(n_embed)\n",
    "        self.mlp = BioMLP(shp=[n_embed, 4*n_embed, n_embed])\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(x)\n",
    "        x = x + self.mlp(x)\n",
    "        return x\n",
    "    \n",
    "    def get_linear_layers(self):\n",
    "        return [self.ln_1, *self.attn.get_linear_layers(), self.ln_2, *self.mlp.get_linear_layers()]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad706adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class BioTransformer(nn.Module):\n",
    "    # Transformer: since our goal is to deal with linear regression, not language, \n",
    "    # we ignore token embeddings and positioanl embeddings. \n",
    "    def __init__(self, in_dim=3, out_dim=3, n_head=2, n_embed=20, n_layer=2, block_size=19):\n",
    "        super().__init__()\n",
    "        self.n_head = n_head\n",
    "        self.n_embed = n_embed\n",
    "        self.n_layer = n_layer\n",
    "        self.l_i = BioLinear(in_dim, n_embed)\n",
    "        self.blocks = nn.ModuleList([BioBlock(n_head=n_head, n_embed=n_embed) for _ in range(n_layer)])\n",
    "        self.ln_f = nn.LayerNorm(n_embed)\n",
    "        #self.ln_f.weight.requires_grad = False\n",
    "        #self.ln_f.bias.requires_grad = False\n",
    "        self.l_f = BioLinear(n_embed, out_dim)\n",
    "        self.in_dim = in_dim\n",
    "        self.out_dim = out_dim\n",
    "        #self.pe = torch.nn.Parameter(torch.normal(0,1,size=(block_size, n_embed)))\n",
    "        \n",
    "        # parameters for the bio-inspired trick\n",
    "        self.l0 = 0.5 # distance between two nearby layers\n",
    "        #self.in_perm = torch.tensor(np.arange(int(self.in_dim/self.l_i.in_fold)), dtype=torch.long)\n",
    "        self.in_perm = nn.Parameter(torch.tensor(np.arange(int(self.in_dim/self.l_i.in_fold)), dtype=torch.float))\n",
    "        #self.out_perm = torch.tensor(np.arange(int(self.out_dim/self.l_f.out_fold)), dtype=torch.long)\n",
    "        self.out_perm = nn.Parameter(torch.tensor(np.arange(int(self.out_dim/self.l_f.out_fold)), dtype=torch.float))\n",
    "        self.top_k = 20\n",
    "        \n",
    "        \n",
    "        self.res_swap = list(np.arange(2*n_layer+1)*3+1)\n",
    "        self.skip_swap = list(np.arange(2*n_layer+1)*3+2)\n",
    "        self.normal_swap = list(np.arange(2*n_layer+2)*3)\n",
    "        \n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x[:,:,self.in_perm.long()]\n",
    "        x = self.l_i(x)\n",
    "        #x = x + self.pe.unsqueeze(dim=0) # positional encoding\n",
    "        for i in range(self.n_layer):\n",
    "            x = self.blocks[i](x)\n",
    "        #x = self.ln_f(x)\n",
    "        y = self.l_f(x)\n",
    "        \n",
    "        out_perm_inv = torch.zeros(self.out_dim, dtype=torch.long)\n",
    "        out_perm_inv[self.out_perm.long()] = torch.arange(self.out_dim)\n",
    "        y = y[:,:,out_perm_inv]\n",
    "        return y\n",
    "    \n",
    "    def get_linear_layers(self):\n",
    "        linear_list = [self.l_i]\n",
    "        for i in range(self.n_layer):\n",
    "            linear_list = [*linear_list, *self.blocks[i].get_linear_layers()]\n",
    "        linear_list.append(self.ln_f)\n",
    "        linear_list.append(self.l_f)\n",
    "        return linear_list\n",
    "    \n",
    "    \n",
    "    def get_cc(self, weight_factor=2.0, bias_penalize=True, ln_penalize=True, no_penalize_last=False):\n",
    "        # compute connection cost\n",
    "        cc = 0\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        for i in range(num_linear):\n",
    "            layer = linears[i]\n",
    "            if isinstance(layer, nn.LayerNorm):\n",
    "                pass\n",
    "                #cc += torch.sum(torch.abs(layer.weight)) + torch.sum(torch.abs(layer.bias))\n",
    "            else:\n",
    "                if i == num_linear - 1 and no_penalize_last:\n",
    "                    weight_factor = 0.\n",
    "                biolinear = linears[i]\n",
    "                dist = torch.abs(biolinear.out_coordinates.unsqueeze(dim=1) - biolinear.in_coordinates.unsqueeze(dim=0))\n",
    "                cc += torch.mean(torch.abs(biolinear.linear.weight)*(weight_factor*dist+self.l0))\n",
    "                if bias_penalize == True:\n",
    "                    cc += torch.mean(torch.abs(biolinear.linear.bias)*(self.l0))\n",
    "        return cc\n",
    "    \n",
    "    def swap_weight(self, weights, j, k, swap_type=\"out\"):\n",
    "        with torch.no_grad():  \n",
    "            if swap_type == \"in\":\n",
    "                temp = weights[:,j].clone()\n",
    "                weights[:,j] = weights[:,k].clone()\n",
    "                weights[:,k] = temp\n",
    "            elif swap_type == \"out\":\n",
    "                temp = weights[j].clone()\n",
    "                weights[j] = weights[k].clone()\n",
    "                weights[k] = temp\n",
    "            else:\n",
    "                raise Exception(\"Swap type {} is not recognized!\".format(swap_type))\n",
    "            \n",
    "    def swap_bias(self, biases, j, k):\n",
    "        with torch.no_grad():  \n",
    "            temp = biases[j].clone()\n",
    "            biases[j] = biases[k].clone()\n",
    "            biases[k] = temp\n",
    "    \n",
    "    def swap(self, i, j, k):\n",
    "        # in the ith layer (of neurons), swap the jth and the kth neuron. \n",
    "        # Note: n layers of weights means n+1 layers of neurons.\n",
    "        # (incoming, outgoing) * weights + biases are swapped. \n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            left = None\n",
    "            right = linears[i]\n",
    "            self.swap_bias(self.in_perm, j, k)\n",
    "        elif i == num_linear:\n",
    "            left = linears[i-1]\n",
    "            right = None\n",
    "            self.swap_bias(self.out_perm, j, k)\n",
    "        else:\n",
    "            left = linears[i-1]\n",
    "            right = linears[i]\n",
    "            \n",
    "        \n",
    "        if left != None:\n",
    "            fold = left.out_fold\n",
    "            fold_dim = int(left.linear.weight.shape[0]/fold)\n",
    "            for l in range(fold):\n",
    "                self.swap_weight(left.linear.weight, j+fold_dim*l, k+fold_dim*l, swap_type=\"out\")\n",
    "                self.swap_bias(left.linear.bias, j+fold_dim*l, k+fold_dim*l)\n",
    "                \n",
    "        if right != None:\n",
    "        \n",
    "            if i in self.normal_swap:\n",
    "                fold = right.in_fold\n",
    "                fold_dim = int(right.linear.weight.shape[1]/fold)\n",
    "                for l in range(fold):\n",
    "                    self.swap_weight(right.linear.weight, j+fold_dim*l, k+fold_dim*l, swap_type=\"in\")\n",
    "\n",
    "            if i in self.res_swap:\n",
    "                rightright = linears[i+1]\n",
    "                fold = rightright.in_fold\n",
    "                fold_dim = int(rightright.linear.weight.shape[1]/fold)\n",
    "                for l in range(fold):\n",
    "                    self.swap_bias(right.weight, j+fold_dim*l, k+fold_dim*l)\n",
    "                    self.swap_bias(right.bias, j+fold_dim*l, k+fold_dim*l)\n",
    "                    self.swap_weight(rightright.linear.weight, j+fold_dim*l, k+fold_dim*l, swap_type=\"in\")\n",
    "\n",
    "            \n",
    "    def get_score(self, i):\n",
    "        \n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i == 0:\n",
    "            left = None\n",
    "            right = linears[i]\n",
    "        elif i == num_linear:\n",
    "            left = linears[i-1]\n",
    "            right = None\n",
    "        else:\n",
    "            left = linears[i-1]\n",
    "            right = linears[i]\n",
    "        \n",
    "        if isinstance(right, nn.LayerNorm):\n",
    "            right = linears[i+1]\n",
    "            \n",
    "        # need to fold attention, fold = 3\n",
    "        score = 0.\n",
    "        if left == None:\n",
    "            pass\n",
    "        else:\n",
    "            fold = left.out_fold\n",
    "            score +=  torch.mean(torch.sum(torch.abs(left.linear.weight), dim=1).reshape(fold, int(left.linear.weight.shape[0]/fold)), dim=0)\n",
    "            \n",
    "        if right == None:\n",
    "            pass\n",
    "        else:\n",
    "            fold2 = right.in_fold\n",
    "            score += torch.mean(torch.sum(torch.abs(right.linear.weight), dim=0).reshape(fold2, int(right.linear.weight.shape[1]/fold2)), dim=0)\n",
    "            \n",
    "        return score\n",
    "    \n",
    "    def get_n_head(self, i):\n",
    "        linears = self.get_linear_layers()\n",
    "        num_layer = len(linears)\n",
    "        if i == 0:\n",
    "            n_head = linears[0].in_head\n",
    "        else:\n",
    "            n_head = linears[i-1].out_head\n",
    "        return n_head\n",
    "\n",
    "    def get_top_id_head(self, i, top_k=20):\n",
    "        \n",
    "        score = 0.\n",
    "        if i == self.res_swap[0]:\n",
    "            for p in self.res_swap:\n",
    "                score += self.get_score(p)\n",
    "        else:\n",
    "            score = self.get_score(i)\n",
    "            \n",
    "        n_head = self.get_n_head(i)\n",
    "        score_head = score.reshape(n_head, int(score.shape[0]/n_head))\n",
    "        score_head = torch.mean(score_head, dim=1)\n",
    "        top_id_head = torch.flip(torch.argsort(score_head),[0])[:top_k]\n",
    "        \n",
    "        return top_id_head # 1D\n",
    "    \n",
    "    \n",
    "    def get_top_id_tail(self, i, top_k=20):\n",
    "        \n",
    "        score = 0.\n",
    "        if i == self.res_swap[0]:\n",
    "            for p in self.res_swap:\n",
    "                score += self.get_score(p)\n",
    "        else:\n",
    "            score = self.get_score(i)\n",
    "        \n",
    "        n_head = self.get_n_head(i)\n",
    "        head_dim = int(score.shape[0]/n_head)\n",
    "        score_head = score.reshape(n_head, head_dim)\n",
    "        \n",
    "        top_id_tail = torch.flip(torch.argsort(score_head, dim=1),[1])[:top_k]\n",
    "        \n",
    "        return top_id_tail # 2D\n",
    "    \n",
    "    \n",
    "    def relocate_ij_head(self, i, j):\n",
    "        # i-th layer, j-th head\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i < num_linear:\n",
    "            if isinstance(linears[i], nn.LayerNorm):\n",
    "                num_neuron = int(linears[i+1].linear.weight.shape[1]/linears[i+1].in_fold)\n",
    "            else:\n",
    "                num_neuron = int(linears[i].linear.weight.shape[1]/linears[i].in_fold)\n",
    "        else:\n",
    "            num_neuron = linears[i-1].linear.weight.shape[0]\n",
    "            \n",
    "        ccs = []\n",
    "                \n",
    "        num_head = self.get_n_head(i)\n",
    "        head_dim = int(num_neuron/num_head)\n",
    "        \n",
    "        for k in range(num_head):\n",
    "            if i != self.res_swap[0]:\n",
    "                self.swap(i,head_dim*j+np.arange(head_dim),head_dim*k+np.arange(head_dim))\n",
    "                \n",
    "            ccs.append(self.get_cc())\n",
    "            \n",
    "            if i != self.res_swap[0]:\n",
    "                self.swap(i,head_dim*j+np.arange(head_dim),head_dim*k+np.arange(head_dim))\n",
    "                \n",
    "        k = torch.argmin(torch.stack(ccs))\n",
    "        \n",
    "        if i != self.res_swap[0]:\n",
    "            self.swap(i,head_dim*j+np.arange(head_dim),head_dim*k+np.arange(head_dim))\n",
    "        \n",
    "    \n",
    "    def relocate_ijk_tail(self, i, j, k):\n",
    "        # i-th layer, j-th head, k-th neuron\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        if i < num_linear:\n",
    "            if isinstance(linears[i], nn.LayerNorm):\n",
    "                num_neuron = int(linears[i+1].linear.weight.shape[1]/linears[i+1].in_fold)\n",
    "            else:\n",
    "                num_neuron = int(linears[i].linear.weight.shape[1]/linears[i].in_fold)\n",
    "        else:\n",
    "            num_neuron = linears[i-1].linear.weight.shape[0]\n",
    "            \n",
    "        ccs = []\n",
    "        \n",
    "        def swap_res(j,k):\n",
    "            for p in self.res_swap:\n",
    "                self.swap(p,j,k)\n",
    "                \n",
    "        num_head = self.get_n_head(i)\n",
    "        head_dim = int(num_neuron/num_head)\n",
    "        \n",
    "        for p in range(head_dim):\n",
    "            if i == self.res_swap[0]:\n",
    "                swap_res(j*head_dim+k,j*head_dim+p)\n",
    "            else:\n",
    "                self.swap(i,j*head_dim+k,j*head_dim+p)\n",
    "                \n",
    "            ccs.append(self.get_cc())\n",
    "            \n",
    "            if i == self.res_swap[0]:\n",
    "                swap_res(j*head_dim+k,j*head_dim+p)\n",
    "            else:\n",
    "                self.swap(i,j*head_dim+k,j*head_dim+p)\n",
    "                \n",
    "        p = torch.argmin(torch.stack(ccs))\n",
    "        \n",
    "        if i == self.res_swap[0]:\n",
    "                swap_res(j*head_dim+k,j*head_dim+p)\n",
    "        else:\n",
    "            self.swap(i,j*head_dim+k,j*head_dim+p)\n",
    "    \n",
    "    def relocate_i(self, i):\n",
    "        # skip swap\n",
    "        if i in self.skip_swap:\n",
    "            return\n",
    "        \n",
    "        # res swap\n",
    "        if i in self.res_swap and i != self.res_swap[0]:\n",
    "            return\n",
    "            \n",
    "        # normal swap + the first res swap\n",
    "        top_id_head = self.get_top_id_head(i, top_k=self.top_k)\n",
    "        for j in top_id_head:\n",
    "            self.relocate_ij_head(i,j)\n",
    "            \n",
    "        top_id_tail = self.get_top_id_tail(i, top_k=self.top_k)\n",
    "        num_head = top_id_tail.shape[0]\n",
    "        for j in range(num_head):\n",
    "            for k in top_id_tail[j]:\n",
    "                self.relocate_ijk_tail(i,j,k)\n",
    "            \n",
    "    def relocate(self):\n",
    "        print('swap')\n",
    "        # Relocate neurons in the whole model\n",
    "        linears = self.get_linear_layers()\n",
    "        num_linear = len(linears)\n",
    "        for i in range(num_linear+1):\n",
    "            #print(i)\n",
    "            self.relocate_i(num_linear-i)\n",
    "            \n",
    "            \n",
    "    def plot(self):\n",
    "        layers = self.get_linear_layers()\n",
    "        linears = [layers[i] for i in [0,2,3,5,6,8,9,11,12,14]]\n",
    "        shp = [2,32,32,32,128,32,32,32,128,32,2]\n",
    "\n",
    "        fig, ax = plt.subplots(figsize=(5,10))\n",
    "        s = 1/(2*max(shp))\n",
    "        for j in range(len(shp)):\n",
    "            N = shp[j]\n",
    "            for i in range(N):\n",
    "                circle = Ellipse((1/(2*N)+i/N, 0.1*j), s, s/15*((len(shp)-1)+0.4), color='black')\n",
    "                ax.add_patch(circle)\n",
    "\n",
    "\n",
    "        plt.ylim(-0.02,0.1*(len(shp)-1)+0.02)\n",
    "        plt.xlim(-0.02,1.02)\n",
    "\n",
    "        #linears = self.linears\n",
    "        ln_id = [1,4,7,10,13]\n",
    "        jj = 0\n",
    "        for ii in range(len(linears)):\n",
    "            biolinear = linears[ii]\n",
    "            p = biolinear.linear.weight.clone()\n",
    "            if ii in [1,3,5,7,9]:\n",
    "                ln = layers[ln_id[jj]]\n",
    "                p = p * ln.weight.clone().unsqueeze(dim=0)\n",
    "                jj += 1\n",
    "            p_shp = p.shape\n",
    "            if ii == 1 or ii == 5:\n",
    "                p_ = p[64:,:]\n",
    "                #p_ = p[8:,:]\n",
    "                p = p/torch.abs(p_).max()\n",
    "            else:\n",
    "                p = p/torch.abs(p[:,:]).max()\n",
    "            fold_num = int(p_shp[1])\n",
    "            if ii == 1 or ii == 5:\n",
    "                p_shp0 = int(p_shp[0]/3)\n",
    "                for i in range(p_shp0):\n",
    "                    for j in range(fold_num):\n",
    "                        plt.plot([1/(2*p_shp0)+i/p_shp0, 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=np.minimum(10*np.abs(p[i+64,j].detach().numpy()), 1), color=\"red\" if p[i,j]>0 else \"blue\")\n",
    "                        #plt.plot([1/(2*p_shp0)+i/p_shp0, 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=np.minimum(10*np.abs(p[i+8,j].detach().numpy()), 1), color=\"red\" if p[i,j]>0 else \"blue\")\n",
    "\n",
    "            else:\n",
    "                for i in range(p_shp[0]):\n",
    "                    for j in range(fold_num):\n",
    "                        plt.plot([1/(2*p_shp[0])+i/p_shp[0], 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=np.minimum(10*np.abs(p[i,j].detach().numpy()), 1), color=\"red\" if p[i,j]>0 else \"blue\")\n",
    "\n",
    "\n",
    "        ax.axis('off')\n",
    "\n",
    "        fontsize=12\n",
    "        plt.text(1.1,0.0, \"input\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.1, \"LN\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.2, \"Atn\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.3, \"Res\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.4, \"MLP hidden\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.5, \"Res\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.6, \"Atn\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.7, \"Res\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.8, \"MLP hidden\", fontsize=fontsize)\n",
    "        plt.text(1.1,0.9, \"Res\", fontsize=fontsize)\n",
    "        plt.text(1.1,1.0, \"output\", fontsize=fontsize)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "570a3b71",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "\n",
    "seed = 1\n",
    "np.random.seed(seed)\n",
    "torch.manual_seed(seed)\n",
    "# create linear regression dataset: (x_1, x2, ... ,x_d) -> y\n",
    "# each sequence has length T (T sample)\n",
    "d = 1\n",
    "T = 2\n",
    "\n",
    "n_w = 10000\n",
    "#n_w = 10\n",
    "w = torch.rand(n_w, d)*2 - 1\n",
    "#torch.normal(0,1,size=(n_w, d))\n",
    "\n",
    "# x has shape (n_w, T, d)\n",
    "#x = torch.normal(0,1,size=(n_w, T, d))\n",
    "x = torch.rand(n_w, T, d)*2 + 1\n",
    "y = torch.einsum('ik,ijk->ij', w, x)\n",
    "y_ = torch.cat([y, torch.zeros([y.shape[0], y.shape[1]*d])], dim=1).reshape(n_w, d+1, T).permute(0,2,1)\n",
    "x_ = torch.cat([torch.zeros([y.shape[0], y.shape[1], 1]), x], dim=2)\n",
    "data = torch.cat([x_, y_], dim=2)\n",
    "data = data.reshape(n_w, -1)\n",
    "inputs = data[:,:(2*T-1)*(d+1)].reshape(n_w, 2*T-1, d+1)\n",
    "labels = data[:,(d+1):].reshape(n_w, 2*T-1, d+1)\n",
    "\n",
    "fraction = 0.8\n",
    "train_num = int(n_w*fraction)\n",
    "test_num = n_w - train_num\n",
    "\n",
    "train_id = np.random.choice(n_w,train_num,replace=False)\n",
    "test_id = np.array(list(set(np.arange(n_w)) - set(train_id)))\n",
    "\n",
    "inputs_train = inputs[train_id].requires_grad_(True)\n",
    "labels_train = labels[train_id]\n",
    "\n",
    "inputs_test = inputs[test_id].requires_grad_(True)\n",
    "labels_test = labels[test_id]\n",
    "\n",
    "### train ###\n",
    "model = BioTransformer(in_dim=d+1, out_dim=d+1, n_head=1, n_embed=32, n_layer=2)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.0)\n",
    "steps = int(40000)\n",
    "log = 100\n",
    "batch_size = 64 #128\n",
    "lamb = 1e-3\n",
    "swap_log = int(1e6) #1000\n",
    "plot_log = 1000\n",
    "\n",
    "\n",
    "for step in range(steps):\n",
    "    \n",
    "    if step == int(steps*1/4):\n",
    "        lamb *= 10\n",
    "    \n",
    "    if step == int(steps*3/4):\n",
    "        lamb *= 10\n",
    "    \n",
    "    \n",
    "    optimizer.zero_grad()\n",
    "    \n",
    "    if batch_size == None:\n",
    "        train_choice = np.arange(train_num)\n",
    "        test_choice = np.arange(test_num)\n",
    "    else:\n",
    "        train_choice = np.random.choice(train_num, batch_size, replace=True)\n",
    "        test_choice = np.random.choice(test_num, batch_size, replace=True)\n",
    "        \n",
    "    pred  = model(inputs_train[train_choice])\n",
    "    diff = pred[:,::2,0] - labels_train[train_choice][:,::2,0]\n",
    "    loss = torch.mean(diff**2)\n",
    "    loss_last = torch.mean(diff[:,-1]**2)\n",
    "    \n",
    "    pred_test  = model(inputs_test[test_choice])\n",
    "    diff_test = pred_test[:,::2,0] - labels_test[test_choice][:,::2,0]\n",
    "    loss_test = torch.mean(diff_test**2)\n",
    "    loss_test_last = torch.mean(diff_test[:,-1]**2)\n",
    "    \n",
    "    cc = model.get_cc(no_penalize_last=False)\n",
    "    total_loss = loss + lamb*cc\n",
    "    total_loss.backward()\n",
    "\n",
    "    \n",
    "    optimizer.step()\n",
    "    \n",
    "    if step % log == 0:\n",
    "        print(\"step = %d | train loss: %.2e | train last: %.2e | test loss %.2e | test last: %.2e | cc: %.2e \"%(step, loss.detach().numpy(), loss_last.detach().numpy(), loss_test.detach().numpy(), loss_test_last.detach().numpy(), cc.detach().numpy()))\n",
    "        \n",
    "    if (step+1) % swap_log == 0:\n",
    "        model.relocate()\n",
    "        \n",
    "    if (step+1) % plot_log == 0:\n",
    "        model.plot()\n",
    "        plt.show()\n",
    "        plt.scatter(pred[:,-1,0].detach().numpy(), labels_train[train_choice][:,-1,0].detach().numpy(), s=5)\n",
    "        plt.show()\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3236b333",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "layers = model.get_linear_layers()\n",
    "linears = [layers[i] for i in [0,2,3,5,6,8,9,11,12,14]]\n",
    "shp = [2,32,32,32,128,32,32,32,128,32,2]\n",
    "#shp = [2,4,4,4,16,4,4,4,16,4,2]\n",
    "#shp = [2,16,16,16,64,16,16,16,64,16,2]\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(5,10))\n",
    "s = 1/(2*max(shp))\n",
    "for j in range(len(shp)):\n",
    "    N = shp[j]\n",
    "    for i in range(N):\n",
    "        circle = Ellipse((1/(2*N)+i/N, 0.1*j), s, s/15*((len(shp)-1)+0.4), color='black')\n",
    "        ax.add_patch(circle)\n",
    "\n",
    "\n",
    "plt.ylim(-0.02,0.1*(len(shp)-1)+0.02)\n",
    "plt.xlim(-0.02,1.02)\n",
    "\n",
    "#linears = self.linears\n",
    "ln_id = [1,4,7,10,13]\n",
    "jj = 0\n",
    "for ii in range(len(linears)):\n",
    "    biolinear = linears[ii]\n",
    "    p = biolinear.linear.weight.clone()\n",
    "    if ii in [1,3,5,7,9]:\n",
    "        ln = layers[ln_id[jj]]\n",
    "        p = p * ln.weight.clone().unsqueeze(dim=0)\n",
    "        jj += 1\n",
    "    p_shp = p.shape\n",
    "    if ii == 1 or ii == 5:\n",
    "        p_ = p[64:,:]\n",
    "        p = p/torch.abs(p_).max()\n",
    "    else:\n",
    "        p = p/torch.abs(p[:,:]).max()\n",
    "    fold_num = int(p_shp[1])\n",
    "    if ii == 1 or ii == 5:\n",
    "        p_shp0 = int(p_shp[0]/3)\n",
    "        for i in range(p_shp0):\n",
    "            for j in range(fold_num):\n",
    "                plt.plot([1/(2*p_shp0)+i/p_shp0, 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=np.minimum(10*np.abs(p[i+64,j].detach().numpy()), 1), color=\"blue\" if p[i,j]>0 else \"red\")\n",
    "\n",
    "    else:\n",
    "        for i in range(p_shp[0]):\n",
    "            for j in range(fold_num):\n",
    "                plt.plot([1/(2*p_shp[0])+i/p_shp[0], 1/(2*fold_num)+j/fold_num], [0.1*(ii+1),0.1*ii], lw=np.minimum(10*np.abs(p[i,j].detach().numpy()), 1), color=\"blue\" if p[i,j]>0 else \"red\")\n",
    "\n",
    "ax.axis('off')\n",
    "\n",
    "fontsize=12\n",
    "shift = 0.03\n",
    "plt.text(1.1,0.0-shift, \"input\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.1-shift, \"Embed\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.2-shift, \"Atn1\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.3-shift, \"Res1\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.4-shift, \"   MLP1 \\n hidden\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.5-shift, \"Res2\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.6-shift, \"Atn2\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.7-shift, \"Res3\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.8-shift, \"   MLP2 \\n hidden\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,0.9-shift, \"Res4\", fontsize=fontsize, rotation=90)\n",
    "plt.text(1.1,1.0-shift, \"output\", fontsize=fontsize, rotation=90)\n",
    "\n",
    "#circle = Ellipse((0.47,0.5), 20*s, 10*s/15*((len(shp)-1)+0.4), color='black', fill=False, linewidth=4)\n",
    "#ax.add_patch(circle)\n",
    "for i in range(32):\n",
    "    plt.text(i/32+0.007, 0.51, i+1, fontsize=6, rotation=90)\n",
    "    plt.text(i/32+0.007, 0.31, i+1, fontsize=6, rotation=90)\n",
    "    plt.text(i/32+0.007, 0.71, i+1, fontsize=6, rotation=90)\n",
    "    plt.text(i/32+0.007, 0.91, i+1, fontsize=6, rotation=90)\n",
    "    \n",
    "plt.text(0.23,-0.02,0, fontsize=fontsize, rotation=90)\n",
    "plt.text(0.73,-0.02,\"x\", fontsize=fontsize, rotation=90)\n",
    "plt.text(0.23,1.02,\"y\", fontsize=fontsize,rotation=90)\n",
    "plt.text(0.73,1.02,0, fontsize=fontsize,rotation=90)\n",
    "\n",
    "#plt.savefig('./results/incontext2/incontext.png', bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ade81fae",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = inputs_train\n",
    "x = x[:,:,model.in_perm.long()]\n",
    "x = model.l_i(x)\n",
    "\n",
    "\n",
    "B, T, C = x.size()\n",
    "# atn1\n",
    "xx = model.blocks[0].attn.l_attn(x)\n",
    "q, k, v = xx[:,:,:C], xx[:,:,C:2*C], xx[:,:,2*C:3*C]\n",
    "n_head = model.blocks[0].attn.n_head\n",
    "assert C % n_head == 0\n",
    "q = q.reshape(B, T, n_head, int(C/n_head))\n",
    "k = k.reshape(B, T, n_head, int(C/n_head))\n",
    "v = v.reshape(B, T, n_head, int(C/n_head))\n",
    "\n",
    "# (causal) self-attention\n",
    "attn = torch.einsum('ijhl,ikhl->ijkh', q, k)/np.sqrt(int(C/n_head))\n",
    "mask = torch.ones(T,T)*float('-inf')\n",
    "mask = torch.tril(mask, diagonal=-1).permute(1,0).unsqueeze(dim=0).unsqueeze(dim=3)\n",
    "attn = attn + mask\n",
    "attn = nn.Softmax(dim=2)(attn)\n",
    "attn = torch.einsum('ijkl,iklh->ijlh', attn, v)\n",
    "attn = attn.reshape(B, T, C)\n",
    "\n",
    "# output projection\n",
    "y = model.blocks[0].attn.l_proj(attn)\n",
    "\n",
    "print(x.shape, y.shape)\n",
    "x = x + y\n",
    "\n",
    "#prob = x.clone()\n",
    "\n",
    "f = lambda x: 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))\n",
    "y = f(model.blocks[0].mlp.linears[0](x))\n",
    "y = model.blocks[0].mlp.linears[1](y)\n",
    "\n",
    "x = x + y\n",
    "\n",
    "prob = x.clone()\n",
    "\n",
    "B, T, C = x.size()\n",
    "# atn1\n",
    "xx = model.blocks[1].attn.l_attn(x)\n",
    "q, k, v = xx[:,:,:C], xx[:,:,C:2*C], xx[:,:,2*C:3*C]\n",
    "n_head = model.blocks[1].attn.n_head\n",
    "assert C % n_head == 0\n",
    "q = q.reshape(B, T, n_head, int(C/n_head))\n",
    "k = k.reshape(B, T, n_head, int(C/n_head))\n",
    "v = v.reshape(B, T, n_head, int(C/n_head))\n",
    "\n",
    "# (causal) self-attention\n",
    "attn = torch.einsum('ijhl,ikhl->ijkh', q, k)/np.sqrt(int(C/n_head))\n",
    "mask = torch.ones(T,T)*float('-inf')\n",
    "mask = torch.tril(mask, diagonal=-1).permute(1,0).unsqueeze(dim=0).unsqueeze(dim=3)\n",
    "attn = attn + mask\n",
    "attn = nn.Softmax(dim=2)(attn)\n",
    "attn = torch.einsum('ijkl,iklh->ijlh', attn, v)\n",
    "attn = attn.reshape(B, T, C)\n",
    "\n",
    "#prob = attn.clone()\n",
    "\n",
    "\n",
    "# output projection\n",
    "y = model.blocks[1].attn.l_proj(attn)\n",
    "\n",
    "x = x + y\n",
    "\n",
    "\n",
    "#prob = x.clone()\n",
    "\n",
    "f = lambda x: 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))\n",
    "y = f(model.blocks[1].mlp.linears[0](x))\n",
    "y = model.blocks[1].mlp.linears[1](y)\n",
    "\n",
    "x = x + y\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "y = model.l_f(x)\n",
    "\n",
    "out_perm_inv = torch.zeros(model.out_dim, dtype=torch.long)\n",
    "out_perm_inv[model.out_perm.long()] = torch.arange(model.out_dim)\n",
    "y = y[:,:,out_perm_inv]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3209ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.get_top_id_tail(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eea3d117",
   "metadata": {},
   "outputs": [],
   "source": [
    "for iid in range(32):\n",
    "    plt.scatter(inputs_train[:,-1,1].detach().numpy(), prob[:,-1,iid].detach().numpy(), s=0.1)\n",
    "    plt.scatter(w[train_id,0].detach().numpy(), prob[:,-1,iid].detach().numpy(), s=0.1)\n",
    "    plt.scatter(w[train_id,0].detach().numpy()*inputs_train[:,-1,1].detach().numpy(), prob[:,-1,iid].detach().numpy(), s=0.1)\n",
    "    print(iid)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "064a41f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "id1 = 8\n",
    "id2 = 11\n",
    "plt.scatter(prob[:,-1,id1].detach().numpy(), prob[:,-1,id2].detach().numpy(), s=0.1, c= w[train_id,0].detach().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd82d7ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6,18))\n",
    "\n",
    "plt.subplot(3,1,1)\n",
    "\n",
    "pred  = model(inputs)\n",
    "\n",
    "plt.scatter(labels[:,-1,0].detach().numpy(),pred[:,-1,0].detach().numpy(), s=0.1)\n",
    "plt.plot([-3,3],[-3,3], ls=\"--\",color=\"red\",alpha=0.7)\n",
    "plt.xlabel('True '+r\"$y$\",fontsize=25)\n",
    "plt.ylabel('Predicted '+r\"$y$\",fontsize=25)\n",
    "\n",
    "plt.subplot(3,1,2)\n",
    "id1 = 8\n",
    "id2 = 9\n",
    "plt.scatter(prob[:,-1,id1].detach().numpy(), prob[:,-1,id2].detach().numpy(), s=0.1, c= w[train_id,0].detach().numpy())\n",
    "cbar = plt.colorbar()\n",
    "plt.xlabel(\"Neuron 9\", fontsize=25)\n",
    "plt.ylabel(\"Neuron 10\", fontsize=25)\n",
    "cbar.ax.set_ylabel(\"weight scalar\",fontsize=25)\n",
    "\n",
    "plt.subplot(3,1,3)\n",
    "id1 = 10\n",
    "id2 = 18\n",
    "plt.scatter(prob[:,-1,id1].detach().numpy(), prob[:,-1,id2].detach().numpy(), s=0.1, c= w[train_id,0].detach().numpy())\n",
    "cbar = plt.colorbar()\n",
    "plt.xlabel(\"Neuron 11\", fontsize=25)\n",
    "plt.ylabel(\"Neuron 19\", fontsize=25)\n",
    "cbar.ax.set_ylabel(\"weight scalar\",fontsize=25)\n",
    "\n",
    "plt.savefig(\"./results/incontext2/weight_neuron.png\", bbox_inches=\"tight\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e477c52e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}