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

v0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing

Browse Source 
AttentionHistory (giant/model/network.py) adds causal self-attention over
the emitted-secondary prefix as the alternative to MarkovHistory, with a
parallel forward() for training and an init_cache()/step() KV-cache path
for sample.py's per-slot AR inference loop, wired into
Stage2Autoregressive via history="attention".

giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled,
mixing ground-truth history with a detached sample_secondaries_ar self-sample
per slot so teacher_forcing="scheduled"/"never" close the train/inference gap
teacher_forcing="always" always avoided; wired into both stage-2 AR trainers.

config.py's validate_config no longer rejects these two previously
unimplemented schema values.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-08-07 13:15:56 +02:00
parent 93b19911f8
commit 200c6d243b
8 changed files with 654 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
giant/config.py
+6 -7
View File
@@ -735,18 +735,17 @@ def validate_config(cfg: dict) -> None:
if _get_path(cfg, "stage2_model.decoder") == "autoregressive":
history = _get_path(cfg, "stage2_model.autoregressive.history")
if history != "markov":
if history not in ("markov", "attention"):
raise ValueError(
f"stage2_model.autoregressive.history = {history!r} is "
"accepted by the schema but not implemented until v0.3.0 "
"step 7 — use 'markov'"
f"stage2_model.autoregressive.history = {history!r} — must "
"be 'markov' or 'attention'"
)
teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing")
if teacher_forcing != "always":
if teacher_forcing not in ("always", "scheduled", "never"):
raise ValueError(
"stage2_model.autoregressive.teacher_forcing = "
f"{teacher_forcing!r} is accepted by the schema but not "
"implemented until v0.3.0 step 7 — use 'always'"
f"{teacher_forcing!r} — must be 'always', 'scheduled' or "
"'never'"
)
giant/model/network.py
+188 -11
View File
@@ -789,9 +789,16 @@ def build_trunk(
class HistoryEncoder(nn.Module):
"""Interface for stage-2 autoregressive per-token history summaries:
`forward(feat, has_prev) -> (B, K, out_dim)`. `MarkovHistory` is the only
implementation until v0.3.0 step 7 (`AttentionHistory`,
`history = "attention"`)."""
`forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
a full (teacher-forced) token sequence — used by training. `MarkovHistory`
and `AttentionHistory` (docs/v0.3.0-design.md §6.2) are the two
implementations. Inference (`giant/sample.py`) generates one token at a
time and cannot afford `forward`'s per-step cost to be O(K) (attention
would then be O(K^2) over a rollout's k_max loop); encoders that need
incremental state for that path additionally implement `init_cache`/
`step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
its per-step cost is already O(1) (it only ever looks at the previous
token, not the full prefix)."""
def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
raise NotImplementedError
@@ -819,6 +826,126 @@ class MarkovHistory(HistoryEncoder):
return self.mlp(x)
class _CausalAttnBlock(nn.Module):
"""One pre-norm causal self-attention block for `AttentionHistory`.
Exposes two forward paths that must agree (see
`test_attention_history_step_matches_forward` in `tests/test_network.py`):
`forward` — the full-sequence, causally-masked pass used for training;
`step` — an incremental pass for inference, given the *pre-attention*
normalized hidden states of every earlier position (`kv_cache`, i.e.
`norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather
than raw `x` is what makes `step` correct: this block's attention needs
exactly that quantity as keys/values, and `LayerNorm` has no cross-position
interaction, so recomputing it per position instead of caching it would
still be correct but pointlessly repeat work. The *next* block's cache is
built from a different sequence (this block's output), so each block owns
an independent cache entry.
"""
def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(
dim, n_heads, dropout=dropout, batch_first=True
)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim)
)
def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
h = self.norm1(x)
attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
x = x + attn_out
x = x + self.mlp(self.norm2(x))
return x
def step(
self, x_new: torch.Tensor, kv_cache: torch.Tensor | None
) -> tuple[torch.Tensor, torch.Tensor]:
"""`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
(first position) or `(B, T, dim)` — `norm1(x)` of every earlier
position at this same block. Returns `(out, new_kv_cache)`, `out`
being this position's block output (`(B, 1, dim)`, to feed the next
block's `step`), `new_kv_cache` the same cache extended by this
position (to reuse at this block's *next* `step` call)."""
h_new = self.norm1(x_new)
kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1)
attn_out, _ = self.attn(h_new, kv, kv, need_weights=False)
x = x_new + attn_out
x = x + self.mlp(self.norm2(x))
return x, kv
class AttentionHistory(HistoryEncoder):
"""Causal self-attention over the emitted-token prefix
(docs/v0.3.0-design.md §6.2) — the more expressive alternative to
`MarkovHistory`'s fixed previous-token-only summary. `feat`/`has_prev`
follow the same shifted-by-one convention `MarkovHistory` and
`Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
own `(energy_fraction, direction, type_representation)`, with a learned
start vector substituted at `has_prev == False` positions (only slot 0 in
practice — see `giant.train._ar_has_prev`). Causal masking then makes
position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens
`0..i-1` — exactly the prefix available when predicting token `i`.
`forward` is the parallel training path (one pass over the whole
teacher-forced sequence); `init_cache`/`step` are the incremental
inference path `giant/sample.py` uses, one new token per call, to avoid
re-encoding the whole prefix from scratch every slot (docs/v0.3.0-design.md
§10's "KV cache" note) — `step` must be called exactly once per slot (its
cache-extension is not idempotent), so a slot's output must be reused for
every model call within that slot (`forward`'s ODE substeps, or a separate
`predict_type` call) rather than re-derived — see
`Stage2Autoregressive.history_step`.
"""
def __init__(
self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2
) -> None:
super().__init__()
self.start = nn.Parameter(torch.zeros(in_dim))
self.in_proj = nn.Linear(in_dim, out_dim)
self.blocks = nn.ModuleList(
[_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)]
)
def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
start = self.start.view(1, 1, -1).expand_as(feat)
x = torch.where(has_prev.unsqueeze(-1), feat, start)
return self.in_proj(x)
def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
B, K, _ = feat.shape
x = self._embed(feat, has_prev)
mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
for block in self.blocks:
x = block(x, mask)
return x
def init_cache(self) -> list[torch.Tensor | None]:
return [None for _ in self.blocks]
def step(
self,
token_feat: torch.Tensor,
has_prev: torch.Tensor,
cache: list[torch.Tensor | None],
) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
"""`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
token's own features (what would be `feat[:, k]` in `forward`).
Advances every block's cache by this position and returns this
position's output (`(B, 1, out_dim)`, the correct history summary for
the NEXT slot) plus the updated cache."""
x = self._embed(token_feat, has_prev)
new_cache: list[torch.Tensor | None] = []
for block, kv in zip(self.blocks, cache):
x, kv_new = block.step(x, kv)
new_cache.append(kv_new)
return x, new_cache
# ---------------------------------------------------------------------------
# Stage models (docs/v0.3.0-design.md §5.3)
# ---------------------------------------------------------------------------
@@ -1082,8 +1209,8 @@ class Stage2OneShot(nn.Module):
class Stage2Autoregressive(nn.Module):
"""Emits secondaries one at a time in descending-energy order
(docs/v0.3.0-design.md §6), instead of `Stage2OneShot`'s simultaneous
k_max-slot prediction. Only `history = "markov"` is implemented (v0.3.0
step 5) — `history = "attention"` raises immediately at construction.
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.
@@ -1122,13 +1249,16 @@ class Stage2Autoregressive(nn.Module):
build_n_sec_head: bool = True,
particle_type_cfg: dict | None = None,
history: str = "markov",
attn_n_heads: int = 4,
attn_n_layers: int = 2,
) -> None:
super().__init__()
if history != "markov":
raise NotImplementedError(
f"stage2_model.autoregressive.history={history!r} is not "
"implemented until v0.3.0 step 7 — use 'markov'"
if history not in ("markov", "attention"):
raise ValueError(
f"stage2_model.autoregressive.history={history!r} — must be "
"'markov' or 'attention'"
)
self.history_kind = history
self.generator_kind = generator
self.noise_dim = noise_dim
self.k_max = k_max
@@ -1149,7 +1279,14 @@ class Stage2Autoregressive(nn.Module):
# 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
self.history_encoder = MarkovHistory(CONT_SLOT_DIM + self.type_dim, history_dim)
hist_in_dim = CONT_SLOT_DIM + self.type_dim
self.history_encoder: HistoryEncoder = (
AttentionHistory(
hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers
)
if history == "attention"
else MarkovHistory(hist_in_dim, history_dim)
)
token_fuse_in = (
cond_out_dim + context_dim + history_dim + 2
) # +2: remaining_frac, slot_idx
@@ -1205,14 +1342,48 @@ class Stage2Autoregressive(nn.Module):
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)
hist = self.history_encoder(history_feat, has_prev)
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): `None` under `history="markov"` (its
per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
`AttentionHistory.init_cache()` under `history="attention"`."""
if isinstance(self.history_encoder, AttentionHistory):
return self.history_encoder.init_cache()
return None
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."""
if isinstance(self.history_encoder, AttentionHistory):
return self.history_encoder.step(token_feat, has_prev, cache)
return self.history_encoder(token_feat, has_prev), cache
def forward(
self,
x_t: torch.Tensor,
@@ -1224,6 +1395,7 @@ class Stage2Autoregressive(nn.Module):
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(
@@ -1234,6 +1406,7 @@ class Stage2Autoregressive(nn.Module):
has_prev,
remaining_frac,
slot_idx,
hist=hist,
)
if self.time_emb is not None:
assert t is not None
@@ -1268,6 +1441,7 @@ class Stage2Autoregressive(nn.Module):
has_prev: torch.Tensor,
remaining_frac: torch.Tensor,
slot_idx: torch.Tensor,
hist: torch.Tensor | None = None,
) -> torch.Tensor:
if self.type_head is None:
raise RuntimeError(
@@ -1285,6 +1459,7 @@ class Stage2Autoregressive(nn.Module):
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)
@@ -1571,6 +1746,8 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
build_n_sec_head=legacy_owner != "stage1",
particle_type_cfg=particle_type_cfg,
history=ar_cfg.get("history", "markov"),
attn_n_heads=ar_cfg.get("attn_n_heads", 4),
attn_n_layers=ar_cfg.get("attn_n_layers", 2),
)
else:
sec_dim = stage2_trunk_sec_dim(
giant/sample.py
+23 -5
View File
@@ -256,17 +256,28 @@ def sample_secondaries_ar(
once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
physics step.
Under `history="attention"` the history encoding is computed once per
slot via `Stage2Autoregressive.history_step` (a KV-cache append, §10)
rather than re-derived by every model call inside that slot — so an ODE
loop's `steps` substeps, and the separate `predict_type` call when the
type slice isn't folded into the trunk output, all reuse the SAME `hist`
tensor for a given `k`. Recomputing per call instead would be merely
wasteful under markov (its per-call cost is already O(1)) but wrong under
attention: `AttentionHistory.step` mutates the cache by appending, so
calling it more than once per slot would double-count that slot's own
(not-yet-existing) predecessor.
The free-running history feature stays UNSNAPPED (mirrors the
established "no snapping" precedent for `particle_type.target =
"physical"` secondaries feeding their own future conditioning):
`"physical"` carries the raw (log_mass, charge) forward as-is;
`"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
one exception — its history slot must be a probability-simplex-shaped
vector (that's what `MarkovHistory` was trained on, `_type_repr`'s
`F.one_hot` ground truth), so it's the hard one-hot of `argmax(logits)`,
not the raw logits themselves. Discretizing further, into a concrete PDG
code, only ever happens once — at secondary-spawn time in
`giant/rollout.py` — never inside this loop.
vector (that's what `MarkovHistory`/`AttentionHistory` were trained on,
`_type_repr`'s `F.one_hot` ground truth), so it's the hard one-hot of
`argmax(logits)`, not the raw logits themselves. Discretizing further,
into a concrete PDG code, only ever happens once — at secondary-spawn
time in `giant/rollout.py` — never inside this loop.
Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
`sample_secondaries`/`sample_secondaries_wgan`'s (see
@@ -291,6 +302,7 @@ def sample_secondaries_ar(
# Running per-token state, threaded from one slot to the next.
prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
remaining = torch.ones(B, device=device)
history_cache = sec_decoder.init_history_cache()
for k in range(k_max):
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
@@ -299,6 +311,9 @@ def sample_secondaries_ar(
slot_idx = torch.full(
(B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
)
hist, history_cache = sec_decoder.history_step(
history_feat, has_prev, history_cache
)
if generator == "wgan":
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -311,6 +326,7 @@ def sample_secondaries_ar(
has_prev,
remaining_frac,
slot_idx,
hist=hist,
)
else:
x = torch.randn(B, 1, token_dim, device=device)
@@ -327,6 +343,7 @@ def sample_secondaries_ar(
remaining_frac,
slot_idx,
t=t,
hist=hist,
)
x = x + v * dt
token = x
@@ -344,6 +361,7 @@ def sample_secondaries_ar(
has_prev,
remaining_frac,
slot_idx,
hist=hist,
).squeeze(1)
sec_cont[:, k] = cont_k
giant/train.py
+195 -14
View File
@@ -27,6 +27,7 @@ from giant.model.schedule import (
flow_matching_loss_secondary_ar,
)
from giant.model.wgan import gradient_penalty, generator_loss
from giant.sample import sample_secondaries_ar
from giant.validate import validate_marginals
_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
@@ -246,6 +247,119 @@ def _assemble_stage2_ar_inputs(
}
def _stage2_tf_prob(
mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
) -> float:
"""P(condition slot k+1 on the TRUE token k rather than the model's own
prediction), for the current epoch (docs/v0.3.0-design.md §3.3
`stage2_model.autoregressive.teacher_forcing`). `"always"`/`"never"` are
the two degenerate constants; `"scheduled"` linearly interpolates
`p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled
sampling (Bengio et al. 2015)."""
if mode == "always":
return 1.0
if mode == "never":
return 0.0
frac = epoch / max(total_epochs - 1, 1)
frac = min(max(frac, 0.0), 1.0)
return p_start + (p_end - p_start) * frac
def _history_repr_from_ar_sample(
sec_cont_pred: torch.Tensor,
sec_type_pred: torch.Tensor,
particle_type_cfg: dict,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""`(fraction, direction, type_repr)` — the same triple `_type_repr` /
`_stick_fraction` derive from ground truth, but from a free-running
`sample_secondaries_ar` self-sample instead, so the two can be mixed
slot-by-slot under scheduled sampling (`_assemble_stage2_ar_inputs_scheduled`).
`target="onehot"` collapses the raw per-slot type logits to a hard
one-hot of `argmax` — `sample_secondaries_ar`'s own history convention
(see its docstring), matching what `MarkovHistory`/`AttentionHistory`
were trained on; the other two targets are already the right
representation."""
fraction = torch.sigmoid(sec_cont_pred[..., 0])
direction = sec_cont_pred[..., 1:4]
if particle_type_cfg.get("target", "physical") == "onehot":
type_dim = sec_type_pred.size(-1)
type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float()
else:
type_repr = sec_type_pred
return fraction, direction, type_repr
def _assemble_stage2_ar_inputs_scheduled(
model: torch.nn.Module,
cond_cont: torch.Tensor,
cond_cat: torch.Tensor,
stage1_ctx: torch.Tensor,
sec_cont: torch.Tensor,
sec_type_idx: torch.Tensor,
n_sec: torch.Tensor,
particle_type_cfg: dict,
cond_enc: torch.nn.Module,
emb_dim: int,
p_tf: float,
sample_steps: int,
) -> dict[str, torch.Tensor]:
"""Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs`
(docs/v0.3.0-design.md §3.3 `teacher_forcing` = "scheduled"/"never"):
each slot's history is the TRUE previous token with probability `p_tf`
(an independent per-example, per-slot Bernoulli draw) and the model's own
free-running prediction otherwise — closing the train/inference gap that
`teacher_forcing="always"` (ground truth throughout training) never sees.
`p_tf >= 1.0` degenerates exactly to `_assemble_stage2_ar_inputs` (and
skips self-sampling entirely), so callers can call this unconditionally.
The free-running estimate is a REAL autoregressive self-sample —
`giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
cheap one-step proxy, so building it costs the same `k_max` (`* steps`
for flow) sequential forwards `sample.py` pays at inference, EVERY batch
this is called on (§6.4's cost note, paid at train time too whenever
teacher_forcing != "always"). Fully detached: gradient only ever flows
through the "real" target path each stage trainer already uses
(`_assemble_stage2_ar_target`), never through this self-sample.
"""
device = sec_cont.device
B, K = sec_cont.shape[0], sec_cont.shape[1]
if p_tf >= 1.0:
return _assemble_stage2_ar_inputs(
sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim
)
was_training = model.training
sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
)
if was_training:
model.train()
fraction_gt = _stick_fraction(sec_cont)
dir_gt = sec_cont[..., 1:4]
type_repr_gt = _type_repr(
sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim
)
fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
sec_cont_pred, sec_type_pred, particle_type_cfg
)
use_gt = torch.rand(B, K, device=device) < p_tf
fraction = torch.where(use_gt, fraction_gt, fraction_pred)
direction = torch.where(use_gt.unsqueeze(-1), dir_gt, dir_pred)
type_repr = torch.where(use_gt.unsqueeze(-1), type_repr_gt, type_repr_pred)
own_feat = torch.cat([fraction.unsqueeze(-1), direction, type_repr], dim=-1)
return {
"history_feat": _shift_prev(own_feat),
"has_prev": _ar_has_prev(K, device).expand(B, -1),
"remaining_frac": _remaining_energy_fraction(fraction),
"slot_idx": (torch.arange(K, device=device).float() / max(K - 1, 1))
.unsqueeze(0)
.expand(B, -1),
}
def _relax_onehot_type_slice(
x_flat: torch.Tensor,
k_max: int,
@@ -362,6 +476,10 @@ class FlowDDPMStageTrainer(StageTrainer):
particle_type_cfg: dict | None = None,
particle_type_emb_dim: int = 16,
decoder: str = "one_shot",
teacher_forcing: str = "always",
tf_p_start: float = 1.0,
tf_p_end: float = 1.0,
ar_sample_steps: int = 10,
) -> None:
if is_stage2 and generator not in ("flow",):
raise NotImplementedError(
@@ -373,6 +491,12 @@ class FlowDDPMStageTrainer(StageTrainer):
self.is_stage2 = is_stage2
self.generator = generator
self.decoder = decoder
self.teacher_forcing = teacher_forcing
self.tf_p_start = tf_p_start
self.tf_p_end = tf_p_end
self.ar_sample_steps = ar_sample_steps
self.total_epochs = epochs
self.steps_per_epoch = max(steps_per_epoch, 1)
self.device = device
self.model = model.to(device)
self.lambda_weight = lambda_weight
@@ -513,7 +637,13 @@ class FlowDDPMStageTrainer(StageTrainer):
l_type = (se * mask).sum() / denom
return l_type, type_acc
def _compute(self, batch: tuple, device: torch.device) -> dict:
def _compute(
self, batch: tuple, device: torch.device, epoch: int | None = None
) -> dict:
"""`epoch=None` (the `val_loss` path) always uses full teacher
forcing (`p_tf=1.0`) regardless of `self.teacher_forcing` — validation
should stay a stable, non-stochastic ground-truth comparison; only
the training `step` path schedules `p_tf` by epoch."""
(
cond_cont,
cond_cat,
@@ -529,12 +659,30 @@ class FlowDDPMStageTrainer(StageTrainer):
x1_s2 = None
ar_inputs = None
if self.is_stage2 and self.decoder == "autoregressive":
ar_inputs = _assemble_stage2_ar_inputs(
p_tf = (
1.0
if epoch is None
else _stage2_tf_prob(
self.teacher_forcing,
self.tf_p_start,
self.tf_p_end,
epoch,
self.total_epochs,
)
)
ar_inputs = _assemble_stage2_ar_inputs_scheduled(
self.model,
cond_cont,
cond_cat,
stage1_ctx,
sec_cont,
sec_type_idx,
n_sec,
self.particle_type_cfg,
self.model.cond_enc,
self.particle_type_emb_dim,
p_tf,
self.ar_sample_steps,
)
x1_s2 = _assemble_stage2_ar_target(
sec_cont,
@@ -613,7 +761,8 @@ class FlowDDPMStageTrainer(StageTrainer):
self.gumbel_tau_start,
self.gumbel_tau_end,
)
out = self._compute(batch, device)
epoch = global_step // self.steps_per_epoch
out = self._compute(batch, device, epoch=epoch)
self.optimizer.zero_grad()
out["total"].backward()
grad_norm = torch.nn.utils.clip_grad_norm_(self.params, 1.0)
@@ -719,10 +868,20 @@ class WGANStageTrainer(StageTrainer):
type_gumbel_tau_start: float = 1.0,
type_gumbel_tau_end: float = 0.1,
decoder: str = "one_shot",
teacher_forcing: str = "always",
tf_p_start: float = 1.0,
tf_p_end: float = 1.0,
ar_sample_steps: int = 10,
) -> None:
self.name = name
self.is_stage2 = is_stage2
self.decoder = decoder
self.teacher_forcing = teacher_forcing
self.tf_p_start = tf_p_start
self.tf_p_end = tf_p_end
self.ar_sample_steps = ar_sample_steps
self.total_epochs = epochs
self.steps_per_epoch = max(steps_per_epoch, 1)
self.device = device
self.model = model.to(device)
self.critic = critic.to(device)
@@ -810,12 +969,27 @@ class WGANStageTrainer(StageTrainer):
return self.critic(x, cond_cont, cond_cat, stage1_ctx)
if self.decoder == "autoregressive":
ar = _assemble_stage2_ar_inputs(
epoch = global_step // self.steps_per_epoch
p_tf = _stage2_tf_prob(
self.teacher_forcing,
self.tf_p_start,
self.tf_p_end,
epoch,
self.total_epochs,
)
ar = _assemble_stage2_ar_inputs_scheduled(
self.model,
cond_cont,
cond_cat,
stage1_ctx,
sec_cont,
sec_type_idx,
n_sec,
self.particle_type_cfg,
self.model.cond_enc,
self.particle_type_emb_dim,
p_tf,
self.ar_sample_steps,
)
real = (
_assemble_stage2_ar_target(
@@ -1051,16 +1225,15 @@ def _build_stage_trainers(
}
particle_type_emb_dim = cfg["conditioning"]["particle"]["emb_dim"]
decoder = stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot"
if is_stage2 and decoder == "autoregressive":
teacher_forcing = (stage_cfg.get("autoregressive") or {}).get(
"teacher_forcing", "always"
)
if teacher_forcing != "always":
raise NotImplementedError(
"stage2_model.autoregressive.teacher_forcing="
f"{teacher_forcing!r} is not implemented until v0.3.0 "
"step 7 — use 'always'"
)
ar_cfg = (stage_cfg.get("autoregressive") or {}) if is_stage2 else {}
teacher_forcing = ar_cfg.get("teacher_forcing", "always")
tf_p_start = ar_cfg.get("tf_p_start", 1.0)
tf_p_end = ar_cfg.get("tf_p_end", 1.0)
# AR self-sampling under scheduled/never teacher forcing reuses
# train.validate_steps as its flow-matching ODE step count — no
# dedicated config key for this (docs/v0.3.0-design.md §3.3 lists
# tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
ar_sample_steps = t.get("validate_steps", 10)
if generator == "wgan":
critic = critics.get(name)
@@ -1090,6 +1263,10 @@ def _build_stage_trainers(
type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
decoder=decoder,
teacher_forcing=teacher_forcing,
tf_p_start=tf_p_start,
tf_p_end=tf_p_end,
ar_sample_steps=ar_sample_steps,
)
else:
ddpm_n_steps = stage_cfg.get("ddpm", {}).get("n_steps", 1000)
@@ -1116,6 +1293,10 @@ def _build_stage_trainers(
particle_type_cfg=particle_type_cfg,
particle_type_emb_dim=particle_type_emb_dim,
decoder=decoder,
teacher_forcing=teacher_forcing,
tf_p_start=tf_p_start,
tf_p_end=tf_p_end,
ar_sample_steps=ar_sample_steps,
)
return trainers
tests/test_config.py
+29 -3
View File
@@ -1,6 +1,8 @@
from datetime import datetime
from pathlib import Path
import pytest
from giant import config as gconfig
_CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"
@@ -645,13 +647,37 @@ def test_validate_config_ar_default_markov_always_passes():
gconfig.validate_config(cfg) # must not raise
def test_validate_config_ar_history_attention_not_implemented():
def test_validate_config_ar_history_attention_passes():
"""v0.3.0 step 7 implements history='attention' — must not raise."""
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.history": "attention",
}
)
gconfig.validate_config(cfg) # must not raise
@pytest.mark.parametrize("teacher_forcing", ["scheduled", "never"])
def test_validate_config_ar_teacher_forcing_scheduled_or_never_passes(teacher_forcing):
"""v0.3.0 step 7 implements teacher_forcing in {'scheduled', 'never'} —
must not raise."""
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.teacher_forcing": teacher_forcing,
}
)
gconfig.validate_config(cfg) # must not raise
def test_validate_config_ar_history_invalid_value_rejected():
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.history": "bogus",
}
)
try:
gconfig.validate_config(cfg)
assert False, "expected ValueError"
@@ -659,11 +685,11 @@ def test_validate_config_ar_history_attention_not_implemented():
assert "history" in str(e)
def test_validate_config_ar_teacher_forcing_scheduled_not_implemented():
def test_validate_config_ar_teacher_forcing_invalid_value_rejected():
cfg = _cfg_with(
**{
"stage2_model.decoder": "autoregressive",
"stage2_model.autoregressive.teacher_forcing": "scheduled",
"stage2_model.autoregressive.teacher_forcing": "bogus",
}
)
try:
tests/test_network.py
+115 -5
View File
@@ -2,6 +2,7 @@ import pytest
import torch
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
from giant.model.network import (
AttentionHistory,
ConditionEncoder,
MarkovHistory,
SinusoidalEmbedding,
@@ -323,6 +324,70 @@ def test_markov_history_uses_start_vector_when_no_prev():
assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
# --- AttentionHistory (docs/v0.3.0-design.md §6.2, v0.3.0 step 7) ----------
def test_attention_history_shape():
hist = AttentionHistory(in_dim=7, out_dim=12, n_heads=2, n_layers=2)
B, K = 3, 5
feat = torch.randn(B, K, 7)
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
out = hist(feat, has_prev)
assert out.shape == (B, K, 12)
def test_attention_history_uses_start_vector_when_no_prev():
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=1)
B, K = 2, 3
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
feat_a = torch.randn(B, K, 4)
feat_b = feat_a.clone()
feat_b[:, 0] = torch.randn(B, 4) * 100
out_a = hist(feat_a, has_prev)
out_b = hist(feat_b, has_prev)
assert torch.allclose(out_a[:, 0], out_b[:, 0], atol=1e-5)
def test_attention_history_is_causal():
"""Position i's output must not depend on feat at positions > i — unlike
MarkovHistory (which only ever looks at position i itself, already
trivially "causal"), this is AttentionHistory's actual contribution:
seeing the full prefix 0..i-1, never anything later."""
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
hist.eval()
B, K = 2, 5
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
feat_a = torch.randn(B, K, 4)
feat_b = feat_a.clone()
feat_b[:, 3:] = torch.randn(B, K - 3, 4) * 100
with torch.no_grad():
out_a = hist(feat_a, has_prev)
out_b = hist(feat_b, has_prev)
assert torch.allclose(out_a[:, :3], out_b[:, :3], atol=1e-5)
def test_attention_history_step_matches_forward():
"""The incremental KV-cache path (`init_cache`/`step`,
`giant/sample.py`'s AR loop) must reproduce `forward`'s parallel-pass
output exactly, one position at a time."""
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
hist.eval()
B, K = 3, 6
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
feat = torch.randn(B, K, 4)
with torch.no_grad():
expected = hist(feat, has_prev)
cache = hist.init_cache()
outs = []
for k in range(K):
out_k, cache = hist.step(feat[:, k : k + 1], has_prev[:, k : k + 1], cache)
outs.append(out_k)
stepped = torch.cat(outs, dim=1)
assert torch.allclose(stepped, expected, atol=1e-5)
# --- Stage2Autoregressive (docs/v0.3.0-design.md §6, v0.3.0 step 5) ---------
@@ -361,16 +426,19 @@ def _ar_inputs(B: int, K: int, hist_dim: int):
return history_feat, has_prev, remaining_frac, slot_idx
def test_stage2_autoregressive_history_attention_raises():
with pytest.raises(NotImplementedError):
_build_stage2_ar("onehot", "wgan", history="attention")
def test_stage2_autoregressive_history_invalid_raises():
with pytest.raises(ValueError):
_build_stage2_ar("onehot", "wgan", history="bogus")
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
@pytest.mark.parametrize("generator", ["wgan", "flow"])
def test_stage2_autoregressive_forward_shape(target, generator):
@pytest.mark.parametrize("history", ["markov", "attention"])
def test_stage2_autoregressive_forward_shape(target, generator, history):
B, K, emb_dim = 4, 5, 6
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
model = _build_stage2_ar(
target, generator, emb_dim=emb_dim, k_max=K, history=history
)
cond_cont = torch.randn(B, COND_DIM)
cond_cat = torch.zeros(B, 2, dtype=torch.long)
stage1_out = torch.randn(B, 9)
@@ -518,3 +586,45 @@ def test_stage2_autoregressive_gradients_flow_onehot():
(flow_out + nsec_out + type_out).backward()
for name, p in model.named_parameters():
assert p.grad is not None, f"no grad for {name}"
def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
"""`init_history_cache`/`history_step` (the incremental path
`giant/sample.py`'s AR loop drives, one slot per call) must reproduce
exactly what one parallel `self.history_encoder(history_feat, has_prev)`
call over the whole shifted sequence would give at each position — the
KV-cache correctness guarantee, exercised through `Stage2Autoregressive`
itself rather than `AttentionHistory` in isolation
(`test_attention_history_step_matches_forward` covers that lower layer)."""
B, K, emb_dim = 3, 6, 6
model = _build_stage2_ar(
"physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
)
model.eval()
type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
hist_in_dim = CONT_SLOT_DIM + type_dim
own_feat = torch.randn(B, K, hist_in_dim) # token i's own raw feature
has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
history_feat = torch.cat(
[torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
)
with torch.no_grad():
expected = model.history_encoder(history_feat, has_prev_full)
cache = model.init_history_cache()
outs = []
prev = torch.zeros(B, 1, hist_in_dim)
for k in range(K):
has_prev_k = torch.full((B, 1), k >= 1, dtype=torch.bool)
hist_k, cache = model.history_step(prev, has_prev_k, cache)
outs.append(hist_k)
prev = own_feat[:, k : k + 1]
stepped = torch.cat(outs, dim=1)
assert torch.allclose(stepped, expected, atol=1e-5)
def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
model = _build_stage2_ar("physical", "wgan", history="markov")
assert model.init_history_cache() is None
tests/test_sample.py
+9 -2
View File
@@ -76,6 +76,7 @@ def _stage2_ar(
pdg: int = 3,
mat: int = 2,
k_max: int = 5,
history: str = "markov",
) -> Stage2Autoregressive:
particle_cfg, material_cfg = _particle_material_cfg(
_conditioning_for(target), emb_dim
@@ -92,6 +93,9 @@ def _stage2_ar(
noise_dim=8,
k_max=k_max,
particle_type_cfg={"target": target},
history=history,
attn_n_heads=2,
attn_n_layers=1,
).eval()
@@ -187,11 +191,14 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
# ── Stage2Autoregressive ─────────────────────────────────────────────────────
@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("generator", ["flow", "wgan"])
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_shapes(target, generator):
def test_sample_secondaries_ar_shapes(target, generator, history):
B, k_max, emb_dim = 4, 5, 6
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
decoder = _stage2_ar(
target, generator, emb_dim=emb_dim, k_max=k_max, history=history
)
cond_cont, cond_cat = _cond(B)
stage1_out = torch.randn(B, X_DIM)
n_sec_pred = torch.randint(0, k_max + 1, (B,))
tests/test_train.py
+89 -6
View File
@@ -2,6 +2,7 @@
import copy
import csv
import math
import tempfile
from pathlib import Path
@@ -29,6 +30,7 @@ from giant.train import (
_relax_onehot_type_slice,
_remaining_energy_fraction,
_shift_prev,
_stage2_tf_prob,
_stick_fraction,
_type_repr,
_wandb_run_config,
@@ -114,6 +116,37 @@ def test_ar_has_prev_false_only_at_slot_zero():
assert has_prev.tolist() == [[False, True, True, True, True]]
# --- _stage2_tf_prob (docs/v0.3.0-design.md §3.3, v0.3.0 step 7) -----------
def test_stage2_tf_prob_always_is_constant_one():
assert _stage2_tf_prob("always", 1.0, 0.0, 0, 10) == 1.0
assert _stage2_tf_prob("always", 1.0, 0.0, 9, 10) == 1.0
def test_stage2_tf_prob_never_is_constant_zero():
assert _stage2_tf_prob("never", 1.0, 1.0, 0, 10) == 0.0
assert _stage2_tf_prob("never", 1.0, 1.0, 9, 10) == 0.0
def test_stage2_tf_prob_scheduled_interpolates_linearly():
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 11) == 1.0
assert abs(_stage2_tf_prob("scheduled", 1.0, 0.0, 5, 11) - 0.5) < 1e-9
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) == 0.0
def test_stage2_tf_prob_scheduled_clamps_beyond_total_epochs():
end = _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11)
beyond = _stage2_tf_prob("scheduled", 1.0, 0.0, 50, 11)
assert beyond == end
def test_stage2_tf_prob_scheduled_handles_single_epoch():
# total_epochs=1 is guarded to a denominator of 1 internally (like
# _gumbel_tau's total_steps=0 guard) — epoch=0 gives zero progress.
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 1) == 1.0
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_type_repr_shapes_and_values(target):
B, K, emb_dim = 3, 4, 6
@@ -592,20 +625,70 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
# --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
def test_build_stage_trainers_rejects_scheduled_teacher_forcing():
@pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
teacher_forcing, history, stage2_generator
):
"""v0.3.0 step 7: history='attention' and teacher_forcing in
{'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
must run and produce a finite loss, for every {history} x
{teacher_forcing} x {generator} combination."""
cfg = _base_cfg()
cfg["stage2_model"]["decoder"] = "autoregressive"
cfg["stage2_model"]["generator"] = stage2_generator
cfg["stage2_model"]["autoregressive"] = {
"history": "markov",
"teacher_forcing": "scheduled",
"history": history,
"teacher_forcing": teacher_forcing,
"tf_p_start": 1.0,
"tf_p_end": 0.0,
"attn_n_heads": 2,
"attn_n_layers": 1,
}
model_config = _model_config(cfg)
models = build_models(model_config)
critics = build_critics(model_config)
with pytest.raises(NotImplementedError):
_build_stage_trainers(
cfg, models, critics, torch.device("cpu"), total_train_batches=4
trainers = _build_stage_trainers(
cfg, models, critics, torch.device("cpu"), total_train_batches=4
)
trainer = trainers["stage2"]
batch = _fake_batches(1, 4)[0]
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
loss_key = "g_loss" if stage2_generator == "wgan" else "loss"
assert math.isfinite(stats[loss_key])
@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
stage2_generator,
):
"""Full `train()` run (not just one `trainer.step()` call) with
history='attention' AND teacher_forcing='scheduled' together — the
combination v0.3.0 step 7 exists to land — must complete and write a
checkpoint + metrics.csv with finite losses throughout."""
cfg = _base_cfg()
cfg["stage2_model"]["decoder"] = "autoregressive"
cfg["stage2_model"]["generator"] = stage2_generator
cfg["stage2_model"]["autoregressive"] = {
"history": "attention",
"teacher_forcing": "scheduled",
"tf_p_start": 1.0,
"tf_p_end": 0.0,
"attn_n_heads": 2,
"attn_n_layers": 1,
}
with tempfile.TemporaryDirectory() as tmp:
out_dir = Path(tmp) / "run"
_run_train(cfg, out_dir)
assert (out_dir / "last.pt").exists()
with open(out_dir / "metrics.csv", newline="") as f:
rows = list(csv.DictReader(f))
assert len(rows) == cfg["train"]["epochs"]
loss_col = (
"stage2_train_g_loss" if stage2_generator == "wgan" else "stage2_train_loss"
)
assert all(math.isfinite(float(r[loss_col])) for r in rows)
def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():