{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "949d02a5", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/mnt/disk1/miniconda3/envs/baodq_hal/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", " from .autonotebook import tqdm as notebook_tqdm\n", "\u001b[32m2026-03-06 19:42:49.736\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msae.autoencoder.decoder\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m141\u001b[0m - \u001b[1mTriton found\u001b[0m\n", "\u001b[32m2026-03-06 19:42:49.737\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36msae.autoencoder.decoder\u001b[0m:\u001b[36m\u001b[0m:\u001b[36m143\u001b[0m - \u001b[1mTriton enabled, using Triton decoder\u001b[0m\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "🚨 `inputs` is part of LlavaModel.forward's signature, but not documented. Make sure to add it to the docstring of the function in /mnt/disk4/baodq/hallucination/model/llava/modeling_llava.py.\n", "🚨 `inputs` is part of LlavaForConditionalGeneration.forward's signature, but not documented. Make sure to add it to the docstring of the function in /mnt/disk4/baodq/hallucination/model/llava/modeling_llava.py.\n" ] } ], "source": [ "from model.blip.hooked_blip import HookedSAEBlipConditionalGeneration\n", "from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration\n", "import torch\n", "from sae.SAE_Tools import *\n", "from sae.SAE_Trainer import DataConfig\n", "from sae.SAE_Blip_Explaining_Utils import *\n", "\n", "from sae.Load_Data import load_lvlm_data" ] }, { "cell_type": "code", "execution_count": 2, "id": "6a045d6b", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Loading checkpoint shards: 100%|██████████| 3/3 [00:07<00:00, 2.57s/it]\n", "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n" ] } ], "source": [ "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n", "dtype = t.bfloat16\n", "\n", "# blip model\n", "# tok_name=\"Salesforce/blip-image-captioning-base\"\n", "# model = HookedSAEBlipConditionalGeneration.from_pretrained(tok_name)\n", "# processor = BlipProcessor.from_pretrained(tok_name)\n", "\n", "\n", "# llava model\n", "tok_name=\"llava-hf/llava-1.5-7b-hf\"\n", "model = HookedSAELlavaConditionalGeneration.from_pretrained(tok_name)\n", "processor = LlavaProcessor.from_pretrained(tok_name)\n", "\n", "\n", "\n", "model = model.to(device, dtype=dtype)\n", "\n", "\n", "# change hf_dataset and path to change dataset\n", "\n", "num_workers=4\n", "hf_dataset=\"yerevann/coco-karpathy\" # yerevann/coco-karpathy\n", "local_train_path=\"COCO-Dataset/train\"\n", "local_val_path=\"COCO-Dataset/val\"\n", "\n", "batch_size=1\n", "max_length=512\n", "filter_seq_length=None # 30 for blip\n", "\n", "data_config = DataConfig(\n", " batch_size=batch_size,\n", " hf_dataset=hf_dataset,\n", " local_train_path=local_train_path,\n", " local_val_path=local_val_path,\n", " num_workers=num_workers,\n", " max_length=max_length, # the processor of blip only allow max tokens (fixed)\n", " processor=tok_name,\n", ")" ] }, { "cell_type": "markdown", "id": "ddda91c3", "metadata": {}, "source": [ "TEXT SAE" ] }, { "cell_type": "code", "execution_count": null, "id": "3e0e09bf", "metadata": {}, "outputs": [], "source": [ "path_list = [\n", " 'cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.0.attention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt', # text\n", " 'cc3m_checkpoints/topk_32.0_32_text_decoder.bert.encoder.layer.11.attention.self.hook_resid_pre_0.001_256_0.03125_42.ckpt', # both\n", "]\n", "\n", "saes = [\n", " load_sae_model(\n", " file_path=sae_path,\n", " model_type=\"blip\",\n", " hook_type=\"text\",\n", " device=device,\n", " ).to(dtype) for sae_path in path_list\n", "]\n", "\n", "train_loader, val_loader = load_lvlm_data(data_config)\n", "\n", "cache_dict, data_toks = cache_sae_lvlm(\n", " saes,\n", " model,\n", " val_loader,\n", " device,\n", " filter_seq_length=filter_seq_length,\n", " stop_at_batch=5,\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "63318051", "metadata": {}, "outputs": [], "source": [ "for key, val in cache_dict.items():\n", " print(key, val[0].shape)\n", " \n", "sae = saes[0]\n", "values, indices = cache_dict[sae.cfg.hook_name]\n", "print(indices[values > 0].unique())" ] }, { "cell_type": "code", "execution_count": null, "id": "2722c70d", "metadata": {}, "outputs": [], "source": [ "FEAT_IDX = 9 # The feature you want to analyze\n", "N_INTERVALS = 5 # How many splits (e.g., High, Med-High, Med-Low, Low)\n", "K_EXAMPLES = 5 # Examples per interval\n", "\n", "# 2. Fetch the data\n", "data = fetch_feature_activation_intervals_blip(\n", " processor=processor,\n", " feat_idx=FEAT_IDX,\n", " values=values, \n", " indices=indices,\n", " toks=data_toks,\n", " n_intervals=N_INTERVALS,\n", " k_per_interval=K_EXAMPLES,\n", " buffer=10 # Context window size\n", ")\n", "\n", "# 3. Generate HTML\n", "html_code = generate_interactive_html(data)\n", "from IPython.display import HTML\n", "display(HTML(html_code))" ] }, { "cell_type": "markdown", "id": "2fe4cb8d", "metadata": {}, "source": [ "VISION SAE" ] }, { "cell_type": "code", "execution_count": 6, "id": "1506f9be", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Loading COCO nocap dataset...\n", "Using Llava processor\n", "Total crops per image: 1\n" ] } ], "source": [ "path_list = [\n", " # 'cc3m_checkpoints/batchtopk_8.0_8_vision_model.encoder.layers.11.hook_resid_post_0.001_128_0.03125_42.ckpt', # image\n", " # \"cc3m_checkpoints/topk_32.0_32_vision_model.encoder.layers.11.hook_resid_post_0.001_256_0.03125_42.ckpt\"\n", " \"cc3m_checkpoints/batchtopk_8.0_8_model.vision_tower.vision_model.encoder.layers.23.hook_resid_post_0.0001_256_0.03125_42.ckpt\",\n", " \"cc3m_checkpoints/batchtopk_8.0_8_model.vision_tower.vision_model.encoder.layers.22.hook_resid_post_0.0001_256_0.03125_42.ckpt\"\n", "]\n", "\n", "saes = [\n", " load_sae_model(\n", " file_path=sae_path,\n", " model_type=\"llava\",\n", " hook_type=\"vision\",\n", " device=device,\n", " ).to(dtype) for sae_path in path_list\n", "]\n", "\n", "nocap_train, nocap_val = load_lvlm_data_nocap(config=data_config)\n", "\n", "\n", "processed_ds = DebatchNoCapDataset(nocap_val, processor=data_config.processor)\n", "\n", "multi_crop_dataset = MultiScaleCropDataset(\n", " original_dataset=processed_ds,\n", " img_size=model.config.vision_config.image_size,\n", " crop_ratios=[1],\n", " stride_ratio=0.5,\n", " resize_to=model.config.vision_config.image_size,\n", ")\n", "\n", "dataloader = DataLoader(multi_crop_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "fdcfcf49", "metadata": {}, "outputs": [], "source": [ "from typing import Literal\n", "def new_cache_vision_sae_lvlm(\n", " saes: List[Any],\n", " model,\n", " dataloader: DataLoader,\n", " device: str,\n", " mode: str | Literal[\"acts_post\", \"acts_pre\"] = \"acts_post\",\n", " filter_seq_length: int | None = None, # filter too long seq\n", " return_toks: bool = True,\n", " stop_at_batch: int | None = None,\n", "): \n", " act_dict = {sae.cfg.hook_name: [] for sae in saes}\n", " def caching_hook(act: Tensor, hook: HookPoint):\n", " # print(act.shape)\n", " act_dict[\".\".join(hook.name.split(\".\")[:-1])].append((act.squeeze(1).detach().cpu(), data_ids[-1]))\n", " \n", " @contextmanager\n", " def _hook_vision_sae():\n", " pass_through_vision_sae_cache = [t.tensor(0)] # placeholder\n", " def hook_fn(act: Tensor, hook: HookPoint, sae_name: str):\n", " if sae_name + \".hook_sae_input\" == hook.name:\n", " pass_through_vision_sae_cache[0] = act\n", " return act.mean(dim=1, keepdim=True)\n", " \n", " elif sae_name + \".hook_sae_output\" == hook.name:\n", " return pass_through_vision_sae_cache[0]\n", " \n", " elif sae_name + \".hook_sae_error\" == hook.name:\n", " act = t.zeros_like(act).mean(dim=1, keepdim=True)\n", " return act\n", " \n", " try:\n", " for vision_sae in saes:\n", " vision_sae.add_hook(\n", " lambda name: True,\n", " partial(hook_fn, sae_name=vision_sae.cfg.hook_name),\n", " dir=\"fwd\",\n", " )\n", " yield\n", " finally:\n", " for vision_sae in saes:\n", " vision_sae.reset_hooks()\n", " \n", " data_toks = []\n", " data_ids = []\n", " with _hook_vision_sae():\n", " with model.saes(saes, use_error_term=True):\n", " with model.hooks(fwd_hooks=[(lambda name: mode in name, caching_hook)]):\n", " for batch_idx_iter, batch in enumerate(tqdm(dataloader, desc=\"Caching SAE acts\")):\n", " if stop_at_batch is not None:\n", " if batch_idx_iter > stop_at_batch:\n", " break\n", " inputs = {\n", " \"pixel_values\": batch[\"pixel_values\"].to(device),\n", " \"input_ids\": batch[\"input_ids\"].to(device),\n", " \"attention_mask\": batch[\"attention_mask\"].to(device),\n", " }\n", " if filter_seq_length is not None and batch[\"input_ids\"].shape[1] > filter_seq_length:\n", " continue\n", " if return_toks:\n", " data_toks.append(batch[\"input_ids\"].cpu())\n", " data_ids.append(batch[\"imgid\"].item())\n", " model(inputs)\n", " \n", " if return_toks:\n", " data_toks = pad_and_concat(data_toks, dim=0, padding_value=0) # [PAD] token \n", " \n", " for key, act_list in act_dict.items():\n", " acts_list = [x[0] for x in act_list]\n", " img_ids = [x[1] for x in act_list]\n", " acts = pad_and_concat(acts_list, dim=0, padding_value=0) # (b, d_sae), pad 0 activation value in batch\n", " # print(acts.shape)\n", " max_k = int((acts > 1e-3).sum(dim=-1).max().item())\n", " values, indices = acts.topk(k=max_k, dim=-1)\n", "\n", " img_ids = t.tensor(img_ids)\n", " act_dict[key] = (values, indices, img_ids)\n", " \n", " return act_dict, data_toks" ] }, { "cell_type": "code", "execution_count": 5, "id": "43ad73e6", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "Caching SAE acts: 0%| | 0/5000 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "import matplotlib.pyplot as plt\n", "img_idx = 20\n", "\n", "x = multi_crop_dataset[img_idx][\"pixel_values\"].detach().cpu() # (3,H,W)\n", "\n", "mean = torch.tensor([0.485, 0.456, 0.406]).view(3,1,1)\n", "std = torch.tensor([0.229, 0.224, 0.225]).view(3,1,1)\n", "\n", "img = (x * std + mean).clamp(0, 1) # back to [0,1]\n", "img = img.permute(1, 2, 0).numpy()\n", "\n", "plt.imshow(img)\n", "plt.axis(\"off\")\n", "plt.show()\n" ] }, { "cell_type": "code", "execution_count": 4, "id": "a2481716", "metadata": {}, "outputs": [], "source": [ "def new_fetch_topk_activating_crops_lvlm(\n", " feat_idx: int,\n", " values: Tensor,\n", " indices: Tensor,\n", " img_ids: Tensor,\n", " dataset: Dataset,\n", " top_k: int,\n", " output_dir: str,\n", " verbose: bool = True,\n", "): \n", " mask = (values > 1e-3) * (indices == feat_idx)\n", " if verbose:\n", " print(\"Feature ID:\", feat_idx, \"Density:\", mask.sum().item() / mask.numel(), \"Max act:\", (values * mask).max().item())\n", " if mask.sum() == 0:\n", " if verbose:\n", " print(\"No activation!\")\n", " return\n", " filtered_values = values * mask.to(values.dtype) # (b, activated_dim)\n", " \n", " sumed_filtered_values = filtered_values.sum(dim=1) # (b)\n", " top_vals, top_indices = sumed_filtered_values.topk(k=top_k) # (topk)\n", " top_acts = filtered_values[top_indices, :] # (topk, activated_dim)\n", " top_img_ids = img_ids[top_indices]\n", " print(\"img_ids:\", top_img_ids)\n", " \n", " visualize_crops(\n", " top_vals = top_vals,\n", " top_indices = top_indices,\n", " top_acts = top_acts,\n", " feat_idx = feat_idx,\n", " multi_crop_dataset = dataset,\n", " output_dir = output_dir,\n", " )" ] }, { "cell_type": "code", "execution_count": 10, "id": "7c16e9a6", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Feature ID: 1903 Density: 0.1245 Max act: 10.125\n", "img_ids: tensor([133928, 472960, 124018, 566049, 196681, 204324, 95692, 539226, 155051,\n", " 199977, 535997, 563295, 573008, 110359, 304924])\n", "Feature act: 7.90625\n", "Feature ID: 3696 Density: 0.114925 Max act: 7.78125\n", "img_ids: tensor([189267, 338325, 263073, 117014, 537864, 259421, 574343, 545407, 509867,\n", " 408535, 296848, 146659, 99242, 40361, 566518])\n", "Feature act: 6.625\n", "Feature ID: 4647 Density: 0.1036 Max act: 7.03125\n", "img_ids: tensor([545407, 189267, 263073, 259421, 509867, 475271, 117014, 251395, 472078,\n", " 338325, 556157, 208135, 17484, 99242, 237735])\n", "Feature act: 6.40625\n", "Feature ID: 3674 Density: 0.0892 Max act: 6.21875\n", "img_ids: tensor([545407, 208135, 259421, 263073, 64157, 392650, 524186, 5477, 81505,\n", " 326308, 297220, 176828, 462931, 117014, 84870])\n", "Feature act: 5.375\n", "Feature ID: 1366 Density: 0.02925 Max act: 6.71875\n", "img_ids: tensor([ 64157, 208135, 323925, 297220, 545407, 5477, 71815, 117014, 326308,\n", " 522198, 524186, 279806, 263073, 462931, 304545])\n", "Feature act: 4.5625\n" ] } ], "source": [ "# Visualize top_k_features for image_idx\n", "top_k_features = 5\n", "for each in range(top_k_features):\n", " new_fetch_topk_activating_crops_lvlm(\n", " feat_idx=indices[img_idx][each].item(),\n", " values=values,\n", " indices=indices,\n", " dataset=multi_crop_dataset,\n", " top_k=15, # top activating crops\n", " output_dir=\"./vision_sae_vis\",\n", " img_ids=img_ids,\n", " )\n", " print(\"Feature act:\", values[img_idx][each].item())" ] }, { "cell_type": "code", "execution_count": 48, "id": "d31b04f3", "metadata": {}, "outputs": [], "source": [ "sae_feature = 3674\n", "raw_image = Image.open(f\"/mnt/disk4/baodq/hallucination/vision_sae_vis/feature_{sae_feature}_top_crops/top_grid.jpg\")\n", "\n", "# 3. Format prompt correctly\n", "prompt = \"USER: \\nASSISTANT:\"\n", "\n", "# 4. Process and Generate\n", "inputs = processor(text=prompt, images=raw_image, return_tensors=\"pt\")\n", "inputs = {k: v.to(device) for k, v in inputs.items()}\n", "\n", "\n", "\n", "_, cache = model.run_with_cache_with_saes(\n", " inputs,\n", " saes=saes,\n", " names_filter=lambda name: (\"vision_model\" in name) and (\"hook_mlp_out\" in name)\n", ")\n" ] }, { "cell_type": "code", "execution_count": 50, "id": "11ee35a0", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "dict_keys([])\n" ] } ], "source": [ "print(cache.keys())\n", "# print(cache[\"model.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_output\"].shape)" ] }, { "cell_type": "code", "execution_count": 8, "id": "a27b43d8", "metadata": {}, "outputs": [ { "ename": "NameError", "evalue": "name 'cache' is not defined", "output_type": "error", "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", "Cell \u001b[0;32mIn[8], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m act \u001b[38;5;241m=\u001b[39m \u001b[43mcache\u001b[49m[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodel.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_acts_post\u001b[39m\u001b[38;5;124m'\u001b[39m]\n\u001b[1;32m 2\u001b[0m sae_feature \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m3674\u001b[39m\n\u001b[1;32m 3\u001b[0m mean \u001b[38;5;241m=\u001b[39m cache[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mmodel.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_acts_post\u001b[39m\u001b[38;5;124m'\u001b[39m][:, :, sae_feature]\u001b[38;5;241m.\u001b[39mmean(dim\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m1\u001b[39m, keepdim\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n", "\u001b[0;31mNameError\u001b[0m: name 'cache' is not defined" ] } ], "source": [ "act = cache['model.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_acts_post']\n", "sae_feature = 3674\n", "mean = cache['model.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_acts_post'][:, :, sae_feature].mean(dim=1, keepdim=False)\n", "feature_act = act[:, 1:, sae_feature]\n", "\n", "mask = ((feature_act < mean).float()).reshape(24, 24).unsqueeze(0).unsqueeze(0)\n", "\n", "\n", "upsampled_mask = F.interpolate(\n", " mask,\n", " size=raw_image.size[::-1], # (H, W)\n", " mode=\"bilinear\",\n", " align_corners=False\n", ")[0,0]\n", "\n", "import torchvision.transforms as T\n", "\n", "img_tensor = T.ToTensor()(raw_image).to(device)\n", "masked_img = img_tensor * upsampled_mask\n", "plt.imshow(masked_img.permute(1,2,0).cpu())\n", "plt.axis(\"off\")" ] }, { "cell_type": "code", "execution_count": null, "id": "31a025da", "metadata": {}, "outputs": [], "source": [ "# complete code for visualization features\n", "\n", "from model.blip.hooked_blip import HookedSAEBlipConditionalGeneration\n", "import torch\n", "from sae.SAE_Tools import *\n", "from sae.SAE_Trainer import DataConfig\n", "from sae.SAE_Blip_Explaining_Utils import *\n", "from sae.Load_Data import load_lvlm_data\n", "\n", "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n", "\n", "model = HookedSAEBlipConditionalGeneration.from_pretrained(\"Salesforce/blip-image-captioning-base\")\n", "model.to(device)\n", "\n", "\n", "# change hf_dataset and path to change dataset\n", "\n", "num_workers=0\n", "hf_dataset=\"yerevann/coco-karpathy\" # yerevann/coco-karpathy\n", "local_train_path=\"COCO-Dataset/train\"\n", "local_val_path=\"COCO-Dataset/val\"\n", "tok_name=\"Salesforce/blip-image-captioning-base\"\n", "batch_size=32\n", "max_length=512\n", "filter_seq_length=30\n", "\n", "processor = BlipProcessor.from_pretrained(\"Salesforce/blip-image-captioning-base\")\n", "\n", "data_config = DataConfig(\n", " batch_size=batch_size,\n", " hf_dataset=hf_dataset,\n", " local_train_path=local_train_path,\n", " local_val_path=local_val_path,\n", " num_workers=num_workers,\n", " max_length=max_length, # the processor of blip only allow max tokens (fixed)\n", " processor=tok_name,\n", ")\n", "\n", "path_list = [\n", " 'cc3m_checkpoints/topk_32.0_32_vision_model.encoder.layers.0.hook_resid_post_0.001_256_0.03125_42.ckpt', # image\n", "]\n", "\n", "saes = [\n", " load_sae_model(\n", " file_path=sae_path,\n", " model=model, # type: ignore\n", " device=device,\n", " ) for sae_path in path_list\n", "]\n", "\n", "nocap_train, nocap_val = load_lvlm_data_nocap(config=data_config)\n", "\n", "processed_ds = DebatchNoCapDataset(nocap_val, processor=data_config.processor)\n", "\n", "multi_crop_dataset = MultiScaleCropDataset(\n", " original_dataset=processed_ds,\n", " img_size=model.config.vision_config.image_size,\n", " crop_ratios=[1],\n", " stride_ratio=0.5,\n", " resize_to=model.config.vision_config.image_size,\n", ")\n", "\n", "dataloader = DataLoader(multi_crop_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)\n", "\n", "cache_dict, _ = cache_vision_sae_lvlm(\n", " saes,\n", " model,\n", " dataloader,\n", " device,\n", " filter_seq_length=filter_seq_length,\n", " stop_at_batch=0,\n", " return_toks=False,\n", ")\n", "\n", "for key, val in cache_dict.items():\n", " print(key, val[0].shape)\n", " \n", "sae = saes[0]\n", "values, indices = cache_dict[sae.cfg.hook_name]\n", "\n", "\n", "fetch_topk_activating_crops_lvlm(\n", " feat_idx=114,\n", " values=values,\n", " indices=indices,\n", " dataset=multi_crop_dataset,\n", " top_k=15,\n", " output_dir=\"./vision_sae_vis\",\n", ")\n", "\n" ] }, { "cell_type": "markdown", "id": "9387a847", "metadata": {}, "source": [ "# Steering Vision Encoder (layer 24) -> not use in model " ] }, { "cell_type": "code", "execution_count": 54, "id": "70b01cd4", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "USER: \n", "Describe this image in detail.\n", "ASSISTANT: The image features a red and white airplane with a red wing, flying through a cloudy sky. The airplane is positioned in the middle of the scene, with its wings spread out. The sky is filled with clouds, creating a dramatic backdrop for the airplane. The scene captures the essence of air travel and the beauty of the sky.\n" ] } ], "source": [ "image = Image.open(\"/mnt/disk4/baodq/hallucination/output.png\")\n", "prompt = \"USER: \\nDescribe this image in detail.\\nASSISTANT:\"\n", "\n", "inputs = processor(text=prompt, images=image, return_tensors=\"pt\")\n", "inputs = {k: v.to(device) for k, v in inputs.items()}\n", "\n", "feature_id = 88\n", "\n", "def ablate_feature_hook(acts, hook):\n", " # acts shape: (batch, seq_len, d_sae)\n", " acts[..., feature_id] = 0\n", " return acts\n", "\n", "with model.saes(saes):\n", " with model.hooks(\n", " fwd_hooks=[\n", " (\n", " \"model.vision_tower.vision_model.encoder.layers.23.hook_resid_post.hook_sae_acts_post\",\n", " ablate_feature_hook\n", " )\n", " ]\n", " ):\n", " output = model.generate(\n", " **inputs,\n", " max_new_tokens=300,\n", " do_sample=False,\n", " )\n", " \n", "generated_text = processor.decode(output[0], skip_special_tokens=True)\n", "print(generated_text)" ] }, { "cell_type": "code", "execution_count": 53, "id": "6b440894", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "USER: \n", "Describe this image in detail.\n", "ASSISTANT: The image features a woman standing in front of a large body of water, possibly the ocean. She is wearing a white shirt and appears to be enjoying her time by the water. The woman is the main focus of the scene, with the vast body of water in the background.\n" ] } ], "source": [ "image = Image.open(\"/mnt/disk4/baodq/hallucination/output.png\")\n", "prompt = \"USER: \\nDescribe this image in detail.\\nASSISTANT:\"\n", "\n", "inputs = processor(text=prompt, images=image, return_tensors=\"pt\")\n", "inputs = {k: v.to(device) for k, v in inputs.items()}\n", "\n", "feature_id = 3674\n", "\n", "def ablate_feature_hook(acts, hook):\n", " # acts shape: (batch, seq_len, d_sae)\n", " acts[..., feature_id] = 300\n", " return acts\n", "\n", "with model.saes(saes):\n", " with model.hooks(\n", " fwd_hooks=[\n", " (\n", " \"model.vision_tower.vision_model.encoder.layers.22.hook_resid_post.hook_sae_acts_post\",\n", " ablate_feature_hook\n", " )\n", " ]\n", " ):\n", " output = model.generate(\n", " **inputs,\n", " max_new_tokens=300,\n", " do_sample=False,\n", " )\n", " \n", "generated_text = processor.decode(output[0], skip_special_tokens=True)\n", "print(generated_text)" ] }, { "cell_type": "markdown", "id": "a35c7478", "metadata": {}, "source": [ "# Steering Code" ] }, { "cell_type": "code", "execution_count": null, "id": "7c812669", "metadata": {}, "outputs": [], "source": [ "from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration\n", "from transformers import LlavaProcessor\n", "from sae_lens import SAE\n", "import torch\n", "\n", "model = HookedSAELlavaConditionalGeneration.from_pretrained(\"gpt2-small\", device=\"cuda\")\n", "sae, _, _ = SAE.from_pretrained(\n", " release=\"gpt2-small-res-jb\",\n", " sae_id=\"blocks.8.hook_resid_pre\",\n", " device=\"cuda\"\n", ")\n", "\n", "# 1. Feature attribution for a specific prediction\n", "prompt = \"USER: \\nDescribe this image in detail.\\nASSISTANT:\"\n", "processor = LlavaProcessor.from_pretrained(\"llava-hf/llava-1.5-7b-hf\")\n", "image = np.zeros((224, 224, 3)) \n", "inputs = processor(images=image, text=prompt, return_tensors=\"pt\")\n", "\n", "_, cache = model.run_with_cache(inputs, \n", " names_filter=lambda name: \"resid\" in name and \"language\" in name))\n", "\n", "hook_name = \"model.language_model.layers.25.hook_resid_post\"\n", "activations = cache[hook_name]\n", "features = sae.encode(activations)\n", "\n", "# Target token\n", "target_token = processor.tokenizer.encode(\" Paris\", add_special_tokens=False)\n", "\n", "# Compute feature contributions to target logit\n", "\n", "# # contribution = feature_activation * sae_decoder_weight * unembedding\n", "# W_dec = sae.W_dec # [d_sae, d_model]\n", "# W_U = model.W_U # [d_model, d_vocab]\n", "# # Feature direction projected to vocabulary\n", "# feature_to_logit = W_dec @ W_U # [d_sae, d_vocab]\n", "# # Contribution of each feature to \"Paris\" at final position\n", "# feature_acts = features[0, -1] # [d_sae]\n", "# contributions = feature_acts * feature_to_logit[:, target_token]\n", "\n", "\n", "\n", "# norm=model.language_model.norm\n", "# lm_head=model.lm_head # [d_model, d_vocab]\n", "# W_dec = sae.W_dec # [d_sae, d_model]\n", "\n", "# with torch.no_grad():\n", " # baseline_logits = lm_head(norm(residual_stream)) # [b, seq, d_vocab]\n", " # modified = residual_stream + feature_activation * W_dec[feature_idx] # [b, seq, d_model]\n", " # modified_logits = lm_head(norm(modified)) # [b, seq, d_vocab]\n", " # contribution = modified_logits[:, -1, target_token] - baseline_logits[:, -1, target_token] # [b] - contribution score for feature to target token\n", "\n", "\n", "\n", "# Top contributing features\n", "top_features = contributions.topk(10)\n", "print(\"Top features contributing to 'Paris':\")\n", "for idx, val in zip(top_features.indices, top_features.values):\n", " print(f\" Feature {idx.item()}: {val.item():.3f}\")\n", "\n", "# 2. Feature steering\n", "def steer_with_feature(feature_idx, strength=5.0):\n", " \"\"\"Add a feature direction to the residual stream.\"\"\"\n", " feature_direction = sae.W_dec[feature_idx] # [d_model]\n", "\n", " def hook(activation, hook_obj):\n", " activation[:, -1, :] += strength * feature_direction\n", " return activation\n", "\n", " output = model.generate(\n", " **inputs,\n", " max_new_tokens=10,\n", " fwd_hooks=[(hook_name, hook)]\n", " )\n", " return processor.tokenizer.decode(output[0], add_special_tokens=False)\n", "\n", "# Try steering with top feature\n", "top_feature_idx = top_features.indices[0].item()\n", "print(f\"\\nSteering with feature {top_feature_idx}:\")\n", "print(steer_with_feature(top_feature_idx, strength=10.0))" ] } ], "metadata": { "kernelspec": { "display_name": "baodq_hal", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.19" } }, "nbformat": 4, "nbformat_minor": 5 }