| # ============================================================================= | |
| # gradient_ascent.sh — bathroom→toilet causal influence map via probe gradient. | |
| # | |
| # For each BATHROOM-ONLY image (bathroom=1 & toilet=0): nudge the residual stream | |
| # at layer l along ĝ = ∂ score_bath_l/∂ h_l (unit gradient ascent in the bathroom | |
| # direction), step scaled by ALPHA*‖h_l‖ so ALPHA is a fraction of the residual norm | |
| # (comparable across layers), propagate, and measure Δ toilet score at every layer | |
| # l' >= l. Averaged → heatmap (intervention layer l × readout l'). | |
| # | |
| # Tunables (override inline, e.g. ALPHA=2 NUM_IMAGES=128 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) | |
| # 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_bathroom/post/seqprobe.pth}" | |
| TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_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}" | |
| # 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). Empty = freely-generated caption. | |
| 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/gradient_ascent_bath2toilet.png}" | |
| OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/gradient_ascent_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}" | |
| --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 | |
| python -m mechanistic_interp.gradient_ascent "${ARGS[@]}" "$@" | |