"""Symbol retrieval modules for the owned DAT backend.""" import math import torch import torch.nn as nn import torch.nn.functional as F from .transformer_components import PositionalEncoding class SymbolicAttentionRetriever(nn.Module): def __init__( self, hidden_dim: int, symbol_dim: int, n_symbols: int, n_heads: int, dropout: float = 0.0, trainable_symbols: bool = True, scale: float | None = None, use_bias: bool = False, ): super().__init__() if hidden_dim % n_heads != 0: raise ValueError(f"hidden_dim ({hidden_dim}) must be divisible by n_heads ({n_heads})") if symbol_dim % n_heads != 0: raise ValueError(f"symbol_dim ({symbol_dim}) must be divisible by n_heads ({n_heads})") if scale is not None and scale <= 0.0: raise ValueError(f"scale must be positive when provided, got {scale}") self.hidden_dim = hidden_dim self.symbol_dim = symbol_dim self.n_symbols = n_symbols self.n_heads = n_heads self.head_dim = hidden_dim // n_heads self.symbol_head_dim = symbol_dim // n_heads self.scale = 1.0 / math.sqrt(self.head_dim) if scale is None else scale self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias) self.template_features = nn.Parameter(torch.empty(n_symbols, hidden_dim)) self.symbol_library = nn.Parameter( torch.empty(n_symbols, symbol_dim), requires_grad=trainable_symbols, ) self.dropout = nn.Dropout(dropout) self._init_weights() def _init_weights(self) -> None: nn.init.normal_(self.template_features, mean=0.0, std=0.9) nn.init.normal_(self.symbol_library, mean=0.0, std=0.9) def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, seq_len, _ = x.shape queries = self.q_proj(x) queries = queries.view(batch_size, seq_len, self.n_heads, self.head_dim) queries = queries.transpose(1, 2) keys = self.template_features.view(self.n_symbols, self.n_heads, self.head_dim) keys = keys.transpose(0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1) values = self.symbol_library.view(self.n_symbols, self.n_heads, self.symbol_head_dim) values = values.transpose(0, 1).unsqueeze(0).expand(batch_size, -1, -1, -1) scores = torch.matmul(queries, keys.transpose(-2, -1)) * self.scale weights = F.softmax(scores, dim=-1) weights = self.dropout(weights) retrieved = torch.matmul(weights, values) retrieved = retrieved.transpose(1, 2).contiguous() return retrieved.view(batch_size, seq_len, self.symbol_dim) class PositionalSymbolRetriever(nn.Module): def __init__( self, symbol_dim: int, max_len: int, sinusoidal: bool = False, ): super().__init__() self.symbol_dim = symbol_dim self.max_len = max_len self.sinusoidal = sinusoidal pe_type = "sinusoidal" if sinusoidal else "learned" self.position_encoder = PositionalEncoding( embedding_dim=symbol_dim, pe_type=pe_type, max_len=max_len, ) def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, seq_len, _ = x.shape pos_info = self.position_encoder.get_positional_info(seq_len, x.device) if pos_info.embeddings is None: raise ValueError(f"Positional symbol retrieval produced no embeddings for seq_len={seq_len}") return pos_info.embeddings.unsqueeze(0).expand(batch_size, -1, -1) class RelativePositionalSymbolRetriever(nn.Module): def __init__( self, symbol_dim: int, max_rel_distance: int, rope: bool = False, theta: float = 10000.0, ): super().__init__() if rope and symbol_dim % 2 != 0: raise ValueError(f"RoPE relative symbols require even symbol_dim, got {symbol_dim}") if theta <= 0.0: raise ValueError(f"theta must be positive, got {theta}") self.symbol_dim = symbol_dim self.max_rel_distance = max_rel_distance self.rope = rope self.theta = theta if rope: self.register_buffer("_rope_relative_cache", torch.empty(0), persistent=False) else: self.position_encoder = PositionalEncoding( embedding_dim=symbol_dim, pe_type="relative", max_len=max_rel_distance * 2, max_rel_pos=max_rel_distance, ) def forward(self, x: torch.Tensor) -> torch.Tensor: seq_len = x.shape[1] if self.rope: return self._rope_relative_symbols(seq_len, x.device, x.dtype) pos_info = self.position_encoder.get_positional_info(seq_len, x.device) if pos_info.rel_embeddings is None: raise ValueError(f"Relative symbol retrieval produced no embeddings for seq_len={seq_len}") return pos_info.rel_embeddings def _rope_relative_symbols( self, seq_len: int, device: torch.device, dtype: torch.dtype, ) -> torch.Tensor: cached = self._rope_relative_cache if ( cached.shape[0] >= seq_len and cached.device == device and cached.dtype == dtype ): return cached[:seq_len, :seq_len, :] positions = torch.arange(seq_len, device=device) distances = positions[None, :] - positions[:, None] if self.max_rel_distance is not None: distances = torch.clamp(distances, -self.max_rel_distance, self.max_rel_distance) inv_freq = PositionalEncoding._rope_inv_freq(self.symbol_dim, self.theta, device=device) phases = distances.to(torch.float32).unsqueeze(-1) * inv_freq symbols = torch.empty(seq_len, seq_len, self.symbol_dim, device=device) symbols[..., 0::2] = torch.cos(phases) symbols[..., 1::2] = torch.sin(phases) self._rope_relative_cache = symbols.to(dtype=dtype) return self._rope_relative_cache class RelationalSymbolicAttentionRetriever(nn.Module): def __init__( self, hidden_dim: int, symbol_dim: int, rel_n_heads: int, symbolic_attn_n_heads: int, n_symbols: int, neighborhood_size: int = 2, include_self: bool = False, normalize_rels: bool = True, dropout: float = 0.0, trainable_symbols: bool = True, rel_scale: float | None = None, symbolic_attn_scale: float | None = None, use_bias: bool = False, ): super().__init__() if hidden_dim % rel_n_heads != 0: raise ValueError( f"hidden_dim ({hidden_dim}) must be divisible by rel_n_heads ({rel_n_heads})" ) if rel_scale is not None and rel_scale <= 0.0: raise ValueError(f"rel_scale must be positive when provided, got {rel_scale}") if symbolic_attn_scale is not None and symbolic_attn_scale <= 0.0: raise ValueError( "symbolic_attn_scale must be positive when provided, " f"got {symbolic_attn_scale}" ) self.hidden_dim = hidden_dim self.symbol_dim = symbol_dim self.rel_n_heads = rel_n_heads self.neighborhood_size = neighborhood_size self.include_self = include_self self.normalize_rels = normalize_rels self.neighborhood_dim = neighborhood_size + (1 if include_self else 0) self.rel_feature_dim = rel_n_heads * self.neighborhood_dim rel_head_dim = hidden_dim // rel_n_heads self.rel_scale = 1.0 / math.sqrt(rel_head_dim) if rel_scale is None else rel_scale self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias) self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=use_bias) self.rel_to_hidden = nn.Linear(self.rel_feature_dim, hidden_dim, bias=True) self.symbolic_attention = SymbolicAttentionRetriever( hidden_dim=hidden_dim, symbol_dim=symbol_dim, n_symbols=n_symbols, n_heads=symbolic_attn_n_heads, dropout=dropout, trainable_symbols=trainable_symbols, scale=symbolic_attn_scale, use_bias=use_bias, ) self._init_weights() def _init_weights(self) -> None: nn.init.xavier_uniform_(self.q_proj.weight) nn.init.xavier_uniform_(self.k_proj.weight) if self.q_proj.bias is not None: nn.init.zeros_(self.q_proj.bias) if self.k_proj.bias is not None: nn.init.zeros_(self.k_proj.bias) nn.init.xavier_uniform_(self.rel_to_hidden.weight) nn.init.zeros_(self.rel_to_hidden.bias) def _compute_neighborhood_indices(self, seq_len: int, device: torch.device) -> torch.Tensor: positions = torch.arange(seq_len, device=device).unsqueeze(1) # Decoder-LM adaptation: DSSL uses bidirectional neighborhoods, but this # owned backend must keep symbol retrieval causal. Keep the old causal # neighborhood order: current/immediate-past positions before older ones. if self.include_self: offsets = torch.arange(0, self.neighborhood_size + 1, device=device).unsqueeze(0) else: offsets = torch.arange(1, self.neighborhood_size + 1, device=device).unsqueeze(0) return (positions - offsets).clamp(0, seq_len - 1) def forward(self, x: torch.Tensor) -> torch.Tensor: batch_size, seq_len, _ = x.shape head_dim = self.hidden_dim // self.rel_n_heads queries = self.q_proj(x).view(batch_size, seq_len, self.rel_n_heads, head_dim) keys = self.k_proj(x).view(batch_size, seq_len, self.rel_n_heads, head_dim) queries = queries.transpose(1, 2) keys = keys.transpose(1, 2) neighbor_indices = self._compute_neighborhood_indices(seq_len, x.device) neighborhood_keys = keys[:, :, neighbor_indices] neighborhood_relations = torch.einsum("bhid,bhijd->bhij", queries, neighborhood_keys) if self.normalize_rels: neighborhood_relations = F.softmax(neighborhood_relations * self.rel_scale, dim=-1) neighborhood_relations = neighborhood_relations.permute(0, 2, 3, 1) neighborhood_relations = neighborhood_relations.contiguous().view( batch_size, seq_len, self.rel_feature_dim, ) relational_features = self.rel_to_hidden(neighborhood_relations) return self.symbolic_attention(relational_features)