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:
@@ -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,))
|
||||
|
||||
Reference in New Issue
Block a user