import argparse
import csv
import io
import math
import subprocess
import sys
from pathlib import Path
from typing import Optional
import matplotlib.pyplot as plt
import numpy as np
def run_example(
repo: Path,
checkpoint: Optional[str],
lengths: list[int],
fractions: list[float],
mode: str,
vocab_size: Optional[int],
embed_dim: Optional[int],
hidden_dim: Optional[int],
output_dim: Optional[int],
num_neighbors: Optional[int],
plasticity: Optional[bool],
) -> list[dict[str, str]]:
example = repo / "target" / "release" / "examples" / "per_position_divergence"
if not example.exists():
print("Building per_position_divergence example...", file=sys.stderr)
subprocess.run(
["cargo", "build", "--example", "per_position_divergence", "--release"],
cwd=repo,
check=True,
)
cmd = [
str(example),
"--lengths",
",".join(str(l) for l in lengths),
"--fractions",
",".join(str(f) for f in fractions),
"--mode",
mode,
]
if checkpoint:
cmd.extend(["--checkpoint", checkpoint])
for flag, value in [
("--vocab-size", vocab_size),
("--embed-dim", embed_dim),
("--hidden-dim", hidden_dim),
("--output-dim", output_dim),
("--num-neighbors", num_neighbors),
]:
if value is not None:
cmd.extend([flag, str(value)])
if plasticity is not None:
cmd.extend(["--plasticity", "true" if plasticity else "false"])
result = subprocess.run(cmd, cwd=repo, check=True, capture_output=True, text=True)
return list(csv.DictReader(io.StringIO(result.stdout)))
def heatmap(ax, values: np.ndarray, fractions: list[float], title: str, cmap: str) -> None:
im = ax.imshow(
values,
aspect="auto",
cmap=cmap,
origin="lower",
extent=[-0.5, values.shape[1] - 0.5, fractions[0], fractions[-1]],
)
ax.set_title(title)
ax.set_xlabel("position")
ax.set_ylabel("geometric-only fraction")
plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04)
def plot(rows: list[dict[str, str]], lengths: list[int], output: Path) -> None:
by_length: dict[int, list[dict[str, str]]] = {l: [] for l in lengths}
for r in rows:
by_length[int(r["length"])].append(r)
n = len(lengths)
cols = 2
fig, axes = plt.subplots(n, cols, figsize=(6 * cols, 3.5 * n), squeeze=False)
for idx, length in enumerate(lengths):
subset = by_length[length]
fractions = sorted({float(r["geometric_fraction"]) for r in subset})
positions = list(range(length))
geo_l2 = np.zeros((len(fractions), length))
mixed_l2 = np.zeros((len(fractions), length))
geo_cos = np.zeros((len(fractions), length))
mixed_cos = np.zeros((len(fractions), length))
frac_index = {f: i for i, f in enumerate(fractions)}
for r in subset:
f = frac_index[float(r["geometric_fraction"])]
p = int(r["position"])
geo_l2[f, p] = float(r["geo_l2"])
mixed_l2[f, p] = float(r["mixed_l2"])
geo_cos[f, p] = float(r["geo_cosine"])
mixed_cos[f, p] = float(r["mixed_cosine"])
ax_geo = axes[idx][0]
ax_mixed = axes[idx][1]
heatmap(ax_geo, mixed_l2, fractions, f"L={length} — mixed hidden L2", "inferno")
if np.max(geo_l2) > 1e-12:
levels = np.linspace(0.0, float(np.max(geo_l2)), 5)
ax_geo.contour(
np.arange(length),
np.array(fractions),
geo_l2,
levels=levels,
colors="cyan",
alpha=0.4,
linewidths=0.8,
)
ax_geo.plot([], [], color="cyan", alpha=0.6, linewidth=0.8, label="geometric L2 contour")
ax_geo.legend(loc="upper right", fontsize=7)
heatmap(ax_mixed, mixed_cos, fractions, f"L={length} — mixed cosine similarity", "viridis")
ax_mixed.set_clim = (-1.0, 1.0)
fig.suptitle("Per-Position Hidden-State Divergence", y=1.0)
fig.tight_layout()
fig.savefig(output, dpi=150, bbox_inches="tight")
print(f"Saved plot to {output}")
def main() -> None:
parser = argparse.ArgumentParser(description="Plot per-position hidden-state divergence")
parser.add_argument("--checkpoint", default=None, help="path to a flatten_params checkpoint")
parser.add_argument("--lengths", default="64", help="comma-separated context lengths")
parser.add_argument(
"--fractions",
default="0.0,0.25,0.5,0.75,1.0",
help="comma-separated geometric-only fractions",
)
parser.add_argument("--mode", default="last-hybrid", choices=["random", "last-hybrid"], help="mask generation mode")
parser.add_argument("--vocab-size", type=int, default=None, help="model vocabulary size")
parser.add_argument("--embed-dim", type=int, default=None, help="attention embedding dimension")
parser.add_argument("--hidden-dim", type=int, default=None, help="MLP hidden dimension")
parser.add_argument("--output-dim", type=int, default=None, help="MLP output dimension")
parser.add_argument("--num-neighbors", type=int, default=None, help="number of geometric neighbours")
parser.add_argument("--plasticity", action="store_true", help="enable learnable edge weights")
parser.add_argument("--csv", default=None, help="use an existing CSV instead of re-running the example")
parser.add_argument("-o", "--output", default="target/per_position_divergence.png", help="output PNG path")
args = parser.parse_args()
repo = Path(__file__).resolve().parent.parent
lengths = [int(x.strip()) for x in args.lengths.split(",")]
fractions = [float(x.strip()) for x in args.fractions.split(",")]
if args.csv:
with open(args.csv, "r", newline="") as f:
rows = list(csv.DictReader(f))
else:
rows = run_example(
repo,
args.checkpoint,
lengths,
fractions,
args.mode,
args.vocab_size,
args.embed_dim,
args.hidden_dim,
args.output_dim,
args.num_neighbors,
args.plasticity if args.plasticity else None,
)
output = Path(args.output)
output.parent.mkdir(parents=True, exist_ok=True)
plot(rows, lengths, output)
if __name__ == "__main__":
main()