"""Plain-PyTorch neural modules for STRATA.""" from __future__ import annotations import math from typing import Literal import torch from torch import nn from torch.nn import functional as F from strata.modeling.config import StrataConfig from strata.modeling.graph_object import GraphObject, apply_graph_object_intervention from strata.modeling.interventions import apply_predicate_memory_intervention from strata.modeling.outputs import GraphObjectBlockOutput, PredicateBlockOutput GraphMode = Literal["causal_lm", "full_graph"] ResidualScale = torch.Tensor | float | int | None class RMSNorm(nn.Module): """Root-mean-square normalization without bias.""" def __init__(self, d_model: int, eps: float = 1e-6) -> None: super().__init__() self.weight = nn.Parameter(torch.ones(d_model)) self.eps = eps def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: variance = hidden_states.pow(2).mean(dim=-1, keepdim=True) hidden_states = hidden_states * torch.rsqrt(variance + self.eps) return hidden_states * self.weight class SwiGLU(nn.Module): """SwiGLU feed-forward block.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.up = nn.Linear(config.d_model, 2 * config.d_ff, bias=False) self.down = nn.Linear(config.d_ff, config.d_model, bias=False) self.dropout = nn.Dropout(config.dropout) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: gate, value = self.up(hidden_states).chunk(2, dim=-1) return self.down(self.dropout(F.silu(gate) * value)) class LocalCausalSelfAttention(nn.Module): """Local causal multi-head attention. The implementation uses dense score tensors masked to a local causal band. It is intentionally simple and deterministic for the foundation codebase; optimized kernels can be introduced later behind the same interface. """ def __init__(self, config: StrataConfig) -> None: super().__init__() self.config = config self.qkv = nn.Linear(config.d_model, 3 * config.d_model, bias=False) self.out = nn.Linear(config.d_model, config.d_model, bias=False) self.dropout = nn.Dropout(config.dropout) def forward( self, hidden_states: torch.Tensor, *, attention_mask: torch.Tensor | None = None, graph_attention_bias: torch.Tensor | None = None, ) -> torch.Tensor: batch_size, seq_len, _ = hidden_states.shape qkv = self.qkv(hidden_states) query, key, value = qkv.chunk(3, dim=-1) query = _split_heads(query, self.config.num_heads) key = _split_heads(key, self.config.num_heads) value = _split_heads(value, self.config.num_heads) scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt( self.config.head_dim ) mask = _local_causal_mask( seq_len, self.config.local_attention_window, device=hidden_states.device, ) scores = scores.masked_fill(~mask.view(1, 1, seq_len, seq_len), -torch.inf) if attention_mask is not None: key_mask = attention_mask.to(torch.bool).view(batch_size, 1, 1, seq_len) scores = scores.masked_fill(~key_mask, -torch.inf) if graph_attention_bias is not None and self.config.use_attention_bias_from_graph: if graph_attention_bias.shape != (batch_size, seq_len, seq_len): raise ValueError( "graph_attention_bias must have shape " f"({batch_size}, {seq_len}, {seq_len}), got " f"{tuple(graph_attention_bias.shape)}" ) scores = scores + graph_attention_bias.unsqueeze(1) weights = torch.softmax(scores, dim=-1) weights = torch.nan_to_num(weights, nan=0.0) weights = self.dropout(weights) output = torch.matmul(weights, value) output = _merge_heads(output) output = self.out(output) if attention_mask is not None: output = output * attention_mask.to(output.dtype).unsqueeze(-1) return output class LexicalValencyProposer(nn.Module): """Predict token-anchored linguistic candidates used by predicate memory.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.node_type = nn.Linear(config.d_model, config.node_type_vocab_size) self.chart_type = nn.Linear(config.d_model, config.chart_type_vocab_size) self.predicate_gate = nn.Linear(config.d_model, 1) self.candidate_key = nn.Linear(config.d_model, config.d_model, bias=False) self.candidate_value = nn.Linear(config.d_model, config.d_model, bias=False) def forward( self, hidden_states: torch.Tensor, *, emit_heads: bool = True ) -> tuple[PredicateBlockOutput, torch.Tensor, torch.Tensor]: # predicate_gate feeds the memory attention (forward path) and is always # computed; node/chart are prediction heads, computed only when emit_heads. predicate_gate = torch.sigmoid(self.predicate_gate(hidden_states)) output = PredicateBlockOutput( node_type_logits=self.node_type(hidden_states) if emit_heads else None, chart_type_logits=self.chart_type(hidden_states) if emit_heads else None, predicate_gate=predicate_gate, ) return output, self.candidate_key(hidden_states), self.candidate_value(hidden_states) class PredicateMemoryAttention(nn.Module): """Causal token-to-predicate memory attention.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.config = config self.query = nn.Linear(config.d_model, config.d_model, bias=False) self.out = nn.Linear(config.d_model, config.d_model, bias=False) self.dropout = nn.Dropout(config.dropout) self.replacement_logit = nn.Parameter(torch.tensor(0.0)) self.edge_type = nn.Linear(config.d_model, config.graph_relation_types) def forward( self, hidden_states: torch.Tensor, *, candidate_key: torch.Tensor, candidate_value: torch.Tensor, predicate_gate: torch.Tensor, attention_mask: torch.Tensor | None, predicate_memory_bias: torch.Tensor | None, mode: GraphMode, return_edge_logits: bool, emit_heads: bool = True, predicate_memory_intervention: str = "none", ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor]: batch_size, seq_len, _ = hidden_states.shape candidate_key, candidate_value, predicate_gate = apply_predicate_memory_intervention( candidate_key=candidate_key, candidate_value=candidate_value, predicate_gate=predicate_gate, attention_mask=attention_mask, intervention=predicate_memory_intervention, ) query = _split_heads(self.query(hidden_states), self.config.num_heads) key = _split_heads(candidate_key, self.config.num_heads) value = _split_heads(candidate_value, self.config.num_heads) scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt( self.config.head_dim ) scores = scores + torch.log(predicate_gate.clamp_min(1e-6)).transpose(1, 2).view( batch_size, 1, 1, seq_len ) if predicate_memory_bias is not None: if predicate_memory_bias.shape != (batch_size, seq_len, seq_len): raise ValueError( "predicate_memory_bias must have shape " f"({batch_size}, {seq_len}, {seq_len}), got " f"{tuple(predicate_memory_bias.shape)}" ) scores = scores + predicate_memory_bias.to(scores.dtype).unsqueeze(1) if mode == "causal_lm": causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=hidden_states.device).tril() scores = scores.masked_fill(~causal.view(1, 1, seq_len, seq_len), -torch.inf) if attention_mask is not None: key_mask = attention_mask.to(torch.bool).view(batch_size, 1, 1, seq_len) scores = scores.masked_fill(~key_mask, -torch.inf) weights = torch.softmax(scores, dim=-1) weights = torch.nan_to_num(weights, nan=0.0) weights = self.dropout(weights) memory_read = self.out(_merge_heads(torch.matmul(weights, value))) if attention_mask is not None: memory_read = memory_read * attention_mask.to(memory_read.dtype).unsqueeze(-1) edge_logits: torch.Tensor | None = None if return_edge_logits and emit_heads: pair_repr = hidden_states.unsqueeze(2) * candidate_key.unsqueeze(1) edge_logits = self.edge_type(pair_repr) if mode == "causal_lm": causal = torch.ones(seq_len, seq_len, dtype=torch.bool, device=hidden_states.device).tril() edge_logits = edge_logits.masked_fill( ~causal.view(1, seq_len, seq_len, 1), -torch.inf ) return memory_read, edge_logits, torch.sigmoid(self.replacement_logit) class BoundedDeductionLayer(nn.Module): """Bounded differentiable closure over predicate-memory states.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.norm = RMSNorm(config.d_model) self.rule_mlp = nn.Sequential( nn.Linear(config.d_model, config.d_ff, bias=False), nn.SiLU(), nn.Linear(config.d_ff, config.d_model, bias=False), ) self.dropout = nn.Dropout(config.dropout) def forward(self, predicate_read: torch.Tensor) -> torch.Tensor: return self.dropout(self.rule_mlp(self.norm(predicate_read))) class GraphObjectMemoryAttention(nn.Module): """Explicit graph-object memory read over token-anchored graph edges.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.config = config self.query = nn.Linear(config.d_model, config.d_model, bias=False) self.key = nn.Linear(config.d_model, config.d_model, bias=False) self.value = nn.Linear(config.d_model, config.d_model, bias=False) self.out = nn.Linear(config.d_model, config.d_model, bias=False) self.node_type = nn.Embedding(config.graph_object_node_types, config.d_model) self.relation_bias = nn.Embedding(config.graph_object_relation_types, config.num_heads) self.relation_value = nn.Embedding(config.graph_object_relation_types, config.d_model) self.use_relation_conditioned_messages = config.graph_object_relation_conditioned_messages self.use_predicate_slot_memory = config.graph_object_predicate_slot_memory if self.use_relation_conditioned_messages: self.message_pair = nn.Linear(2 * config.d_model, config.d_model, bias=False) self.message_gamma = nn.Embedding(config.graph_object_relation_types, config.d_model) self.message_beta = nn.Embedding(config.graph_object_relation_types, config.d_model) if self.use_predicate_slot_memory: self.slot_event = nn.Linear(config.d_model, config.d_model, bias=False) self.slot_filler = nn.Linear(config.d_model, config.d_model, bias=False) self.slot_role_key = nn.Embedding(config.graph_object_relation_types, config.d_model) self.slot_role_gamma = nn.Embedding(config.graph_object_relation_types, config.d_model) self.slot_role_beta = nn.Embedding(config.graph_object_relation_types, config.d_model) self.relation_aux = nn.Linear(config.d_model, config.graph_object_relation_types, bias=False) self.src_aux = nn.Linear(config.d_model, config.d_model, bias=False) self.dst_aux = nn.Linear(config.d_model, config.d_model, bias=False) self.gate_logit = nn.Parameter(torch.tensor(float(config.graph_object_gate_init))) self.dropout = nn.Dropout(config.dropout) def forward( self, hidden_states: torch.Tensor, *, graph_object: GraphObject | None, attention_mask: torch.Tensor | None, mode: GraphMode, graph_object_intervention: str = "none", return_aux_logits: bool = False, ) -> tuple[torch.Tensor, GraphObjectBlockOutput | None]: if graph_object is None: return torch.zeros_like(hidden_states), None graph_object = apply_graph_object_intervention( graph_object, intervention=graph_object_intervention, relation_vocab_size=self.config.graph_object_relation_types, ) if graph_object is None: return torch.zeros_like(hidden_states), None batch_size, seq_len, _ = hidden_states.shape node_type = graph_object["node_type"] node_mask = graph_object["node_mask"].to(torch.bool) edge_src = graph_object["edge_src"] edge_dst = graph_object["edge_dst"] edge_rel = graph_object["edge_rel"] edge_slot_mask = graph_object["edge_slot_mask"].to(torch.bool) if node_type.shape != (batch_size, seq_len): raise ValueError(f"graph node_type must have shape {(batch_size, seq_len)}, got {tuple(node_type.shape)}") if edge_src.ndim != 2 or edge_src.shape != edge_dst.shape or edge_src.shape != edge_rel.shape: raise ValueError("graph edge_src/edge_dst/edge_rel must have matching shape [batch, edges]") if edge_src.shape[0] != batch_size: raise ValueError(f"graph edge batch size {edge_src.shape[0]} != hidden batch size {batch_size}") safe_node_type = node_type.clamp_min(0).clamp_max(self.config.graph_object_node_types - 1) node_repr = hidden_states + self.node_type(safe_node_type) * node_mask.to(hidden_states.dtype).unsqueeze(-1) query = _split_heads(self.query(hidden_states), self.config.num_heads) edge_count = edge_src.shape[1] if edge_count == 0: return torch.zeros_like(hidden_states), None safe_src = edge_src.clamp(0, seq_len - 1) safe_dst = edge_dst.clamp(0, seq_len - 1) gather = lambda x, idx: x.gather(1, idx.unsqueeze(-1).expand(batch_size, edge_count, x.shape[-1])) src_repr = gather(node_repr, safe_src) dst_repr = gather(node_repr, safe_dst) safe_edge_type = edge_rel.clamp_min(0).clamp_max(self.config.graph_object_relation_types - 1) rel_repr = self.relation_value(safe_edge_type) if self.use_predicate_slot_memory: # Relation labels are addresses, not decorations: the key names the # event-role slot and the value writes the endpoint filler through a # role-specific projection. Untyped collapse therefore destroys the # ARG0/ARG1 address distinction even when endpoints/topology remain. role_key = self.slot_role_key(safe_edge_type) relation_gamma = torch.tanh(self.slot_role_gamma(safe_edge_type)) relation_beta = self.slot_role_beta(safe_edge_type) slot_key_repr = self.slot_event(dst_repr) + role_key slot_value_repr = self.slot_filler(src_repr) * (1.0 + relation_gamma) + relation_beta slot_repr = slot_key_repr + slot_value_repr elif self.use_relation_conditioned_messages: pair_repr = torch.cat([src_repr, dst_repr], dim=-1) base_message = self.message_pair(pair_repr) relation_gamma = torch.tanh(self.message_gamma(safe_edge_type)) relation_beta = self.message_beta(safe_edge_type) slot_repr = base_message * (1.0 + relation_gamma) + relation_beta + rel_repr slot_key_repr = slot_repr slot_value_repr = slot_repr else: slot_repr = src_repr + dst_repr + rel_repr slot_key_repr = slot_repr slot_value_repr = slot_repr aux_output: GraphObjectBlockOutput | None = None if return_aux_logits: src_query = self.src_aux(slot_repr) dst_query = self.dst_aux(slot_repr) token_keys = hidden_states.transpose(1, 2) aux_output = GraphObjectBlockOutput( relation_logits=self.relation_aux(slot_repr), src_logits=torch.matmul(src_query, token_keys) / math.sqrt(self.config.d_model), dst_logits=torch.matmul(dst_query, token_keys) / math.sqrt(self.config.d_model), ) key = _split_heads(self.key(slot_key_repr), self.config.num_heads) value = _split_heads(self.value(slot_value_repr), self.config.num_heads) scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.config.head_dim) rel_bias = self.relation_bias(safe_edge_type).transpose(1, 2).unsqueeze(2) scores = scores + rel_bias valid_edge = edge_slot_mask.view(batch_size, 1, edge_count).expand(batch_size, seq_len, edge_count) if mode == "causal_lm": positions = torch.arange(seq_len, device=hidden_states.device).view(1, seq_len, 1) endpoints_visible = (safe_src.view(batch_size, 1, edge_count) <= positions) & ( safe_dst.view(batch_size, 1, edge_count) <= positions ) valid_edge = valid_edge & endpoints_visible if attention_mask is not None: query_mask = attention_mask.to(torch.bool).view(batch_size, seq_len, 1) src_valid = attention_mask.gather(1, safe_src).to(torch.bool).view(batch_size, 1, edge_count) dst_valid = attention_mask.gather(1, safe_dst).to(torch.bool).view(batch_size, 1, edge_count) valid_edge = valid_edge & query_mask & src_valid & dst_valid scores = scores.masked_fill(~valid_edge.view(batch_size, 1, seq_len, edge_count), -torch.inf) weights = torch.softmax(scores, dim=-1) weights = torch.nan_to_num(weights, nan=0.0) weights = self.dropout(weights) edge_read = torch.matmul(weights, value) update = self.out(_merge_heads(edge_read)) update = update * torch.sigmoid(self.gate_logit) if attention_mask is not None: update = update * attention_mask.to(update.dtype).unsqueeze(-1) return update, aux_output class PredicateBlock(nn.Module): """Lexical proposal, predicate memory read, and bounded closure.""" def __init__(self, config: StrataConfig) -> None: super().__init__() self.norm = RMSNorm(config.d_model) self.proposer = LexicalValencyProposer(config) self.memory = PredicateMemoryAttention(config) self.deduction = BoundedDeductionLayer(config) self.dropout = nn.Dropout(config.dropout) def forward( self, hidden_states: torch.Tensor, *, attention_mask: torch.Tensor | None, predicate_memory_bias: torch.Tensor | None, mode: GraphMode, return_edge_logits: bool, emit_heads: bool = True, predicate_memory_intervention: str = "none", predicate_memory_residual_scale: ResidualScale = None, ) -> tuple[torch.Tensor, PredicateBlockOutput, torch.Tensor]: normalized = self.norm(hidden_states) proposed, candidate_key, candidate_value = self.proposer(normalized, emit_heads=emit_heads) memory_read, edge_logits, replacement_gate = self.memory( normalized, candidate_key=candidate_key, candidate_value=candidate_value, predicate_gate=proposed.predicate_gate, attention_mask=attention_mask, predicate_memory_bias=predicate_memory_bias, mode=mode, return_edge_logits=return_edge_logits, emit_heads=emit_heads, predicate_memory_intervention=predicate_memory_intervention, ) proposed.edge_logits = edge_logits closure = self.deduction(memory_read) update = replacement_gate * memory_read + (1.0 - replacement_gate) * closure update = _apply_residual_scale(update, predicate_memory_residual_scale, name="predicate_memory_residual_scale") return hidden_states + self.dropout(update), proposed, replacement_gate class StrataDecoderBlock(nn.Module): """Local sequence computation plus optional STRATA predicate block.""" def __init__(self, config: StrataConfig, *, has_predicate_block: bool) -> None: super().__init__() self.attn_norm = RMSNorm(config.d_model) self.attention = LocalCausalSelfAttention(config) self.ffn_norm = RMSNorm(config.d_model) self.ffn = SwiGLU(config) self.predicate = PredicateBlock(config) if has_predicate_block else None self.graph_object = GraphObjectMemoryAttention(config) if has_predicate_block and config.use_graph_object_memory else None self.dropout = nn.Dropout(config.dropout) def forward( self, hidden_states: torch.Tensor, *, attention_mask: torch.Tensor | None, graph_attention_bias: torch.Tensor | None, predicate_memory_bias: torch.Tensor | None, mode: GraphMode, return_edge_logits: bool, emit_heads: bool = True, predicate_memory_intervention: str = "none", predicate_memory_residual_scale: ResidualScale = None, graph_object: GraphObject | None = None, graph_object_intervention: str = "none", graph_object_residual_scale: ResidualScale = None, return_graph_object_logits: bool = False, ) -> tuple[torch.Tensor, PredicateBlockOutput | None, torch.Tensor | None, GraphObjectBlockOutput | None]: hidden_states = hidden_states + self.dropout( self.attention( self.attn_norm(hidden_states), attention_mask=attention_mask, graph_attention_bias=graph_attention_bias, ) ) hidden_states = hidden_states + self.dropout(self.ffn(self.ffn_norm(hidden_states))) if self.predicate is None: return hidden_states, None, None, None hidden_states, predicate_output, replacement_gate = self.predicate( hidden_states, attention_mask=attention_mask, predicate_memory_bias=predicate_memory_bias, mode=mode, return_edge_logits=return_edge_logits, emit_heads=emit_heads, predicate_memory_intervention=predicate_memory_intervention, predicate_memory_residual_scale=predicate_memory_residual_scale, ) graph_object_output: GraphObjectBlockOutput | None = None if self.graph_object is not None and graph_object is not None: graph_update, graph_object_output = self.graph_object( hidden_states, graph_object=graph_object, attention_mask=attention_mask, mode=mode, graph_object_intervention=graph_object_intervention, return_aux_logits=return_graph_object_logits, ) graph_update = _apply_residual_scale(graph_update, graph_object_residual_scale, name="graph_object_residual_scale") hidden_states = hidden_states + self.dropout(graph_update) return hidden_states, predicate_output, replacement_gate, graph_object_output def _apply_residual_scale(update: torch.Tensor, scale: ResidualScale, *, name: str) -> torch.Tensor: if scale is None: return update if isinstance(scale, (float, int)): return update * float(scale) if scale.ndim == 0: return update * scale.to(update.dtype) batch_size = update.shape[0] if scale.shape == (batch_size,): shaped = scale.view(batch_size, 1, 1) elif scale.shape == (batch_size, 1): shaped = scale.view(batch_size, 1, 1) elif scale.shape == (batch_size, 1, 1): shaped = scale else: raise ValueError( f"{name} must be scalar or have shape ({batch_size},), " f"({batch_size}, 1), or ({batch_size}, 1, 1); got {tuple(scale.shape)}" ) return update * shaped.to(update.dtype) def _split_heads(tensor: torch.Tensor, num_heads: int) -> torch.Tensor: batch_size, seq_len, d_model = tensor.shape head_dim = d_model // num_heads return tensor.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2) def _merge_heads(tensor: torch.Tensor) -> torch.Tensor: batch_size, num_heads, seq_len, head_dim = tensor.shape return tensor.transpose(1, 2).contiguous().view(batch_size, seq_len, num_heads * head_dim) def _local_causal_mask(seq_len: int, window: int, *, device: torch.device) -> torch.Tensor: positions = torch.arange(seq_len, device=device) query = positions.view(seq_len, 1) key = positions.view(1, seq_len) return (key <= query) & ((query - key) < window)