mechanistic_interp/integrated_gradient_reverse.py

"""
Integrated-gradient causal influence map: toilet → bathroom across layers.

Idea
────
Identical to ``gradient_ascent.py`` EXCEPT the steered direction at layer ``l`` is
the **integrated gradient** of the *toilet* probe along the straight-line path
from a fixed *base* residual to the toilet image's residual — instead of the raw
local probe gradient at the image point.

  base  = mean residual stream of NEGATIVE-label validation images (the attention
          probe's validation negatives, from --neg_jsonl[--baseline_split], capped
          at --baseline_num=1000). One (d_model,) mean vector per layer, computed
          over all caption tokens of those negatives (forwarded with the SAME
          --forced_text as the toilet images), broadcast over the caption block.
          This is the IG baseline x' (a "no-concept" reference point).

Per image (run on TOILET-ONLY images: toilet=1 & bathroom=0):
  1. Generate a caption, forced-forward, capture resid_post at all layers; slice to
     the generated-caption tokens (the positions the probe pools over).
  2. score_toilet_l(h_l): toilet probe at layer l over the caption sequence (scalar).
  3. Integrated gradient of score_toilet_l from base b_l to image x_l (caption block):
        IG_l = (x_l - b_l) ⊙ (1/m) Σ_{k=1..m} ∂score_toilet_l/∂h |_{b_l + (k/m)(x_l - b_l)}
     This attributes, per residual coordinate, how much moving base→image along the
     toilet direction raises the toilet readout (Σ IG_l ≈ score(x_l) − score(b_l)).
     Normalized by its total magnitude: ĝ = IG_l / ‖IG_l‖ (single Frobenius norm).
  4. Intervene (per alpha): h_l' = h_l + alpha * ‖h_l‖ * ĝ  (steer along the unit
     integrated-gradient toilet direction, step scaled by the caption-block residual
     norm ‖h_l‖ so alpha is a *fraction* of residual magnitude — comparable across layers).
     PLUS one extra "full" panel: h_l' = h_l + (x_l - b_l)  — one full base→image
     displacement pushed further in the toilet direction.
  5. Forward with h_l patched; read bathroom score at every layer l' >= l:
        delta[l, l'] = score_bathroom_l'(intervened) - score_bathroom_l'(baseline)   (NOT /alpha)
  6. Average over images → heatmap (intervention layer l × readout layer l').

A positive band above the diagonal = pushing toilet-ness at l causally raises the
bathroom readout downstream → evidence for a toilet→bathroom mechanism.

Saves: heatmap PNG + the raw (L, L) matrix as JSON.
"""

import argparse
import json
import os
import random

import numpy as np
import torch as t
from PIL import Image
from transformers import AutoConfig, LlavaProcessor

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from mechanistic_interp.sequence_probe import sequence_layer_probes_from_checkpoint
from hallucination.mechanistic_interp.compare_baselines import apply_efuf_edit, apply_nullu_edit

_HOOK_SUFFIX = {"pre": "hook_resid_pre", "mid": "hook_resid_mid", "post": "hook_resid_post"}


# ── Model-variant loading (base / lora / nullu / efuf), per attribution_patching ──
def _decoder_layers(model):
    lm = getattr(model, "language_model", None) or getattr(getattr(model, "model", None), "language_model", None)
    inner = getattr(lm, "model", None)
    if inner is not None and hasattr(inner, "layers"):
        return inner.layers
    return lm.layers


def load_variant_model(args, dtype, device):
    """Build the chosen model variant. NOTE: the steering/readout PROBES are loaded
    separately (from base-trained checkpoints) — only the MODEL changes here."""
    v = args.variant
    if v == "lora":
        from model.llava.hooked_lora_llava import HookedLoRALlava
        model = HookedLoRALlava.from_pretrained(
            args.model_name, torch_dtype=dtype, device_map={"": device}).eval()
        lora_dir = args.lora_path if os.path.isdir(args.lora_path) else os.path.dirname(args.lora_path)
        model.load_lora_adapter(lora_dir, merge=True)
        print(f"[int-grad] model=lora (merged {lora_dir})")
    else:
        model = HookedSAELlavaConditionalGeneration.from_pretrained(
            args.model_name, torch_dtype=dtype, device_map={"": device}).eval()
        if v == "efuf":
            n = apply_efuf_edit(model, args.efuf_path); print(f"[int-grad] model=efuf ({n} proj tensors, {args.efuf_path})")
        elif v == "nullu":
            n = apply_nullu_edit(model, _decoder_layers(model), args.nullu_path, args.nullu_lowest, args.nullu_highest)
            print(f"[int-grad] model=nullu ({n} layers [{args.nullu_lowest},{args.nullu_highest}), {args.nullu_path})")
        else:
            print(f"[int-grad] model=base")
    return model


def hp_name(layer: int, hook_type: str) -> str:
    return f"model.language_model.layers.{layer}.{_HOOK_SUFFIX[hook_type]}"


def caption_slice(attn_mask, asst_text, processor, max_seq_tokens):
    """Return (start, end) indices of the generated-caption tokens in the sequence.

    Caption tokens sit flush at the tail of the real (unpadded) region; we keep the
    leading min(cap_len, max_seq_tokens) of them — matching training's a[:max_seq_tokens].
    """
    seq_len = int(attn_mask.sum().item())
    cap_len = len(processor.tokenizer(asst_text, add_special_tokens=False)["input_ids"])
    if cap_len <= 0:
        return None
    start = seq_len - cap_len
    end = start + min(cap_len, max_seq_tokens)
    if start < 0:
        return None
    return start, end


@t.no_grad()
def generate_caption(model, processor, image, question, device, max_new_tokens):
    prompt = f"USER: <image>\n{question}\nASSISTANT:"
    inp = processor(images=[image], text=[prompt], return_tensors="pt").to(device)
    out = model.generate(**inp, do_sample=False, num_beams=1, use_cache=True,
                         max_new_tokens=max_new_tokens)
    cap = processor.batch_decode(out, skip_special_tokens=True)[0]
    return cap.split("ASSISTANT:")[-1].strip()


def probe_logit(probe_module, layer, seq_feats):
    """Scalar logit for one image from one layer's attention probe. seq_feats: (1,T,d)."""
    kpm = t.zeros(seq_feats.shape[:2], dtype=t.bool, device=seq_feats.device)  # no padding
    return probe_module.probes[probe_module._idx[layer]](seq_feats, kpm).squeeze(0)


def probe_score(probe_module, layer, seq_feats):
    """Bounded probe score σ(logit) = P(concept present) ∈ [0,1] (logits are unbounded)."""
    return t.sigmoid(probe_logit(probe_module, layer, seq_feats))


def caption_block_acts(model, processor, img, args, device, hps):
    """One image → its caption-block residual block per hook point.

    Runs the SAME pipeline as the main loop (forced_text if set, else generated
    caption), forced-forwards, and slices each hook's activations to the caption
    tokens. Returns ({hp: (1, T, d_model)}, T) or (None, 0) if the slice is empty.
    """
    if args.forced_text:
        asst = args.forced_text
    else:
        asst = generate_caption(model, processor, img, args.question, device, args.max_new_tokens)
    forced = f"USER: <image>\n{args.question}\nASSISTANT: {asst}"
    fwd = processor(images=[img], text=[forced], return_tensors="pt").to(device)
    sl = caption_slice(fwd["attention_mask"], asst, processor, args.max_seq_tokens)
    if sl is None:
        return None, 0
    s0, s1 = sl
    acts = {}

    def make_cap(name):
        def _fn(act, hook):
            acts[name] = act[:, s0:s1]
        return _fn

    with t.no_grad():
        model.run_with_hooks(fwd, fwd_hooks=[(hp, make_cap(hp)) for hp in hps])
    return acts, (s1 - s0)


def compute_baseline(model, processor, neg_stems, stem_to_file, args, device, hps, layers):
    """IG baseline: mean residual over caption tokens of NEGATIVE-label validation images.

    For each negative image, forward via the same pipeline as the main loop, slice the
    caption block, and accumulate a running per-layer sum over tokens. Returns
    {layer: (d_model,)} mean vectors (the no-concept reference point x' for IG).
    """
    sums = {l: None for l in layers}
    n_tok = 0
    n_img = 0
    for j, stem in enumerate(neg_stems):
        try:
            img = Image.open(stem_to_file[stem]).convert("RGB")
        except Exception:
            continue
        acts, T = caption_block_acts(model, processor, img, args, device, hps)
        if acts is None or T == 0:
            continue
        for l in layers:
            block = acts[hps[l]][0].float().sum(dim=0)   # (d_model,) sum over caption tokens
            sums[l] = block if sums[l] is None else sums[l] + block
        n_tok += T
        n_img += 1
        if (j + 1) % 50 == 0 or j == 0:
            print(f"[int-grad] baseline: {n_img} negatives processed ({n_tok} tokens)")
    if n_tok == 0:
        raise SystemExit("Baseline: no negative images processed — check --neg_jsonl / image folder.")
    print(f"[int-grad] baseline built from {n_img} negative images, {n_tok} caption tokens")
    return {l: (sums[l] / n_tok) for l in layers}


def panel_label(key):
    """Display label for a steering panel key (float alpha or the 'full' sentinel)."""
    return "full (x−b)" if key == "full" else f"α = {key}"


def plot_heatmaps(deltas, panel_keys, out, title, mode="meannorm", cbar_label="Δ toilet score σ",
                  steer_name="bathroom", readout_name="toilet"):
    """Heatmap grid of {key: (L,L) Δ matrix} in one of two modes:
      'normal'   — each panel on its OWN diverging scale + own colorbar (true per-panel magnitude).
      'meannorm' — every cell ÷ mean|Δ| over all panels; ONE shared, robust (p99) scale
                   so panels are comparable and colour reads as '× typical effect'.
    Panel keys are float alphas plus the 'full' (x−b) endpoint panel.
    """
    nA = len(panel_keys)
    ncol = min(3, nA)
    nrow = -(-nA // ncol)
    fig, axes = plt.subplots(nrow, ncol, figsize=(5.0 * ncol, 4.5 * nrow), squeeze=False)
    if mode == "meannorm":
        allabs = np.abs(np.concatenate([deltas[a][np.isfinite(deltas[a])].ravel() for a in panel_keys]))
        mean_mag = float(allabs.mean()) + 1e-12
        vmax = float(np.percentile(allabs / mean_mag, 99))
        print(f"[int-grad] mean|Δ|={mean_mag:.5f}; shared p99 cap={vmax:.2f}× mean "
              f"(max={allabs.max()/mean_mag:.1f}× mean)")
        im = None
        for k, a in enumerate(panel_keys):
            ax = axes[k // ncol][k % ncol]
            im = ax.imshow(deltas[a] / mean_mag, origin="upper", cmap="RdBu_r",
                           vmin=-vmax, vmax=vmax, aspect="auto")
            ax.set_title(panel_label(a))
            ax.set_xlabel(f"readout layer l′ ({readout_name})")
            ax.set_ylabel(f"intervention layer l ({steer_name})")
        for k in range(nA, nrow * ncol):
            axes[k // ncol][k % ncol].axis("off")
        cbar = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.02)
        cbar.set_label(f"{cbar_label} / mean|Δ|  (× typical effect)")
    else:  # normal — per-panel scale
        for k, a in enumerate(panel_keys):
            ax = axes[k // ncol][k % ncol]
            d = deltas[a]
            vmax = float(np.nanmax(np.abs(d)))
            im = ax.imshow(d, origin="upper", cmap="RdBu_r", vmin=-vmax, vmax=vmax, aspect="auto")
            ax.set_title(panel_label(a))
            ax.set_xlabel(f"readout layer l′ ({readout_name})")
            ax.set_ylabel(f"intervention layer l ({steer_name})")
            fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, label=cbar_label)
        for k in range(nA, nrow * ncol):
            axes[k // ncol][k % ncol].axis("off")
        fig.tight_layout()
    fig.suptitle(title, fontsize=12, y=1.03)
    os.makedirs(os.path.dirname(out), exist_ok=True)
    fig.savefig(out, dpi=150, bbox_inches="tight")
    plt.close(fig)
    print(f"[int-grad] saved {out}")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--device_id", type=int, default=0)
    ap.add_argument("--dtype", default="bfloat16", choices=["float32", "float16", "bfloat16"])

    # Model variant — which residual stream to steer/measure on. Probes stay base-trained.
    ap.add_argument("--variant", default="base", choices=["base", "lora", "nullu", "efuf"],
                    help="Model variant the integrated gradient runs on (probes still load from base).")
    ap.add_argument("--lora_path", default="/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter",
                    help="LoRA(ours) adapter dir (variant=lora).")
    ap.add_argument("--efuf_path", default="/data/caotue/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_paper_10ep/epoch_002.pth",
                    help="EFUF .pth checkpoint (variant=efuf).")
    ap.add_argument("--nullu_path", default="/data/caotue/nullu/edited_models/LLaVA-7B-top4-0-32-bathroom_toilet",
                    help="Nullu edited-model HF dir (variant=nullu).")
    ap.add_argument("--nullu_lowest", type=int, default=8)
    ap.add_argument("--nullu_highest", type=int, default=32)

    # Probes — 4variant-trained checkpoints (the readout/steer directions). NOT the variant.
    # REVERSED: steer = toilet probe, readout = bathroom probe.
    ap.add_argument("--bath_probe", default="/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth",
                    help="4variant-trained SequenceLayerProbes for the steered direction (toilet in reversed mode).")
    ap.add_argument("--toilet_probe", default="/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth",
                    help="4variant-trained SequenceLayerProbes for the measured readout (bathroom in reversed mode).")

    ap.add_argument("--image_folder", default="/data/caotue/CC3M-Dataset/cc3m_images")
    ap.add_argument("--samples_json", default="mechanistic_interp/toilet_bathroom/samples.json",
                    help="Sample file with per-image base_mentions_object flags.")
    ap.add_argument("--base_prompt", default="Describe this image.",
                    help="Prompt whose base_mentions_object flag selects images.")
    ap.add_argument("--base_mentions", default="any", choices=["false", "true", "any"],
                    help="Filter toilet-only images by base_mentions_object: 'any' = no filter (default for reversed).")
    ap.add_argument("--id_col", default="image_id")
    ap.add_argument("--toilet_col", default="toilet")
    ap.add_argument("--question", default="Describe this image.")

    # Alternative image selection from an HF dataset (for relations with NO samples.json).
    # When --hf_dataset is set, scene-only images = rows with scene_col==1 & object_col==0
    # (the direct analog of bathroom-only), and --samples_json / --base_mentions are ignored.
    ap.add_argument("--hf_dataset", default=None,
                    help="If set, pick scene-only images (--scene_col==1 & --object_col==0) "
                         "from this HF dataset instead of --samples_json.")
    ap.add_argument("--split", default="validation",
                    help="HF split used for --hf_dataset selection (default validation).")
    ap.add_argument("--scene_col", default=None,
                    help="Steer concept column (present==1) in --hf_dataset.")
    ap.add_argument("--object_col", default=None,
                    help="Readout concept column (absent==0) in --hf_dataset.")
    ap.add_argument("--steer_name", default="toilet",
                    help="Display name of the steered concept (plot labels/title).")
    ap.add_argument("--readout_name", default="bathroom",
                    help="Display name of the readout concept (plot labels/title).")
    ap.add_argument("--forced_text", default="In this image, there is a",
                    help="If set, skip generation and force the ASSISTANT answer to exactly "
                         "this string (e.g. 'This image features a bathroom with a'). The probe "
                         "then reads these forced answer tokens — a controlled, constant context "
                         "across all images (and for the baseline negatives). If unset, use the "
                         "model's freely-generated caption.")
    ap.add_argument("--num_images", type=int, default=100,
                    help="Random-subsample cap (0 = ALL). Default 100 (shuffled by --seed).")
    ap.add_argument("--seed", type=int, default=0, help="Seed for the random image subsample.")

    # ── Integrated-gradient baseline (mean NEGATIVE-label validation residual) ──
    ap.add_argument("--neg_jsonl", default="mechanistic_interp/neg_cc3m_5k.json",
                    help='JSON {"train":[ids],"validation":[ids]} of negative image stems. '
                         "The IG baseline is the mean caption-token residual over the "
                         "--baseline_split negatives (the attention probe's validation negatives).")
    ap.add_argument("--baseline_split", default="validation", choices=["train", "validation"],
                    help="Which --neg_jsonl split to build the baseline from (default validation).")
    ap.add_argument("--baseline_num", type=int, default=1000,
                    help="Cap on negative images used for the baseline mean (0 = all in split).")
    ap.add_argument("--ig_steps", type=int, default=32,
                    help="Riemann steps for the integrated-gradient path base→image.")

    ap.add_argument("--alphas", type=float, nargs="+", default=[0.01, 0.02, 0.04, 0.08, 0.1],
                    help="Steering step sizes to sweep. h_l' = h_l + alpha*‖h_l‖*ĝ, where ĝ is "
                         "the unit integrated-gradient direction, so alpha is a FRACTION of the "
                         "caption-block residual norm (alpha=1 ⇒ step magnitude == ‖h_l‖). An extra "
                         "'full' panel (h_l' = h_l + (x−b)) is always added. The IG direction is "
                         "computed once per layer; only the patched forward repeats per panel.")
    ap.add_argument("--hook_type", default="post", choices=["pre", "mid", "post"])
    ap.add_argument("--plot_mode", default="meannorm", choices=["normal", "meannorm"],
                    help="normal = each panel its own colorbar; meannorm = ÷mean|Δ|, shared p99 scale.")
    ap.add_argument("--max_new_tokens", type=int, default=64)
    ap.add_argument("--max_seq_tokens", type=int, default=64)

    ap.add_argument("--out", default="mechanistic_interp/graph/integrated_gradient_toilet2bath.png")
    ap.add_argument("--out_json", default="mechanistic_interp/graph/integrated_gradient_toilet2bath.json")
    args = ap.parse_args()

    dtype = {"float32": t.float32, "float16": t.float16, "bfloat16": t.bfloat16}[args.dtype]
    device = f"cuda:{args.device_id}" if t.cuda.is_available() else "cpu"

    n_layers = AutoConfig.from_pretrained(args.model_name).text_config.num_hidden_layers
    layers = list(range(n_layers))

    print(f"[int-grad] device={device} dtype={dtype} layers={n_layers} alphas={args.alphas} ig_steps={args.ig_steps}")
    print(f"[int-grad] steer={args.steer_name} ({args.bath_probe})")
    print(f"[int-grad] readout={args.readout_name} ({args.toilet_probe})")
    if args.forced_text:
        print(f"[int-grad] FORCED answer: '{args.forced_text}'")
    else:
        print(f"[int-grad] answer: freely-generated caption")

    # ── Select "scene-only" images = steer concept present, readout concept absent.
    #    Two backends:
    #      (a) --hf_dataset set: rows with scene_col==1 & object_col==0 (no samples.json
    #          / base_mentions filter — for relations that lack one).
    #      (b) else --samples_json: category==bathroom_only & toilet==0, filtered by
    #          base_mentions_object per --base_mentions (flagship bathroom→toilet study).
    stem_to_file = {}
    for root, _, files in os.walk(args.image_folder):
        for fn in files:
            if fn.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
                stem_to_file[os.path.splitext(fn)[0]] = os.path.join(root, fn)
    ids = []
    if args.hf_dataset:
        if not args.scene_col or not args.object_col:
            raise SystemExit("--hf_dataset requires --scene_col and --object_col.")
        from datasets import load_dataset
        ds = load_dataset(args.hf_dataset, split=args.split)
        for row in ds:
            if row[args.scene_col] == 1 and row[args.object_col] == 0:
                stem = os.path.splitext(os.path.basename(str(row[args.id_col])))[0]
                if stem in stem_to_file:
                    ids.append(stem)
        sel_desc = f"{args.scene_col}=1 & {args.object_col}=0 from {args.hf_dataset}[{args.split}]"
    else:
        want = {"false": False, "true": True, "any": None}[args.base_mentions]
        data = json.load(open(args.samples_json))
        for it in data:
            if it.get("category") != "toilet_only" or it.get("toilet") != 1:
                continue
            pr = it.get("prompt_results", {}).get(args.base_prompt, {})
            if want is not None and pr.get("base_mentions_object") is not want:
                continue
            stem = os.path.splitext(os.path.basename(it[args.id_col]))[0]
            if stem in stem_to_file:
                ids.append(stem)
        halluc = {"false": "NON-hallucinating", "true": "HALLUCINATING", "any": "all"}[args.base_mentions]
        sel_desc = f"{halluc} toilet-only"
    if args.num_images and args.num_images > 0:
        random.Random(args.seed).shuffle(ids)          # random subsample
        ids = ids[:args.num_images]
    print(f"[int-grad] scene-only images ({sel_desc}): {len(ids)} (variant={args.variant}, "
          f"probes=base, alphas={args.alphas})")

    # ── Model (variant) + probes (base-trained) ──────────────────────────────
    model = load_variant_model(args, dtype, device)
    processor = LlavaProcessor.from_pretrained(args.model_name)
    bath = sequence_layer_probes_from_checkpoint(args.bath_probe, device)
    toilet = sequence_layer_probes_from_checkpoint(args.toilet_probe, device)
    bath.eval(); toilet.eval()

    hps = [hp_name(l, args.hook_type) for l in layers]

    # ── IG baseline: mean residual of NEGATIVE-label validation images per layer ──
    with open(args.neg_jsonl) as f:
        neg_split = json.load(f)
    neg_stems = [os.path.splitext(os.path.basename(str(s)))[0]
                 for s in neg_split.get(args.baseline_split, [])]
    neg_stems = [s for s in neg_stems if s in stem_to_file]
    if args.baseline_num and args.baseline_num > 0:
        neg_stems = neg_stems[:args.baseline_num]
    print(f"[int-grad] baseline negatives ({args.baseline_split}, on disk): {len(neg_stems)}")
    base_mean = compute_baseline(model, processor, neg_stems, stem_to_file,
                                 args, device, hps, layers)   # {l: (d_model,)}

    # Steering panels: the alpha sweep (unit IG direction, alpha*‖h‖ step) + a "full"
    # endpoint panel (h' = h + (x−b), one full base→image displacement). Keys index
    # delta_sum / deltas; the 'full' sentinel sorts last in the plot grid.
    panel_keys = list(args.alphas) + ["full"]

    # delta[key][l, l'] accumulator over the upper triangle (l' >= l); NaN below diagonal.
    tri = np.triu(np.ones((n_layers, n_layers))) > 0
    delta_sum = {k: np.where(tri, 0.0, np.nan) for k in panel_keys}
    count = 0

    def make_cap(name, store):
        def _fn(act, hook):
            store[name] = act
        return _fn

    for i, stem in enumerate(ids):
        try:
            img = Image.open(stem_to_file[stem]).convert("RGB")
        except Exception:
            continue
        if args.forced_text:
            asst = args.forced_text
        else:
            asst = generate_caption(model, processor, img, args.question, device, args.max_new_tokens)
        forced = f"USER: <image>\n{args.question}\nASSISTANT: {asst}"
        fwd = processor(images=[img], text=[forced], return_tensors="pt").to(device)

        sl = caption_slice(fwd["attention_mask"], asst, processor, args.max_seq_tokens)
        if sl is None:
            continue
        s0, s1 = sl

        # 1. Baseline forward: capture resid_post at all layers.
        base_acts = {}
        with t.no_grad():
            model.run_with_hooks(fwd, fwd_hooks=[(hp, make_cap(hp, base_acts)) for hp in hps])

        # Baseline toilet scores per layer (bounded σ(logit) ∈ [0,1]).
        toi_base = {}
        with t.no_grad():
            for l in layers:
                feats = base_acts[hps[l]][:, s0:s1].float()
                toi_base[l] = float(probe_score(toilet, l, feats).item())

        # 2. For each intervention layer l: integrated-gradient steering in bathroom dir,
        #    then propagate. ĝ and (x−b) depend only on l (not the panel) → compute once.
        for l in layers:
            h_l = base_acts[hps[l]]                       # (1, S, d_model)
            x = h_l[:, s0:s1].float().detach()            # (1, T, d_model) image residual (IG input)
            b = base_mean[l].view(1, 1, -1).expand_as(x)  # (1, T, d_model) baseline (IG reference)

            # Integrated gradient: avg probe gradient along the path base→image, ⊙ (x − b).
            # Local to the probe — no model backward, so all m steps are cheap.
            diff = (x - b)
            grad_sum = t.zeros_like(x)
            for k in range(1, args.ig_steps + 1):
                ak = k / args.ig_steps
                interp = (b + ak * diff).detach().requires_grad_(True)
                score = probe_logit(bath, l, interp)       # scalar bathroom logit
                gk, = t.autograd.grad(score, interp)
                grad_sum = grad_sum + gk.detach()
            ig = diff * (grad_sum / args.ig_steps)         # (1, T, d_model) integrated gradient
            g = ig / (ig.norm() + 1e-8)                    # unit steering direction (whole-tensor norm)
            g = g.to(h_l.dtype)

            h_norm = h_l[:, s0:s1].float().norm().to(h_l.dtype)   # ‖h_l‖ over caption block; makes alpha a fraction of residual magnitude (comparable across layers)
            full_step = diff.to(h_l.dtype)                        # the 'full' panel step: h' = h + (x−b)

            # Per-panel steering vector applied to the caption block.
            steps = {a: (a * h_norm * g) for a in args.alphas}
            steps["full"] = full_step

            for key in panel_keys:
                patched = h_l.clone()
                patched[:, s0:s1] = patched[:, s0:s1] + steps[key]

                # l' == l: no propagation needed. Δ = score(steered) − score(baseline).
                with t.no_grad():
                    toi_int_l = float(probe_score(toilet, l, patched[:, s0:s1].float()).item())
                delta_sum[key][l, l] += (toi_int_l - toi_base[l])

                # l' > l: patch layer l, capture downstream.
                if l < n_layers - 1:
                    down = {}
                    def patch_fn(act, hook, p=patched):
                        return p
                    hooks = [(hps[l], patch_fn)] + \
                            [(hps[lp], make_cap(hps[lp], down)) for lp in range(l + 1, n_layers)]
                    with t.no_grad():
                        model.run_with_hooks(fwd, fwd_hooks=hooks)
                    with t.no_grad():
                        for lp in range(l + 1, n_layers):
                            feats = down[hps[lp]][:, s0:s1].float()
                            toi_int = float(probe_score(toilet, lp, feats).item())
                            delta_sum[key][l, lp] += (toi_int - toi_base[lp])

        count += 1
        if (i + 1) % 10 == 0 or i == 0:
            print(f"[int-grad] processed {count}/{len(ids)}")

    if count == 0:
        raise SystemExit("No images processed.")
    deltas = {k: delta_sum[k] / count for k in panel_keys}
    print(f"[int-grad] averaged over {count} images")

    # ── Save JSON (all panels) ──────────────────────────────────────────────────
    os.makedirs(os.path.dirname(args.out_json), exist_ok=True)
    with open(args.out_json, "w") as f:
        json.dump({"alphas": args.alphas, "panels": [str(k) for k in panel_keys],
                   "n_images": count, "hook_type": args.hook_type,
                   "ig_steps": args.ig_steps, "baseline_split": args.baseline_split,
                   "baseline_num": len(neg_stems),
                   "delta": {str(k): np.where(np.isnan(d), None, d).tolist()
                             for k, d in deltas.items()}}, f)
    print(f"[int-grad] saved {args.out_json}")

    # ── Heatmap (plot_mode = normal | meannorm) ─────────────────────────────────
    sel_tag = (f"{args.scene_col}=1&{args.object_col}=0" if args.hf_dataset
               else f"base_mentions={args.base_mentions}")
    plot_heatmaps(
        deltas, panel_keys, args.out,
        title=f"{args.steer_name}→{args.readout_name} integrated-gradient influence ({args.plot_mode}) — "
              f"model={args.variant}, {count} imgs [{sel_tag}], {args.hook_type}",
        mode=args.plot_mode,
        cbar_label=f"Δ {args.readout_name} score σ",
        steer_name=args.steer_name, readout_name=args.readout_name)


if __name__ == "__main__":
    main()