File size: 6,788 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 | #!/bin/bash
# =============================================================================
# integrated_gradient.sh — bathroom→toilet causal influence map via INTEGRATED
# GRADIENT of the bathroom probe (replaces gradient_ascent.sh's local-grad steer).
#
# Steered direction at layer l = integrated gradient of score_bath_l along the
# straight path base→image, where base = mean caption-token residual of the
# NEGATIVE-label validation images (NEG_JSONL[BASELINE_SPLIT], capped BASELINE_NUM,
# forwarded with the same FORCED_TEXT). Direction normalized to unit; step is
# ALPHA*‖h_l‖*ĝ (ALPHA a fraction of the residual norm) — plus a 'full' panel
# h'=h+(x−b). Δ toilet score read at every layer l'>=l. Averaged → heatmap.
#
# Tunables (override inline, e.g. NUM_IMAGES=128 IG_STEPS=64 bash <script>):
# DEVICE_ID GPU id (default 0)
# DTYPE bfloat16 | float16 | float32 (default bfloat16)
# BATH_PROBE steered-direction SequenceLayerProbes checkpoint (e.g. scene)
# TOILET_PROBE measured-readout SequenceLayerProbes checkpoint (e.g. object)
# IMAGE_FOLDER folder of images
# SAMPLES_JSON samples.json with base_mentions_object flags (empty to disable)
# BASE_PROMPT prompt whose base_mentions_object selects non-hallucinating imgs
# ── Alternative image selection from an HF dataset (no samples.json needed) ──
# HF_DATASET if set, pick scene-only images (SCENE_COL=1 & OBJECT_COL=0) from it
# SPLIT HF split for selection (default validation)
# SCENE_COL steer concept column (present) e.g. dining_room
# OBJECT_COL readout concept column (absent) e.g. plate
# STEER_NAME display name of steered concept (plots)
# READOUT_NAME display name of readout concept (plots)
# ── Integrated-gradient baseline (mean NEGATIVE-label validation residual) ──
# NEG_JSONL {"train":[ids],"validation":[ids]} negative stems
# BASELINE_SPLIT which NEG_JSONL split for the baseline (default validation)
# BASELINE_NUM cap on negatives for baseline mean (default 1000; 0=all)
# IG_STEPS Riemann steps for the IG path base→image (default 32)
# QUESTION prompt question
# NUM_IMAGES cap on images; 0 = ALL non-hallucinating bathroom-only (default 0)
# ALPHAS steps as FRACTION of ‖h_l‖ (h'=h+ALPHA*‖h_l‖*ĝ) (default "0.05 0.1 0.2 0.4 0.8")
# HOOK_TYPE pre | mid | post (default post)
# MAX_NEW_TOKENS caption generation length (default 64)
# MAX_SEQ_TOKENS caption tokens kept for the probe (default 64)
# OUT / OUT_JSON output heatmap PNG / matrix JSON
# =============================================================================
export HF_HOME="/data/caotue/hf_cache"
export HF_DATASETS_CACHE="/data/caotue/hf_cache/datasets"
export TORCH_HOME="/data/caotue/torch_cache"
export TMPDIR="/data/caotue/tmp"
DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/data/caotue/CC3M-Dataset/cc3m_images}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
BASE_PROMPT="${BASE_PROMPT:-Describe this image.}"
# HF-dataset image selection (alternative to SAMPLES_JSON). Empty = use SAMPLES_JSON.
HF_DATASET="${HF_DATASET:-}"
SPLIT="${SPLIT:-validation}"
SCENE_COL="${SCENE_COL:-}"
OBJECT_COL="${OBJECT_COL:-}"
STEER_NAME="${STEER_NAME:-bathroom}"
READOUT_NAME="${READOUT_NAME:-toilet}"
# Integrated-gradient baseline (mean NEGATIVE-label validation residual).
NEG_JSONL="${NEG_JSONL:-mechanistic_interp/neg_cc3m_5k.json}"
BASELINE_SPLIT="${BASELINE_SPLIT:-validation}"
BASELINE_NUM="${BASELINE_NUM:-1000}"
IG_STEPS="${IG_STEPS:-32}"
# MEAN_STEER=1 → IG input = mean residual over ALL selected bathroom-only images
# (image-independent steering direction). Empty/0 = per-image (default).
MEAN_STEER="${MEAN_STEER:-0}"
# BASE_MENTIONS: false = base does NOT hallucinate toilet (+alpha induction);
# true = base DOES hallucinate toilet (-alpha suppression/necessity); any = no filter.
BASE_MENTIONS="${BASE_MENTIONS:-false}"
QUESTION="${QUESTION:-Describe this image.}"
# FORCED_TEXT: if set, skip generation and force the ASSISTANT answer to this exact
# string (controlled, constant context across images AND baseline negatives).
FORCED_TEXT="${FORCED_TEXT:-This image features a bathroom with a}"
# NUM_IMAGES=0 → ALL non-hallucinating bathroom-only images (base says no toilet).
NUM_IMAGES="${NUM_IMAGES:-0}"
ALPHAS="${ALPHAS:-0.05 0.1 0.2 0.4 0.8}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/integrated_gradient_bath2toilet.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/integrated_gradient_bath2toilet.json}"
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
ARGS=(
--device_id "${DEVICE_ID}"
--dtype "${DTYPE}"
--bath_probe "${BATH_PROBE}"
--toilet_probe "${TOILET_PROBE}"
--image_folder "${IMAGE_FOLDER}"
--base_prompt "${BASE_PROMPT}"
--base_mentions "${BASE_MENTIONS}"
--question "${QUESTION}"
--num_images "${NUM_IMAGES}"
--neg_jsonl "${NEG_JSONL}"
--baseline_split "${BASELINE_SPLIT}"
--baseline_num "${BASELINE_NUM}"
--ig_steps "${IG_STEPS}"
--alphas ${ALPHAS}
--hook_type "${HOOK_TYPE}"
--max_new_tokens "${MAX_NEW_TOKENS}"
--max_seq_tokens "${MAX_SEQ_TOKENS}"
--steer_name "${STEER_NAME}"
--readout_name "${READOUT_NAME}"
--out "${OUT}"
--out_json "${OUT_JSON}"
)
# Image selection: HF dataset (scene-only via columns) takes priority over samples.json.
if [ -n "${HF_DATASET}" ]; then
ARGS+=(--hf_dataset "${HF_DATASET}" --split "${SPLIT}"
--scene_col "${SCENE_COL}" --object_col "${OBJECT_COL}")
elif [ -n "${SAMPLES_JSON}" ]; then
ARGS+=(--samples_json "${SAMPLES_JSON}")
fi
# Only pass --forced_text when non-empty (empty = use generated captions).
if [ -n "${FORCED_TEXT}" ]; then
ARGS+=(--forced_text "${FORCED_TEXT}")
fi
# Mean-steer toggle (image-independent IG direction over all bathroom-only images).
if [ "${MEAN_STEER}" = "1" ] || [ "${MEAN_STEER}" = "true" ]; then
ARGS+=(--mean_steer)
fi
python -m mechanistic_interp.integrated_gradient "${ARGS[@]}" "$@"
|