topk_sae.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
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)