File size: 4,527 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 | """Re-render a gradient-ascent heatmap JSON with a SHARED color scale across alphas.
The original per-panel scaling normalizes each alpha to its own max, so a large-alpha
panel (bigger real changes) gets a stretched scale and looks PALER than a small-alpha
panel — visually backwards. This re-plots all panels on ONE shared, robust
(percentile-clipped) diverging scale so color magnitude is comparable across alpha:
bigger change → more saturated.
Usage:
python -m mechanistic_interp.scripts.replot_gradient_ascent \
--json mechanistic_interp/graph/gradient_ascent_bath2toilet.json \
--out mechanistic_interp/graph/gradient_ascent_bath2toilet_shared.png \
[--pct 99] [--drop_early 0]
"""
import argparse
import json
import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--json", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--pct", type=float, default=99.0,
help="Percentile of |Δ| (over ALL panels) used as the shared vmax — "
"robust to the early-layer outlier. 100 = plain max.")
ap.add_argument("--normalize", action="store_true",
help="Alias for --norm_by max (divide by largest |Δ|, scale ±1).")
ap.add_argument("--norm_by", choices=["none", "max", "mean"], default="none",
help="Divide every cell by max|Δ| or mean|Δ| (over ALL panels). "
"'mean' = color is in multiples of the typical effect (crisp); "
"'max' = biggest cell = ±1 (artifact-dominated, blurry).")
ap.add_argument("--drop_early", type=int, default=0,
help="Blank the first N intervention layers (rows) — the off-manifold "
"artifact — from both the scale and the display.")
ap.add_argument("--title", default=None)
args = ap.parse_args()
d = json.load(open(args.json))
alphas = d["alphas"]
mats = {}
for a in alphas:
M = np.array([[np.nan if v is None else v for v in row] for row in d["delta"][str(a)]])
if args.drop_early > 0:
M[: args.drop_early, :] = np.nan
mats[a] = M
allvals = np.concatenate([m[np.isfinite(m)].ravel() for m in mats.values()])
absvals = np.abs(allvals)
norm_by = "max" if args.normalize else args.norm_by
if norm_by == "max":
norm = float(absvals.max())
for a in mats:
mats[a] = mats[a] / norm
vmax = 1.0
cbar_label = "Δ toilet σ / max|Δ| (±1)"
print(f"normalized by max|Δ| = {norm:.5f} → scale [-1, 1]")
elif norm_by == "mean":
norm = float(absvals.mean())
for a in mats:
mats[a] = mats[a] / norm
znorm = np.abs(np.concatenate([m[np.isfinite(m)].ravel() for m in mats.values()]))
vmax = float(np.percentile(znorm, args.pct)) # robust crisp cap, in mean-units
cbar_label = "Δ toilet σ / mean|Δ| (× typical effect)"
print(f"normalized by mean|Δ| = {norm:.5f}; crisp cap p{args.pct} = {vmax:.2f}× mean "
f"(max = {absvals.max()/norm:.1f}× mean)")
else:
vmax = float(np.percentile(absvals, args.pct))
cbar_label = "Δ toilet score σ (shared scale)"
print(f"shared vmax (p{args.pct}) = {vmax:.4f} (raw max |Δ| = {absvals.max():.4f})")
n = len(alphas)
ncol = min(3, n)
nrow = -(-n // ncol)
fig, axes = plt.subplots(nrow, ncol, figsize=(4.8 * ncol, 4.3 * nrow), squeeze=False)
im = None
for k, a in enumerate(alphas):
ax = axes[k // ncol][k % ncol]
im = ax.imshow(mats[a], origin="upper", cmap="RdBu_r",
vmin=-vmax, vmax=vmax, aspect="auto")
ax.set_title(f"α = {a}")
ax.set_xlabel("readout layer l′ (toilet)")
ax.set_ylabel("intervention layer l (bathroom)")
for k in range(n, nrow * ncol):
axes[k // ncol][k % ncol].axis("off")
# One shared colorbar for all panels.
cbar = fig.colorbar(im, ax=axes, fraction=0.025, pad=0.02)
cbar.set_label(cbar_label)
title = args.title or (f"{d.get('n_images','?')} images — shared color scale "
f"(p{args.pct} clip, drop_early={args.drop_early})")
fig.suptitle(title, fontsize=12)
fig.savefig(args.out, dpi=150, bbox_inches="tight")
print(f"saved {args.out}")
if __name__ == "__main__":
main()
|