lars
878e9ddca3
CI / Format (ruff format) (push) Failing after 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 33s
CI / Type check (ty) (push) Successful in 37s
CI / Format (ruff format) (pull_request) Failing after 37s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 37s
CI / Tests (pull_request) Successful in 2m49s
CI / Tests (push) Successful in 2m55s
The design doc and its followups doc are no longer needed as a live reference now that the v0.3.0 redesign is implemented — comments and docstrings across the codebase cited it extensively (file path, "design doc §X.Y", "decision N", or bare "§X.Y" section numbers) as design rationale. Removed docs/ and edited every citing comment/docstring to drop the now-dangling reference while keeping the substantive explanation next to it. CLAUDE.md's v0.3.0 roadmap bullet loses its trailing pointer to the deleted file. Verified: no remaining "docs/v0.3.0", "design doc", "decision N", or "§N.N" references (repo-wide grep); ruff and ty clean; full test suite on the heaviest-touched modules (network, sample, rollout, migration, config, train) passes. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1105 lines
36 KiB
Python
1105 lines
36 KiB
Python
"""Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0
|
|
"step 1" commit (eb6dd27), i.e. the last commit before the step-2
|
|
composable-parts decomposition.
|
|
|
|
This is a deliberate verbatim copy, not an import of the live module — the
|
|
whole point is that this file's classes keep behaving exactly as v0.2 did
|
|
even after `giant/model/network.py` itself is rewritten, so
|
|
`tests/test_migration_v02_v03.py` has a stable "old" side to diff the new
|
|
`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical
|
|
acceptance test). Do not edit this file to track future
|
|
`network.py` changes — it exists specifically to stop tracking them.
|
|
"""
|
|
|
|
import inspect
|
|
import math
|
|
import re
|
|
from collections.abc import Sequence
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
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):
|
|
def __init__(self, dim: int) -> None:
|
|
super().__init__()
|
|
assert dim % 2 == 0, "dim must be even"
|
|
half = dim // 2
|
|
freqs = torch.exp(
|
|
-math.log(10000)
|
|
* torch.arange(half, dtype=torch.float32)
|
|
/ max(half - 1, 1)
|
|
)
|
|
self.register_buffer("freqs", freqs)
|
|
|
|
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
|
t = t.reshape(-1, 1).float()
|
|
args = t * self.freqs.unsqueeze(0) # (B, half)
|
|
return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim)
|
|
|
|
|
|
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,
|
|
mat_vocab: int,
|
|
cont_dim: int = COND_DIM,
|
|
emb_dim: int = 16,
|
|
out_dim: int = 128,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
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(),
|
|
nn.Linear(out_dim, out_dim),
|
|
)
|
|
|
|
def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
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)
|
|
|
|
|
|
class ResBlock(nn.Module):
|
|
def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> None:
|
|
super().__init__()
|
|
self.norm = nn.LayerNorm(dim)
|
|
self.linear1 = nn.Linear(dim, dim)
|
|
self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
|
|
self.act = nn.SiLU()
|
|
self.dropout = nn.Dropout(dropout)
|
|
self.linear2 = nn.Linear(dim, dim)
|
|
|
|
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
|
h = self.norm(x)
|
|
h = self.linear1(h) + self.cond_proj(cond)
|
|
h = self.act(h)
|
|
h = self.dropout(h)
|
|
h = self.linear2(h)
|
|
return x + h
|
|
|
|
|
|
class DenoisingMLP(nn.Module):
|
|
"""Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits.
|
|
|
|
The n_sec head runs on the condition encoding only (no diffusion noise),
|
|
so it can be called at inference time independently via `predict_n_sec`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
time_dim: int = 64,
|
|
cond_out_dim: int = 128,
|
|
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)
|
|
self.cond_enc = ConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
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)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(hidden_dim, x_dim)
|
|
# Predicts n_sec as classification over {0, 1, ..., k_max}.
|
|
# Applied to the condition encoding (not the diffused latent).
|
|
self.n_sec_head = nn.Sequential(
|
|
nn.Linear(cond_out_dim, hidden_dim // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_dim // 2, k_max + 1),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x_t: torch.Tensor,
|
|
t: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
t_emb = self.time_emb(t) # (B, time_dim)
|
|
c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim)
|
|
cond = torch.cat([t_emb, c_emb], dim=-1)
|
|
x = self.input_proj(x_t)
|
|
for block in self.blocks:
|
|
x = block(x, cond)
|
|
return self.out_proj(x)
|
|
|
|
def predict_n_sec(
|
|
self,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Return n_sec logits (B, K_MAX+1) from conditioning alone."""
|
|
c_emb = self.cond_enc(cond_cont, cond_cat)
|
|
return self.n_sec_head(c_emb)
|
|
|
|
|
|
class SecondaryConditionEncoder(nn.Module):
|
|
"""Encodes pre-step conditioning + Stage-1 output for the secondary decoder."""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
emb_dim: int = 16,
|
|
cond_out_dim: int = 128,
|
|
stage1_dim: int = X_DIM,
|
|
stage1_proj_dim: int = 64,
|
|
out_dim: int = 128,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.base = ConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
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
|
|
self.fuse = nn.Sequential(
|
|
nn.Linear(fused_dim, out_dim),
|
|
nn.SiLU(),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
base = self.base(cond_cont, cond_cat) # (B, cond_out_dim)
|
|
s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim)
|
|
return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim)
|
|
|
|
|
|
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, 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__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
time_dim: int = 64,
|
|
cond_out_dim: int = 128,
|
|
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)
|
|
self.cond_enc = SecondaryConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
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)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(hidden_dim, sec_dim)
|
|
|
|
def forward(
|
|
self,
|
|
x_t: torch.Tensor,
|
|
t: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
t_emb = self.time_emb(t)
|
|
c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
|
|
cond = torch.cat([t_emb, c_emb], dim=-1)
|
|
x = self.input_proj(x_t)
|
|
for block in self.blocks:
|
|
x = block(x, cond)
|
|
return self.out_proj(x)
|
|
|
|
|
|
class WGANGenerator(nn.Module):
|
|
"""Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time.
|
|
|
|
Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the
|
|
input is a noise vector `z` (not a diffused/interpolated `x_t`) and the
|
|
ResBlocks condition on the condition encoding alone (no time embedding to
|
|
concatenate) — see `giant/model/wgan.py` for the adversarial losses, and
|
|
`giant.sample.sample_wgan` for single-pass sampling.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
cond_out_dim: int = 128,
|
|
x_dim: int = X_DIM,
|
|
noise_dim: int = 64,
|
|
dropout: float = 0.1,
|
|
k_max: int = K_MAX,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.noise_dim = noise_dim
|
|
self.cond_enc = ConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
out_dim=cond_out_dim,
|
|
conditioning=conditioning,
|
|
)
|
|
self.input_proj = nn.Linear(noise_dim, hidden_dim)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(hidden_dim, x_dim)
|
|
self.n_sec_head = nn.Sequential(
|
|
nn.Linear(cond_out_dim, hidden_dim // 2),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_dim // 2, k_max + 1),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
z: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
cond = self.cond_enc(cond_cont, cond_cat)
|
|
x = self.input_proj(z)
|
|
for block in self.blocks:
|
|
x = block(x, cond)
|
|
return self.out_proj(x)
|
|
|
|
def predict_n_sec(
|
|
self,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Return n_sec logits (B, K_MAX+1) from conditioning alone."""
|
|
c_emb = self.cond_enc(cond_cont, cond_cat)
|
|
return self.n_sec_head(c_emb)
|
|
|
|
|
|
class Critic(nn.Module):
|
|
"""Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`.
|
|
|
|
Kept structurally parallel to `WGANGenerator` (own condition encoder —
|
|
separate weights from the generator's, standard GAN practice) but has no
|
|
n_sec head: n_sec is never adversarial, it stays a plain classifier on
|
|
the generator side.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
cond_out_dim: int = 128,
|
|
x_dim: int = X_DIM,
|
|
dropout: float = 0.1,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.cond_enc = ConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
out_dim=cond_out_dim,
|
|
conditioning=conditioning,
|
|
)
|
|
self.input_proj = nn.Linear(x_dim, hidden_dim)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_norm = nn.LayerNorm(hidden_dim)
|
|
self.out_proj = nn.Linear(hidden_dim, 1)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
cond = self.cond_enc(cond_cont, cond_cat)
|
|
h = self.input_proj(x)
|
|
for block in self.blocks:
|
|
h = block(h, cond)
|
|
return self.out_proj(self.out_norm(h)).squeeze(-1)
|
|
|
|
|
|
class WGANSecondaryGenerator(nn.Module):
|
|
"""Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots.
|
|
|
|
Mirrors `SecondaryDecoder` minus the time embedding, the same way
|
|
`WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`,
|
|
conditions on `SecondaryConditionEncoder`'s output alone.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
cond_out_dim: int = 128,
|
|
stage1_proj_dim: int = 64,
|
|
sec_dim: int = SEC_DIM,
|
|
noise_dim: int = 64,
|
|
dropout: float = 0.1,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.noise_dim = noise_dim
|
|
self.cond_enc = SecondaryConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
cond_out_dim=cond_out_dim,
|
|
stage1_proj_dim=stage1_proj_dim,
|
|
out_dim=cond_out_dim,
|
|
conditioning=conditioning,
|
|
)
|
|
self.input_proj = nn.Linear(noise_dim, hidden_dim)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(hidden_dim, sec_dim)
|
|
|
|
def forward(
|
|
self,
|
|
z: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
|
|
x = self.input_proj(z)
|
|
for block in self.blocks:
|
|
x = block(x, cond)
|
|
return self.out_proj(x)
|
|
|
|
|
|
class SecondaryCritic(nn.Module):
|
|
"""Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots."""
|
|
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
hidden_dim: int = 256,
|
|
n_blocks: int = 6,
|
|
emb_dim: int = 16,
|
|
cond_out_dim: int = 128,
|
|
stage1_proj_dim: int = 64,
|
|
sec_dim: int = SEC_DIM,
|
|
dropout: float = 0.1,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.cond_enc = SecondaryConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
cond_out_dim=cond_out_dim,
|
|
stage1_proj_dim=stage1_proj_dim,
|
|
out_dim=cond_out_dim,
|
|
conditioning=conditioning,
|
|
)
|
|
self.input_proj = nn.Linear(sec_dim, hidden_dim)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_norm = nn.LayerNorm(hidden_dim)
|
|
self.out_proj = nn.Linear(hidden_dim, 1)
|
|
|
|
def forward(
|
|
self,
|
|
x: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
cond = self.cond_enc(cond_cont, cond_cat, stage1_out)
|
|
h = self.input_proj(x)
|
|
for block in self.blocks:
|
|
h = block(h, cond)
|
|
return self.out_proj(self.out_norm(h)).squeeze(-1)
|
|
|
|
|
|
class Router(nn.Module):
|
|
"""Contract for a pluggable mixture-of-experts routing axis."""
|
|
|
|
def __init__(self, n_experts: int) -> None:
|
|
super().__init__()
|
|
self.n_experts = n_experts
|
|
self.gumbel = False
|
|
self.gumbel_tau = 1.0
|
|
|
|
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
raise NotImplementedError
|
|
|
|
def combine_weights(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
|
) -> torch.Tensor:
|
|
probs = self.gate(cond_cont, cond_cat)
|
|
if not (self.gumbel and self.training):
|
|
return probs
|
|
log_probs = torch.log(probs.clamp_min(1e-8))
|
|
return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
|
|
|
|
def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
return self.gate(cond_cont, cond_cat).argmax(dim=-1)
|
|
|
|
def balance_loss(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
|
) -> torch.Tensor:
|
|
importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,)
|
|
return (importance.std() / (importance.mean() + 1e-8)) ** 2
|
|
|
|
def classify_loss(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
|
|
) -> torch.Tensor:
|
|
return torch.zeros((), device=cond_cont.device)
|
|
|
|
def entropy_loss(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
|
) -> torch.Tensor:
|
|
norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
|
|
return norm_entropy
|
|
|
|
def gate_stats(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
gate = self.gate(cond_cont, cond_cat) # (B, n_experts)
|
|
row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,)
|
|
norm_entropy = row_entropy.mean() / math.log(self.n_experts)
|
|
importance = gate.sum(dim=0) # (n_experts,)
|
|
return norm_entropy, importance
|
|
|
|
|
|
ROUTER_REGISTRY: dict[str, type[Router]] = {}
|
|
|
|
|
|
def register_router(name: str):
|
|
def decorator(cls: type[Router]) -> type[Router]:
|
|
ROUTER_REGISTRY[name] = cls
|
|
return cls
|
|
|
|
return decorator
|
|
|
|
|
|
def build_router(name: str, n_experts: int, **kwargs) -> Router:
|
|
if name not in ROUTER_REGISTRY:
|
|
raise ValueError(
|
|
f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
|
|
)
|
|
cls = ROUTER_REGISTRY[name]
|
|
accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
|
|
filtered = {k: v for k, v in kwargs.items() if k in accepted}
|
|
return cls(n_experts=n_experts, **filtered)
|
|
|
|
|
|
def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
|
|
return lo + (hi - lo) * torch.sigmoid(raw)
|
|
|
|
|
|
def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
|
|
p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
|
|
return math.log(p / (1 - p))
|
|
|
|
|
|
@register_router("energy")
|
|
class EnergyRouter(Router):
|
|
def __init__(
|
|
self,
|
|
n_experts: int = 4,
|
|
temperature: float = 0.5,
|
|
learn_centers: bool = True,
|
|
energy_idx: int = 3,
|
|
centers_init: Sequence[float] | None = None,
|
|
learn_width: bool = False,
|
|
learn_temperature: bool = False,
|
|
width_min_ratio: float = 0.1,
|
|
width_max_ratio: float = 10.0,
|
|
) -> None:
|
|
super().__init__(n_experts)
|
|
if learn_width and learn_temperature:
|
|
raise ValueError("learn_width and learn_temperature are mutually exclusive")
|
|
self.temperature = temperature
|
|
self.energy_idx = energy_idx
|
|
self.learn_width = learn_width
|
|
self.learn_temperature = learn_temperature
|
|
if learn_width or learn_temperature:
|
|
if not (width_min_ratio < 1.0 < width_max_ratio):
|
|
raise ValueError(
|
|
f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
|
|
f"({width_max_ratio}) must bracket 1.0"
|
|
)
|
|
self._width_lo = width_min_ratio * temperature
|
|
self._width_hi = width_max_ratio * temperature
|
|
raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
|
|
if learn_width:
|
|
self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
|
|
else:
|
|
self.raw_temperature = nn.Parameter(torch.tensor(raw0))
|
|
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:
|
|
self.register_buffer("centers", centers)
|
|
|
|
def effective_width(self) -> torch.Tensor | float:
|
|
if self.learn_width:
|
|
return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
|
|
if self.learn_temperature:
|
|
return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
|
|
return self.temperature
|
|
|
|
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1)
|
|
d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts)
|
|
return torch.softmax(-d2 / self.effective_width(), dim=-1)
|
|
|
|
|
|
@register_router("pdg")
|
|
class PdgRouter(Router):
|
|
def __init__(
|
|
self,
|
|
n_experts: int,
|
|
pdg_vocab: int,
|
|
emb_dim: int = 8,
|
|
temperature: float = 0.5,
|
|
learn_centers: bool = True,
|
|
) -> None:
|
|
super().__init__(n_experts)
|
|
self.temperature = temperature
|
|
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
|
|
centers = torch.randn(n_experts, emb_dim) * 0.1
|
|
if learn_centers:
|
|
self.centers = nn.Parameter(centers)
|
|
else:
|
|
self.register_buffer("centers", centers)
|
|
|
|
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim)
|
|
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
|
|
-1
|
|
) # (B, n_experts)
|
|
return torch.softmax(-d2 / self.temperature, dim=-1)
|
|
|
|
|
|
@register_router("process")
|
|
class ProcessRouter(Router):
|
|
def __init__(
|
|
self,
|
|
n_experts: int,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
emb_dim: int = 8,
|
|
hidden_dim: int = 64,
|
|
) -> None:
|
|
super().__init__(n_experts)
|
|
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
|
|
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
|
|
self.classifier = nn.Sequential(
|
|
nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
|
|
nn.SiLU(),
|
|
nn.Linear(hidden_dim, n_experts),
|
|
)
|
|
|
|
def logits(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])
|
|
h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
|
|
return self.classifier(h)
|
|
|
|
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
|
|
|
|
def classify_loss(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
|
|
) -> torch.Tensor:
|
|
return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
|
|
|
|
|
|
class ComposedRouter(Router):
|
|
def __init__(self, routers: list[Router]) -> None:
|
|
if not routers:
|
|
raise ValueError("ComposedRouter needs at least one sub-router")
|
|
n_experts = 1
|
|
for r in routers:
|
|
n_experts *= r.n_experts
|
|
super().__init__(n_experts)
|
|
self.routers = nn.ModuleList(routers)
|
|
|
|
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
|
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
|
|
for router in self.routers[1:]:
|
|
g = router.gate(cond_cont, cond_cat) # (B, n_i)
|
|
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
|
|
1
|
|
) # (B, prod so far)
|
|
return joint
|
|
|
|
def classify_loss(
|
|
self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
|
|
) -> torch.Tensor:
|
|
total = torch.zeros((), device=cond_cont.device)
|
|
for router in self.routers:
|
|
total = total + router.classify_loss(cond_cont, cond_cat, labels)
|
|
return total
|
|
|
|
|
|
def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
|
|
routers = [
|
|
build_router(
|
|
spec["type"],
|
|
spec["n_experts"],
|
|
**{
|
|
**shared_kwargs,
|
|
**{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
|
|
},
|
|
)
|
|
for spec in specs
|
|
]
|
|
return ComposedRouter(routers)
|
|
|
|
|
|
class ExpertTrunk(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_dim: int,
|
|
hidden_dim: int,
|
|
n_blocks: int,
|
|
merged_cond_dim: int,
|
|
dropout: float = 0.1,
|
|
) -> None:
|
|
super().__init__()
|
|
self.input_proj = nn.Linear(in_dim, hidden_dim)
|
|
self.blocks = nn.ModuleList(
|
|
[
|
|
ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
|
|
for _ in range(n_blocks)
|
|
]
|
|
)
|
|
self.out_proj = nn.Linear(hidden_dim, in_dim)
|
|
|
|
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
|
x = self.input_proj(x)
|
|
for block in self.blocks:
|
|
x = block(x, cond)
|
|
return self.out_proj(x)
|
|
|
|
|
|
def _route_forward(
|
|
experts: nn.ModuleList,
|
|
router: Router,
|
|
x: torch.Tensor,
|
|
cond: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
training: bool,
|
|
) -> torch.Tensor:
|
|
if training:
|
|
weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts)
|
|
out = torch.zeros_like(x)
|
|
for i, expert in enumerate(experts):
|
|
out = out + weights[:, i : i + 1] * expert(x, cond)
|
|
return out
|
|
|
|
idx = router.top1(cond_cont, cond_cat) # (B,)
|
|
out = torch.zeros_like(x)
|
|
for i, expert in enumerate(experts):
|
|
mask = idx == i
|
|
if mask.any():
|
|
out[mask] = expert(x[mask], cond[mask])
|
|
return out
|
|
|
|
|
|
class RoutedDenoisingMLP(nn.Module):
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
router: Router,
|
|
expert_hidden_dim: int = 128,
|
|
expert_n_blocks: int = 3,
|
|
emb_dim: int = EMB_DIM,
|
|
time_dim: int = 64,
|
|
cond_out_dim: int = 128,
|
|
x_dim: int = X_DIM,
|
|
dropout: float = 0.1,
|
|
k_max: int = K_MAX,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.router = router
|
|
self.time_emb = SinusoidalEmbedding(time_dim)
|
|
self.cond_enc = ConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
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(
|
|
[
|
|
ExpertTrunk(
|
|
x_dim,
|
|
expert_hidden_dim,
|
|
expert_n_blocks,
|
|
merged_cond_dim,
|
|
dropout=dropout,
|
|
)
|
|
for _ in range(router.n_experts)
|
|
]
|
|
)
|
|
self.n_sec_head = nn.Sequential(
|
|
nn.Linear(cond_out_dim, cond_out_dim),
|
|
nn.SiLU(),
|
|
nn.Linear(cond_out_dim, k_max + 1),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x_t: torch.Tensor,
|
|
t: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
t_emb = self.time_emb(t)
|
|
c_emb = self.cond_enc(cond_cont, cond_cat)
|
|
cond = torch.cat([t_emb, c_emb], dim=-1)
|
|
return _route_forward(
|
|
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
|
|
)
|
|
|
|
def predict_n_sec(
|
|
self,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
c_emb = self.cond_enc(cond_cont, cond_cat)
|
|
return self.n_sec_head(c_emb)
|
|
|
|
|
|
class RoutedSecondaryDecoder(nn.Module):
|
|
def __init__(
|
|
self,
|
|
pdg_vocab: int,
|
|
mat_vocab: int,
|
|
router: Router,
|
|
expert_hidden_dim: int = 128,
|
|
expert_n_blocks: int = 3,
|
|
emb_dim: int = EMB_DIM,
|
|
time_dim: int = 64,
|
|
cond_out_dim: int = 128,
|
|
stage1_proj_dim: int = 64,
|
|
sec_dim: int = SEC_DIM,
|
|
dropout: float = 0.1,
|
|
conditioning: str = "embedding",
|
|
) -> None:
|
|
super().__init__()
|
|
self.router = router
|
|
self.time_emb = SinusoidalEmbedding(time_dim)
|
|
self.cond_enc = SecondaryConditionEncoder(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
emb_dim=emb_dim,
|
|
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(
|
|
[
|
|
ExpertTrunk(
|
|
sec_dim,
|
|
expert_hidden_dim,
|
|
expert_n_blocks,
|
|
merged_cond_dim,
|
|
dropout=dropout,
|
|
)
|
|
for _ in range(router.n_experts)
|
|
]
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
x_t: torch.Tensor,
|
|
t: torch.Tensor,
|
|
cond_cont: torch.Tensor,
|
|
cond_cat: torch.Tensor,
|
|
stage1_out: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
t_emb = self.time_emb(t)
|
|
c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
|
|
cond = torch.cat([t_emb, c_emb], dim=-1)
|
|
return _route_forward(
|
|
self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
|
|
)
|
|
|
|
|
|
_STAGE1_MODEL_KEYS = {
|
|
"pdg_vocab",
|
|
"mat_vocab",
|
|
"hidden_dim",
|
|
"n_blocks",
|
|
"emb_dim",
|
|
"dropout",
|
|
"k_max",
|
|
"conditioning",
|
|
}
|
|
_SEC_DECODER_MODEL_KEYS = {
|
|
"pdg_vocab",
|
|
"mat_vocab",
|
|
"hidden_dim",
|
|
"n_blocks",
|
|
"emb_dim",
|
|
"dropout",
|
|
"conditioning",
|
|
}
|
|
_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"}
|
|
_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"}
|
|
_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"}
|
|
|
|
|
|
_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
|
|
|
|
|
|
def _parse_composed_axes(router_cfg: dict) -> list[dict]:
|
|
axes: dict[int, dict] = {}
|
|
for key, value in router_cfg.items():
|
|
m = _AXIS_KEY_RE.match(key)
|
|
if m is None:
|
|
continue
|
|
idx, field = int(m.group(1)), m.group(2)
|
|
axes.setdefault(idx, {})[field] = value
|
|
missing = set(range(len(axes))) - axes.keys()
|
|
if missing:
|
|
raise ValueError(f"composed router config has gaps at axis indices {missing}")
|
|
return [axes[i] for i in range(len(axes))]
|
|
|
|
|
|
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
|
|
|
|
|
|
def _check_router_conditioning_compat(
|
|
router_types: list[str], conditioning: str
|
|
) -> None:
|
|
bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
|
|
if bad and conditioning == "physical":
|
|
raise ValueError(
|
|
f"router type(s) {bad} always use a training-vocab PDG embedding, "
|
|
"which is incompatible with conditioning='physical' (whose whole "
|
|
"point is generalizing beyond that vocab) — pick a different "
|
|
"router type (e.g. 'energy') or use conditioning='embedding'."
|
|
)
|
|
|
|
|
|
def _build_router_from_cfg(
|
|
router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
|
|
) -> Router:
|
|
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
|
|
if router_cfg["type"] == "composed":
|
|
axes = _parse_composed_axes(router_cfg)
|
|
_check_router_conditioning_compat([a["type"] for a in axes], conditioning)
|
|
router = build_composed_router(axes, **shared_vocab)
|
|
router.gumbel = bool(router_cfg.get("gumbel", False))
|
|
return router
|
|
_check_router_conditioning_compat([router_cfg["type"]], conditioning)
|
|
router_kwargs = {
|
|
k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
|
|
}
|
|
router_kwargs.setdefault("pdg_vocab", pdg_vocab)
|
|
router_kwargs.setdefault("mat_vocab", mat_vocab)
|
|
router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
|
|
router.gumbel = bool(router_cfg.get("gumbel", False))
|
|
return router
|
|
|
|
|
|
def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
|
|
if model_config.get("mode") == "wgan":
|
|
stage1 = WGANGenerator(
|
|
**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
|
|
)
|
|
sec_decoder = WGANSecondaryGenerator(
|
|
**{
|
|
k: v
|
|
for k, v in model_config.items()
|
|
if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
|
|
}
|
|
)
|
|
return stage1, sec_decoder
|
|
|
|
router_cfg = model_config.get("router")
|
|
if router_cfg and router_cfg.get("enabled"):
|
|
pdg_vocab = model_config["pdg_vocab"]
|
|
mat_vocab = model_config["mat_vocab"]
|
|
shared = dict(
|
|
pdg_vocab=pdg_vocab,
|
|
mat_vocab=mat_vocab,
|
|
expert_hidden_dim=model_config.get("expert_hidden_dim")
|
|
or model_config.get("hidden_dim", 128),
|
|
expert_n_blocks=model_config.get("expert_n_blocks")
|
|
or model_config.get("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"),
|
|
)
|
|
conditioning = shared["conditioning"]
|
|
stage1 = RoutedDenoisingMLP(
|
|
router=_build_router_from_cfg(
|
|
router_cfg, pdg_vocab, mat_vocab, conditioning
|
|
),
|
|
k_max=model_config.get("k_max", K_MAX),
|
|
**shared,
|
|
)
|
|
sec_decoder = RoutedSecondaryDecoder(
|
|
router=_build_router_from_cfg(
|
|
router_cfg, pdg_vocab, mat_vocab, conditioning
|
|
),
|
|
**shared,
|
|
)
|
|
return stage1, sec_decoder
|
|
|
|
stage1 = DenoisingMLP(
|
|
**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
|
|
)
|
|
sec_decoder = SecondaryDecoder(
|
|
**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
|
|
)
|
|
return stage1, sec_decoder
|
|
|
|
|
|
def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
|
|
critic = Critic(
|
|
**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
|
|
)
|
|
sec_critic = SecondaryCritic(
|
|
**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
|
|
)
|
|
return critic, sec_critic
|