CLEF-Screen: JointSchemaHead redo FT (seed1, 3 epochs) from jev-retrain MI355X run 2026-11-04
2b605e8 verified | """The checkpoint's joint schema head (joint_schema_model.py `JointSchemaHead.forward`, one record) in NumPy float32. | |
| Weights: joint_head.safetensors (bfloat16, 122 tensors, widened to float32 exactly; `load_state_dict(strict=True)` names) | |
| and joint_head_config.json {hidden_size 4096, width 1024, routing_layers 2, layers 4, heads 16, feedforward 4096}. | |
| Inputs: the row graph's hidden states at the record's T real positions (after the final RMSNorm), the record's ids | |
| (processor form) and spans, and the lm_head table (`lm_head.weight` rows, for the lexical vectors). | |
| x = hidden_norm(hidden) LayerNorm(4096); memory = memory_projection(x) [T, 1024] | |
| global = x[T - 1]; q_i = mean of x over question i's span; ctx_o = mean of x over option o's span | |
| lexical_o = mean of the lm_head rows of option o's span ids (not normalized) | |
| queries_o = option_context_projection(ctx_o) + option_lexical_projection(lexical_o) | |
| + option_question_projection(q_i) i = the question of option o | |
| 2 x evidence routing: queries += MHA(query_norm(queries), memory_norm(memory)); queries += FFN(feedforward_norm(.)) | |
| (pre-LN, 16 heads of 64, FFN = Linear 1024->4096, GELU, Linear 4096->1024) | |
| per question: w = softmax over its options of routed . question_projection(q_i) / 32; summary_i = sum w * routed | |
| fields = question_projection(q) + option_summary_norm(summary) + global_projection(global) + type_embedding(type) | |
| 4 x TransformerDecoderLayer (norm_first, GELU): self-attention over the questions, cross-attention to the memory | |
| (not normalized), FFN; fields = field_norm(fields) | |
| prior_o = e^min(prior_logit_scale, ln 100) * cos(lexical_o, q_i + global) | |
| joint_o = e^min(joint_logit_scale, ln 100) * cos(field_i, option_norm(routed_o)) + residual_scorer([f, o, f*o, |f-o|]) | |
| logit_o = prior_o + sigmoid(residual_gate) * joint_o | |
| Every tensor is float32 as in the author's float32 module; LayerNorm eps 1e-5; attention = softmax(q k^T / 8) v per | |
| head in float32; GELU = the exact erf form (evaluated in float64, rounded to float32); F.normalize eps 1e-12 and | |
| F.cosine_similarity eps 1e-8 (each side divided by max(norm, eps), then the dot product, as torch computes it). | |
| Dropout is the identity (eval). The per-question softmax of the logits is the response's (response.py).""" | |
| from __future__ import annotations | |
| import json | |
| import math | |
| from pathlib import Path | |
| import numpy as np | |
| from ._erf import gelu | |
| from .tables import load_bf16_tensors | |
| CONFIG = {"hidden_size": 4096, "width": 1024, "routing_layers": 2, "layers": 4, "heads": 16, "feedforward": 4096} | |
| LN_EPS = 1e-5 | |
| F32 = np.float32 | |
| def layer_norm(x: np.ndarray, w: np.ndarray, b: np.ndarray, eps: float = LN_EPS) -> np.ndarray: | |
| mean = x.mean(axis=-1, keepdims=True, dtype=F32) | |
| xc = x - mean | |
| var = (xc * xc).mean(axis=-1, keepdims=True, dtype=F32) | |
| return (xc * (F32(1.0) / np.sqrt(var + F32(eps)))) * w + b | |
| def normalize(x: np.ndarray, eps: float) -> np.ndarray: | |
| """x / max(||x||_2, eps) over the last axis (F.normalize).""" | |
| n = np.sqrt((x * x).sum(axis=-1, keepdims=True, dtype=F32)) | |
| return x / np.maximum(n, F32(eps)) | |
| def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray: | |
| e = np.exp(x - x.max(axis=axis, keepdims=True)) | |
| return e / e.sum(axis=axis, keepdims=True, dtype=F32) | |
| class _Linear: | |
| def __init__(self, w: np.ndarray, b: np.ndarray | None = None): | |
| self.wt = np.ascontiguousarray(w.T, dtype=F32) # [in, out] | |
| self.b = None if b is None else np.ascontiguousarray(b, dtype=F32) | |
| def __call__(self, x: np.ndarray) -> np.ndarray: | |
| y = x @ self.wt | |
| return y if self.b is None else y + self.b | |
| class _LayerNorm: | |
| def __init__(self, w: np.ndarray, b: np.ndarray): | |
| self.w, self.b = w.astype(F32), b.astype(F32) | |
| def __call__(self, x: np.ndarray) -> np.ndarray: | |
| return layer_norm(x, self.w, self.b) | |
| class _Attention: | |
| """nn.MultiheadAttention (batch_first, eval, no masks) for one sequence: q_in [Lq, E], kv_in [Lk, E] -> [Lq, E].""" | |
| def __init__(self, t: dict, prefix: str, heads: int): | |
| w, b = t[prefix + "in_proj_weight"], t[prefix + "in_proj_bias"] | |
| e = w.shape[1] | |
| self.q = _Linear(w[:e], b[:e]) | |
| self.k = _Linear(w[e:2 * e], b[e:2 * e]) | |
| self.v = _Linear(w[2 * e:], b[2 * e:]) | |
| self.out = _Linear(t[prefix + "out_proj.weight"], t[prefix + "out_proj.bias"]) | |
| self.heads, self.dim = heads, e // heads | |
| self.scale = F32(1.0 / math.sqrt(self.dim)) | |
| def __call__(self, q_in: np.ndarray, kv_in: np.ndarray) -> np.ndarray: | |
| h, d = self.heads, self.dim | |
| q = self.q(q_in).reshape(q_in.shape[0], h, d).transpose(1, 0, 2) # [H, Lq, D] | |
| k = self.k(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2) # [H, Lk, D] | |
| v = self.v(kv_in).reshape(kv_in.shape[0], h, d).transpose(1, 0, 2) | |
| p = softmax((q @ k.transpose(0, 2, 1)) * self.scale) # [H, Lq, Lk] | |
| o = (p @ v).transpose(1, 0, 2).reshape(q_in.shape[0], h * d) | |
| return self.out(o) | |
| class _Tracked(dict): | |
| """The weights by name, recording which names were read (load_state_dict(strict=True): none missing, none unused).""" | |
| def __init__(self, *a): | |
| super().__init__(*a) | |
| self.read = set() | |
| def __getitem__(self, key): | |
| self.read.add(key) | |
| return super().__getitem__(key) | |
| class JointHead: | |
| """The author's JointSchemaHead for one record. `scores(hidden, ids, questions, lm_head)` -> one float32 logit | |
| array per question, in the author's option order.""" | |
| def __init__(self, tensors: dict[str, np.ndarray], config: dict | None = None): | |
| cfg = dict(CONFIG if config is None else config) | |
| if cfg != CONFIG: | |
| raise ValueError(f"head config {cfg} is not the release's {CONFIG}") | |
| self.config = cfg | |
| t = _Tracked({k: np.asarray(v, dtype=F32) for k, v in tensors.items()}) | |
| heads = cfg["heads"] | |
| ln = lambda p: _LayerNorm(t[p + ".weight"], t[p + ".bias"]) # noqa: E731 | |
| lin = lambda p, bias=True: _Linear(t[p + ".weight"], t[p + ".bias"] if bias else None) # noqa: E731 | |
| self.hidden_norm = ln("hidden_norm") | |
| self.memory_projection = lin("memory_projection", False) | |
| self.question_projection = lin("question_projection", False) | |
| self.option_question_projection = lin("option_question_projection", False) | |
| self.global_projection = lin("global_projection", False) | |
| self.option_context_projection = lin("option_context_projection", False) | |
| self.option_lexical_projection = lin("option_lexical_projection", False) | |
| self.type_embedding = t["type_embedding.weight"] | |
| self.evidence = [{"query_norm": ln(f"evidence_layers.{i}.query_norm"), | |
| "memory_norm": ln(f"evidence_layers.{i}.memory_norm"), | |
| "attention": _Attention(t, f"evidence_layers.{i}.attention.", heads), | |
| "feedforward_norm": ln(f"evidence_layers.{i}.feedforward_norm"), | |
| "ff0": lin(f"evidence_layers.{i}.feedforward.0"), | |
| "ff3": lin(f"evidence_layers.{i}.feedforward.3")} | |
| for i in range(cfg["routing_layers"])] | |
| self.option_summary_norm = ln("option_summary_norm") | |
| self.layers = [{"norm1": ln(f"layers.{i}.norm1"), "norm2": ln(f"layers.{i}.norm2"), | |
| "norm3": ln(f"layers.{i}.norm3"), | |
| "self_attn": _Attention(t, f"layers.{i}.self_attn.", heads), | |
| "multihead_attn": _Attention(t, f"layers.{i}.multihead_attn.", heads), | |
| "linear1": lin(f"layers.{i}.linear1"), "linear2": lin(f"layers.{i}.linear2")} | |
| for i in range(cfg["layers"])] | |
| self.field_norm = ln("field_norm") | |
| self.option_norm = ln("option_norm") | |
| self.scorer0 = lin("residual_scorer.0") | |
| self.scorer3 = lin("residual_scorer.3") | |
| cap = F32(math.log(100.0)) | |
| self.prior_scale = np.exp(np.minimum(t["prior_logit_scale"].reshape(()), cap)).astype(F32) | |
| self.joint_scale = np.exp(np.minimum(t["joint_logit_scale"].reshape(()), cap)).astype(F32) | |
| g = t["residual_gate"].reshape(()) | |
| self.gate = (F32(1.0) / (F32(1.0) + np.exp(-g))).astype(F32) | |
| unused = set(t) - t.read # a missing name already raised KeyError when it was read | |
| if unused: | |
| raise ValueError(f"head weights not used: {sorted(unused)}") | |
| self.n_tensors = len(t) | |
| self.n_params = int(sum(v.size for v in dict.values(t))) | |
| def from_files(cls, weights: str | Path, config: str | Path | None = None) -> "JointHead": | |
| cfg = json.loads(Path(config).read_text()) if config is not None else None | |
| return cls(load_bf16_tensors(weights), cfg) | |
| def lexical_rows(lm_head, ids, questions) -> list[np.ndarray]: | |
| """Per question, [n_options, 4096]: the mean of the lm_head rows of each option span's ids (float32).""" | |
| ids = np.asarray(ids, dtype=np.int64) | |
| out = [] | |
| for q in questions: | |
| rows = [lm_head.rows(ids[a:b]).mean(axis=0, dtype=F32) for a, b in q.option_spans] | |
| out.append(np.stack(rows).astype(F32)) | |
| return out | |
| def scores(self, hidden: np.ndarray, ids, questions, lm_head=None, lexical: list | None = None) -> list[np.ndarray]: | |
| """hidden [T, 4096] (the record's real positions), ids (len T), questions (request.EncodedQuestion or | |
| anything with question_span / option_spans / type_id) -> [float32 logits per question]. | |
| `lexical` (one [n_options, 4096] array per question) replaces the lm_head gather (controls).""" | |
| hidden = np.asarray(hidden, dtype=F32) | |
| if hidden.ndim != 2 or hidden.shape[0] != len(ids): | |
| raise ValueError(f"hidden {hidden.shape} for {len(ids)} ids") | |
| if lexical is None: | |
| lexical = self.lexical_rows(lm_head, ids, questions) | |
| x = self.hidden_norm(hidden) # [T, 4096] | |
| memory = self.memory_projection(x) # [T, 1024] | |
| global_vector = x[-1] | |
| question_vectors = np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in (q.question_span for q in questions)]) | |
| type_rows = self.type_embedding[[q.type_id for q in questions]] | |
| counts = [len(q.option_spans) for q in questions] | |
| context = [np.stack([x[a:b].mean(axis=0, dtype=F32) for a, b in q.option_spans]) for q in questions] | |
| queries = np.concatenate([ | |
| self.option_context_projection(ctx) + self.option_lexical_projection(lex) | |
| + self.option_question_projection(question_vectors[i])[None] | |
| for i, (ctx, lex) in enumerate(zip(context, lexical))]) # [O, 1024] | |
| for layer in self.evidence: | |
| routed = layer["attention"](layer["query_norm"](queries), layer["memory_norm"](memory)) | |
| queries = queries + routed | |
| queries = queries + layer["ff3"](gelu(layer["ff0"](layer["feedforward_norm"](queries)))) | |
| split = np.split(queries, np.cumsum(counts)[:-1]) | |
| base = self.question_projection(question_vectors) # [Q, 1024] | |
| summaries = [] | |
| for field, options in zip(base, split): | |
| w = softmax((options @ field) / F32(math.sqrt(options.shape[-1])), axis=0) | |
| summaries.append((w[:, None] * options).sum(axis=0, dtype=F32)) | |
| fields = (base + self.option_summary_norm(np.stack(summaries)) + self.global_projection(global_vector)[None] | |
| + type_rows) | |
| for layer in self.layers: | |
| f1 = layer["norm1"](fields) | |
| fields = fields + layer["self_attn"](f1, f1) | |
| fields = fields + layer["multihead_attn"](layer["norm2"](fields), memory) | |
| fields = fields + layer["linear2"](gelu(layer["linear1"](layer["norm3"](fields)))) | |
| fields = self.field_norm(fields) | |
| out = [] | |
| for i, (field, lex, routed) in enumerate(zip(fields, lexical, split)): | |
| anchor = normalize(question_vectors[i] + global_vector, 1e-12) | |
| prior = self.prior_scale * (normalize(lex, 1e-12) @ anchor) | |
| options = self.option_norm(routed) | |
| f = np.broadcast_to(field, options.shape) | |
| cosine = (normalize(f, 1e-8) * normalize(options, 1e-8)).sum(axis=-1, dtype=F32) | |
| features = np.concatenate([f, options, f * options, np.abs(f - options)], axis=-1) | |
| residual = self.scorer3(gelu(self.scorer0(features)))[:, 0] | |
| joint = self.joint_scale * cosine + residual | |
| out.append((prior + self.gate * joint).astype(F32)) | |
| return out | |