File size: 6,290 Bytes
a2ffd07 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import torch
from torch import Tensor
from .decoder import decode
from .types import Stats, TopK, SAEOut
from .Standard import SAE_Template, SAEConfig
from .SAE_Wrapper import _disable_hooks
from typing import Any, Tuple, Dict, List, Optional
from .Utils import standardize, unit_norm_decoder
class BatchTopKSAE(SAE_Template):
threshold: torch.Tensor
"""The inference threshold for JumpReLU activation function."""
def __init__(
self,
d_in: int,
d_sae: int,
hook_names: list[str],
k: int,
dead_steps_threshold: int,
dead_threshold: float = 1e-3,
auxk: int | None = 256,
standardize: bool = True,
threshold_beta: float = 0.99,
**kwargs,
) -> None:
"""
Args:
d_in (int): The number of inputs.
d_sae (int): The number of latents.
k (int): The number of largest latents to keep.
dead_steps_threshold (int): The number of steps after which a latent is
flagged as dead during training.
dead_threshold (float): The threshold for a latent to be considered
activated. Defaults to 1e-3.
auxk (int | None): The number of dead latents with which to model the
reconstruction error. Defaults to 256.
standardize (bool): Whether to standardize the inputs. Defaults to True.
"""
super().__init__(
d_in=d_in,
d_sae=d_sae,
hook_names=hook_names,
auxk=auxk,
dead_steps_threshold=dead_steps_threshold,
dead_threshold=dead_threshold,
standardize=standardize,
)
self.register_buffer("threshold", torch.tensor(dead_threshold, dtype=torch.float32))
self.threshold_beta = threshold_beta
self.k = k
def encode(
self, inputs: torch.Tensor, use_threshold: bool = True,
) -> tuple[Tensor, TopK, TopK | None, Stats | None, torch.Tensor]:
inputs = self.hook_sae_input(inputs)
stats = None
if self.cfg.standardize:
inputs, stats = standardize(inputs)
# Keep a reference to the latents before the TopK activation function
hidden_pre: Tensor = self.hook_sae_acts_pre(self.encoder.forward(inputs - self.pre_encoder_bias))
if use_threshold:
latents = self.hook_sae_acts_post(hidden_pre * (torch.relu(hidden_pre) > self.threshold))
else:
# Perform batchtopk
num_toks = hidden_pre.numel() // self.cfg.d_sae
flatten_values, flatten_indices = torch.topk(
torch.relu(hidden_pre.flatten()),
k=self.k * num_toks,
sorted=False
)
latents = torch.zeros_like(hidden_pre.flatten())
latents = self.hook_sae_acts_post(
latents.scatter_(-1, flatten_indices, flatten_values).reshape(hidden_pre.shape)
)
self.update_threshold(latents)
max_k = torch.max(torch.sum(latents > self.cfg.dead_threshold, dim=-1)).item()
values, indices = torch.topk(
latents,
k=max_k, # type: ignore
sorted=False
)
topk = TopK(values, indices)
self.update_last_nonzero(topk, inputs.device)
dead, auxk = self.compute_dead_latents_and_auxk(hidden_pre)
return latents, topk, auxk, stats, dead
def decode(self, latents: Tensor, stats: Stats | None = None) -> torch.Tensor:
recons = (latents @ self.decoder.weight.T) + self.pre_encoder_bias
if stats is not None:
recons = recons * stats.std + stats.mean
return self.hook_sae_recons(recons)
def forward_training(self, inputs: torch.Tensor) -> SAEOut:
latents, topk, auxk, stats, dead = self.encode(inputs, use_threshold=False)
recons = self.decode(latents, stats)
auxk_recons = None
if auxk is not None:
auxk_latents = torch.zeros_like(latents)
auxk_latents.scatter_(
-1,
auxk.indices,
torch.relu(auxk.values),
)
auxk_recons = self.decode(auxk_latents)
return SAEOut(topk, recons, auxk, auxk_recons, dead, 0)
def update_threshold(self, latents: Tensor):
device_type = "cuda" if latents.is_cuda else "cpu"
with torch.autocast(device_type=device_type, enabled=False), torch.no_grad():
active = latents[latents > self.cfg.dead_threshold]
if active.size(0) == 0:
min_activation = self.cfg.dead_threshold
else:
min_activation = active.min().detach().to(dtype=torch.float32)
self.threshold = (self.threshold_beta * self.threshold) + (
(1 - self.threshold_beta) * min_activation
)
def update_last_nonzero(self, topk: TopK, device: torch.device) -> None:
# Update the number of steps since the latents have activated
last_nonzero = torch.zeros_like(self.last_nonzero, device=device)
last_nonzero.scatter_add_(
dim=0,
index=topk.indices.reshape(-1),
src=(
topk.values > self.cfg.dead_threshold
).to(last_nonzero.dtype).reshape(-1),
)
self.last_nonzero *= 1 - last_nonzero.clamp(max=1)
self.last_nonzero += 1
def compute_dead_latents_and_auxk(self, latents: Tensor) -> tuple[Tensor, TopK | None]:
# Mask the latents flagged as dead during training
dead_mask = self.last_nonzero >= self.cfg.dead_steps_threshold
latents = latents * dead_mask
# Compute the fraction of dead latents
dead = torch.sum(dead_mask, dtype=torch.float32).detach() / self.cfg.d_sae
# If auxk is not None, find the auxk largest dead latents
auxk = None
if self.cfg.auxk is not None:
values_auxk , indices_auxk = torch.topk(
latents,
k=self.cfg.auxk,
sorted=False
)
auxk = TopK(values_auxk , indices_auxk)
return dead, auxk |