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

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:
2026-08-07 13:15:56 +02:00
parent 93b19911f8
commit 200c6d243b
8 changed files with 654 additions and 53 deletions
+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,))