File size: 5,545 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 | #!/bin/bash
# =============================================================================
# mlp_trace.sh — MLP-layer leak diagnostic on the *suppressed* cohort.
#
# Free-gen from PROMPT under base, then under ΔW@MLP_{L_int} for each
# L_int in TARGET_LAYERS. Per condition, per layer L, captures the SAE
# feature aggregate at {hook_resid_mid, hook_resid_post} over the
# generated tokens → 4 scalars per layer:
#
# {obj1}@mid(L) {obj2}@mid(L) {obj1}@post(L) {obj2}@post(L)
#
# Cohort = samples where base says {object1} and LoRA does not
# (filter_samples with gen_mode='prefix' against SAMPLES_JSON).
#
# Outputs:
# ${OUT_DIR}/per_sample/{image_id}.pt — per-sample scalars + gen text
# ${OUT_DIR}/summary.json — config + leak-sweep numbers
# ${GRAPH_DIR}/four_line_base.png — 4 lines × N layers (base)
# ${GRAPH_DIR}/four_line_dw_L{L:02d}.png — same under ΔW@MLP_L
# ${GRAPH_DIR}/leak_layer_sweep.png — mean {obj2}@post(L=L_int)
# base vs ΔW@MLP_{L_int}
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
# Object names for the dataset
OBJECT1="${OBJECT1:-toilet}"
OBJECT2="${OBJECT2:-bathroom}"
# Shortcut names for concept references (e.g., 'bath', 'toi', 'kit', 'oven')
OBJECT1_SHORT="${OBJECT1_SHORT:-toi}"
OBJECT2_SHORT="${OBJECT2_SHORT:-bath}"
# Four feature JSONs (select_features.py format), one per (concept, site).
OBJ1_MID_FEATURES_JSON="${OBJ1_MID_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_toilet_mid_top10.json}"
OBJ1_POST_FEATURES_JSON="${OBJ1_POST_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_toilet_top10.json}"
OBJ2_MID_FEATURES_JSON="${OBJ2_MID_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_bathroom_mid_top10.json}"
OBJ2_POST_FEATURES_JSON="${OBJ2_POST_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_bathroom_post_top10.json}"
SAMPLES_JSON="${SAMPLES_JSON:-/data/caotue/hallucination/mechanistic_interp/toilet_bathroom/samples.json}"
SAE_CKPT="${SAE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/last.ckpt}"
ADAPTER_PATH="${ADAPTER_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_model.safetensors}"
ADAPTER_CFG="${ADAPTER_CFG:-/data/caotue/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_config.json}"
HF_DATASET="${HF_DATASET:-pbcong/bathroom-toilet}"
HF_SPLIT="${HF_SPLIT:-validation}"
ID_COL="${ID_COL:-image_id}"
PROMPT="${PROMPT:-Describe this image.}"
# Assistant-side prefill appended after 'ASSISTANT:' (e.g. 'there is a
# bathroom'). Empty = no prefill. Aggregation window covers
# (prefix tokens + GEN_TOKENS).
PREFIX="${PREFIX:-This image features}"
GEN_TOKENS="${GEN_TOKENS:-256}"
# CATEGORY: one of any | {object1}_only | {object2}_only | {object1}_{object2}
CATEGORY="${CATEGORY:-bathroom_only}"
MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
DEVICE="${DEVICE:-cuda:5}"
# DTYPE: one of float32 | float16 | bfloat16
DTYPE="${DTYPE:-bfloat16}"
N_SAMPLES="${N_SAMPLES:-0}" # 0 = all filter-passing samples
# POOL: one of max | mean (pool over the K generated-token positions)
POOL="${POOL:-max}"
# FEATURE_AGG: one of max | mean (aggregate over per-layer feature ids per concept)
FEATURE_AGG="${FEATURE_AGG:-mean}"
# TARGET_LAYERS: 'all' = every layer | '' = base only | '5,12,18' = list
TARGET_LAYERS="${TARGET_LAYERS:-all}"
# PLOT_ONLY: if 1, skip tracing and load existing .pt files for plotting only
PLOT_ONLY="${PLOT_ONLY:-0}"
RUN_TAG="${RUN_TAG:-${HF_SPLIT}/${POOL}-${FEATURE_AGG}/${CATEGORY}/top10}"
OUT_DIR="${OUT_DIR:-mechanistic_interp/mlp_trace/${RUN_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/mlp_trace/${RUN_TAG}/graphs}"
ARGS=(
-m mechanistic_interp.mlp_trace_features
--object1 "${OBJECT1}"
--object2 "${OBJECT2}"
--obj1_mid_features_json "${OBJ1_MID_FEATURES_JSON}"
--obj1_post_features_json "${OBJ1_POST_FEATURES_JSON}"
--obj2_mid_features_json "${OBJ2_MID_FEATURES_JSON}"
--obj2_post_features_json "${OBJ2_POST_FEATURES_JSON}"
--samples_json "${SAMPLES_JSON}"
--prompt "${PROMPT}"
--prefix "${PREFIX}"
--gen_tokens "${GEN_TOKENS}"
--category "${CATEGORY}"
--hf_dataset "${HF_DATASET}"
--hf_split "${HF_SPLIT}"
--id_col "${ID_COL}"
--adapter_path "${ADAPTER_PATH}"
--adapter_cfg "${ADAPTER_CFG}"
--sae_ckpt "${SAE_CKPT}"
--model_name "${MODEL_NAME}"
--device "${DEVICE}"
--dtype "${DTYPE}"
--n_samples "${N_SAMPLES}"
--pool "${POOL}"
--feature_agg "${FEATURE_AGG}"
--target_layers "${TARGET_LAYERS}"
--out_dir "${OUT_DIR}"
--graph_dir "${GRAPH_DIR}"
)
# --plot_only: skip tracing, load existing .pt files for plotting only
if [[ "${PLOT_ONLY:-}" == "1" ]]; then
ARGS+=(--plot_only)
echo "Plot-only mode: loading existing per-sample .pt files"
fi
echo "Running: python ${ARGS[*]}"
python "${ARGS[@]}" |