#!/bin/bash # Top-k d_sae features -> toilet SAE features (per-feature gradient attribution). # # For a chosen readout layer L_inspect with toilet feature set T, find WHICH # features anywhere in the d_sae at the layers BEFORE L_inspect most affect T: # M_{L_inspect} = mean_tokens( sum_{f in T} z[L_inspect][.., f] ) # a[L_patch][f] = sum_tokens z_clean[L_patch][t,f] * dM_{L_inspect}/dz[L_patch][t,f] # normalised by the inspect-layer metric M_{L_inspect} (fraction of the toilet # readout the feature explains) so inspect layers at different depths compare # (raw |a| grows ~exponentially with depth: ~500 at L_inspect=1 vs ~4e7 at 31): # score[L_inspect][L_patch][f] = a[L_inspect][L_patch][f] / M_{L_inspect} # # Output: a JSON (OUT) holding, per L_inspect, the top-K features per upstream # layer AND a GLOBAL cross-layer ranking; plus an INTERACTIVE Plotly HTML (GRAPH) # with a dropdown over L_inspect showing a horizontal bar chart of the global # top-N features ranked ACROSS all upstream layers (one comparable ranking — a # heatmap is NOT used, since its per-row top-k can't compare ranks between layers). # Each bar is labelled "L·#", coloured by signed normalised attr # (red promotes T, blue suppresses T), hover reveals the raw value. # # Tunables (override via env): # DEVICE target GPU (default cuda:2) # DTYPE bfloat16 | float16 | float32 (default bfloat16) # SAE_CKPT BatchTopK SAE checkpoint (threshold buffer required) # TOILET_FEATS JSON {layer_: {features: [...]}} - readout set T # K features kept per upstream layer in the JSON (default 10) # TOP_N features in the global cross-layer ranking / bars (default 30) # VERIFY 1 = re-rank the shortlist by FAITHFUL ablation ΔM (comparable # across layers); 0 = first-order score only (default 0). First- # order inflates early layers by orders of magnitude, so it is # NOT comparable across L_patch - use VERIFY=1 to compare layers. # Ablates ONLY the top_n shortlist, never the whole d_sae. # VERIFY_BATCH features ablated per forward when VERIFY=1 (default 8) # RANK_BY abs | pos | neg (default abs) # abs = strongest either way, pos = promoters, neg = suppressors # INSPECT_LAYERS space-separated readout layers (default: all with T; e.g. "24 31") # OUT_TAG extra suffix on output filenames (e.g. top10) # SAMPLES samples.json to filter images (base mentioned object, LoRA # suppressed it). Set SAMPLES="" to use the IMAGES list below. # SAMPLES_PROMPT prompt_results key driving the filter (default "Describe this image.") # IMAGE_ROOT directory holding .jpg for SAMPLES filtering. # PROMPT text fed to the model for the forward/backward pass. # # Usage: # bash mechanistic_interp/scripts/sae_attribution_patching.sh # INSPECT_LAYERS="24 31" K=15 bash mechanistic_interp/scripts/sae_attribution_patching.sh # RANK_BY=pos bash mechanistic_interp/scripts/sae_attribution_patching.sh set -euo pipefail REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" MODEL_NAME="llava-hf/llava-1.5-7b-hf" SAE_CKPT="${SAE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/last.ckpt}" TOILET_FEATS="${TOILET_FEATS:-${REPO_ROOT}/mechanistic_interp/probes/selected_features_f1_post_toilet_top10.json}" DEVICE="${DEVICE:-cuda:5}" DTYPE="${DTYPE:-bfloat16}" K="${K:-20}" TOP_N="${TOP_N:-30}" VERIFY="${VERIFY:-1}" VERIFY_BATCH="${VERIFY_BATCH:-16}" RANK_BY="${RANK_BY:-pos}" INSPECT_LAYERS="${INSPECT_LAYERS:-}" OUT_TAG="${OUT_TAG:-}" TAG_SUFFIX="${OUT_TAG:+_${OUT_TAG}}" OUT="${REPO_ROOT}/mechanistic_interp/graph/d_sae_correlation_patching${TAG_SUFFIX}.json" GRAPH="${REPO_ROOT}/mechanistic_interp/graph/d_sae_correlation_patching${TAG_SUFFIX}.html" SAMPLES="${SAMPLES:-${REPO_ROOT}/mechanistic_interp/toilet_bathroom/samples.json}" SAMPLES_PROMPT="${SAMPLES_PROMPT:-Describe this image.}" IMAGE_ROOT="${IMAGE_ROOT:-/data/caotue/CC3M-Dataset/cc3m_images/train}" PROMPT="${PROMPT:-USER: Describe this image. ASSISTANT: In this image, there is a bathroom and a}" # Fallback image when SAMPLES="" and no --images is passed. IMAGES=( /data/caotue/CC3M-Dataset/cc3m_images/train/001607409.jpg ) # ------------------------------------------------------------------ # mkdir -p "$(dirname "${OUT}")" echo "========================================================" echo " top-k d_sae features -> toilet features (per-feature attribution)" echo "========================================================" echo " Toilet feats : ${TOILET_FEATS}" echo " SAE ckpt : ${SAE_CKPT}" echo " Device / dtype : ${DEVICE} / ${DTYPE}" echo " k / top_n : ${K} / ${TOP_N}" echo " rank_by : ${RANK_BY}" echo " verify : ${VERIFY} (faithful ΔM ablation; batch ${VERIFY_BATCH})" echo " Inspect layers : ${INSPECT_LAYERS:-}" echo " Prompt : ${PROMPT}" echo " Out JSON : ${OUT}" echo " Out graph : ${GRAPH}" echo "========================================================" cd "${REPO_ROOT}" export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}" ARGS=( --toilet_feats "${TOILET_FEATS}" --sae_ckpt "${SAE_CKPT}" --model_name "${MODEL_NAME}" --device "${DEVICE}" --dtype "${DTYPE}" --prompt "${PROMPT}" --k "${K}" --top_n "${TOP_N}" --verify_batch "${VERIFY_BATCH}" --rank_by "${RANK_BY}" --out "${OUT}" --graph "${GRAPH}" ) [ -n "${INSPECT_LAYERS}" ] && ARGS+=(--inspect_layers ${INSPECT_LAYERS}) [ "${VERIFY}" = "1" ] && ARGS+=(--verify) # Image source precedence: explicit --images on the CLI > SAMPLES filter > IMAGES list. if [[ "$*" == *"--images"* ]]; then python -m mechanistic_interp.d_sae_correlation_patching "${ARGS[@]}" "$@" elif [ -n "${SAMPLES}" ]; then python -m mechanistic_interp.d_sae_correlation_patching \ --samples "${SAMPLES}" \ --samples_prompt "${SAMPLES_PROMPT}" \ --image_root "${IMAGE_ROOT}" \ "${ARGS[@]}" \ "$@" else python -m mechanistic_interp.d_sae_correlation_patching \ --images "${IMAGES[@]}" \ "${ARGS[@]}" \ "$@" fi