import torch as t from torch import Tensor from .Graph_Template import Graph, GraphName, Node, Index from typing import List, Tuple, Dict, Union, Callable, Any from tqdm import tqdm from transformer_lens.hook_points import HookPoint from transformer_lens import ( utils, HookedTransformer, ActivationCache, ) from copy import deepcopy from collections import OrderedDict, defaultdict from warnings import warn from contextlib import contextmanager from model.hooked_blip import HookedSAEBlipConditionalGeneration from functools import partial from itertools import product from collections import OrderedDict from .Graph_utils import nested_dict_to_string from .Sparse_act import SparseAct from torch import sparse_coo_tensor import einops from .vision_sae_wrapper import * def cache_to_sparseact( cache: Dict[str, Tensor] | ActivationCache, act_name: str, res_name: str | None = None, resc_name: str | None = None, ) -> SparseAct: return SparseAct( act=cache[act_name], res=cache[res_name] if res_name is not None else None, resc=cache[resc_name] if resc_name is not None else None, ) def interpolate(start: Tensor, end: Tensor, frac: float,) -> Tensor: assert 0 <= frac <= 1, "frac must be in [0, 1]" return start * frac + end * (1 - frac) def create_list(dict: Dict[str, Tensor], name: str) -> Dict[str, Tensor]: if name not in dict: dict[name] = 0 # type: ignore return dict class ConnectionNode(Node): def __init__(self, name: str): self._name = name @property def name(self) -> str: return self._name def __eq__(self, other) -> bool: if isinstance(other, ConnectionNode): return self._name == other.name elif isinstance(other, str): return self._name == other else: raise NotImplementedError("other is not an instance of ConnectionNode or str") def __repr__(self) -> str: return self.name def __hash__(self) -> int: return hash(self.name) class ConnectionIndex(Index): def __init__( self, list_index: tuple[int|None, ...] | None = None, ): if list_index is None: self.list_index = (None,) else: for index in list_index: assert type(index) == int or index == None, "index is not an instance of int or None" self.list_index = list_index @property def as_index(self) -> Tuple[int | slice, ...]: return tuple(slice(None) if x is None else x for x in self.list_index) # for indexing def __repr__(self) -> str: ret = "[" for idx, x in enumerate(self.list_index): if idx > 0: ret += ", " if x is None: ret += ":" elif type(x) == int: ret += str(x) else: raise NotImplementedError(x) ret += "]" return ret def __eq__(self, other) -> bool: if isinstance(other, ConnectionIndex): return self.list_index == other.list_index elif isinstance(other, tuple): return self.list_index == other else: raise NotImplementedError("other is not an instance of ConnectionIndex or tuple") def __hash__(self) -> int: return hash(self.list_index) class FeatureIndex(ConnectionIndex): def __init__( self, idx: Tuple[int, ...] | List[int], length = 2, ): self.idx = tuple(idx) super().__init__(return_idx(idx, length)) class FeatureErrorIndex(ConnectionIndex): def __init__( self, idx: Tuple[int, ...] | List[int], length = 2, ): self.idx = tuple(idx) super().__init__(return_idx(idx, length)) class ErrorIndex(ConnectionIndex): def __init__( self, idx: Tuple[int, ...] | List[int], length = 1, ): self.idx = tuple(idx) super().__init__(return_idx(idx, length)) def sae_hook_name(sae_name: str) -> str: return f'{sae_name}.hook_sae_acts_post' # e.g. ...hook_resid_pre.hook_sae_acts_post def error_term_name(sae_name: str) -> str: return f'{sae_name}.hook_sae_error' def output_hook_name(sae_name: str) -> str: return f'{sae_name}.hook_sae_output' def input_hook_name(sae_name: str) -> str: return f'{sae_name}.hook_sae_input' def recons_hook_name(sae_name: str) -> str: return f'{sae_name}.hook_sae_recons' def revert_hook_name(sae_name: str) -> str: return ".".join(sae_name.split(".")[:-1]) # e.g. ...hook_resid_pre.hook_sae_acts_post -> ...hook_resid_pre def return_idx(idx: Tuple | List, length: int = 2) -> Tuple: assert len(idx) <= length, "The length of idx must be less than or equal to length" return tuple([None for _ in range(length - len(idx))] + list(idx)) class Feature_Graph_Blip(Graph): def __init__( self, model: HookedSAEBlipConditionalGeneration, text_saes: Dict[int, List[Tuple[str, Any]]], # {layer: list[{hook_position: HookedSAE}]}, can define granularity here vision_saes: Dict[int, List[Tuple[str, Any]]], use_error_term: bool = False, ): ''' ''' self.model = model self.cfg = model.cfg self.use_error_term = use_error_term self.device = self.cfg.device self.n_layers = self.cfg.n_layers # self.n_heads = self.cfg.n_heads assert len(text_saes) == self.n_layers, "please provide SAEs" self.text_saes = text_saes self.vision_saes = vision_saes self.dict_vision_saes = self.extract_saes(vision_saes) self.dict_text_saes = self.extract_saes(text_saes) self.dict_saes: Dict[str, Any] = self.dict_vision_saes | self.dict_text_saes self.reset_graph() def graph_type(self) -> str: return GraphName.feature_graph def extract_saes(self, saes: Dict[int, List[Tuple[str, Any]]]) -> Dict[str, Any]: dict_saes = {} for layer in range(self.n_layers): for hook_position, sae in saes[layer]: dict_saes[hook_position] = sae return dict_saes def build_default_connection( self, seq_length: int, token_wise: bool = False, inter: bool = True, intra: bool = False, ) -> Dict: ''' Sample tokens to get the shape of the activations, necessary to build the graph of features If token_wise, the each node is a feature of a token, else, the graph is built feature-wise ''' self.reset_graph() # vision for layer in range(self.n_layers): # Setup self.connection[layer] = OrderedDict() for hook_position, _ in self.vision_saes[layer]: self.connection[layer][(ConnectionNode(hook_position), ConnectionIndex())] = [] # Inter-layer connections if inter: for i, (hook_position_end, _) in reversed(list(enumerate(self.vision_saes[layer]))): for j, (hook_position_start, _) in enumerate(self.vision_saes[layer]): if i > j: self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) # Intra-layer connections if intra: for prev_layer in range(layer): for hook_position_end, _ in self.vision_saes[layer]: for hook_position_start, _ in self.vision_saes[prev_layer]: self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) # Sequential connection layer-1 --> layer else: if layer == 0: continue for hook_position_start, _ in self.vision_saes[layer-1]: # connect to the first node of the next layer hook_position_end = self.vision_saes[layer][0][0] self.connection[layer][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) # text for layer in range(self.n_layers): # Setup self.connection[layer + self.n_layers] = OrderedDict() # text has layer of: layer + n_layers for hook_position, _ in self.text_saes[layer]: self.connection[layer + self.n_layers][(ConnectionNode(hook_position), ConnectionIndex())] = [] # Inter-layer connections if inter: for i, (hook_position_end, _) in reversed(list(enumerate(self.text_saes[layer]))): for j, (hook_position_start, _) in enumerate(self.text_saes[layer]): if i > j: self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) # Intra-layer connections if intra: # intra connection with vision for hook_position_end, _ in self.text_saes[layer]: if layer != 0 or ("crossattention.self.hook_attn_out" in hook_position_end or "mlp" in hook_position_end): for vision_layer in range(self.n_layers): for hook_position_start, _ in self.vision_saes[vision_layer]: self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) for prev_layer in range(layer): for hook_position_end, _ in self.text_saes[layer]: for hook_position_start, _ in self.text_saes[prev_layer]: self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) # Sequential connection layer-1 --> layer else: if layer == 0: # connect with the last vision layer hook_position_start = self.vision_saes[self.n_layers-1][-1][0] hook_position_end = self.text_saes[0][0][0] if "crossattention.self.hook_attn_out" in hook_position_end or "mlp" in hook_position_end: self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) else: for hook_position_start, _ in self.text_saes[layer-1]: # connect to the first node of the prev layer hook_position_end = self.text_saes[layer][0][0] self.connection[layer + self.n_layers][(ConnectionNode(hook_position_end), ConnectionIndex())].append( (ConnectionNode(hook_position_start), ConnectionIndex()) ) self._build_nodes(seq_length, token_wise) self._build_graphs() return self.connection def _build_nodes(self, seq_length: int, token_wise: bool) -> None: self.seq_length = seq_length self.token_wise = token_wise # text for sae_name, sae in self.dict_text_saes.items(): node_shape = (self.seq_length, sae.cfg.d_sae) self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct( t.ones(node_shape, requires_grad=False).to(self.device), None, t.ones(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None, ) self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct( t.zeros(node_shape, requires_grad=False).to(self.device), None, t.zeros(self.seq_length, requires_grad=False).to(self.device) if self.use_error_term else None, ) # vision for sae_name, sae in self.dict_vision_saes.items(): node_shape = (1, sae.cfg.d_sae) # HARDCODE self.nodes[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct( t.ones(node_shape, requires_grad=False).to(self.device), None, t.ones(1, requires_grad=False).to(self.device) if self.use_error_term else None, ) self.node_scores[(ConnectionNode(sae_name), ConnectionIndex())] = SparseAct( t.zeros(node_shape, requires_grad=False).to(self.device), None, t.zeros(1, requires_grad=False).to(self.device) if self.use_error_term else None, ) def _build_graphs(self) -> None: assert self.nodes, "Please build the nodes first" for _, connections in self.connection.items(): for hook_position_end, list_hook_positions_start in connections.items(): self.edges[hook_position_end] = OrderedDict() self.edge_scores[hook_position_end] = OrderedDict() for hook_position_start in list_hook_positions_start: self.edges[hook_position_end][hook_position_start] = t.zeros([0]) # place holder self.edge_scores[hook_position_end][hook_position_start] = t.zeros([0]) # place holder def reset_graph(self) -> None: self.connection: OrderedDict[int, OrderedDict[Tuple[Node, Index], List[Tuple[Node, Index]]]] = OrderedDict() # {layer: {hook_position_end: [hook_position_start]}} self.nodes: Dict[Tuple[Node, Index], SparseAct] = {} # {hook_position: SparseAct[act: Tensor['seq d_sae'], resc: Tensor['seq']]} self.node_scores: Dict[Tuple[Node, Index], SparseAct] = {} self.edges: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict() # {hook_position_end: {hook_position_start: Tensor}} self.edge_scores: OrderedDict[Tuple[Node, Index], OrderedDict[Tuple[Node, Index], Tensor]] = OrderedDict() self.seq_length: int | None = None self.token_wise: bool | None = None def _check_graph(self) -> None: assert len(self.connection) > 0, "Graph is empty" assert len(self.nodes) > 0, "Nodes is empty" assert len(self.node_scores) > 0, "Node scores is empty" assert len(self.edges) > 0, "Edges is empty" assert len(self.edge_scores) > 0, "Edge scores is empty" assert self.seq_length is not None assert self.token_wise is not None def __repr__(self) -> str: self._check_graph() return nested_dict_to_string(self.connection, indent=4) def _check_feat_idx(self, feat_idx: Tuple | List) -> None: if self.token_wise: assert len(feat_idx) == 2, "Please provide the correct index" else: assert len(feat_idx) == 1, "Please provide the correct index" def _check_error_idx(self, error_idx: Tuple | List) -> None: if self.token_wise: assert len(error_idx) == 1, "Please provide the correct index" else: assert len(error_idx) == 0, "Please provide the correct index" def _active_nodes( self, sae_name: Node, node_idx: Index, reverse: bool = False, ) -> Tuple[List, ...]: self._check_graph() if self.token_wise: node = self.nodes[(sae_name, node_idx)] # act: (seq, d_sae), resc: (seq) else: node = self.nodes[(sae_name, node_idx)].mean(dim=0) # act: (d_sae), resc: () - COLAPSED TENSOR if reverse: active_node = (node.act == 0).nonzero().tolist() active_error = (node.resc == 0).nonzero().tolist() if self.use_error_term else [] # type: ignore else: active_node = node.act.nonzero().tolist() active_error = node.resc.nonzero().tolist() if self.use_error_term else [] # type: ignore return active_node, active_error def active_nodes( self, sae_name: Node, node_idx: Index, reverse: bool = False, ) -> List[Tuple[Node, Index]]: active_node, active_error = self._active_nodes(sae_name, node_idx, reverse) feat_list = [ (sae_name, FeatureIndex(idx)) for idx in active_node ] error_list = [ (sae_name, ErrorIndex(idx)) for idx in active_error ] return feat_list + error_list # type: ignore def add_node( self, sae_name: Node, node_idx: Index, ) -> None: self._check_graph() pass # no implementation needed def add_edge( self, sae_name_start: Node, node_idx_start: Index, sae_name_end: Node, node_idx_end: Index, ) -> None: self._check_graph() pass # no implementation needed def delete_node( self, sae_name: Node, node_idx: Index, ) -> None: self._check_graph() pass # no implementation needed def delete_edge( self, sae_name_start: Node, node_idx_start: Index, sae_name_end: Node, node_idx_end: Index, ) -> None: self._check_graph() pass # no implementation needed def _check_node_shape(self, sae_name: Node, sparseact: SparseAct): seq = self.seq_length if "vision" not in sae_name.name else 1 try: if self.token_wise: assert sparseact.act.shape == t.Size([seq, self.dict_saes[sae_name.name].cfg.d_sae]) # type: ignore else: assert sparseact.act.shape == t.Size([self.dict_saes[sae_name.name].cfg.d_sae]) if self.use_error_term: if self.token_wise: assert sparseact.resc.shape == t.Size([seq]) # type: ignore else: assert sparseact.resc.numel() == 1 # type: ignore except: raise ValueError("Wrong input shape for nodes.") def update_node( self, sae_name: Node, node_idx: Index, value_and_mask: Tuple[SparseAct, SparseAct], ) -> None: self._check_graph() value, mask = value_and_mask assert isinstance(value, SparseAct), "value is not an instance of SparseAct" assert isinstance(mask, SparseAct), "mask is not an instance of SparseAct" self._check_node_shape(sae_name, value) if not self.token_wise: seq = self.seq_length if "vision" not in sae_name.name else 1 # convert to the right format of (seq, d_sae) and (seq) value_act = einops.repeat(value.act, 'd_sae -> seq d_sae', seq=seq) mask_act = einops.repeat(mask.act, 'd_sae -> seq d_sae', seq=seq) value_resc = einops.repeat(value.resc, ' -> seq', seq=seq) if self.use_error_term else None mask_resc = einops.repeat(mask.resc, ' -> seq', seq=seq) if self.use_error_term else None value = SparseAct(value_act, None, value_resc) mask = SparseAct(mask_act, None, mask_resc) self.nodes[(sae_name, node_idx)] = mask.to(t.float32) self.node_scores[(sae_name, node_idx)] = value def update_edge( self, sae_name_start: Node, node_idx_start: Index, sae_name_end: Node, node_idx_end: Index, value_and_mask: Tuple[Tensor, Tensor], ) -> None: self._check_graph() value, mask = value_and_mask assert isinstance(value, Tensor), "value is not an instance of Tensor" assert isinstance(mask, Tensor), "mask is not an instance of Tensor" self.edges[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = mask self.edge_scores[(sae_name_end, node_idx_end)][(sae_name_start, node_idx_start)] = value def find_deleted_nodes(self, reverse = False) -> List[Tuple[Node, Index]]: self._check_graph() all_list = [] for sae_name, node_idx in self.nodes.keys(): inactive_node, inactive_error = self._active_nodes(sae_name, node_idx, not reverse) feat_list = [ (sae_name, FeatureIndex(idx)) for idx in inactive_node ] error_list = [ (sae_name, ErrorIndex(idx)) for idx in inactive_error ] all_list += feat_list + error_list return all_list def find_deleted_edges(self, reverse = False) -> List[Tuple[Node, Index, Node, Index]]: self._check_graph() all_list = [] for end_node, end_idx in self.edges.keys(): for start_node, start_idx in self.edges[(end_node, end_idx)].keys(): edge = self.edges[(end_node, end_idx)][(start_node, start_idx)] if not reverse: inactive_node = (edge.values() == 0) # (num_active, seq, d_sae+1) or (num_active, d_sae+1) else: inactive_node = edge.values() d_sae_end = self.dict_saes[end_node.name].cfg.d_sae d_sae_start = self.dict_saes[start_node.name].cfg.d_sae active_idx = [index for _, index in self.active_nodes(start_node, start_idx)] revised_active_nodes = [ index.idx if isinstance(index, FeatureIndex) else index.idx + (d_sae_start,) # type: ignore for index in active_idx ] for num_active_idx in range(edge.indices().shape[1]): end_node_idx = edge.indices()[:, num_active_idx].tolist() # check if the end_node_idx is the error term check_error = True if self.use_error_term and end_node_idx[-1] == d_sae_end else False # create the end_index class to be used in the tuple end_index = FeatureIndex(end_node_idx) if not check_error else ErrorIndex(end_node_idx[:-1]) for i, idx in enumerate(revised_active_nodes): if inactive_node[(num_active_idx,) + idx] > 0: all_list.append( ( start_node, active_idx[i], end_node, end_index, ) ) return all_list def iterate_nodes(self) -> List[Tuple[Node, Index]]: self._check_graph() return list(self.nodes.keys()) def iterate_edges(self) -> List[Tuple[Node, Index, Node, Index]]: self._check_graph() all_edges = [] for end, edges in self.edges.items(): for start, edge in edges.items(): all_edges.append(start + end) return all_edges def add_single_feature( self, sae_name: Node, node_idx: Index, feat_idx: Tuple[int, ...] | List[int] | FeatureIndex, ) -> None: self._check_graph() self._set_feat_value(sae_name, node_idx, feat_idx, 1) def delete_single_feature( self, sae_name: Node, node_idx: Index, feat_idx: Tuple[int, ...] | List[int] | FeatureIndex, ) -> None: self._check_graph() self._set_feat_value(sae_name, node_idx, feat_idx, 0) def update_single_feature( self, sae_name: Node, node_idx: Index, feat_idx: Tuple[int, ...] | List[int] | FeatureIndex, value: int | float | Tensor, ) -> None: self._check_graph() assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor." self._set_feat_value(sae_name, node_idx, feat_idx, value) def add_single_error( self, sae_name: Node, node_idx: Index, error_idx: Tuple[int, ...] | List[int] | ErrorIndex, ) -> None: self._check_graph() self._set_error_value(sae_name, node_idx, error_idx, 1) def delete_single_error( self, sae_name: Node, node_idx: Index, error_idx: Tuple[int, ...] | List[int] | ErrorIndex, ) -> None: self._check_graph() self._set_error_value(sae_name, node_idx, error_idx, 0) def update_single_error( self, sae_name: Node, node_idx: Index, error_idx: Tuple[int, ...] | List[int] | ErrorIndex, value: int | float | Tensor, ) -> None: self._check_graph() assert isinstance(value, (int, float, Tensor)), "The value must be int, float or Tensor." self._set_error_value(sae_name, node_idx, error_idx, value) def _set_feat_value( self, sae_name: Node, node_idx: Index, feat_idx: Tuple[int, ...] | List[int] | FeatureIndex, value: int | float | Tensor, ) -> None: self._check_graph() if isinstance(feat_idx, (tuple, list)): self._check_feat_idx(feat_idx) self.nodes[(sae_name, node_idx)].act[return_idx(feat_idx)] = value elif isinstance(feat_idx, FeatureIndex): self.nodes[(sae_name, node_idx)].act[feat_idx.as_index] = value else: raise ValueError(f"feat_idx of type {type(feat_idx)} is not supported.") def _set_error_value( self, sae_name: Node, node_idx: Index, error_idx: Tuple[int, ...] | List[int] | ErrorIndex, value: int | float | Tensor, ) -> None: if not self.use_error_term: # nothing to delete return if isinstance(error_idx, (tuple, list)): self._check_error_idx(error_idx) self.nodes[(sae_name, node_idx)].resc[return_idx(error_idx, 1)] = value # type: ignore elif isinstance(error_idx, ErrorIndex): self.nodes[(sae_name, node_idx)].resc[error_idx.as_index] = value # type: ignore else: raise ValueError(f"error_idx of type {type(error_idx)} is not supported.") def remove_error_term(self) -> None: if self.use_error_term: for sae_name, node_idx in self.nodes.keys(): if self.nodes[(sae_name, node_idx)].resc is not None: seq = self.seq_length if "vision" not in sae_name.name else 1 self.nodes[(sae_name, node_idx)].resc = t.zeros(seq, requires_grad=False).to(self.device) # type: ignore def forward( self, inputs: Dict[str, Tensor], corrupt_cache: ActivationCache | Dict[str, Tensor] | None, patch_deleted_comp: bool = False, **kwargs, ) -> Tuple[Tensor, Dict[str, SparseAct]]: ''' Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation ''' self._check_graph() self.model.reset_hooks() self.model_setup() fwd_cache = {} def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor: if hook.name == sae_hook_name(sae_name): if patch_deleted_comp and corrupt_cache is not None: act_mask = self.nodes[(sae_name, (None,))].act == 0 # type: ignore act[:, act_mask] = corrupt_cache[sae_hook_name(sae_name)][:, act_mask] fwd_cache[sae_hook_name(sae_name)] = act.detach() elif hook.name == error_term_name(sae_name) and self.use_error_term: if patch_deleted_comp and corrupt_cache is not None: resc_mask = self.nodes[(sae_name, (None,))].resc == 0 # type: ignore act[:, resc_mask] = corrupt_cache[error_term_name(sae_name)][:, resc_mask] fwd_cache[error_term_name(sae_name)] = act.detach() return act with t.no_grad() and self._hook_vision_sae(): with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term): with self.model.hooks( fwd_hooks=[ (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) for sae_name in self.dict_saes.keys() ], ): logits = self.model(inputs) cache = {} for sae_name in self.dict_saes.keys(): cache[sae_name] = SparseAct( act=fwd_cache[sae_hook_name(sae_name)], res=fwd_cache[error_term_name(sae_name)] if self.use_error_term else None, ) for sae in self.dict_saes.values(): sae.reset_hooks() self.model.reset_hooks() return logits, cache def __call__(self, *args, **kwargs) -> Tuple[Tensor, Dict[str, SparseAct]]: return self.forward(*args, **kwargs) def forward_backward_gradient( self, inputs: Dict[str, Tensor], corrupt_cache: ActivationCache | Dict[str, Tensor], metric: Callable[[Tensor], Tensor], retain_graph: bool = False, mode: str = 'node', gradient_mode: str = 'standard', pass_through_grad: bool = False, verbose: bool = False, **kwargs, ) -> Tuple[ Dict[Tuple[Node, Index], SparseAct], # node effects Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects ]: if mode == 'node': if verbose: print("Calculating node gradients...") if gradient_mode == "standard": node_grads, clean_cache = self._gradient_wrt_nodes( inputs, metric, retain_graph, pass_through_grad, **kwargs ) elif gradient_mode == "ig": node_grads, clean_cache = self._gradient_wrt_nodes_ig( inputs, corrupt_cache, metric, retain_graph, verbose, **kwargs ) else: raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported") node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache) return node_effect, {} elif mode == 'edge': if kwargs.get('node_grads', None) is None or kwargs.get('node_effect', None) is None: # run the _gradient_wrt_nodes to get node_grads, and use node_effect to prune out unimportant nodes if verbose: print("Calculating node gradients...") if gradient_mode == "standard": node_grads, clean_cache = self._gradient_wrt_nodes( inputs, metric, retain_graph, pass_through_grad, **kwargs ) elif gradient_mode == "ig": node_grads, clean_cache = self._gradient_wrt_nodes_ig( inputs, corrupt_cache, metric, retain_graph, verbose, **kwargs ) else: raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported") node_effect = self._attrib_effect_node(node_grads, corrupt_cache, clean_cache) else: node_grads: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_grads') # type: ignore node_effect: Dict[Tuple[Node, Index], SparseAct] = kwargs.get('node_effect') # type: ignore # Pruning if kwargs.get('prune', False): if verbose: print("Pruning nodes...") self._prune_nodes(node_effect, verbose, **kwargs) if kwargs.get('gradient_only', False): if verbose: print("Returning edge gradients only...") for name, sparse_act in node_grads.items(): node_grads[name] = sparse_act.to_sparse_like_self(t.ones_like(sparse_act.to_tensor())) del sparse_act edge_grads, _ = self._gradient_wrt_edges( inputs, corrupt_cache, node_grads, verbose, **kwargs ) return node_effect, edge_grads else: raise NotImplementedError(f"mode {mode} is not supported") def _attrib_effect_node( self, grads: Dict, corrupt_cache: ActivationCache | Dict[str, Tensor], clean_cache: Dict[str, SparseAct], ) -> Dict: attrib_effect = {} aggregate_dim = [0] if self.token_wise else [0, 1] for (node, index), grad in grads.items(): corrupt_sparse_act = cache_to_sparseact( corrupt_cache, sae_hook_name(node.name), error_term_name(node.name) if self.use_error_term else None, ) attrib_effect[(node, index)] = ( # act: (seq, d_sae), resc: (seq) || act: (d_sae), resc: () grad @ (corrupt_sparse_act - clean_cache[node.name]) ).sum(aggregate_dim) return attrib_effect def _gradient_wrt_nodes( self, inputs: Dict[str, Tensor], metric: Callable[[Tensor], Tensor], retain_graph: bool = False, pass_through_grad: bool = False, verbose: bool = False, **kwargs, ) -> Tuple[ Dict[Tuple[Node, Index], SparseAct], # node effects Dict[str, SparseAct] ]: ''' Forward pass of the graph with clean tokens, if the edge exists, replace the activation with corrupted activation Backward pass on the graph wrt the metric Return the gradients wrt nodes, edges, and activation cache ''' self._check_graph() self.model_setup() self.model.reset_hooks() for _, sae in self.dict_saes.items(): sae.reset_hooks() bwd_cache = {} pass_through_cache = {} def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None: if hook.name == sae_hook_name(sae_name): bwd_cache[sae_hook_name(sae_name)] = grad.detach() elif hook.name == output_hook_name(sae_name): if self.use_error_term: # we have: output = recon + stop_grad(error_term) # so, the TRUE error_grad (if not stop_grad) is output_grad # due to: output = recon + error_term if "vision" not in sae_name: bwd_cache[error_term_name(sae_name)] = grad.detach() else: bwd_cache[error_term_name(sae_name)] = t.zeros_like(grad).mean(dim=1, keepdim=True) if pass_through_grad: pass_through_cache[output_hook_name(sae_name)] = grad.detach() elif hook.name == input_hook_name(sae_name): if "vision" in sae_name: grad.zero_() else: if pass_through_grad: # we have to modify inplace instead of grad = ... and then return grad # because, returning a tensor in bwd pass hook is buggy somehow grad.copy_(pass_through_cache[output_hook_name(sae_name)]) fwd_cache = {} def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str) -> Tensor: if hook.name == sae_hook_name(sae_name): fwd_cache[sae_hook_name(sae_name)] = act.detach() elif error_term_name(sae_name) == hook.name and self.use_error_term: fwd_cache[error_term_name(sae_name)] = act.detach() return act with t.set_grad_enabled(True): with self._detach_error_term(True) and self._hook_vision_sae(): with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term): with self.model.hooks( fwd_hooks=[ (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name)) for sae_name in self.dict_saes.keys() ], bwd_hooks=[ (lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_bwd, sae_name=sae_name)) for sae_name in self.dict_saes.keys() ], ): metric(self.model(inputs)).backward(retain_graph=retain_graph) node_grads = {} for node, index in self.nodes.keys(): node_grads[(node, index)] = cache_to_sparseact( bwd_cache, sae_hook_name(node.name), error_term_name(node.name) if self.use_error_term else None, ) cache = {} for sae_name in self.dict_saes.keys(): cache[sae_name] = cache_to_sparseact( fwd_cache, sae_hook_name(sae_name), error_term_name(sae_name) if self.use_error_term else None, ) self.model.reset_hooks() for sae in self.dict_saes.values(): sae.reset_hooks() return node_grads, cache def _gradient_wrt_nodes_ig( self, inputs: Dict[str, Tensor], corrupt_cache: ActivationCache | Dict[str, Tensor], metric: Callable[[Tensor], Tensor], retain_graph: bool = False, verbose: bool = False, **kwargs, ) -> Tuple[ Dict[Tuple[Node, Index], SparseAct], # node effects Dict[str, SparseAct] ]: steps = kwargs.get('steps', 10) self._check_graph() self.model_setup() self.model.reset_hooks() for _, sae in self.dict_saes.items(): sae.reset_hooks() bwd_cache: Dict[str, Tensor] = {} def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str) -> None: if hook.name == sae_hook_name(sae_name): create_list(bwd_cache, sae_hook_name(sae_name)) bwd_cache[sae_hook_name(sae_name)] += grad.detach() elif hook.name == output_hook_name(sae_name): if self.use_error_term: # we have: output = recon + stop_grad(error_term) # so, the TRUE error_grad (if not stop_grad) is output_grad # due to: output = recon + error_term create_list(bwd_cache, error_term_name(sae_name)) if "vision" not in sae_name: bwd_cache[error_term_name(sae_name)] += grad.detach() else: bwd_cache[error_term_name(sae_name)] += t.zeros_like(grad).mean(dim=1, keepdim=True) fwd_cache = {} def hook_sae_fwd(act: Tensor, hook: HookPoint, sae_name: str, target_name: str, frac: float) -> Tensor: if hook.name == sae_hook_name(sae_name): # interpolate for integrated gradients if hook.name == sae_hook_name(target_name): act = interpolate( corrupt_cache[sae_hook_name(sae_name)], act, frac, ) fwd_cache[sae_hook_name(sae_name)] = act.detach() elif error_term_name(sae_name) == hook.name and self.use_error_term: # interpolate for integrated gradients if hook.name == error_term_name(target_name): act = interpolate( corrupt_cache[error_term_name(sae_name)], act, frac, ) fwd_cache[error_term_name(sae_name)] = act.detach() return act with t.set_grad_enabled(True): with self._detach_error_term(True) and self._hook_vision_sae(): with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term): for target_name in self.dict_saes.keys(): for step in range(steps): frac = step / steps with self.model.hooks( fwd_hooks=[ ( lambda name, sae_name=sae_name: sae_name in name, partial(hook_sae_fwd, sae_name=sae_name, target_name=target_name, frac=frac) ) for sae_name in self.dict_saes.keys() ], bwd_hooks=[ ( lambda name, target_name=target_name: target_name in name, partial(hook_sae_bwd, sae_name=target_name) ) ], ): metric(self.model(inputs)).backward(retain_graph=retain_graph) # average the gradients for key in bwd_cache.keys(): bwd_cache[key] /= steps node_grads = {} for node, index in self.nodes.keys(): node_grads[(node, index)] = cache_to_sparseact( bwd_cache, sae_hook_name(node.name), error_term_name(node.name) if self.use_error_term else None, ) cache = {} for sae_name in self.dict_saes.keys(): cache[sae_name] = cache_to_sparseact( fwd_cache, sae_hook_name(sae_name), error_term_name(sae_name) if self.use_error_term else None, ) self.model.reset_hooks() for sae in self.dict_saes.values(): sae.reset_hooks() return node_grads, cache def _gradient_wrt_edges( self, inputs: Dict[str, Tensor], corrupt_cache: ActivationCache | Dict[str, Tensor], node_grads: Dict[Tuple[Node, Index], SparseAct], verbose: bool = False, **kwargs, ) -> Tuple[ Dict[Tuple[Node, Index, Node, Index], Tensor], # edge effects Dict[str, SparseAct] ]: self._check_graph() self.model_setup() self.model.reset_hooks() for _, sae in self.dict_saes.items(): sae.reset_hooks() gradient_mode = kwargs.get('edge_gradient_mode', 'gradient') _, clean_cache = self.forward(inputs, corrupt_cache=None) edge_grads: Dict[Tuple[Node, Index, Node, Index], Tensor] = {} for layer, connection in tqdm(self.connection.items(), disable=not verbose): if verbose: print(f"Layer {layer}:") for hook_position_end, list_hook_positions_start in tqdm(connection.items(), disable=not verbose): assert hook_position_end in node_grads, f"Node gradient of {hook_position_end} is not provided." for hook_position_start in list_hook_positions_start: corrupt_sparse_act = cache_to_sparseact( corrupt_cache, sae_hook_name(hook_position_start[0].name), error_term_name(hook_position_start[0].name) if self.use_error_term else None, ) right_vec = corrupt_sparse_act - clean_cache[hook_position_start[0].name] if gradient_mode == "gradient": edge_grads[hook_position_start + hook_position_end] = self._edge_attribution( inputs, hook_position_end, hook_position_start, node_grads[hook_position_end], right_vec, layer, **kwargs, ) else: raise NotImplementedError(f"gradient_mode {gradient_mode} is not supported") return edge_grads, clean_cache def _edge_attribution( self, inputs: Dict[str, Tensor], hook_position_end: Tuple[Node, Index], hook_position_start: Tuple[Node, Index], leftvec: SparseAct, rightvec: SparseAct, layer: int, **kwargs, ) -> Tensor: def hook_sae_bwd(grad: Tensor, hook: HookPoint, sae_name: str, bwd_cache: Dict) -> None: if hook.name == sae_hook_name(hook_position_start[0].name): # only store the gradient at the start position bwd_cache[sae_hook_name(hook_position_start[0].name)] = grad.detach() elif hook.name == output_hook_name(hook_position_start[0].name): # we have: output = recon + stop_grad(error_term) # so, the TRUE error_grad (if not stop_grad) is output_grad # due to: output = recon + error_term if self.use_error_term: # only store the gradient at the start position if "vision" not in sae_name: bwd_cache[error_term_name(hook_position_start[0].name)] = grad.detach() else: bwd_cache[error_term_name(hook_position_start[0].name)] = t.zeros_like(grad).mean(dim=1, keepdim=True) # IMPORTANT NOTE for reproducibility: # we zero grad of intermediate components # but zero grad the output hook of SAE will ALSO ZERO GRAD the resid mid / post --> no downstream grads # so, we instead zero grad of the input hook of SAE elif hook.name == input_hook_name(sae_name): if hook.name != input_hook_name(hook_position_end[0].name): grad.zero_() def hook_sae_fwd(act: Tensor, hook: HookPoint, to_bwd_cache: Dict) -> Tensor: if hook.name == sae_hook_name(hook_position_end[0].name): # store activation for backward later to_bwd_cache[sae_hook_name(hook_position_end[0].name)] = act elif error_term_name(hook_position_end[0].name) == hook.name and self.use_error_term: # store activation for backward later to_bwd_cache[error_term_name(hook_position_end[0].name)] = act return act d_sae_end = self.dict_saes[hook_position_end[0].name].cfg.d_sae d_sae_start = self.dict_saes[hook_position_start[0].name].cfg.d_sae to_bwd_cache = {} bwd_cache = {} edge_effect = OrderedDict() with t.set_grad_enabled(True): with self._detach_error_term(False, hook_position_end[0].name) and self._hook_vision_sae(): with self.model.saes(saes=self._saes_to_list(), use_error_term=self.use_error_term): with self.model.hooks( fwd_hooks=[ (lambda name: True, partial( hook_sae_fwd, to_bwd_cache=to_bwd_cache, )) ], bwd_hooks=[ (lambda name, sae_name=sae_name: sae_name in name, partial( hook_sae_bwd, sae_name=sae_name, bwd_cache=bwd_cache, )) for sae_name in self.dict_saes.keys() ], ): self.model.forward(inputs) aggregate_dim = [0] if self.token_wise else [0, 1] to_bwd = ( # (seq, d_sae+1) || (d_sae+1) cache_to_sparseact( to_bwd_cache, sae_hook_name(hook_position_end[0].name), error_term_name(hook_position_end[0].name) if self.use_error_term else None, ) @ leftvec.detach() ).sum(aggregate_dim).to_tensor() del to_bwd_cache for active_idx, (end_node, end_index) in enumerate(self.active_nodes(*hook_position_end)): if isinstance(end_index, ErrorIndex): # the last index is error: shape (d_sae+1) so last index is d_sae to_bwd[end_index.idx + (d_sae_end,)].backward(retain_graph=True) index = t.tensor(list(end_index.idx + (d_sae_end,)), device=self.device) elif isinstance(end_index, FeatureIndex): to_bwd[end_index.idx].backward(retain_graph=True) index = t.tensor(list(end_index.idx), device=self.device) else: raise ValueError(f"end_index of type {type(end_index)} is not supported.") ''' edge_effect shape (seq, d_sae+1, seq, d_sae+1) or (d_sae+1, d_sae+1) in sparse_coo tensor the sparse_coo will have the shape: --> indices of shape (2, num_active) or (1, num_active) --> values of shape (num_active, seq, d_sae+1) or (num_active, d_sae+1) ''' edge_effect[active_idx] = ( index, ( # (seq, d_sae+1) || (d_sae+1) cache_to_sparseact( bwd_cache, sae_hook_name(hook_position_start[0].name), error_term_name(hook_position_start[0].name) if self.use_error_term else None, ) @ rightvec ).sum(aggregate_dim).to_tensor() ) del bwd_cache seq_start = rightvec.act.shape[1] # type: ignore seq_end = leftvec.act.shape[1] # type: ignore num_end = d_sae_end num_start = d_sae_start if self.use_error_term: num_end += 1 num_start += 1 if len(edge_effect.keys()) != 0: indices = t.stack([val[0] for val in edge_effect.values()], dim=0).T # shape (2, num_active) or (1, num_active) values = t.stack([val[1] for val in edge_effect.values()], dim=0) # shape (num_active, seq, d_sae+1) or (num_active, d_sae+1) # if no active nodes, return empty tensor else: indices = t.empty((2, 0) if self.token_wise else (1, 0), dtype=t.long).to(self.device) values = t.empty((0, seq_start, num_start) if self.token_wise else (0, num_start), dtype=t.float).to(self.device) if self.token_wise: return t.sparse_coo_tensor(indices, values, size=(seq_end, num_end, seq_start, num_start)).coalesce() else: return t.sparse_coo_tensor(indices, values, size=(num_end, num_start)).coalesce() def _prune_nodes( self, node_effect: Dict[Tuple[Node, Index], SparseAct], verbose: bool = False, **kwargs, ) -> None: if kwargs.get("node_threshold", None) is None: raise ValueError("Please provide the node_threshold") else: node_threshold = kwargs.get("node_threshold") assert isinstance(node_threshold, (float, int, Tensor)), "node_threshold must be a int, float, or Tensor" for node, index in tqdm(node_effect.keys(), disable=not verbose): if kwargs.get("reverse_prune_node", False): effect_feat_mask = node_effect[(node, index)].act.abs() < node_threshold else: effect_feat_mask = node_effect[(node, index)].act.abs() > node_threshold if verbose: num_pruned = (~effect_feat_mask).sum().item() print(f"{(node, index)}: pruned ({num_pruned / effect_feat_mask.numel() * 100:.5f}%) features") if self.use_error_term: if kwargs.get("reverse_prune_node", False): effect_error_mask = node_effect[(node, index)].resc.abs() < node_threshold # type: ignore else: effect_error_mask = node_effect[(node, index)].resc.abs() > node_threshold # type: ignore if verbose: num_pruned = (~effect_error_mask).sum().item() print(f"{(node, index)}: pruned ({num_pruned / effect_error_mask.numel() * 100:.5f}%) errors") # replace the nodes with mask to prune mask = SparseAct( act=effect_feat_mask, resc=effect_error_mask if self.use_error_term else None, # type: ignore ) self.nodes[(node, index)] = self.nodes[(node, index)] * mask def model_setup(self): pass def run_model( self, inputs: Dict[str, Tensor], use_error_term: bool | None = None, ) -> Tuple[Tensor, ActivationCache]: with self._hook_vision_sae(): out = self.model.run_with_cache_with_saes( inputs, saes=self._saes_to_list(), use_error_term=self.use_error_term if use_error_term is None else use_error_term, names_filter=lambda name: "sae" in name, ) return out # type: ignore def _saes_to_list(self) -> List[Any]: return [sae for _, sae in self.dict_saes.items()] @contextmanager def _detach_error_term(self, detach: bool, sae_name: str | None = None): orig_detach_error_term = {} orig_disable_error_grad = {} try: for name, sae in self.dict_saes.items(): if sae_name is None or sae_name == name: orig_detach_error_term[name] = sae.detach_error_term sae.detach_error_term = detach else: orig_detach_error_term[name] = sae.detach_error_term sae.detach_error_term = True # default detach error term orig_disable_error_grad[name] = sae.disable_error_grad sae.disable_error_grad = False # allow grad flows through error hook yield finally: for name, sae in self.dict_saes.items(): sae.detach_error_term = orig_detach_error_term[name] sae.disable_error_grad = orig_disable_error_grad[name] @contextmanager def _hook_vision_sae(self): pass_through_vision_sae_cache = [t.tensor(0)] # placeholder def hook_fn(act: Tensor, hook: HookPoint, sae_name: str): if input_hook_name(sae_name) == hook.name: pass_through_vision_sae_cache[0] = act # nodetach allow gradient to flow return act.mean(dim=1, keepdim=True) elif output_hook_name(sae_name) == hook.name: return pass_through_vision_sae_cache[0] + (act - act.detach()) / 2 # allow gradient to flow through sae elif error_term_name(sae_name) == hook.name: act = t.zeros_like(act).mean(dim=1, keepdim=True) return act try: for sae_name, vision_sae in self.dict_vision_saes.items(): vision_sae.add_hook( lambda name: True, partial(hook_fn, sae_name=sae_name), dir="fwd", ) yield finally: for vision_sae in self.dict_vision_saes.values(): vision_sae.reset_hooks()