wip

notebooks/11d_sae_no_imp.ipynb

@@ -174,27 +174,7 @@
    "execution_count": 4,
    "metadata": {},
    "outputs": [],
-   "source": [
-    "# # insample_datasets = list(set(ds_val['ds_string_base']))\n",
-    "# # outsample_datasets = list(set(ds_ood['ds_string_base']))\n",
-    "# # print(insample_datasets, outsample_datasets)\n",
-    "# from src.datasets.act_dm import ActivationDataModule, SharedDataset\n",
-    "\n",
-    "\n",
-    "# class ActivationDataModule2(ActivationDataModule):\n",
-    "#     def to_tds(self, ds, name):\n",
-    "#         \"\"\"huggingface dataset to pytorch.\"\"\"\n",
-    "#         h = self.hparams\n",
-    "#         # 4x faster if we make it a tensor ourselves\n",
-    "#         ds = ds.with_format(None)\n",
-    "#         tds = torch.utils.data.TensorDataset(\n",
-    "#             torch.FloatTensor(ds['X'][..., 0]), torch.FloatTensor(ds['y']))\n",
-    "        \n",
-    "#         # this shared dataset is 10x faster with multiple workers\n",
-    "#         if h.num_workers>0: \n",
-    "#             tds = SharedDataset(tds, f\"{self.hparams.name}_{name}\") \n",
-    "#         return tds"
-   ]
+   "source": []
   },
   {
    "cell_type": "code",
@@ -246,7 +226,7 @@
     "    for layer in layers_names:\n",
     "        # Stack the base and adapter representations as a 4th dim\n",
     "        X1 = [ds[f'end_residual_{layer}_base'], ds[f'end_residual_{layer}_adapt']]\n",
-    "        X1 = rearrange(X1, 'versions b l f  -> b l f versions')[..., 0]\n",
+    "        X1 = rearrange(X1, 'versions b l f  -> b l f versions')[..., 0]  # NOTE: here we take only the BASE version!\n",
     "        data.append(X1)\n",
     "    \n",
     "    # concat layers\n",

src/helpers/lightning.py

@@ -53,7 +53,7 @@ def plot_hist(df_hist, allowlist=None, logy=False):
     for suffix in suffixes:
         if allowlist and suffix not in allowlist:
             continue
-        plt.figure(figsize=(8, 3))
+        plt.figure(figsize=(5, 2))
         df_hist[[c for c in df_hist.columns if c.endswith(suffix) and '/' in c]].plot(title=suffix, style='.', logy=logy, ax=plt.gca())
         plt.title(suffix)   
         plt.show()

src/probes/importance_matrix.py

@@ -36,8 +36,8 @@ def get_importance_matrix(saved_adaptop_file, layers=['fc1', 'Wqkv']):
     importance_matrix = importance_matrix + 1
 
     # square to make it positive
-    importance_matrix = importance_matrix.clamp(0, None) 
-    importance_matrix -= importance_matrix.mean() - 1
+    # importance_matrix = importance_matrix.clamp(0, None) 
+    # importance_matrix -= importance_matrix.mean() - 1
 
     return importance_matrix
 

src/vae/sae2.py

@@ -44,16 +44,20 @@ class NormedLinear(nn.Linear):
             F.normalize(self.weight, dim=self.norm_dim, out=self.weight)
 
 class NormedLinears(nn.Module):
-    def __init__(self, n_instances: int, n_input_ae: int, n_output: int, weight_norm_dim=None):
+    def __init__(self, n_instances: int, n_input_ae: int, n_output: int, weight_norm_dim=None, act: Callable = nn.ReLU()):
         super().__init__()
         self.linears = nn.ModuleList(NormedLinear(n_input_ae, n_output, norm_dim=weight_norm_dim) for _ in range(n_instances))
+        self.act = act
 
     def weight_norm(self) -> None:
         for m in self.linears:
             m.weight_norm()
 
     def forward(self, x: Tensor) -> Tensor:
-        return t.stack([m(x[:, i]) for i, m in enumerate(self.linears)], dim=1)
+        x = t.stack([m(x[:, i]) for i, m in enumerate(self.linears)], dim=1)
+        if self.act is not None:
+            x = self.act(x)
+        return x
     
 
 class AffineInstanceNorm1d(nn.BatchNorm1d):
@@ -93,36 +97,52 @@ class AutoEncoder(nn.Module):
         # instead of a tied bias, we use a batch norm type module to track and adjust for the training mean and std. We also have an inverse function to undo the normalization.
         self.norm = Affines(cfg.n_instances, cfg.n_input_ae)
 
-        self.encoder = []
-        for i in range(cfg.depth):
-            self.encoder.append(NormedLinears(cfg.n_instances, cfg.n_input_ae, cfg.n_hidden_ae))
-            self.encoder.append(nn.ReLU())
+        self.encoder = [
+            NormedLinears(cfg.n_instances, cfg.n_input_ae, cfg.n_hidden_ae)
+        ]
+        for i in range(1, cfg.depth):
+            self.encoder.append(NormedLinears(cfg.n_instances, cfg.n_hidden_ae, cfg.n_hidden_ae))
         self.encoder = nn.Sequential(*self.encoder)
 
         self.decoder = []
-        for i in range(cfg.depth):
-            self.decoder.append(NormedLinears(cfg.n_instances, cfg.n_hidden_ae, cfg.n_input_ae, weight_norm_dim=0))
-            if i<cfg.depth-1:
-                self.decoder.append(nn.ReLU())
+        for i in range(cfg.depth-1):
+            self.decoder.append(NormedLinears(cfg.n_instances, cfg.n_hidden_ae, cfg.n_hidden_ae, weight_norm_dim=0))
+        self.decoder.append(NormedLinears(cfg.n_instances, cfg.n_hidden_ae, cfg.n_input_ae, weight_norm_dim=0, act=None))
+
         self.decoder = nn.Sequential(*self.decoder)
 
     def forward(self, h: Float[Tensor, "batch_size n_instances n_hidden"]):
 
         # Compute activations
+        depth = self.cfg.depth
         h_cent = self.norm(h)
-        acts = self.encoder(h_cent)
-        h_reconstructed = self.norm.inv(self.decoder(acts))
+        
+        # We're trying gradual sparsity. Where the middle layer has the full sparsity, but the outer layers have less. This should encourage the model to learn increasing sparse representations.
+        l1_loss = 0
+        for i, m in enumerate(self.encoder):
+            h_cent = m(h_cent)
+            l1_loss += h_cent.abs().mean(2).sum(1) / (depth-i)**2 # shape [batch_size n_instances n_latent]
+        acts = latent = h_cent
+
+        # acts = self.encoder(h_cent)
+
+        for i, m in enumerate(self.decoder):
+            acts = m(acts)
+            l1_loss += acts.abs().mean(2).sum(1) / (i+1)**2 # shape [batch_size n_instances n_latent]
+        h_reconstructed = self.norm.inv(acts)
+
+        # h_reconstructed = self.norm.inv(self.decoder(acts))
 
         # 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).sum(2).sum(1) # shape [batch_size n_instances features]
-        l1_loss = acts.abs().sum(2).sum(1) # shape [batch_size n_instances n_latent]
+        l2_loss = h_err.pow(2).mean(2).sum(1) # shape [batch_size n_instances features]
+        # l1_loss = acts.abs().mean(2).sum(1) # shape [batch_size n_instances n_latent]
         loss = (self.cfg.l1_coeff * l1_loss + l2_loss).mean(0) # scalar
 
-        return l1_loss, l2_loss, loss, acts, h_reconstructed
+        return l1_loss, l2_loss, loss, latent, h_reconstructed
 
     @t.no_grad()
     def normalize_decoder(self) -> None: