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
def run_sweep(
repo: Path,
checkpoint: Optional[str],
lengths: list[int],
target_ms: 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" / "tiebreak_sweep"
if not example.exists():
print("Building tiebreak_sweep example...", file=sys.stderr)
subprocess.run(
["cargo", "build", "--example", "tiebreak_sweep", "--release"],
cwd=repo,
check=True,
)
cmd = [
str(example),
"--lengths",
",".join(str(l) for l in lengths),
"--target-ms",
str(target_ms),
"--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 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 = math.ceil(math.sqrt(n))
rows_grid = math.ceil(n / cols)
fig, axes = plt.subplots(rows_grid, cols, figsize=(5 * cols, 4 * rows_grid), squeeze=False)
for idx, length in enumerate(lengths):
ax = axes[idx // cols][idx % cols]
subset = by_length[length]
frac = [float(r["geometric_fraction"]) for r in subset]
hybrid_ns = [float(r["hybrid_ns"]) for r in subset]
tiebreak_ns = [float(r["tiebreak_ns"]) for r in subset]
geometric_ns = [float(r["geometric_ns"]) for r in subset]
geo_mse = [max(1e-16, float(r["geo_mse"])) for r in subset]
tb_mse = [max(1e-16, float(r["tb_mse"])) for r in subset]
speedup = [h / t for h, t in zip(hybrid_ns, tiebreak_ns)]
geo_speedup = hybrid_ns[0] / geometric_ns[0]
ax.scatter(tb_mse, speedup, c=frac, cmap="viridis", edgecolors="black", linewidths=0.5, zorder=3)
ax.plot(tb_mse, speedup, "-", color="gray", alpha=0.4, zorder=2)
ax.scatter(
[geo_mse[0]],
[geo_speedup],
marker="s",
s=80,
color="C1",
edgecolors="black",
linewidths=0.5,
label=f"pure geometric ({geo_speedup:.1f}×)",
zorder=4,
)
ax.set_xscale("log")
ax.set_xlabel("logit MSE vs hybrid")
ax.set_ylabel("speedup over hybrid")
ax.set_title(f"L = {length}")
ax.grid(True, alpha=0.3, which="both")
ax.legend(loc="best")
for idx in range(n, cols * rows_grid):
axes[idx // cols][idx % cols].set_visible(False)
cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
sm = plt.cm.ScalarMappable(cmap="viridis", norm=plt.Normalize(0.0, 1.0))
sm.set_array([])
fig.colorbar(sm, cax=cbar_ax, label="geometric-only fraction")
fig.suptitle("Approximate Transformer Math — Tiebreak Pareto Sweep", y=1.02)
fig.tight_layout(rect=[0, 0, 0.9, 1])
fig.savefig(output, dpi=150, bbox_inches="tight")
print(f"Saved plot to {output}")
def main() -> None:
parser = argparse.ArgumentParser(description="Plot tiebreak sweep across lengths")
parser.add_argument("--checkpoint", default=None, help="path to a flatten_params checkpoint")
parser.add_argument(
"--lengths",
default="64,128,256,512",
help="comma-separated context lengths",
)
parser.add_argument("--target-ms", type=float, default=50.0, help="target benchmark time per mode")
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 sweep")
parser.add_argument("-o", "--output", default="target/tiebreak_sweep.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(",")]
if args.csv:
with open(args.csv, "r", newline="") as f:
rows = list(csv.DictReader(f))
else:
rows = run_sweep(
repo,
args.checkpoint,
lengths,
args.target_ms,
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()