Create topk_sae.py

topk_sae.py

@@ -0,0 +1,261 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+import numpy as np
+from torch.utils.data import DataLoader, TensorDataset
+from tqdm import tqdm
+import wandb
+import os
+import glob
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class FastAutoencoder(nn.Module):
+    def __init__(self, n_dirs: int, d_model: int, k: int, auxk: int, multik: int, dead_steps_threshold: int = 266):
+        super().__init__()
+        self.n_dirs = n_dirs
+        self.d_model = d_model
+        self.k = k
+        self.auxk = auxk
+        self.multik = multik
+        self.dead_steps_threshold = dead_steps_threshold
+
+        self.encoder = nn.Linear(d_model, n_dirs, bias=False)
+        self.decoder = nn.Linear(n_dirs, d_model, bias=False)
+
+        self.pre_bias = nn.Parameter(torch.zeros(d_model))
+        self.latent_bias = nn.Parameter(torch.zeros(n_dirs))
+
+        self.stats_last_nonzero = torch.zeros(n_dirs, dtype=torch.long, device=device)
+
+    def forward(self, x):
+        x = x - self.pre_bias
+        latents_pre_act = self.encoder(x) + self.latent_bias
+
+        # Main top-k selection
+        topk_values, topk_indices = torch.topk(latents_pre_act, k=self.k, dim=-1)
+        topk_values = F.relu(topk_values)
+        multik_values, multik_indices = torch.topk(latents_pre_act, k=4*self.k, dim=-1)
+        multik_values = F.relu(multik_values)
+
+        latents = torch.zeros_like(latents_pre_act)
+        latents.scatter_(-1, topk_indices, topk_values)
+        multik_latents = torch.zeros_like(latents_pre_act)
+        multik_latents.scatter_(-1, multik_indices, multik_values)
+
+        # Update stats_last_nonzero
+        self.stats_last_nonzero += 1
+        self.stats_last_nonzero.scatter_(0, topk_indices.unique(), 0)
+
+        recons = self.decoder(latents) + self.pre_bias
+        multik_recons = self.decoder(multik_latents) + self.pre_bias
+
+        # AuxK
+        if self.auxk is not None:
+            # Create dead latents mask
+            dead_mask = (self.stats_last_nonzero > self.dead_steps_threshold).float()
+            
+            # Apply mask to latents_pre_act
+            dead_latents_pre_act = latents_pre_act * dead_mask
+            
+            # Select top-k_aux from dead latents
+            auxk_values, auxk_indices = torch.topk(dead_latents_pre_act, k=self.auxk, dim=-1)
+            auxk_values = F.relu(auxk_values)
+
+        else:
+            auxk_values, auxk_indices = None, None
+
+        return recons, {
+            "topk_indices": topk_indices,
+            "topk_values": topk_values,
+            "multik_indices": multik_indices,
+            "multik_values": multik_values,
+            "multik_recons": multik_recons,
+            "auxk_indices": auxk_indices,
+            "auxk_values": auxk_values,
+            "latents_pre_act": latents_pre_act,
+            "latents_post_act": latents,
+        }
+
+    def decode_sparse(self, indices, values):
+        latents = torch.zeros(self.n_dirs, device=indices.device)
+        latents.scatter_(-1, indices, values)
+        return self.decoder(latents) + self.pre_bias
+
+    # def decode_sparse(self, indices, values):
+    #     latents = torch.zeros(1, self.n_dirs, device=indices.device, dtype=torch.float32)
+    #     latents.scatter_(-1, indices.unsqueeze(0), values.unsqueeze(0))
+    #     return self.decoder(latents.squeeze(0)) + self.pre_bias
+    
+    def print_tensor_info(self, tensor, name):
+        print(f"{name} - Shape: {tensor.shape}, Dtype: {tensor.dtype}, Device: {tensor.device}")
+
+    def decode_clamp(self, latents, clamp):
+        topk_values, topk_indices = torch.topk(latents, k = 64, dim=-1)
+        topk_values = F.relu(topk_values)
+        latents = torch.zeros_like(latents)
+        latents.scatter_(-1, topk_indices, topk_values)
+        # multiply latents by clamp, which is 1D but has has the same size as each latent vector
+        latents = latents * clamp
+        
+        return self.decoder(latents) + self.pre_bias
+    
+    def decode_at_k(self, latents, k):
+        topk_values, topk_indices = torch.topk(latents, k=k, dim=-1)
+        topk_values = F.relu(topk_values)
+        latents = torch.zeros_like(latents)
+        latents.scatter_(-1, topk_indices, topk_values)
+        
+        return self.decoder(latents) + self.pre_bias
+
+def unit_norm_decoder_(autoencoder: FastAutoencoder) -> None:
+    with torch.no_grad():
+        autoencoder.decoder.weight.div_(autoencoder.decoder.weight.norm(dim=0, keepdim=True))
+
+def unit_norm_decoder_grad_adjustment_(autoencoder: FastAutoencoder) -> None:
+    if autoencoder.decoder.weight.grad is not None:
+        with torch.no_grad():
+            proj = torch.sum(autoencoder.decoder.weight * autoencoder.decoder.weight.grad, dim=0, keepdim=True)
+            autoencoder.decoder.weight.grad.sub_(proj * autoencoder.decoder.weight)
+
+def mse(output, target):
+    return F.mse_loss(output, target)
+
+def normalized_mse(recon, xs):
+    return mse(recon, xs) / mse(xs.mean(dim=0, keepdim=True).expand_as(xs), xs)
+
+def loss_fn(ae, x, recons, info, auxk_coef, multik_coef):
+    recons_loss = normalized_mse(recons, x)
+    recons_loss += multik_coef * normalized_mse(info["multik_recons"], x)
+    
+    if ae.auxk is not None:
+        e = x - recons.detach()  # reconstruction error
+        auxk_latents = torch.zeros_like(info["latents_pre_act"])
+        auxk_latents.scatter_(-1, info["auxk_indices"], info["auxk_values"])
+        e_hat = ae.decoder(auxk_latents)  # reconstruction of error using dead latents
+        auxk_loss = normalized_mse(e_hat, e)
+        total_loss = recons_loss + auxk_coef * auxk_loss
+    else:
+        auxk_loss = torch.tensor(0.0, device=device)
+        total_loss = recons_loss
+    
+    return total_loss, recons_loss, auxk_loss
+
+def init_from_data_(ae, data_sample):
+    # set pre_bias to median of data
+    ae.pre_bias.data = torch.median(data_sample, dim=0).values
+    nn.init.xavier_uniform_(ae.decoder.weight)
+
+    # decoder is unit norm
+    unit_norm_decoder_(ae)
+
+    # encoder as transpose of decoder
+    ae.encoder.weight.data = ae.decoder.weight.t().clone()
+
+    nn.init.zeros_(ae.latent_bias)
+
+def train(ae, train_loader, optimizer, epochs, k, auxk_coef, multik_coef, clip_grad=None, save_dir="../models", model_name=""):
+    os.makedirs(save_dir, exist_ok=True)
+    step = 0
+    num_batches = len(train_loader)
+    for epoch in range(epochs):
+        ae.train()
+        total_loss = 0
+        for batch in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}"):
+            optimizer.zero_grad()
+            x = batch[0].to(device)
+            recons, info = ae(x)
+            loss, recons_loss, auxk_loss = loss_fn(ae, x, recons, info, auxk_coef, multik_coef)
+            loss.backward()
+            step += 1
+
+            # calculate proportion of dead latents (not fired in last num_batches = 1 epoch)
+            dead_latents_prop = (ae.stats_last_nonzero > num_batches).float().mean().item()
+
+            wandb.log({
+                "total_loss": loss.item(),
+                "reconstruction_loss": recons_loss.item(),
+                "auxiliary_loss": auxk_loss.item(),
+                "dead_latents_proportion": dead_latents_prop,
+                "l0_norm": k,
+                "step": step
+            })
+            
+            unit_norm_decoder_grad_adjustment_(ae)
+            
+            if clip_grad is not None:
+                torch.nn.utils.clip_grad_norm_(ae.parameters(), clip_grad)
+            
+            optimizer.step()
+            unit_norm_decoder_(ae)
+            
+            total_loss += loss.item()
+        
+        avg_loss = total_loss / len(train_loader)
+        print(f"Epoch {epoch+1}, Average Loss: {avg_loss:.4f}")
+
+        # Delete previous model saves for this configuration
+        for old_model in glob.glob(os.path.join(save_dir, f"{model_name}_epoch_*.pth")):
+            os.remove(old_model)
+
+        # Save new model
+        save_path = os.path.join(save_dir, f"{model_name}_epoch_{epoch+1}.pth")
+        torch.save(ae.state_dict(), save_path)
+        print(f"Model saved to {save_path}")
+
+def main():
+    d_model = 1536
+    n_dirs = 3072 #9216
+    k = 64 #64
+    auxk = k*2 #256
+    multik = 128
+    batch_size = 1024
+    lr = 1e-4
+    auxk_coef = 1/32
+    clip_grad = 1.0
+    multik_coef = 0 # turn it off
+
+    csLG = False
+
+    # Create model name
+    model_name = f"{k}_{n_dirs}_{auxk}_auxk" if not csLG else f"{k}_{n_dirs}_{auxk}_auxk_csLG"
+    epochs = 50 if not csLG else 137
+
+    wandb.init(project="saerch", name=model_name, config={
+        "n_dirs": n_dirs,
+        "d_model": d_model,
+        "k": k,
+        "auxk": auxk,
+        "batch_size": batch_size,
+        "lr": lr,
+        "epochs": epochs,
+        "auxk_coef": auxk_coef,
+        "multik_coef": multik_coef,
+        "clip_grad": clip_grad,
+        "device": device.type
+    })
+
+    if not csLG:
+        data = np.load("../data/vector_store_astroPH/abstract_embeddings.npy")
+        print("Doing astro.ph...")
+    else:
+        data = np.load("../data/vector_store_csLG/abstract_embeddings.npy")
+        print("Doing csLG...")
+    data_tensor = torch.from_numpy(data).float()
+    # Print shape
+    print(f"Data shape: {data_tensor.shape}")
+    dataset = TensorDataset(data_tensor)
+    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+    ae = FastAutoencoder(n_dirs, d_model, k, auxk, multik).to(device)
+    init_from_data_(ae, data_tensor[:10000].to(device))
+
+    optimizer = optim.Adam(ae.parameters(), lr=lr)
+
+    train(ae, train_loader, optimizer, epochs, k, auxk_coef, multik_coef, clip_grad=clip_grad, model_name=model_name)
+
+    wandb.finish()
+
+if __name__ == "__main__":
+    main()