lars/giant
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
d0cbcbce802eb3fdf93023550cc9922692eacb82
giant/giant/model/models.py
T
lars c1c4957e2f Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40) 
Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:

- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
  P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
  (n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
  stop head instead of a pre-resolved count: each row stops the first slot
  its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
  default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
  batch loop breaks once every row has stopped, so cost scales with the
  realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
  (the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
  rollout.py/cli.py/validate.py now derive the realized count from
  sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
  rather than resolving it up front — a no-op reordering under every other
  n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
  the per-slot target/mask (one slot wider than the existing token-content
  sec_mask, since the stop slot itself needs supervision) and
  StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
  exactly like _n_sec_loss gates on n_sec_head. Wired into both the
  flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
  stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
  head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
  and n_sec.owner="stage2") instead of always rejecting it.

Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 10:44:14 +02:00

759 lines
32 KiB
Python
Raw Blame History

"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`,
`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""
import torch
import torch.nn as nn
from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
from giant.model.encoders import ConditionEncoder
from giant.model.history import HistoryEncoder, build_history
from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head
from giant.model.objectives import build_objective
from giant.model.routers import Router
from giant.model.trunks import build_trunk
# ---------------------------------------------------------------------------
# Stage models
# ---------------------------------------------------------------------------
def resolve_type_n_classes(particle_type_cfg: ParticleTypeConfig, particle_emb_dim: int) -> int:
"""Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
place of a bare `conditioning.particle.emb_dim` read. Under
`target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
the secondary-species vocabulary size from the unrelated
physical-conditioning MLP's output width. Under `target = "embedding"`
(or `"physical"`, which ignores this value entirely) `n_classes` doesn't
apply — the width stays `conditioning.particle.emb_dim`, the embedding
table's own dimensionality (`validate_config` requires
`conditioning.particle.type = "embedding"` here)."""
if particle_type_cfg.target == "onehot":
return particle_type_cfg.n_classes or particle_emb_dim
return particle_emb_dim
def stage2_type_dim(particle_type_cfg: ParticleTypeConfig, emb_dim: int) -> int:
"""Width of a single secondary slot's type slice —
`PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
`emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
this many classes/dims wide — callers resolve `emb_dim` via
`resolve_type_n_classes` first)."""
return PARTICLE_PHYS_DIM if particle_type_cfg.target == "physical" else emb_dim
def stage2_trunk_sec_dim(particle_type_cfg: ParticleTypeConfig, generator: str, k_max: int, emb_dim: int) -> int:
"""`Stage2OneShot`'s trunk output width.
`target = "physical"` is untouched from v0.2/today:
`k_max * SEC_SLOT_DIM`, the type slice folded into the same
flow-matched/WGAN vector as the continuous stick/dir slots.
`target` in `("onehot", "embedding")`: under an objective with
`folds_type_slice` (currently just wgan) the type slice is still folded
in (adversarial for onehot via ST-Gumbel, already-continuous for
embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide:
`k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice
isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job
instead — so the trunk only covers `k_max * CONT_SLOT_DIM`.
"""
if particle_type_cfg.target == "physical":
return k_max * SEC_SLOT_DIM
if build_objective(generator).folds_type_slice:
return k_max * (CONT_SLOT_DIM + emb_dim)
return k_max * CONT_SLOT_DIM
class StageModel(nn.Module):
"""Base owning the scaffolding common to `Stage1Model`, `Stage2OneShot`,
`Stage2Autoregressive` (gitea #39): build-or-share `cond_enc`,
`particle_type_cfg` normalisation, and — via `_build_trunk_and_heads`,
called by each subclass's `__init__` once its own conditioning-assembly
modules exist — the objective/time-embedding/trunk construction and the
`n_sec_head`/`type_head` classifier heads. A subclass supplies only its
own conditioning assembly (`Stage1Model` uses `cond_enc` directly;
`Stage2OneShot`/`Stage2Autoregressive` add a context-fusion path) and its
trunk's output width.
`cond_enc`, if given, is used in place of building a fresh
`ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
constructs one shared instance and passes it to both stages, halving the
conditioning parameter count and forcing a common representation."""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
cond_out_dim: int,
generator: str,
noise_dim: int,
k_max: int | None = None,
particle_type_cfg: ParticleTypeConfig | None = None,
cond_enc: ConditionEncoder | None = None,
) -> None:
super().__init__()
self.generator_kind = generator
self.noise_dim = noise_dim
self.k_max = k_max
# `ParticleTypeConfig()`'s own dataclass default is target="onehot"
# (the config.toml default when [stage2_model.particle_type] is
# omitted) — a different question from "nobody passed anything to
# this constructor", which direct/test construction relies on
# defaulting to "physical" (build_models/build_critics always pass
# particle_type_cfg explicitly, so this sentinel is never hit there).
self.particle_type_cfg = (
particle_type_cfg if particle_type_cfg is not None else ParticleTypeConfig(target="physical")
)
self.type_dim = stage2_type_dim(
self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg.emb_dim)
)
self.cond_enc = (
cond_enc
if cond_enc is not None
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
)
def _build_trunk_and_heads(
self,
*,
trunk_out_dim: int,
hidden_dim: int,
n_res_blocks: int,
cond_out_dim: int,
time_dim: int,
router: Router | None,
trunk_type: str,
block_conditioning: str,
dropout: float,
n_sec_head_k_max: int | None,
n_sec_head_cfg: dict | None,
type_head_out_dim: int | None,
type_head_cfg: dict | None,
build_stop_head: bool = False,
stop_head_cfg: dict | None = None,
) -> None:
"""Builds `self.time_emb`, `self.trunk`, `self.n_sec_head`,
`self.type_head`, `self.stop_head`. Called by a subclass's `__init__`
after it has set up its own conditioning-assembly modules —
`merged_cond_dim` below must match the width that assembly
(`_cond_embed`/`_base_cond`/`_token_cond`, or plain `cond_enc` for
`Stage1Model`) actually produces.
`n_sec_head` is built iff `n_sec_head_k_max is not None` (output
width `n_sec_head_k_max + 1`) — `Stage1Model` passes this only for a
migrated v0.2 checkpoint, `Stage2OneShot`/`Stage2Autoregressive` pass
it whenever `build_n_sec_head=True`. `type_head` is built iff
`type_head_out_dim is not None` (the caller — only the two Stage2
classes — passes `None` exactly when `particle_type_cfg.target ==
"physical"`) *and* the objective doesn't fold the type slice into its
own trunk output (checked here, since `objective` is already needed
for the trunk itself). `stop_head` is built iff `build_stop_head` —
only `Stage2Autoregressive` ever passes `True` (`n_sec.mode ==
"stop_token"`, mutually exclusive with `n_sec_head`), a single
`cond_out_dim -> 1` logit per call, same `HeadConfig` shape rules as
the other two heads.
"""
objective = build_objective(self.generator_kind)
has_time = objective.needs_time
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
in_dim = objective.trunk_in_dim(trunk_out_dim, self.noise_dim)
self.trunk = build_trunk(
router,
trunk_type,
in_dim,
trunk_out_dim,
hidden_dim,
n_res_blocks,
merged_cond_dim,
dropout,
block_conditioning,
)
self.n_sec_head = None
if n_sec_head_k_max is not None:
head_cfg = HeadConfig.from_dict(n_sec_head_cfg)
hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
self.n_sec_head = build_mlp_head(cond_out_dim, n_sec_head_k_max + 1, hidden, head_cfg.depth)
self.type_head = None
if type_head_out_dim is not None and not objective.folds_type_slice:
head_cfg = HeadConfig.from_dict(type_head_cfg)
hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
self.type_head = build_mlp_head(cond_out_dim, type_head_out_dim, hidden, head_cfg.depth)
self.stop_head = None
if build_stop_head:
head_cfg = HeadConfig.from_dict(stop_head_cfg)
hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
def _require_n_sec_head(self) -> None:
if self.n_sec_head is None:
raise RuntimeError(
f"this {type(self).__name__} has no n_sec_head — it belongs to "
"a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
"stage1.predict_n_sec(cond_cont, cond_cat) instead"
)
def _require_type_head(self) -> None:
if self.type_head is None:
raise RuntimeError(
f"this {type(self).__name__} has no type_head — either "
"particle_type.target='physical' (the type slice is part of "
"forward()'s own output) or generator='wgan' (the WGAN "
"trainer reads the type slice out of forward()'s output "
"directly instead)"
)
def _require_stop_head(self) -> None:
if self.stop_head is None:
raise RuntimeError(
f"this {type(self).__name__} has no stop_head — only a "
"Stage2Autoregressive built with stage2_model.n_sec.mode = "
"'stop_token' owns one"
)
class Stage1Model(StageModel):
"""Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
move it to stage 2, except for a migrated v0.2 checkpoint
(`n_sec_head_k_max` given), where it stays attached here
since that's where its weights live and what conditioning it was trained
against (see `_migrate_legacy_model_config`)."""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
hidden_dim: int = 256,
n_res_blocks: int = 6,
cond_out_dim: int = 128,
x_dim: int = X_DIM,
dropout: float = 0.0,
generator: str = "flow",
time_dim: int = 64,
noise_dim: int = 64,
router: Router | None = None,
trunk_type: str = "resmlp",
block_conditioning: str = "add",
n_sec_head_k_max: int | None = None,
cond_enc: ConditionEncoder | None = None,
n_sec_head_cfg: dict | None = None,
) -> None:
super().__init__(
pdg_vocab,
mat_vocab,
particle_cfg,
material_cfg,
cond_out_dim=cond_out_dim,
generator=generator,
noise_dim=noise_dim,
cond_enc=cond_enc,
)
self._build_trunk_and_heads(
trunk_out_dim=x_dim,
hidden_dim=hidden_dim,
n_res_blocks=n_res_blocks,
cond_out_dim=cond_out_dim,
time_dim=time_dim,
router=router,
trunk_type=trunk_type,
block_conditioning=block_conditioning,
dropout=dropout,
n_sec_head_k_max=n_sec_head_k_max,
n_sec_head_cfg=n_sec_head_cfg,
type_head_out_dim=None,
type_head_cfg=None,
)
def forward(
self,
x_t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
t: torch.Tensor | None = None,
) -> torch.Tensor:
c_emb = self.cond_enc(cond_cont, cond_cat)
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
return self.trunk(x_t, cond, cond_cont, cond_cat)
def _require_n_sec_head(self) -> None:
"""Overrides `StageModel`'s guard — a `Stage1Model` with no
`n_sec_head` points the caller to stage 2 (n_sec's default owner),
not to `stage1` as the base's message would."""
if self.n_sec_head is None:
raise RuntimeError(
"this Stage1Model has no n_sec_head — n_sec now lives on "
"stage 2 by default; this method only exists "
"for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
)
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. Only
valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
configs predict n_sec from Stage2OneShot instead."""
self._require_n_sec_head()
assert self.n_sec_head is not None
c_emb = self.cond_enc(cond_cont, cond_cat)
return self.n_sec_head(c_emb)
class Stage2OneShot(StageModel):
"""Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
step 4/5, not implemented yet).
Owns `n_sec_head` by default unless `build_n_sec_head=False`
(a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
Stage1Model — see `_migrate_legacy_model_config`).
`particle_type_cfg.target` (default `"physical"`) selects the
secondary-type mechanism: `"physical"` keeps the type slice folded into
the trunk's own
flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
the caller via `stage2_trunk_sec_dim` — already reflects this). Under
`"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that
doesn't fold the type slice (flow/ddpm), the type
slice is predicted by a separate `type_head` instead (same shape pattern
as `n_sec_head`) — `sec_dim` then covers only the continuous
stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
Under a folding objective (wgan) the type slice stays folded into `sec_dim`
(just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
"""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
hidden_dim: int = 256,
n_res_blocks: int = 6,
cond_out_dim: int = 128,
context_dim: int = 64,
sec_dim: int = SEC_DIM,
x_dim: int = X_DIM,
dropout: float = 0.0,
generator: str = "wgan",
time_dim: int = 64,
noise_dim: int = 64,
k_max: int = K_MAX,
router: Router | None = None,
trunk_type: str = "resmlp",
block_conditioning: str = "add",
build_n_sec_head: bool = True,
particle_type_cfg: ParticleTypeConfig | None = None,
cond_enc: ConditionEncoder | None = None,
n_sec_head_cfg: dict | None = None,
type_head_cfg: dict | None = None,
) -> None:
super().__init__(
pdg_vocab,
mat_vocab,
particle_cfg,
material_cfg,
cond_out_dim=cond_out_dim,
generator=generator,
noise_dim=noise_dim,
k_max=k_max,
particle_type_cfg=particle_type_cfg,
cond_enc=cond_enc,
)
self.context_adapter = ContextAdapter(x_dim, context_dim)
self.fuse = nn.Sequential(
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
nn.SiLU(),
)
target = self.particle_type_cfg.target
type_head_out_dim = None if target == "physical" else k_max * self.type_dim
self._build_trunk_and_heads(
trunk_out_dim=sec_dim,
hidden_dim=hidden_dim,
n_res_blocks=n_res_blocks,
cond_out_dim=cond_out_dim,
time_dim=time_dim,
router=router,
trunk_type=trunk_type,
block_conditioning=block_conditioning,
dropout=dropout,
n_sec_head_k_max=k_max if build_n_sec_head else None,
n_sec_head_cfg=n_sec_head_cfg,
type_head_out_dim=type_head_out_dim,
type_head_cfg=type_head_cfg,
)
def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
base = self.cond_enc(cond_cont, cond_cat)
ctx = self.context_adapter(stage1_out)
return self.fuse(torch.cat([base, ctx], dim=-1))
def forward(
self,
x_t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
t: torch.Tensor | None = None,
) -> torch.Tensor:
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
return self.trunk(x_t, cond, cond_cont, cond_cat)
def predict_n_sec(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
) -> torch.Tensor:
self._require_n_sec_head()
assert self.n_sec_head is not None
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
return self.n_sec_head(c_emb)
def predict_type(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
) -> torch.Tensor:
"""`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
vectors (`target="embedding"`) — only under `generator in ("flow",
"ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
own output instead (see class docstring)."""
self._require_type_head()
assert self.type_head is not None
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
return self.type_head(c_emb).view(-1, self.k_max, self.type_dim)
class Stage2Autoregressive(StageModel):
"""Emits secondaries one at a time in descending-energy order, instead
of `Stage2OneShot`'s simultaneous
k_max-slot prediction. `history` selects `MarkovHistory` or
`AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
`teacher_forcing` handling lives entirely in the trainer
(`giant/train.py`), since it only affects how training inputs are
assembled, not this module's architecture.
Under teacher forcing every token's conditioning is built from ground
truth, so a whole K-token sequence trains in one parallel batched pass:
`forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
to `k_max`) — this also means a future one-token-at-a-time inference loop
(`K=1` per call, step 6) needs no interface change here.
Two independent conditioning paths, mirroring `Stage2OneShot`'s
`_cond_embed` but split in two: `_base_cond` (`cond_enc` +
`context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
on token position; `_token_cond` additionally fuses in the history
encoding and two running scalars (remaining energy-budget fraction,
normalized slot index), and feeds `forward`/`predict_type`/`predict_stop`/
the trunk.
`n_sec.mode = "stop_token"` (`build_stop_head=True`) replaces
`predict_n_sec`'s one-shot classifier with `predict_stop`'s per-token EOS
logit instead — the two heads are mutually exclusive (`build_n_sec_head`
is `False` whenever this is `True`, see `giant.model.builders`).
"""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
hidden_dim: int = 256,
n_res_blocks: int = 6,
cond_out_dim: int = 128,
context_dim: int = 64,
x_dim: int = X_DIM,
dropout: float = 0.0,
generator: str = "wgan",
time_dim: int = 64,
noise_dim: int = 64,
k_max: int = K_MAX,
router: Router | None = None,
trunk_type: str = "resmlp",
block_conditioning: str = "add",
build_n_sec_head: bool = True,
particle_type_cfg: ParticleTypeConfig | None = None,
history: str = "markov",
attn_n_heads: int = 4,
attn_n_layers: int = 2,
cond_enc: ConditionEncoder | None = None,
n_sec_head_cfg: dict | None = None,
type_head_cfg: dict | None = None,
build_stop_head: bool = False,
stop_sampling: str = "greedy",
stop_head_cfg: dict | None = None,
) -> None:
super().__init__(
pdg_vocab,
mat_vocab,
particle_cfg,
material_cfg,
cond_out_dim=cond_out_dim,
generator=generator,
noise_dim=noise_dim,
k_max=k_max,
particle_type_cfg=particle_type_cfg,
cond_enc=cond_enc,
)
self.history_kind = history
self.stop_sampling = stop_sampling
self.context_adapter = ContextAdapter(x_dim, context_dim)
self.base_fuse = nn.Sequential(
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
nn.SiLU(),
)
# Reuses conditioning.out_dim for the history encoder's own output
# width — there's no dedicated stage2_model.autoregressive key for
# this, a reasonable default rather than a design-doc-specified value.
history_dim = cond_out_dim
hist_in_dim = CONT_SLOT_DIM + self.type_dim
self.history_encoder: HistoryEncoder = build_history(
history, hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
)
token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx
self.token_fuse = nn.Sequential(
nn.Linear(token_fuse_in, cond_out_dim),
nn.SiLU(),
)
# `self.type_dim` (set by StageModel.__init__) doubles as the raw
# `emb_dim` `stage2_trunk_sec_dim` wants: for a non-"physical" target
# `stage2_type_dim` already resolved `type_dim` to exactly that value;
# for "physical" the emb_dim argument goes unused anyway.
token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, self.type_dim)
target = self.particle_type_cfg.target
type_head_out_dim = None if target == "physical" else self.type_dim
self._build_trunk_and_heads(
trunk_out_dim=token_dim,
hidden_dim=hidden_dim,
n_res_blocks=n_res_blocks,
cond_out_dim=cond_out_dim,
time_dim=time_dim,
router=router,
trunk_type=trunk_type,
block_conditioning=block_conditioning,
dropout=dropout,
n_sec_head_k_max=k_max if build_n_sec_head else None,
n_sec_head_cfg=n_sec_head_cfg,
type_head_out_dim=type_head_out_dim,
type_head_cfg=type_head_cfg,
build_stop_head=build_stop_head,
stop_head_cfg=stop_head_cfg,
)
def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
base = self.cond_enc(cond_cont, cond_cat)
ctx = self.context_adapter(stage1_out)
return self.base_fuse(torch.cat([base, ctx], dim=-1))
def _token_cond(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
history_feat: torch.Tensor,
has_prev: torch.Tensor,
remaining_frac: torch.Tensor,
slot_idx: torch.Tensor,
hist: torch.Tensor | None = None,
) -> torch.Tensor:
"""`hist`, if given, overrides recomputing `self.history_encoder`
from `history_feat`/`has_prev` — the inference-time KV-cache path
(`Stage2Autoregressive.history_step`) precomputes it once per slot and
passes it in here so a slot's (possibly several) model calls — an ODE
loop's substeps, or a separate `predict_type` call — read the same
cached history instead of each re-deriving (and, under attention,
re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
K = history_feat.size(1)
base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
if hist is None:
hist = self.history_encoder(history_feat, has_prev)
scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
def init_history_cache(self):
"""Inference-only incremental-decoding state for `self.history_encoder`
(`giant/sample.py`'s AR loop) — whatever `self.history_encoder.init_cache()`
returns for the configured `history` type: `None` under `history="markov"`
(its per-step cost is already O(1) — see `HistoryEncoder`'s docstring),
or `AttentionHistory.init_cache()`'s real per-block KV cache under
`history="attention"`."""
return self.history_encoder.init_cache()
def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
"""One inference slot's worth of history encoding: advances `cache`
(from `init_history_cache`, or a previous `history_step` call) by
`token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
token, same convention `giant.sample.sample_secondaries_ar` already
threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
this slot's history summary (pass it as `_token_cond`'s `hist=` to
every model call made for this slot), `new_cache` is what to pass into
the *next* slot's `history_step`. Must be called exactly once per
slot — see `AttentionHistory.step`'s docstring."""
return self.history_encoder.step(token_feat, has_prev, cache)
def forward(
self,
x_t: torch.Tensor,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
history_feat: torch.Tensor,
has_prev: torch.Tensor,
remaining_frac: torch.Tensor,
slot_idx: torch.Tensor,
t: torch.Tensor | None = None,
hist: torch.Tensor | None = None,
) -> torch.Tensor:
B, K = x_t.shape[0], x_t.shape[1]
c_emb = self._token_cond(
cond_cont,
cond_cat,
stage1_out,
history_feat,
has_prev,
remaining_frac,
slot_idx,
hist=hist,
)
if self.time_emb is not None:
assert t is not None
t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
cond = torch.cat([t_emb, c_emb], dim=-1)
else:
cond = c_emb
x_flat = x_t.reshape(B * K, -1)
cond_flat = cond.reshape(B * K, -1)
cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
return out.view(B, K, -1)
def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
self._require_n_sec_head()
assert self.n_sec_head is not None
return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
def predict_type(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
history_feat: torch.Tensor,
has_prev: torch.Tensor,
remaining_frac: torch.Tensor,
slot_idx: torch.Tensor,
hist: torch.Tensor | None = None,
) -> torch.Tensor:
self._require_type_head()
assert self.type_head is not None
c_emb = self._token_cond(
cond_cont,
cond_cat,
stage1_out,
history_feat,
has_prev,
remaining_frac,
slot_idx,
hist=hist,
)
B, K, _ = c_emb.shape
return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
def predict_stop(
self,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_out: torch.Tensor,
history_feat: torch.Tensor,
has_prev: torch.Tensor,
remaining_frac: torch.Tensor,
slot_idx: torch.Tensor,
hist: torch.Tensor | None = None,
) -> torch.Tensor:
"""`(B, K)` raw stop logits — `n_sec.mode = "stop_token"` only.
Evaluated on slot `k`'s own conditioning (which carries slot `k-1`'s
history, same as `predict_type`), so this is `P(n_sec == k |
prefix)`: a high logit at slot `k` means "stop before generating a
token here" — the caller (`giant.sample.sample_secondaries_ar`)
checks it before spending a model call on that slot's token."""
self._require_stop_head()
assert self.stop_head is not None
c_emb = self._token_cond(
cond_cont,
cond_cat,
stage1_out,
history_feat,
has_prev,
remaining_frac,
slot_idx,
hist=hist,
)
B, K, _ = c_emb.shape
return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
class CriticModel(nn.Module):
"""Generator-agnostic WGAN-GP critic body: a scalar realism score, for
either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
`Stage2OneShot`). Used only when that stage's `generator == "wgan"`."""
def __init__(
self,
pdg_vocab: int,
mat_vocab: int,
particle_cfg: ConditioningAxisConfig,
material_cfg: ConditioningAxisConfig,
in_dim: int,
hidden_dim: int = 256,
n_res_blocks: int = 6,
cond_out_dim: int = 128,
dropout: float = 0.0,
stage: str = "stage1",
context_dim: int = 64,
context_in_dim: int = X_DIM,
) -> None:
super().__init__()
if stage not in ("stage1", "stage2"):
raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
self.stage = stage
self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
if stage == "stage2":
self.context_adapter = ContextAdapter(context_in_dim, context_dim)
self.fuse = nn.Sequential(
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
nn.SiLU(),
)
self.input_proj = nn.Linear(in_dim, hidden_dim)
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_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 | None = None,
) -> torch.Tensor:
base = self.cond_enc(cond_cont, cond_cat)
if self.stage == "stage2":
ctx = self.context_adapter(stage1_out)
cond = self.fuse(torch.cat([base, ctx], dim=-1))
else:
cond = base
h = self.input_proj(x)
for block in self.blocks:
h = block(h, cond)
return self.out_proj(self.out_norm(h)).squeeze(-1)
Reference in New Issue View Git Blame Copy Permalink