wip

notebooks/11b_vae_w_direct.ipynb

@@ -2,9 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 107,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
+     ]
+    }
+   ],
    "source": [
     "import os\n",
     "import numpy as np\n",
@@ -80,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 108,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -125,9 +134,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 109,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(PosixPath('/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/lightning_logs/version_24/hidden_states/.ds/ds_valtest_8b8fd6070504d5ef'),\n",
+       " PosixPath('/media/wassname/SGIronWolf/projects5/elk/sgd_probes_are_lie_detectors/notebooks/lightning_logs/version_24/hidden_states/.ds/ds_OOD_a41d3a61513ade30'))"
+      ]
+     },
+     "execution_count": 109,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "# load hidden state from a previously loaded adapter\n",
     "# the columns with _base are from the base model, and adapt from adapter\n",
@@ -139,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 110,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -150,7 +171,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 111,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -207,9 +228,52 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 112,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "select rows are 74.39% based on knowledge\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2107d9b230ac4d7cb6f84c90d1c57690",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Map:   0%|          | 0/615 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "select rows are 74.39% based on knowledge\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ba4b4630e534440db8ffe9b480a344f1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Map:   0%|          | 0/615 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
    "source": [
     "dm = load_file_to_dm(f1_val)\n",
     "dm_ood = load_file_to_dm(f1_ood)"
@@ -247,7 +311,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 96,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -481,7 +545,7 @@
     "    weight_decay=wd,\n",
     "    hs=64,\n",
     "    depth=3,\n",
-    "    dropout=1,\n",
+    "    dropout=0,\n",
     "    n_latent=16,\n",
     "    l1_coeff=l1_coeff, \n",
     "    importance_matrix=importance_matrix,\n",
@@ -544,7 +608,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 97,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1488,7 +1552,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 98,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1520,7 +1584,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 100,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1572,7 +1636,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 101,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -1660,7 +1724,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 103,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1673,7 +1737,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 104,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1682,7 +1746,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 105,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -1694,7 +1758,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 106,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {

research_log.md

@@ -737,3 +737,4 @@ https://github.com/ai-safety-foundation/sparse_autoencoder
 
 TODO:
 - try SAE
+Perhaps consider 

src/vae/sae.py

@@ -0,0 +1,76 @@
+# from https://github.com/callummcdougall/sae-exercises-mats/blob/116ecf3f8f7ffcd66cf628518009f81989e57bac/solutions.py#L692
+import torch as t
+from torch import nn, Tensor
+from torch.nn import functional as F
+from dataclasses import dataclass
+import einops
+
+from jaxtyping import Float
+from typing import Optional, Union, Callable
+
+@dataclass
+class AutoEncoderConfig:
+    # We optimize n_instances models in a single training loop to let us sweep over
+    # sparsity or importance curves  efficiently. You should treat `n_instances` as
+    # kinda like a batch dimension, but one which is built into our training setup.
+    n_instances: int
+    # this is the hidden states, and the latent size
+    n_input_ae: int
+    n_hidden_ae: int
+    l1_coeff: float = 0.5
+    tied_weights: bool = False
+
+
+class AutoEncoder(nn.Module):
+    W_enc: Float[Tensor, "n_instances n_input_ae n_hidden_ae"]
+    W_dec: Float[Tensor, "n_instances n_hidden_ae n_input_ae"]
+    b_enc: Float[Tensor, "n_instances n_hidden_ae"]
+    b_dec: Float[Tensor, "n_instances n_input_ae"]
+
+    def __init__(self, cfg: AutoEncoderConfig, importance_matrix: Float[Tensor, "n_instances n_input_ae"] = None):
+        super().__init__()
+        self.cfg = cfg
+        self.importance_matrix = importance_matrix
+
+        self.W_enc = nn.Parameter(nn.init.xavier_normal_(t.empty((cfg.n_instances, cfg.n_input_ae, cfg.n_hidden_ae))))
+        if not(cfg.tied_weights):
+            self.W_dec = nn.Parameter(nn.init.xavier_normal_(t.empty((cfg.n_instances, cfg.n_hidden_ae, cfg.n_input_ae))))
+        self.b_enc = nn.Parameter(t.zeros(cfg.n_instances, cfg.n_hidden_ae))
+        self.b_dec = nn.Parameter(t.zeros(cfg.n_instances, cfg.n_input_ae))
+
+    def forward(self, h: Float[Tensor, "batch_size n_instances n_hidden"]):
+
+        # Compute activations
+        h_cent = h - self.b_dec
+        acts = einops.einsum(
+            h_cent, self.W_enc,
+            "batch_size n_instances n_input_ae, n_instances n_input_ae n_hidden_ae -> batch_size n_instances n_hidden_ae"
+        )
+        acts = F.relu(acts + self.b_enc)
+
+        # Compute reconstructed input
+        h_reconstructed = einops.einsum(
+            acts, (self.W_enc.transpose(-1, -2) if self.cfg.tied_weights else self.W_dec),
+            "batch_size n_instances n_hidden_ae, n_instances n_hidden_ae n_input_ae -> batch_size n_instances n_input_ae"
+        ) + self.b_dec
+
+        # Compute loss, return values
+        h_err = h_reconstructed - h
+        if self.importance_matrix is not None:
+            importance_matrix = self.importance_matrix[None, : ].to(h_err.device)
+            h_err = h_err * importance_matrix
+        l2_loss = h_err.pow(2).mean(-1) # shape [batch_size n_instances]
+        l1_loss = acts.abs().sum(-1) # shape [batch_size n_instances]
+        loss = (self.cfg.l1_coeff * l1_loss + l2_loss).mean(0).sum() # scalar
+
+        return l1_loss, l2_loss, loss, acts, h_reconstructed
+
+    @t.no_grad()
+    def normalize_decoder(self) -> None:
+        '''
+        Normalizes the decoder weights to have unit norm. If using tied weights, we we assume W_enc is used for both.
+        '''
+        if self.cfg.tied_weights:
+            self.W_enc.data = self.W_enc.data / self.W_enc.data.norm(dim=1, keepdim=True)
+        else:
+            self.W_dec.data = self.W_dec.data / self.W_dec.data.norm(dim=2, keepdim=True)