From ee29b9a30394aae3837ee22e89c5a2ee40ce931d Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 27 Jul 2026 13:11:16 +0200 Subject: [PATCH] router: seed EnergyRouter centers from data quantiles instead of a fixed linspace The 2026-07-22 rollout benchmark's router_gating diagnostic showed the 10-expert EnergyRouter's default linspace(-2, 2, n_experts) init assumes a roughly uniform z-normalized energy distribution, leaving experts heavily overlapping instead of partitioning the range. Add an optional centers_init kwarg (backward compatible, defaults to the old linspace) and have giant train estimate it from a reservoir sample of the real energy column, collected during the existing normalizer-fitting pass. --- giant/data/transforms.py | 48 ++++++++++++++++++++++++++++++++++++++++ giant/model/network.py | 24 +++++++++++++++----- giant/pipeline.py | 28 ++++++++++++++++++++++- tests/test_router.py | 36 ++++++++++++++++++++++++++++++ 4 files changed, 130 insertions(+), 6 deletions(-) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index d83148c..a4a5222 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -228,6 +228,54 @@ class _WelfordAccumulator: return norm +class _ReservoirSampler: + """Uniform random sample of a fixed capacity drawn from a data stream. + + Algorithm R (Vitter 1985), vectorized per chunk so it stays cheap over + hundreds of millions of rows: use to get a representative subsample of + a column for a distribution estimate (e.g. quantiles) without + materializing the full column. + + sampler = _ReservoirSampler(capacity=100_000) + for chunk in data: + sampler.update(chunk) + sample = sampler.sample + """ + + def __init__(self, capacity: int, seed: int = 0) -> None: + self.capacity = capacity + self.n_seen = 0 + self._rng = np.random.default_rng(seed) + self._reservoir = np.empty(0, dtype=np.float64) + + def update(self, values: np.ndarray) -> None: + values = np.asarray(values, dtype=np.float64).reshape(-1) + if values.size == 0: + return + n_before = self.n_seen + if n_before < self.capacity: + take = min(values.size, self.capacity - n_before) + self._reservoir = np.concatenate([self._reservoir, values[:take]]) + values = values[take:] + n_before += take + self.n_seen = n_before + values.size + if values.size == 0 or self.capacity == 0: + return + # remaining elements are past the fill phase: element at 1-based + # stream position j replaces a uniformly random reservoir slot with + # probability capacity/j, which yields a uniform sample overall. + positions = n_before + np.arange(1, values.size + 1) + accept = self._rng.random(values.size) < (self.capacity / positions) + accept_idx = np.nonzero(accept)[0] + if accept_idx.size > 0: + slots = self._rng.integers(0, self.capacity, size=accept_idx.size) + self._reservoir[slots] = values[accept_idx] + + @property + def sample(self) -> np.ndarray: + return self._reservoir.astype(np.float32) + + def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray: """World-frame unit vector pointing from pre_pos to post_pos. diff --git a/giant/model/network.py b/giant/model/network.py index 5596a17..5345747 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1,6 +1,7 @@ import inspect import math import re +from collections.abc import Sequence import torch import torch.nn as nn @@ -600,10 +601,14 @@ class EnergyRouter(Router): """Soft turn-on gate over normalized pre-step log-energy. Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or - fixed) 1-D centers, initialized spread across [-2, 2] — roughly the - z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, - differentiable in e; as tau -> 0 this hardens to nearest-center - (Voronoi) selection, which is exactly what `top1` uses at eval. + fixed) 1-D centers. By default initialized spread evenly across + [-2, 2] — an assumed-uniform z-normalized energy range that may not + match the true (often skewed) distribution and can leave experts + overlapping instead of partitioning the range; pass `centers_init` to + seed them from data (e.g. energy quantiles) instead. + `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as + tau -> 0 this hardens to nearest-center (Voronoi) selection, which is + exactly what `top1` uses at eval. """ def __init__( @@ -612,11 +617,20 @@ class EnergyRouter(Router): temperature: float = 0.5, learn_centers: bool = True, energy_idx: int = 3, + centers_init: Sequence[float] | None = None, ) -> None: super().__init__(n_experts) self.temperature = temperature self.energy_idx = energy_idx - centers = torch.linspace(-2.0, 2.0, n_experts) + if centers_init is None: + centers = torch.linspace(-2.0, 2.0, n_experts) + else: + if len(centers_init) != n_experts: + raise ValueError( + f"centers_init has {len(centers_init)} values, " + f"expected n_experts={n_experts}" + ) + centers = torch.tensor(list(centers_init), dtype=torch.float32) if learn_centers: self.centers = nn.Parameter(centers) else: diff --git a/giant/pipeline.py b/giant/pipeline.py index c06a0d0..6518059 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -20,7 +20,7 @@ from giant.data.loader import ( build_index_maps_from_files, build_process_map_from_files, ) -from giant.data.transforms import build_features, _WelfordAccumulator +from giant.data.transforms import build_features, _WelfordAccumulator, _ReservoirSampler from giant.data.dataset import make_event_split, StreamingStepsDataset from giant.model.network import build_models, build_critics from giant.train import train as run_training @@ -83,6 +83,18 @@ def run_train_job( cond_acc = _WelfordAccumulator(COND_DIM) tgt_acc = _WelfordAccumulator(X_DIM) sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM) + # EnergyRouter's default center spread (linspace over [-2, 2]) assumes + # the z-normalized energy column is roughly uniform, which real energy + # spectra rarely are — collect a reservoir sample here (reusing this + # same pass, not a second scan) so centers can instead be seeded from + # actual data quantiles below. + energy_router_active = ( + router_cfg.get("enabled") and router_cfg.get("type") == "energy" + ) + energy_idx = router_cfg.get("energy_idx", 3) + energy_sampler = ( + _ReservoirSampler(capacity=100_000) if energy_router_active else None + ) for path in files: for chunk in iter_file_chunks(path): mask = np.isin(chunk["event_id"], events_arr) @@ -99,6 +111,8 @@ def run_train_job( ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) + if energy_sampler is not None: + energy_sampler.update(cond_cont[:, energy_idx]) sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] sec_phys = sec_cont[:, :, 4:6][sec_valid] if len(sec_phys) > 0: @@ -107,6 +121,18 @@ def run_train_job( tgt_norm = tgt_acc.to_normalizer() sec_phys_norm = sec_phys_acc.to_normalizer() + if energy_sampler is not None and energy_sampler.n_seen > 0: + assert cond_norm.mean is not None and cond_norm.std is not None + normalized_sample = ( + energy_sampler.sample - cond_norm.mean[energy_idx] + ) / cond_norm.std[energy_idx] + quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"]) + centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32) + router_cfg["centers_init"] = centers_init.tolist() + echo( + f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}" + ) + train_ds = StreamingStepsDataset( files=files, split_events=train_events, diff --git a/tests/test_router.py b/tests/test_router.py index 8f3fa2a..62a8c86 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -97,6 +97,42 @@ def test_build_router_ignores_unrecognized_kwargs(): assert router.temperature == 0.3 +def test_energy_router_default_centers_are_linspace(): + router = EnergyRouter(n_experts=4) + torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4)) + + +def test_energy_router_centers_init_overrides_default(): + centers_init = [-1.0, 0.0, 0.5, 3.0] + router = EnergyRouter(n_experts=4, centers_init=centers_init) + torch.testing.assert_close(router.centers, torch.tensor(centers_init)) + + +def test_energy_router_centers_init_wrong_length_raises(): + try: + EnergyRouter(n_experts=4, centers_init=[0.0, 1.0]) + except ValueError: + return + raise AssertionError("expected ValueError for centers_init length mismatch") + + +def test_energy_router_centers_init_respects_learn_centers_flag(): + learned = EnergyRouter( + n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True + ) + fixed = EnergyRouter( + n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False + ) + assert isinstance(learned.centers, torch.nn.Parameter) + assert not isinstance(fixed.centers, torch.nn.Parameter) + + +def test_build_router_threads_centers_init_through_energy_router(): + centers_init = [-1.5, -0.5, 0.5, 1.5] + router = build_router("energy", 4, centers_init=centers_init) + torch.testing.assert_close(router.centers, torch.tensor(centers_init)) + + def test_build_router_unknown_type_raises(): try: build_router("nonexistent", 4)