File size: 4,657 Bytes
026f63e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """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)) # cls / question / state
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))
# continuous 0..1 score for `model.score(state, question)` (trained with BCE when rows carry `score`)
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(), # logit; sigmoid -> 0..1
"embedding": cls, "choice_mask": choice_mask}
|