lars/giant
1
0
Fork 0
Code Issues 2 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Condition on material/particle physical properties instead of learned embeddings

Browse Source 
Adds model.conditioning = "physical" | "embedding": physical mode routes
particle mass/charge and material Z_eff/A_eff/density/X0/lambda_int through
small MLPs to replace the learned PDG/material embedding tables, so the
surrogate generalizes to PDG codes/materials outside the training vocab
instead of memorizing it. "embedding" stays available as the comparison
baseline (old checkpoints without the key default to it).

Stage 2 now regresses a secondary's mass/charge directly against a fixed
physics-derived target instead of a learned/snapped embedding, and uses no
snapping at inference — the model's raw predicted (mass, charge) is the
secondary's physical identity, including for its own further rollout steps.
A separate reporting-only nearest-known-PDG lookup (never fed back into the
model) populates output pdg columns / the embedding-mode rollout fallback.

giant/materials.py's table is populated with Geant4's own built-in NIST
constants (Z_eff, A_eff, density, X0, lambda_int), extracted directly from
the Geant4 11.4.1 build vendored in minicalosim via G4NistManager rather
than hand-typed literature values. G4_LYSO is left unfilled: confirmed (both
by runtime lookup and by searching minicalosim's history) that it's never
actually a constructed Geant4 material there, only documentation/UI color-map
text.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-17 15:12:54 +02:00
parent a6bb142a40
commit 68fb99bed8
23 changed files with 1252 additions and 304 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/model/network.py
+85 -19
View File
@@ -6,7 +6,16 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM
from giant.constants import (
COND_DIM,
COND_DIM_BASE,
EMB_DIM,
K_MAX,
MATERIAL_PHYS_DIM,
PARTICLE_PHYS_DIM,
SEC_DIM,
X_DIM,
)
class SinusoidalEmbedding(nn.Module):
@@ -28,6 +37,23 @@ class SinusoidalEmbedding(nn.Module):
class ConditionEncoder(nn.Module):
"""Fuses continuous conditioning with particle/material identity.
Two mutually exclusive ways to turn (pdg, material) identity into the
two `emb_dim`-wide vectors concatenated with the base continuous
conditioning before the fusion MLP:
- "embedding": a learned `nn.Embedding` lookup table per axis, indexed
by `cond_cat`'s dense training-vocab index. Memorizes the training
menu; the original Phase-2 design.
- "physical": a small MLP per axis, mapping the axis's raw physical
properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
giant.data.transforms.build_features) to an `emb_dim`-wide vector —
a drop-in replacement for the embedding lookup, computable for any
PDG code / material name rather than only ones seen in training.
Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the
fusion MLP, so only how the two vectors are produced differs.
"""
def __init__(
self,
pdg_vocab: int,
@@ -35,11 +61,27 @@ class ConditionEncoder(nn.Module):
cont_dim: int = COND_DIM,
emb_dim: int = 16,
out_dim: int = 128,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
in_dim = cont_dim + 2 * emb_dim
if conditioning not in ("embedding", "physical"):
raise ValueError(f"unknown conditioning mode {conditioning!r}")
self.conditioning = conditioning
if conditioning == "embedding":
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
else:
self.particle_mlp = nn.Sequential(
nn.Linear(PARTICLE_PHYS_DIM, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
self.material_mlp = nn.Sequential(
nn.Linear(MATERIAL_PHYS_DIM, emb_dim),
nn.SiLU(),
nn.Linear(emb_dim, emb_dim),
)
in_dim = COND_DIM_BASE + 2 * emb_dim
self.mlp = nn.Sequential(
nn.Linear(in_dim, out_dim),
nn.SiLU(),
@@ -47,9 +89,17 @@ class ConditionEncoder(nn.Module):
)
def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
pdg_e = self.pdg_emb(cond_cat[:, 0])
mat_e = self.mat_emb(cond_cat[:, 1])
x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
if self.conditioning == "embedding":
pdg_e = self.pdg_emb(cond_cat[:, 0])
mat_e = self.mat_emb(cond_cat[:, 1])
else:
particle_phys = cond_cont[
:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
]
material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
pdg_e = self.particle_mlp(particle_phys)
mat_e = self.material_mlp(material_phys)
x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
return self.mlp(x)
@@ -91,6 +141,7 @@ class DenoisingMLP(nn.Module):
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
@@ -99,6 +150,7 @@ class DenoisingMLP(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(x_dim, hidden_dim)
@@ -141,10 +193,6 @@ class DenoisingMLP(nn.Module):
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
def pdg_embedding_weight(self) -> torch.Tensor:
"""Return the PDG embedding table weights for secondary type targets."""
return self.cond_enc.pdg_emb.weight
class SecondaryConditionEncoder(nn.Module):
"""Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""
@@ -158,6 +206,7 @@ class SecondaryConditionEncoder(nn.Module):
stage1_dim: int = X_DIM,
stage1_proj_dim: int = 64,
out_dim: int = 128,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.base = ConditionEncoder(
@@ -165,6 +214,7 @@ class SecondaryConditionEncoder(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim)
fused_dim = cond_out_dim + stage1_proj_dim
@@ -187,8 +237,12 @@ class SecondaryConditionEncoder(nn.Module):
class SecondaryDecoder(nn.Module):
"""Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously.
Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one
secondary ordered by descending energy. Padded slots are masked from loss.
Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for
one secondary ordered by descending energy — mass/charge are the
secondary's predicted physical identity, regressed directly against real
physics targets (see giant.data.transforms.encode_secondaries), used
as-is with no snapping to a discrete PDG code. Padded slots are masked
from loss.
"""
def __init__(
@@ -203,6 +257,7 @@ class SecondaryDecoder(nn.Module):
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.time_emb = SinusoidalEmbedding(time_dim)
@@ -213,6 +268,7 @@ class SecondaryDecoder(nn.Module):
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.input_proj = nn.Linear(sec_dim, hidden_dim)
@@ -587,8 +643,8 @@ class RoutedDenoisingMLP(nn.Module):
Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all
tiny) across experts and routes only the trunk (where the FLOPs are).
Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as
`DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes.
Same `forward`/`predict_n_sec` signatures as `DenoisingMLP`, so
sample.py/rollout.py/validate.py need no changes.
"""
def __init__(
@@ -604,6 +660,7 @@ class RoutedDenoisingMLP(nn.Module):
x_dim: int = X_DIM,
dropout: float = 0.1,
k_max: int = K_MAX,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.router = router
@@ -613,6 +670,7 @@ class RoutedDenoisingMLP(nn.Module):
mat_vocab=mat_vocab,
emb_dim=emb_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
@@ -656,10 +714,6 @@ class RoutedDenoisingMLP(nn.Module):
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
def pdg_embedding_weight(self) -> torch.Tensor:
"""Return the PDG embedding table weights for secondary type targets."""
return self.cond_enc.pdg_emb.weight
class RoutedSecondaryDecoder(nn.Module):
"""Routed drop-in for `SecondaryDecoder`.
@@ -682,6 +736,7 @@ class RoutedSecondaryDecoder(nn.Module):
stage1_proj_dim: int = 64,
sec_dim: int = SEC_DIM,
dropout: float = 0.1,
conditioning: str = "embedding",
) -> None:
super().__init__()
self.router = router
@@ -693,6 +748,7 @@ class RoutedSecondaryDecoder(nn.Module):
cond_out_dim=cond_out_dim,
stage1_proj_dim=stage1_proj_dim,
out_dim=cond_out_dim,
conditioning=conditioning,
)
merged_cond_dim = time_dim + cond_out_dim
self.experts = nn.ModuleList(
@@ -732,6 +788,7 @@ _STAGE1_MODEL_KEYS = {
"emb_dim",
"dropout",
"k_max",
"conditioning",
}
_SEC_DECODER_MODEL_KEYS = {
"pdg_vocab",
@@ -740,6 +797,7 @@ _SEC_DECODER_MODEL_KEYS = {
"n_blocks",
"emb_dim",
"dropout",
"conditioning",
}
@@ -798,6 +856,13 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
is truthy; a missing/absent "router" key (pre-routing checkpoints)
falls back to the monolithic pair unchanged, so this is a drop-in
replacement for the ad-hoc constructions it replaces.
`model_config.get("conditioning", "embedding")` — old checkpoints have no
"conditioning" key and must keep loading with their original embedding
tables, so the default here is "embedding", not the training-time
default (which is "physical" — see giant.config.DEFAULT_CONFIG). Read
once and passed to both stage1/sec_decoder, so they structurally always
share one mode.
"""
router_cfg = model_config.get("router")
if router_cfg and router_cfg.get("enabled"):
@@ -810,6 +875,7 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
expert_n_blocks=model_config.get("expert_n_blocks", 3),
emb_dim=model_config.get("emb_dim", EMB_DIM),
dropout=model_config.get("dropout", 0.1),
conditioning=model_config.get("conditioning", "embedding"),
)
stage1 = RoutedDenoisingMLP(
router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),