| """MiniJEV Core: the domain-agnostic decision network. |
| |
| It only sees vectors in a shared latent space of width `d`: |
| |
| state_tokens (B, S, d) + state_mask (B, S) from any StateEncoder (grid, text, numeric, ...) |
| question_tokens (B, Q, d) + question_mask (B, Q) from the question encoder |
| choice_emb (B, N, d) + choice_mask (B, N) from the ChoiceEncoder (N is dynamic) |
| |
| and returns one logit per choice (+ value, scalar score and the CLS embedding). |
| It knows nothing about Snake, text, bytes or tokenizers. |
| |
| [CLS] + question tokens + state tokens ──► Transformer ──► context H |
| │ |
| scorer "cross": choices cross-attend H, then self-attend each other ─► Linear ─► logit_i |
| scorer "bi": logit_i = MLP([h, c_i, h*c_i]) (light cross scorer on pooled h = CLS) |
| scorer "fixed_head": Linear(CLS) -> max_choices logits (ablation: ignores choice meaning) |
| |
| "cross" and "bi" are permutation-equivariant over choices (no positional information). |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| from torch import nn |
|
|
| from .layers import DecisionBlock, EncoderBlock |
|
|
| NEG = -1e9 |
| SCORERS = ("cross", "bi", "fixed_head") |
|
|
|
|
| class MiniJEVCore(nn.Module): |
| def __init__(self, d: int, n_layers: int, n_heads: int = 4, mlp_ratio: float = 4.0, dropout: float = 0.0, |
| n_decision_layers: int = 1, scorer: str = "cross", max_choices: int = 16): |
| super().__init__() |
| if scorer not in SCORERS: |
| raise ValueError(f"scorer must be one of {SCORERS}, got {scorer!r}") |
| self.d, self.scorer = d, scorer |
| self.cls = nn.Parameter(torch.zeros(1, 1, d)) |
| self.type_emb = nn.Parameter(torch.zeros(3, d)) |
| nn.init.normal_(self.type_emb, std=0.02) |
| self.blocks = nn.ModuleList(EncoderBlock(d, n_heads, mlp_ratio, dropout) for _ in range(n_layers)) |
| self.ln_ctx = nn.LayerNorm(d) |
| if scorer == "cross": |
| self.decision = nn.ModuleList(DecisionBlock(d, n_heads, mlp_ratio, dropout) |
| for _ in range(n_decision_layers)) |
| self.choice_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 1)) |
| elif scorer == "bi": |
| self.ln_h, self.ln_c = nn.LayerNorm(d), nn.LayerNorm(d) |
| self.choice_head = nn.Sequential(nn.Linear(3 * d, d), nn.GELU(), nn.Linear(d, 1)) |
| else: |
| self.choice_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, max_choices)) |
| self.value_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1)) |
| |
| self.scalar_head = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, 1)) |
|
|
| def encode_context(self, state_tokens, state_mask, question_tokens, question_mask): |
| B = question_tokens.shape[0] |
| parts = [self.cls.expand(B, -1, -1) + self.type_emb[0], question_tokens + self.type_emb[1]] |
| masks = [torch.ones(B, 1, dtype=torch.bool, device=question_tokens.device), question_mask] |
| if state_tokens is not None: |
| parts.append(state_tokens + self.type_emb[2]) |
| masks.append(state_mask) |
| x, mask = torch.cat(parts, 1), torch.cat(masks, 1) |
| for blk in self.blocks: |
| x = blk(x, key_mask=mask) |
| return self.ln_ctx(x), mask |
|
|
| def score_choices(self, ctx, ctx_mask, choice_emb, choice_mask): |
| if self.scorer == "cross": |
| c = choice_emb |
| for blk in self.decision: |
| c = blk(c, ctx, ctx_mask, choice_mask) |
| logits = self.choice_head(c).squeeze(-1) |
| elif self.scorer == "bi": |
| h = self.ln_h(ctx[:, :1]).expand(-1, choice_emb.shape[1], -1) |
| c = self.ln_c(choice_emb) |
| logits = self.choice_head(torch.cat([h, c, h * c], -1)).squeeze(-1) |
| else: |
| logits = self.choice_head(ctx[:, 0])[:, : choice_emb.shape[1]] |
| return logits.float().masked_fill(~choice_mask, NEG) |
|
|
| def forward(self, state_tokens, state_mask, question_tokens, question_mask, choice_emb, choice_mask): |
| ctx, ctx_mask = self.encode_context(state_tokens, state_mask, question_tokens, question_mask) |
| cls = ctx[:, 0] |
| return {"logits": self.score_choices(ctx, ctx_mask, choice_emb, choice_mask), |
| "value": self.value_head(cls).squeeze(-1).float(), |
| "score": self.scalar_head(cls).squeeze(-1).float(), |
| "embedding": cls, "choice_mask": choice_mask} |
|
|
