From 0d2967fa6bc97fd703bc07c260b1973ca753ab02 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Thu, 23 Jul 2026 14:15:39 +0200 Subject: [PATCH] Add router gating diagnostic for MoE checkpoints plot_router_gating visualizes soft expert gate weights vs. a continuous routing axis (e.g. pre-step energy), binned into equal-population quantiles and stacked to show the router's soft decision boundaries. Wired into rollout_validation.ipynb as a new notebook-only section that loads a checkpoint's Router directly, since gate weights aren't present in rollout/predict parquet output. Co-Authored-By: Claude Sonnet 5 --- analysis/rollout_validation.ipynb | 136 ++++++++++++++++++++++++++---- giant/analysis.py | 63 ++++++++++++++ 2 files changed, 183 insertions(+), 16 deletions(-) diff --git a/analysis/rollout_validation.ipynb b/analysis/rollout_validation.ipynb index cbb7dbf..c4f5506 100644 --- a/analysis/rollout_validation.ipynb +++ b/analysis/rollout_validation.ipynb @@ -19,18 +19,18 @@ "source": [ "# GIANT rollout-vs-truth validation notebook\n", "\n", - "Diagnostics for a full autoregressive `giant rollout` shower, compared against a held-out ground-truth steps file (the same schema `giant train` consumes — see `giant.data.loader.load_steps`) rather than one-step-ahead `giant predict` output.\n", + "Diagnostics for a full autoregressive `giant rollout` shower, compared against a held-out ground-truth steps file (the same schema `giant train` consumes \u2014 see `giant.data.loader.load_steps`) rather than one-step-ahead `giant predict` output.\n", "\n", - "This is the sibling of `validation.ipynb`: that notebook checks whether one-step generation (conditioned on the *real* preceding state, every row) reproduces real marginals/correlations/shower observables. This one checks the thing that actually matters for deployment — whether a shower **rolled out autoregressively from the model's own outputs** still looks physical, which is where covariate shift (small per-step errors compounding across a track) would show up.\n", + "This is the sibling of `validation.ipynb`: that notebook checks whether one-step generation (conditioned on the *real* preceding state, every row) reproduces real marginals/correlations/shower observables. This one checks the thing that actually matters for deployment \u2014 whether a shower **rolled out autoregressively from the model's own outputs** still looks physical, which is where covariate shift (small per-step errors compounding across a track) would show up.\n", "\n", - "Built on `RolloutVsTruth`, which treats the rollout file as \"generated\" and the truth file as \"real\". Unlike the paired predict-parquet `source` (`pred_*`/`true_*` columns of the same row), the two files here are **independent, unpaired datasets** — a rollout doesn't replay real events row-for-row, so real/generated may have different lengths and there's no per-row correspondence. Everything below only ever compares real-vs-generated *distributions*, never individual paired rows, and every check still streams (no `SampleCollection`, no full-file materialization) — see `giant.analysis`'s module docstring for the `RolloutVsTruth` mechanics.\n", + "Built on `RolloutVsTruth`, which treats the rollout file as \"generated\" and the truth file as \"real\". Unlike the paired predict-parquet `source` (`pred_*`/`true_*` columns of the same row), the two files here are **independent, unpaired datasets** \u2014 a rollout doesn't replay real events row-for-row, so real/generated may have different lengths and there's no per-row correspondence. Everything below only ever compares real-vs-generated *distributions*, never individual paired rows, and every check still streams (no `SampleCollection`, no full-file materialization) \u2014 see `giant.analysis`'s module docstring for the `RolloutVsTruth` mechanics.\n", "\n", - "Same four tiers as `validation.ipynb`, all built on the same functions — pass a `RolloutVsTruth` in place of the predict-parquet path/LazyFrame everywhere:\n", + "Same four tiers as `validation.ipynb`, all built on the same functions \u2014 pass a `RolloutVsTruth` in place of the predict-parquet path/LazyFrame everywhere:\n", "\n", - "1. **stratified marginals** — per-dimension real-vs-generated, sliced by pdg/material/energy\n", - "2. **joint structure** — correlation matrices, physically-coupled pairwise plots, direction alignment\n", - "3. **physical constraints** — unit-norm directions, non-negative step_length/delta_e/edep (checked on the rollout's own output — with autoregression, a constraint violation early in a track can compound into later steps, unlike one-step-ahead validation)\n", - "4. **event-level (shower) observables** — total/mean/median energy and length per event, longitudinal/transverse profiles, shower-max depth, computed directly from the rollout shower against the truth file's own events (`compute_rollout_vs_truth_observables_pl`, the Tier 4 counterpart to `RolloutVsTruth`)" + "1. **stratified marginals** \u2014 per-dimension real-vs-generated, sliced by pdg/material/energy\n", + "2. **joint structure** \u2014 correlation matrices, physically-coupled pairwise plots, direction alignment\n", + "3. **physical constraints** \u2014 unit-norm directions, non-negative step_length/delta_e/edep (checked on the rollout's own output \u2014 with autoregression, a constraint violation early in a track can compound into later steps, unlike one-step-ahead validation)\n", + "4. **event-level (shower) observables** \u2014 total/mean/median energy and length per event, longitudinal/transverse profiles, shower-max depth, computed directly from the rollout shower against the truth file's own events (`compute_rollout_vs_truth_observables_pl`, the Tier 4 counterpart to `RolloutVsTruth`)" ] }, { @@ -47,7 +47,7 @@ " \"/ceph/lbogner/geant_steps/predictions/9e76bc2c-f4ef-4488-9f62-b6d14e1f298e.parquet\"\n", ")\n", "# Any held-out file sharing giant train's input schema (real miniCaloSim\n", - "# steps) — e.g. the val split the rollout's seed events were drawn from.\n", + "# steps) \u2014 e.g. the val split the rollout's seed events were drawn from.\n", "TRUTH_FILE = (\n", " \"/ceph/lbogner/geant_steps/processed/steps/gen3/schema2/pbwo4/shard-009.parquet\"\n", ")\n", @@ -55,7 +55,7 @@ "# sample_frac subsamples each side of the Tier 1-3 checks independently\n", "# (kept memory-bounded for large files); defaults to every row. Tier 4\n", "# (compute_rollout_vs_truth_observables_pl, below) always streams every row\n", - "# regardless — per-event sums would be silently corrupted by row subsampling.\n", + "# regardless \u2014 per-event sums would be silently corrupted by row subsampling.\n", "SOURCE = RolloutVsTruth(rollout=ROLLOUT_FILE, truth=TRUTH_FILE)" ] }, @@ -138,7 +138,7 @@ "outputs": [], "source": [ "# Real vs. generated Pearson correlation matrices (+ their difference) over\n", - "# the 9 raw target dims — catches a model that decorrelates targets that are\n", + "# the 9 raw target dims \u2014 catches a model that decorrelates targets that are\n", "# physically coupled even when every individual marginal looks clean.\n", "_ = plot_correlation_matrices(SOURCE)" ] @@ -150,7 +150,7 @@ "metadata": {}, "outputs": [], "source": [ - "# Scatter for physically-coupled pairs (step_length/delta_e/edep) — the\n", + "# Scatter for physically-coupled pairs (step_length/delta_e/edep) \u2014 the\n", "# joint-structure check correlation matrices alone can't fully capture.\n", "_ = plot_pairwise(SOURCE, n_sample=10000)" ] @@ -162,7 +162,7 @@ "metadata": {}, "outputs": [], "source": [ - "# cos(angle) between post_dir and travel_dir — coupled through the\n", + "# cos(angle) between post_dir and travel_dir \u2014 coupled through the\n", "# scattering physics, so this is another joint-structure check.\n", "_ = plot_direction_alignment(SOURCE)" ] @@ -174,7 +174,7 @@ "source": [ "## Tier 3: physical constraints\n", "\n", - "Unit-norm direction vectors, non-negative step_length/delta_e/edep. `constraint_report_pl`/`plot_constraint_violations` only ever check the *generated* side (here the rollout output) — under autoregression a violation isn't just a one-off artifact, it can feed the next step's conditioning, so this is worth watching more closely here than in one-step-ahead validation." + "Unit-norm direction vectors, non-negative step_length/delta_e/edep. `constraint_report_pl`/`plot_constraint_violations` only ever check the *generated* side (here the rollout output) \u2014 under autoregression a violation isn't just a one-off artifact, it can feed the next step's conditioning, so this is worth watching more closely here than in one-step-ahead validation." ] }, { @@ -194,7 +194,7 @@ "source": [ "## Tier 4: event-level (shower) observables\n", "\n", - "Built on `compute_rollout_vs_truth_observables_pl`, not `compute_event_observables_pl` — the rollout file carries its own `track_id`/`termination_reason` columns the event-level aggregation needs, and the shower here already *is* a full autoregressive rollout rather than one-step generations re-aggregated by event. Entry axis/point and per-event totals are computed separately per side (rollout and truth events are unrelated), but depth/transverse bin edges are shared across both so the profiles below overlay on one binning.\n", + "Built on `compute_rollout_vs_truth_observables_pl`, not `compute_event_observables_pl` \u2014 the rollout file carries its own `track_id`/`termination_reason` columns the event-level aggregation needs, and the shower here already *is* a full autoregressive rollout rather than one-step generations re-aggregated by event. Entry axis/point and per-event totals are computed separately per side (rollout and truth events are unrelated), but depth/transverse bin edges are shared across both so the profiles below overlay on one binning.\n", "\n", "Returns the same `EventObservables` `compute_event_observables_pl` does, so every plot function from `validation.ipynb` works unchanged here too." ] @@ -292,7 +292,111 @@ "source": [ "---\n", "\n", - "For the dataset-wide breakdown of which particle species contributed how much of the total energy/length (`pdg_contribution_table_pl`), see `validation.ipynb` — it needs the paired predict schema, which this rollout-vs-truth comparison doesn't have." + "For the dataset-wide breakdown of which particle species contributed how much of the total energy/length (`pdg_contribution_table_pl`), see `validation.ipynb` \u2014 it needs the paired predict schema, which this rollout-vs-truth comparison doesn't have." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Router gating showcase (MoE)\n", + "\n", + "Every other section above is file-only \u2014 it reads `ROLLOUT_FILE` and never touches\n", + "a checkpoint (see `giant.analysis`'s module docstring). This section is the one\n", + "deliberate exception: soft gate weights only exist inside the trained `Router`,\n", + "not in the rollout parquet, so this loads the checkpoint that produced\n", + "`ROLLOUT_FILE` and calls `model.router.gate(...)` directly on that shower's\n", + "pre-step conditioning.\n", + "\n", + "`model.router` is Stage 1's router; Stage 2 (`sec_decoder.router`) is a separate,\n", + "independently trained `Router` instance over the same axis (see\n", + "`giant.model.network.build_models`) and isn't shown here.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "import polars as pl\n", + "import torch\n", + "\n", + "from giant.analysis import plot_router_gating\n", + "from giant.data.transforms import Normalizer, build_cond_features\n", + "from giant.model.network import build_models\n", + "\n", + "# Checkpoint that produced ROLLOUT_FILE (needs `model.router` enabled at\n", + "# train time, i.e. trained with `--router` / `model.router.enabled = true`).\n", + "CHECKPOINT = \"/ceph/lbogner/geant_steps/checkpoints/REPLACE_ME/best.pt\"\n", + "\n", + "ckpt = torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=False)\n", + "model_cfg = ckpt[\"model_config\"]\n", + "conditioning = model_cfg.get(\"conditioning\", \"embedding\")\n", + "pdg_map = {int(k): v for k, v in ckpt[\"pdg_map\"].items()}\n", + "mat_map = {str(k): v for k, v in ckpt[\"mat_map\"].items()}\n", + "cond_norm = Normalizer.from_dict(ckpt[\"normalizer\"][\"cond\"])\n", + "\n", + "model, _sec_decoder = build_models(model_cfg)\n", + "model.load_state_dict(ckpt[\"model\"])\n", + "model.eval()\n", + "\n", + "if not hasattr(model, \"router\"):\n", + " raise RuntimeError(\n", + " f\"{CHECKPOINT} has no router \u2014 it was trained with model.router.enabled=False\"\n", + " )\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Pre-step conditioning for every row of the rollout shower, reconstructed\n", + "# the same way `giant predict`/`giant rollout` do (giant.data.transforms).\n", + "cols = [\"pdg\", \"pre_x\", \"pre_y\", \"pre_z\", \"pre_E\", \"pre_dx\", \"pre_dy\", \"pre_dz\",\n", + " \"material\", \"layer_id\"]\n", + "df = pl.read_parquet(ROLLOUT_FILE, columns=cols)\n", + "\n", + "# Rows whose pdg/material fell outside the training vocab can't be encoded\n", + "# (mirrors the pdg_mask filtering in `giant predict`'s CLI path).\n", + "known = df[\"pdg\"].map_elements(lambda p: int(p) in pdg_map, return_dtype=pl.Boolean) & df[\n", + " \"material\"\n", + "].map_elements(lambda m: str(m) in mat_map, return_dtype=pl.Boolean)\n", + "n_dropped = (~known).sum()\n", + "if n_dropped:\n", + " print(f\"dropping {n_dropped}/{len(df)} rows with unknown pdg/material\")\n", + "df = df.filter(known)\n", + "\n", + "data = {\n", + " \"pre_pos\": df.select(\"pre_x\", \"pre_y\", \"pre_z\").to_numpy().astype(np.float32),\n", + " \"pre_E\": df[\"pre_E\"].to_numpy().astype(np.float32),\n", + " \"pre_dir\": df.select(\"pre_dx\", \"pre_dy\", \"pre_dz\").to_numpy().astype(np.float32),\n", + " \"layer_id\": df[\"layer_id\"].to_numpy(),\n", + " \"pdg\": df[\"pdg\"].to_numpy(),\n", + " \"material\": df[\"material\"].to_numpy(),\n", + "}\n", + "cond_cont, cond_cat = build_cond_features(\n", + " data, pdg_map, mat_map, cond_norm, conditioning=conditioning\n", + ")\n", + "cc = torch.from_numpy(cond_cont).float()\n", + "ck = torch.from_numpy(cond_cat).long()\n", + "\n", + "with torch.no_grad():\n", + " gate_weights = model.router.gate(cc, ck).numpy() # (N, n_experts), rows sum to 1\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# EnergyRouter gates on pre-step energy, so that's the natural x-axis here \u2014\n", + "# swap for a categorical plot if this checkpoint used a different router type.\n", + "_ = plot_router_gating(data[\"pre_E\"], gate_weights, x_label=\"pre_E\")\n" ] } ], diff --git a/giant/analysis.py b/giant/analysis.py index 2e1354c..ec2fe40 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -2294,3 +2294,66 @@ def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6): "length traveled share by particle type", max_slices, ) + + +def plot_router_gating( + x: np.ndarray, + gate_weights: np.ndarray, + x_label: str = "energy", + log_x: bool = True, + n_bins: int = 40, + figsize: tuple[float, float] = (7, 4), +): + """Soft mixture-of-experts gate weight vs. a continuous routing axis. + + Unlike everything else in this module, this doesn't stream from a + predict/rollout file — `gate_weights` (N, n_experts, rows already summing + to 1, `giant.model.network.Router.gate`'s contract) has to come from a + live `Router.gate(cond_cont, cond_cat)` call against a loaded checkpoint, + which is a deliberate exception to this module's file-only-diagnostics + design (see the module docstring); that on-the-fly step belongs in the + calling notebook, not here. + + `x` is binned into `n_bins` equal-population (quantile) bins rather than + equal-width ones, since routing axes like energy are usually heavy-tailed + and equal-width bins would leave the upper end almost empty. One line per + expert, mean gate weight per bin stacked as filled areas — since rows of + `gate_weights` are a partition of unity, the stack always fills exactly + to 1, and the visible crossover bands are the router's soft decision + boundaries (where two experts' means cross ~0.5). + """ + x_arr = np.asarray(x) + n_experts = gate_weights.shape[1] + order = np.argsort(x_arr) + x_sorted = x_arr[order] + gw_sorted = gate_weights[order] + + edges = np.quantile(x_sorted, np.linspace(0, 1, n_bins + 1)) + edges[-1] = np.nextafter(edges[-1], np.inf) # include the max value + bin_idx = np.clip(np.digitize(x_sorted, edges[1:-1]), 0, n_bins - 1) + + centers = np.full(n_bins, np.nan) + means = np.full((n_bins, n_experts), np.nan) + for b in range(n_bins): + mask = bin_idx == b + if mask.any(): + centers[b] = x_sorted[mask].mean() + means[b] = gw_sorted[mask].mean(axis=0) + + valid = ~np.isnan(centers) + centers, means = centers[valid], means[valid] + + fig, ax = plt.subplots(figsize=figsize) + cum = np.zeros(len(centers)) + for i in range(n_experts): + ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}") + cum = cum + means[:, i] + if log_x: + ax.set_xscale("log") + ax.set_xlabel(x_label) + ax.set_ylabel("mean gate weight") + ax.set_ylim(0, 1) + ax.set_title("soft router gating") + ax.legend(fontsize=8, ncol=min(n_experts, 4)) + fig.tight_layout() + return fig