Publish final EMA checkpoint

README.md

@@ -1,3 +1,56 @@
 ---
 license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- babylm
+- babylm-2026
+- causal-lm
+- dual-attention-transformer
+- nextlat
+- ema
+- custom-code
 ---
+
+# DAT Strict NextLat Final
+
+DAT Strict NextLat Final is a causal Dual Attention Transformer language model trained from scratch for the BabyLM 2026 Strict track. It combines self-attention, relational attention, position-relative symbols, and a NextLat auxiliary training objective. The repository contains the tokenizer and custom HuggingFace Transformers code required by the Auto classes.
+
+## Architecture
+
+- 16 transformer layers
+- Hidden dimension 1,024
+- 12 self-attention heads and 4 relational-attention heads
+- Rotary token positional encoding
+- Shared position-relative symbol retrieval without relative-symbol RoPE
+- RCA relational attention
+- SwiGLU feed-forward layers
+- Vocabulary size 16,384
+- Maximum model sequence length 514
+- Untied token embedding and language-model head
+
+## Training
+
+The model was trained for ten passes through the BabyLM 2026 Strict corpus under the 100 million word data limit. Training used batches of 42 sequences of 512 tokens, seed 1, Muon for hidden matrix parameters, and LambW for the remaining parameters. The Muon learning rate was 0.02, the auxiliary learning rate was 0.001, weight decay was 0.01, and the schedule used one percent warmup followed by cosine decay to ten percent of the initial rate.
+
+NextLat used horizon 1 with cross-entropy weight 0.0, KL weight 0.5, and MSE weight 1.0. An exponential moving average with decay 0.999 was maintained during training. The published final model and checkpoint revisions use the EMA weights. No teacher model, synthetic data augmentation, or multimodal input was used.
+
+## Checkpoints
+
+The final EMA model is stored on `main`. Intermediate EMA checkpoints use the official BabyLM Strict revision names from `chck_1M` through `chck_1000M`. These names represent cumulative words processed across repeated passes through the corpus.
+
+## Loading
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo_id = "abe123/babylm-dat-strict-nextlat-final"
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)
+```
+
+Use an explicit immutable commit SHA for reproducible evaluation.
+
+## Intended use and limitations
+
+This checkpoint is intended for BabyLM evaluation and language-model research. It uses custom model code and therefore requires `trust_remote_code=True` when loaded through HuggingFace Auto classes.

attention_config.py

@@ -0,0 +1,5 @@
+"""Shared attention configuration constants."""
+
+SUPPORTED_SEQUENCE_BOUNDARY_POLICIES = frozenset(
+    {"eos_document", "none", "segment_document"}
+)

config.json

@@ -0,0 +1,67 @@
+{
+  "architectures": [
+    "DatForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dat.DatConfig",
+    "AutoModel": "modeling_dat.DatModel",
+    "AutoModelForCausalLM": "modeling_dat.DatForCausalLM",
+    "AutoModelForMaskedLM": "modeling_dat.DatForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_dat.DatForSequenceClassification"
+  },
+  "bos_token_id": 1,
+  "dff_factor": 4,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "ffn_activation": "swiglu",
+  "ffn_hidden_dim_mode": "dff_factor",
+  "hidden_dim": 1024,
+  "init_range": 0.15,
+  "init_scheme": "normal_0_02_scaled_projection",
+  "max_position_embeddings": 514,
+  "max_rel_pos": 512,
+  "max_seq_len": 514,
+  "mlm_head_enabled": false,
+  "model_type": "dat",
+  "n_heads_ra": 4,
+  "n_heads_sa": 12,
+  "n_layers": 16,
+  "n_symbols": null,
+  "norm_first": true,
+  "norm_type": "layernorm",
+  "pad_token_id": 3,
+  "pe_type": "rope",
+  "positional_symbols_sinusoidal": false,
+  "ra_n_relations": null,
+  "ra_rel_activation": "identity",
+  "ra_symmetric_rels": false,
+  "ra_type": "rca",
+  "relative_symbols_rope": false,
+  "relsymbolic_dropout": 0.0,
+  "relsymbolic_include_self": false,
+  "relsymbolic_neighborhood_size": 2,
+  "relsymbolic_normalize_rels": true,
+  "relsymbolic_rel_n_heads": 4,
+  "relsymbolic_rel_scale": null,
+  "relsymbolic_symbolic_attn_n_heads": 4,
+  "relsymbolic_symbolic_attn_scale": null,
+  "relsymbolic_trainable_symbols": true,
+  "relsymbolic_use_bias": false,
+  "rope_theta": 10000.0,
+  "segment_boundary_token_id": null,
+  "sequence_boundary_policy": "eos_document",
+  "share_attn_params": false,
+  "shared_symbol_retriever": true,
+  "symbol_dim": null,
+  "symbol_retrieval": "relative",
+  "symbolic_attn_n_heads": null,
+  "symbolic_use_bias": false,
+  "tie_lm_head": false,
+  "tie_word_embeddings": false,
+  "transformers_version": "5.9.0",
+  "use_bias_ffn": true,
+  "use_bias_out": false,
+  "use_bias_qkv": false,
+  "vocab_size": 16384
+}

configuration_dat.py

@@ -0,0 +1,198 @@
+"""HuggingFace PretrainedConfig for the dual-attention (DAT) decoder LM.
+
+This is the source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a
+generated HF repository (alongside modeling_dat.py and the flattened model
+source) so the model can be loaded with
+``AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True)``.
+
+The config simply carries every field of models.attention.dat.config.DatLMConfig
+so that modeling_dat.py can rebuild the exact DatLMConfig at load time.
+"""
+
+from transformers import PretrainedConfig
+
+
+# Every field needed to rebuild DatLMConfig. New exported fields are appended
+# after existing constructor parameters to preserve positional compatibility.
+# modeling_dat.py rebuilds DatLMConfig as
+# DatLMConfig(**{f: getattr(config, f) for f in DAT_LM_FIELDS}).
+DAT_LM_FIELDS = (
+    "vocab_size",
+    "max_seq_len",
+    "pe_type",
+    "hidden_dim",
+    "n_heads_sa",
+    "n_heads_ra",
+    "n_layers",
+    "dropout",
+    "dff_factor",
+    "ffn_hidden_dim_mode",
+    "ffn_activation",
+    "rope_theta",
+    "max_rel_pos",
+    "init_range",
+    "init_scheme",
+    "norm_type",
+    "norm_first",
+    "use_bias_qkv",
+    "use_bias_out",
+    "use_bias_ffn",
+    "tie_lm_head",
+    "symbol_dim",
+    "n_symbols",
+    "symbolic_attn_n_heads",
+    "symbol_retrieval",
+    "symbolic_use_bias",
+    "shared_symbol_retriever",
+    "share_attn_params",
+    "positional_symbols_sinusoidal",
+    "relative_symbols_rope",
+    "relsymbolic_rel_n_heads",
+    "relsymbolic_symbolic_attn_n_heads",
+    "relsymbolic_neighborhood_size",
+    "relsymbolic_include_self",
+    "relsymbolic_normalize_rels",
+    "relsymbolic_trainable_symbols",
+    "relsymbolic_dropout",
+    "relsymbolic_rel_scale",
+    "relsymbolic_symbolic_attn_scale",
+    "relsymbolic_use_bias",
+    "ra_type",
+    "ra_n_relations",
+    "ra_rel_activation",
+    "ra_symmetric_rels",
+    "sequence_boundary_policy",
+    "segment_boundary_token_id",
+    "pad_token_id",
+    "bos_token_id",
+    "eos_token_id",
+    "mlm_head_enabled",
+)
+
+
+class DatConfig(PretrainedConfig):
+    model_type = "dat"
+    # Expose the model width under HF's conventional name so generic eval
+    # harnesses (e.g. babylm-eval's finetuning head, which reads
+    # config.hidden_size) can size their classification head. The model's own
+    # field is hidden_dim; this maps the alias onto it for both get and set.
+    attribute_map = {"hidden_size": "hidden_dim"}
+
+    def __init__(
+        self,
+        vocab_size: int = 16384,
+        max_seq_len: int = 513,
+        pe_type: str = "rope",
+        hidden_dim: int = 256,
+        n_heads_sa: int = 2,
+        n_heads_ra: int = 2,
+        n_layers: int = 4,
+        dropout: float = 0.0,
+        dff_factor: int = 4,
+        ffn_hidden_dim_mode: str = "dff_factor",
+        ffn_activation: str = "gelu",
+        rope_theta: float = 10000.0,
+        max_rel_pos: int | None = None,
+        init_range: float = 0.15,
+        init_scheme: str = "xavier_uniform",
+        norm_type: str = "rmsnorm",
+        norm_first: bool = True,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+        use_bias_ffn: bool = True,
+        tie_lm_head: bool = True,
+        symbol_dim: int | None = None,
+        n_symbols: int | None = None,
+        symbolic_attn_n_heads: int | None = None,
+        symbol_retrieval: str = "symbolic",
+        symbolic_use_bias: bool = False,
+        shared_symbol_retriever: bool = True,
+        share_attn_params: bool = False,
+        positional_symbols_sinusoidal: bool = False,
+        relative_symbols_rope: bool = False,
+        relsymbolic_rel_n_heads: int = 4,
+        relsymbolic_symbolic_attn_n_heads: int = 4,
+        relsymbolic_neighborhood_size: int = 2,
+        relsymbolic_include_self: bool = False,
+        relsymbolic_normalize_rels: bool = True,
+        relsymbolic_trainable_symbols: bool = True,
+        relsymbolic_dropout: float = 0.0,
+        relsymbolic_rel_scale: float | None = None,
+        relsymbolic_symbolic_attn_scale: float | None = None,
+        relsymbolic_use_bias: bool = False,
+        ra_type: str = "ra",
+        ra_n_relations: int | None = None,
+        ra_rel_activation: str = "identity",
+        ra_symmetric_rels: bool = False,
+        sequence_boundary_policy: str = "eos_document",
+        segment_boundary_token_id: int | None = None,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        mlm_head_enabled: bool = False,
+        **kwargs,
+    ) -> None:
+        self.vocab_size = vocab_size
+        self.max_seq_len = max_seq_len
+        self.pe_type = pe_type
+        self.hidden_dim = hidden_dim
+        self.n_heads_sa = n_heads_sa
+        self.n_heads_ra = n_heads_ra
+        self.n_layers = n_layers
+        self.dropout = dropout
+        self.dff_factor = dff_factor
+        self.ffn_hidden_dim_mode = ffn_hidden_dim_mode
+        self.ffn_activation = ffn_activation
+        self.rope_theta = rope_theta
+        self.max_rel_pos = max_rel_pos
+        self.init_range = init_range
+        self.init_scheme = init_scheme
+        self.norm_type = norm_type
+        self.norm_first = norm_first
+        self.use_bias_qkv = use_bias_qkv
+        self.use_bias_out = use_bias_out
+        self.use_bias_ffn = use_bias_ffn
+        self.tie_lm_head = tie_lm_head
+        self.symbol_dim = symbol_dim
+        self.n_symbols = n_symbols
+        self.symbolic_attn_n_heads = symbolic_attn_n_heads
+        self.symbol_retrieval = symbol_retrieval
+        self.symbolic_use_bias = symbolic_use_bias
+        self.shared_symbol_retriever = shared_symbol_retriever
+        self.share_attn_params = share_attn_params
+        self.positional_symbols_sinusoidal = positional_symbols_sinusoidal
+        self.relative_symbols_rope = relative_symbols_rope
+        self.relsymbolic_rel_n_heads = relsymbolic_rel_n_heads
+        self.relsymbolic_symbolic_attn_n_heads = relsymbolic_symbolic_attn_n_heads
+        self.relsymbolic_neighborhood_size = relsymbolic_neighborhood_size
+        self.relsymbolic_include_self = relsymbolic_include_self
+        self.relsymbolic_normalize_rels = relsymbolic_normalize_rels
+        self.relsymbolic_trainable_symbols = relsymbolic_trainable_symbols
+        self.relsymbolic_dropout = relsymbolic_dropout
+        self.relsymbolic_rel_scale = relsymbolic_rel_scale
+        self.relsymbolic_symbolic_attn_scale = relsymbolic_symbolic_attn_scale
+        self.relsymbolic_use_bias = relsymbolic_use_bias
+        self.ra_type = ra_type
+        self.ra_n_relations = ra_n_relations
+        self.ra_rel_activation = ra_rel_activation
+        self.ra_symmetric_rels = ra_symmetric_rels
+        self.sequence_boundary_policy = sequence_boundary_policy
+        self.segment_boundary_token_id = segment_boundary_token_id
+        self.mlm_head_enabled = mlm_head_enabled
+        # HF base also exposes max_position_embeddings for generation utilities.
+        self.max_position_embeddings = max_seq_len
+        # tie_lm_head is the source of truth for weight tying; keep HF's
+        # tie_word_embeddings in lockstep so from_pretrained never re-ties an
+        # untied head (or fails to tie a tied one). Drop any value coming in via
+        # kwargs (e.g. a serialized config.json) so the two cannot disagree.
+        # NOTE: tie_lm_head controls only lm_head <-> token_embeddings tying.
+        # The mlm_head.linear_out <-> lm_head tying is handled separately by
+        # the get_expanded_tied_weights_keys override in modeling_dat.py.
+        kwargs.pop("tie_word_embeddings", None)
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_lm_head,
+            **kwargs,
+        )

dat_config.py

@@ -0,0 +1,283 @@
+"""Configuration for the owned dual-attention decoder LM."""
+
+from dataclasses import dataclass
+
+from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES
+
+
+SUPPORTED_PE_TYPES = {"sinusoidal", "learned", "relative", "rope", "none"}
+SUPPORTED_SYMBOL_RETRIEVAL = {"symbolic", "positional", "relative", "relsymbolic"}
+SUPPORTED_RA_TYPES = {"ra", "rca", "disrca"}
+SUPPORTED_RA_ACTIVATIONS = {"softmax", "identity", "relu", "tanh", "sigmoid", "gelu"}
+SUPPORTED_FFN_ACTIVATIONS = {"gelu", "relu", "swiglu", "identity"}
+SUPPORTED_FFN_HIDDEN_DIM_MODES = {"dff_factor", "swiglu_parameter_matched"}
+SUPPORTED_INIT_SCHEMES = {"xavier_uniform", "normal_0_02_scaled_projection"}
+
+
+@dataclass(frozen=True)
+class DatLMConfig:
+    vocab_size: int
+    max_seq_len: int
+    pe_type: str = "rope"
+    hidden_dim: int = 256
+    n_heads_sa: int = 2
+    n_heads_ra: int = 2
+    n_layers: int = 4
+    dropout: float = 0.0
+    dff_factor: int = 4
+    ffn_hidden_dim_mode: str = "dff_factor"
+    ffn_activation: str = "gelu"
+    rope_theta: float = 10000.0
+    max_rel_pos: int | None = None
+    sequence_boundary_policy: str = "eos_document"
+    segment_boundary_token_id: int | None = None
+    init_range: float = 0.15
+    init_scheme: str = "xavier_uniform"
+    norm_type: str = "rmsnorm"
+    norm_first: bool = True
+    use_bias_qkv: bool = False
+    use_bias_out: bool = True
+    use_bias_ffn: bool = True
+    tie_lm_head: bool = True
+    symbol_dim: int | None = None
+    n_symbols: int | None = None
+    symbolic_attn_n_heads: int | None = None
+    symbol_retrieval: str = "symbolic"
+    symbolic_use_bias: bool = False
+    shared_symbol_retriever: bool = True
+    share_attn_params: bool = False
+    positional_symbols_sinusoidal: bool = False
+    relative_symbols_rope: bool = False
+    relsymbolic_rel_n_heads: int = 4
+    relsymbolic_symbolic_attn_n_heads: int = 4
+    relsymbolic_neighborhood_size: int = 2
+    relsymbolic_include_self: bool = False
+    relsymbolic_normalize_rels: bool = True
+    relsymbolic_trainable_symbols: bool = True
+    relsymbolic_dropout: float = 0.0
+    relsymbolic_rel_scale: float | None = None
+    relsymbolic_symbolic_attn_scale: float | None = None
+    relsymbolic_use_bias: bool = False
+    ra_type: str = "ra"
+    ra_n_relations: int | None = None
+    ra_rel_activation: str = "identity"
+    ra_symmetric_rels: bool = False
+    pad_token_id: int = 0
+    bos_token_id: int = 1
+    eos_token_id: int = 2
+    mlm_head_enabled: bool = False
+
+    def __post_init__(self) -> None:
+        if self.vocab_size <= 0:
+            raise ValueError(f"vocab_size must be positive, got {self.vocab_size}")
+        if self.max_seq_len <= 0:
+            raise ValueError(f"max_seq_len must be positive, got {self.max_seq_len}")
+        if self.pe_type not in SUPPORTED_PE_TYPES:
+            raise ValueError(f"Unsupported pe_type: {self.pe_type}")
+        if self.hidden_dim <= 0:
+            raise ValueError(f"hidden_dim must be positive, got {self.hidden_dim}")
+        if self.n_heads_sa <= 0:
+            raise ValueError(f"n_heads_sa must be positive for DAT, got {self.n_heads_sa}")
+        if self.n_heads_ra <= 0:
+            raise ValueError(f"n_heads_ra must be positive for DAT, got {self.n_heads_ra}")
+        total_heads = self.total_n_heads
+        if self.hidden_dim % total_heads != 0:
+            raise ValueError(
+                f"hidden_dim ({self.hidden_dim}) must be divisible by total DAT heads "
+                f"({total_heads} = {self.n_heads_sa} SA + {self.n_heads_ra} RA)"
+            )
+        if self.n_layers <= 0:
+            raise ValueError(f"n_layers must be positive, got {self.n_layers}")
+        if not 0.0 <= self.dropout < 1.0:
+            raise ValueError(f"dropout must be in [0.0, 1.0), got {self.dropout}")
+        if self.dff_factor <= 0:
+            raise ValueError(f"dff_factor must be positive, got {self.dff_factor}")
+        if self.ffn_hidden_dim_mode not in SUPPORTED_FFN_HIDDEN_DIM_MODES:
+            raise ValueError(f"Unsupported ffn_hidden_dim_mode for DAT: {self.ffn_hidden_dim_mode}")
+        if self.ffn_activation not in SUPPORTED_FFN_ACTIVATIONS:
+            raise ValueError(f"Unsupported ffn_activation for DAT: {self.ffn_activation}")
+        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched" and self.ffn_activation != "swiglu":
+            raise ValueError(
+                "ffn_hidden_dim_mode='swiglu_parameter_matched' requires "
+                f"ffn_activation='swiglu', got {self.ffn_activation}"
+            )
+        if self.rope_theta <= 0.0:
+            raise ValueError(f"rope_theta must be positive, got {self.rope_theta}")
+        if self.max_rel_pos is not None and self.max_rel_pos <= 0:
+            raise ValueError(f"max_rel_pos must be positive when provided, got {self.max_rel_pos}")
+        if self.sequence_boundary_policy not in SUPPORTED_SEQUENCE_BOUNDARY_POLICIES:
+            raise ValueError(
+                "Unsupported sequence_boundary_policy for DAT: "
+                f"{self.sequence_boundary_policy}"
+            )
+        if self.sequence_boundary_policy == "segment_document":
+            if self.segment_boundary_token_id is None:
+                raise ValueError(
+                    "segment_boundary_token_id is required when "
+                    "sequence_boundary_policy='segment_document'"
+                )
+        if self.init_range <= 0.0:
+            raise ValueError(f"init_range must be positive, got {self.init_range}")
+        if self.init_scheme not in SUPPORTED_INIT_SCHEMES:
+            raise ValueError(f"Unsupported init_scheme for DAT: {self.init_scheme}")
+        if self.norm_type not in {"layernorm", "rmsnorm"}:
+            raise ValueError(
+                f"norm_type must be 'layernorm' or 'rmsnorm', got {self.norm_type}"
+            )
+        if self.pe_type == "rope" and self.head_dim % 2 != 0:
+            raise ValueError(
+                "RoPE requires even DAT head_dim, got "
+                f"{self.head_dim} from hidden_dim={self.hidden_dim}, total_heads={total_heads}"
+            )
+        if self.pe_type == "sinusoidal" and self.hidden_dim % 2 != 0:
+            raise ValueError(f"Sinusoidal encoding requires even hidden_dim, got {self.hidden_dim}")
+        if self.share_attn_params and self.n_heads_sa != self.n_heads_ra:
+            raise ValueError(
+                "share_attn_params=True requires n_heads_sa == n_heads_ra, "
+                f"got {self.n_heads_sa} and {self.n_heads_ra}"
+            )
+        if self.symbol_dim is not None and self.symbol_dim <= 0:
+            raise ValueError(f"symbol_dim must be positive when provided, got {self.symbol_dim}")
+        if self.symbolic_attn_n_heads is not None and self.symbolic_attn_n_heads <= 0:
+            raise ValueError(
+                "symbolic_attn_n_heads must be positive when provided, "
+                f"got {self.symbolic_attn_n_heads}"
+            )
+        if self.symbol_retrieval not in SUPPORTED_SYMBOL_RETRIEVAL:
+            raise ValueError(f"Unsupported symbol_retrieval for DAT: {self.symbol_retrieval}")
+        if self.symbol_retrieval == "symbolic":
+            symbolic_heads = self.resolved_symbolic_attn_n_heads
+            if self.hidden_dim % symbolic_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by symbolic_attn_n_heads "
+                    f"({symbolic_heads}) for symbolic retrieval"
+                )
+            if self.resolved_symbol_dim % symbolic_heads != 0:
+                raise ValueError(
+                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by symbolic_attn_n_heads "
+                    f"({symbolic_heads}) for symbolic retrieval"
+                )
+        if (
+            self.symbol_retrieval == "positional"
+            and self.positional_symbols_sinusoidal
+            and self.resolved_symbol_dim % 2 != 0
+        ):
+            raise ValueError(
+                "Sinusoidal positional symbols require even symbol_dim, "
+                f"got {self.resolved_symbol_dim}"
+            )
+        if self.relative_symbols_rope:
+            if self.symbol_retrieval != "relative":
+                raise ValueError(
+                    "relative_symbols_rope=True requires symbol_retrieval='relative', "
+                    f"got {self.symbol_retrieval!r}"
+                )
+            if self.resolved_symbol_dim % 2 != 0:
+                raise ValueError(
+                    "RoPE relative symbols require even symbol_dim, "
+                    f"got {self.resolved_symbol_dim}"
+                )
+        if self.positional_symbols_sinusoidal and self.symbol_retrieval != "positional":
+            raise ValueError(
+                "positional_symbols_sinusoidal=True requires symbol_retrieval='positional', "
+                f"got {self.symbol_retrieval!r}"
+            )
+        if self.resolved_n_symbols <= 0:
+            raise ValueError(f"resolved_n_symbols must be positive, got {self.resolved_n_symbols}")
+        if self.symbol_retrieval == "relsymbolic":
+            if self.relsymbolic_rel_n_heads <= 0:
+                raise ValueError(
+                    f"relsymbolic_rel_n_heads must be positive, got {self.relsymbolic_rel_n_heads}"
+                )
+            if self.hidden_dim % self.relsymbolic_rel_n_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
+                    f"relsymbolic_rel_n_heads ({self.relsymbolic_rel_n_heads})"
+                )
+            if self.relsymbolic_symbolic_attn_n_heads <= 0:
+                raise ValueError(
+                    "relsymbolic_symbolic_attn_n_heads must be positive, "
+                    f"got {self.relsymbolic_symbolic_attn_n_heads}"
+                )
+            if self.hidden_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
+                raise ValueError(
+                    f"hidden_dim ({self.hidden_dim}) must be divisible by "
+                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
+                )
+            if self.resolved_symbol_dim % self.relsymbolic_symbolic_attn_n_heads != 0:
+                raise ValueError(
+                    f"symbol_dim ({self.resolved_symbol_dim}) must be divisible by "
+                    f"relsymbolic_symbolic_attn_n_heads ({self.relsymbolic_symbolic_attn_n_heads})"
+                )
+            if self.relsymbolic_neighborhood_size <= 0:
+                raise ValueError(
+                    "relsymbolic_neighborhood_size must be positive, "
+                    f"got {self.relsymbolic_neighborhood_size}"
+                )
+            if not 0.0 <= self.relsymbolic_dropout < 1.0:
+                raise ValueError(
+                    f"relsymbolic_dropout must be in [0.0, 1.0), got {self.relsymbolic_dropout}"
+                )
+            if self.relsymbolic_rel_scale is not None and self.relsymbolic_rel_scale <= 0.0:
+                raise ValueError(
+                    "relsymbolic_rel_scale must be positive when provided, "
+                    f"got {self.relsymbolic_rel_scale}"
+                )
+            if (
+                self.relsymbolic_symbolic_attn_scale is not None
+                and self.relsymbolic_symbolic_attn_scale <= 0.0
+            ):
+                raise ValueError(
+                    "relsymbolic_symbolic_attn_scale must be positive when provided, "
+                    f"got {self.relsymbolic_symbolic_attn_scale}"
+                )
+        if self.ra_type not in SUPPORTED_RA_TYPES:
+            raise ValueError(f"Unsupported ra_type for DAT: {self.ra_type}")
+        if self.ra_n_relations is not None and self.ra_n_relations <= 0:
+            raise ValueError(
+                f"ra_n_relations must be positive when provided, got {self.ra_n_relations}"
+            )
+        if self.ra_type != "ra" and self.ra_n_relations is not None:
+            raise ValueError(f"ra_n_relations applies only to ra_type='ra', got {self.ra_type}")
+        if self.ra_type != "ra" and self.ra_symmetric_rels:
+            raise ValueError(f"ra_symmetric_rels applies only to ra_type='ra', got {self.ra_type}")
+        n_relations = self.resolved_ra_n_relations
+        if self.ra_type == "ra" and (self.head_dim * self.n_heads_ra) % n_relations != 0:
+            raise ValueError(
+                f"head_dim * n_heads_ra ({self.head_dim * self.n_heads_ra}) must be "
+                f"divisible by ra_n_relations ({n_relations})"
+            )
+        if self.ra_rel_activation not in SUPPORTED_RA_ACTIVATIONS:
+            raise ValueError(
+                f"Unsupported ra_rel_activation for DAT: {self.ra_rel_activation}"
+            )
+
+    @property
+    def total_n_heads(self) -> int:
+        return self.n_heads_sa + self.n_heads_ra
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_dim // self.total_n_heads
+
+    @property
+    def resolved_symbol_dim(self) -> int:
+        return self.hidden_dim if self.symbol_dim is None else self.symbol_dim
+
+    @property
+    def resolved_symbolic_attn_n_heads(self) -> int:
+        return self.total_n_heads if self.symbolic_attn_n_heads is None else self.symbolic_attn_n_heads
+
+    @property
+    def resolved_ffn_hidden_dim(self) -> int:
+        if self.ffn_hidden_dim_mode == "swiglu_parameter_matched":
+            return int(8 / 3 * self.hidden_dim)
+        return self.hidden_dim * self.dff_factor
+
+    @property
+    def resolved_n_symbols(self) -> int:
+        return self.max_seq_len if self.n_symbols is None else self.n_symbols
+
+    @property
+    def resolved_ra_n_relations(self) -> int:
+        return self.n_heads_ra if self.ra_n_relations is None else self.ra_n_relations

dat_core.py

@@ -0,0 +1,328 @@
+"""Dual-attention primitives for decoder language models."""
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import PositionalInfo
+from .transformer_core import MultiHeadAttentionBase
+
+
+def _activate_scores(scores: torch.Tensor, activation: str) -> torch.Tensor:
+    if activation == "softmax":
+        return torch.softmax(scores, dim=-1)
+    if activation == "identity":
+        return scores
+    if activation == "relu":
+        return torch.relu(scores)
+    if activation == "tanh":
+        return torch.tanh(scores)
+    if activation == "sigmoid":
+        return torch.sigmoid(scores)
+    if activation == "gelu":
+        return torch.nn.functional.gelu(scores)
+    raise ValueError(f"Unsupported attention activation: {activation}")
+
+
+class RelationalAttentionBase(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        n_relations: Optional[int] = None,
+        rel_activation: str = "identity",
+        symmetric_rels: bool = False,
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // total_n_heads
+        output_dim = n_heads * head_dim
+        super().__init__(
+            query_dim=hidden_dim,
+            output_dim=output_dim,
+            key_dim=hidden_dim,
+            value_dim=symbol_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.symbol_dim = symbol_dim
+        self.rel_activation = rel_activation
+        self.symmetric_rels = symmetric_rels
+        self.use_relative_positional_symbols = use_relative_positional_symbols
+        self.n_relations = n_heads if n_relations is None else n_relations
+        total_rel_dim = self.head_dim * n_heads
+        if total_rel_dim % self.n_relations != 0:
+            raise ValueError(
+                f"head_dim * n_heads ({total_rel_dim}) must be divisible by n_relations "
+                f"({self.n_relations})"
+            )
+        self.rel_proj_dim = total_rel_dim // self.n_relations
+        self.rel_scale = 1.0 / math.sqrt(self.rel_proj_dim)
+        rel_total_dim = self.n_relations * self.rel_proj_dim
+        self.wq_rel = nn.Linear(hidden_dim, rel_total_dim, bias=False)
+        self.wk_rel = self.wq_rel if symmetric_rels else nn.Linear(
+            hidden_dim,
+            rel_total_dim,
+            bias=False,
+        )
+        nn.init.xavier_uniform_(self.wq_rel.weight)
+        if self.wk_rel is not self.wq_rel:
+            nn.init.xavier_uniform_(self.wk_rel.weight)
+
+    def _compute_base_attention(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q = self._reshape_for_multihead(self.q_proj(x), batch_size, seq_len)
+        k = self._reshape_for_multihead(self.k_proj(x), batch_size, seq_len)
+        scores = self._compute_attn_scores(q, k, pos_info)
+        return self._apply_activation_and_mask(scores, mask)
+
+    def compute_relational_scores(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        return_as: str,
+    ) -> torch.Tensor:
+        batch_size, seq_len, _ = x.shape
+        q_rel = self.wq_rel(x).view(batch_size, seq_len, self.n_relations, self.rel_proj_dim)
+        k_rel = self.wk_rel(x).view(batch_size, seq_len, self.n_relations, self.rel_proj_dim)
+        q_rel = q_rel.transpose(1, 2)
+        k_rel = k_rel.transpose(1, 2)
+        relations = torch.matmul(q_rel, k_rel.transpose(-2, -1)) * self.rel_scale
+        processed_mask = self._process_mask(mask)
+        if self.rel_activation == "softmax" and processed_mask is not None:
+            relations.masked_fill_(~processed_mask, torch.finfo(relations.dtype).min)
+        relations = _activate_scores(relations, self.rel_activation)
+        if processed_mask is not None:
+            relations = relations.masked_fill(~processed_mask, 0.0)
+        if return_as == "vectors":
+            return relations.permute(0, 2, 3, 1)
+        if return_as == "scores":
+            return relations
+        raise ValueError(f"return_as must be 'vectors' or 'scores', got {return_as}")
+
+    def combine_attention_and_relations(
+        self,
+        attn_weights: torch.Tensor,
+        relation_scores: torch.Tensor,
+    ) -> torch.Tensor:
+        return attn_weights * relation_scores
+
+    def _process_symbols_with_attention(
+        self,
+        symbols: torch.Tensor,
+        attn_weights: torch.Tensor,
+    ) -> torch.Tensor:
+        batch_size, _, seq_len_q, seq_len_k = attn_weights.shape
+        values = self.v_proj(symbols)
+        if self.use_relative_positional_symbols:
+            values = values.view(seq_len_q, seq_len_k, self.n_heads, self.head_dim)
+            return torch.einsum("bhij,ijhd->bihd", attn_weights, values)
+
+        values = values.view(batch_size, seq_len_k, self.n_heads, self.head_dim)
+        values = values.transpose(1, 2)
+        output = torch.matmul(attn_weights, values)
+        return output.transpose(1, 2)
+
+    def _apply_output_projection(self, output: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len = output.shape[:2]
+        output = output.contiguous().view(batch_size, seq_len, self.output_dim)
+        output = self.o_proj(output)
+        output = self.dropout(output)
+        return output
+
+    def _validate_symbols(self, x: torch.Tensor, symbols: torch.Tensor) -> None:
+        if self.use_relative_positional_symbols:
+            seq_len = x.shape[1]
+            expected_shape = (seq_len, seq_len, self.symbol_dim)
+            if tuple(symbols.shape) != expected_shape:
+                raise ValueError(
+                    f"Relative symbols must have shape {expected_shape}, got {tuple(symbols.shape)}"
+                )
+
+
+class RelationalAttention(RelationalAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        n_relations: int,
+        dropout: float = 0.0,
+        rel_activation: str = "identity",
+        symmetric_rels: bool = False,
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__(
+            hidden_dim=hidden_dim,
+            symbol_dim=symbol_dim,
+            n_heads=n_heads,
+            total_n_heads=total_n_heads,
+            dropout=dropout,
+            n_relations=n_relations,
+            rel_activation=rel_activation,
+            symmetric_rels=symmetric_rels,
+            use_relative_positional_symbols=use_relative_positional_symbols,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.wr_proj = nn.Parameter(torch.empty(n_heads, self.head_dim, n_relations))
+        nn.init.xavier_uniform_(self.wr_proj)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        self._validate_symbols(x, symbols)
+        attn_weights = self._compute_base_attention(x, mask, pos_info)
+        attn_weights = self.attn_dropout(attn_weights)
+        relation_vectors = self.compute_relational_scores(x, mask, return_as="vectors")
+        attended_symbols = self._process_symbols_with_attention(symbols, attn_weights)
+        projected_relations = torch.einsum(
+            "bhij,bijr,hdr->bihd",
+            attn_weights,
+            relation_vectors,
+            self.wr_proj,
+        )
+        output = self._apply_output_projection(attended_symbols + projected_relations)
+        self.last_attn_weights = attn_weights.detach()
+        return output, {"attention": attn_weights, "relations": relation_vectors}
+
+
+class RelationalCrossAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        activation: str = "identity",
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // total_n_heads
+        super().__init__(
+            query_dim=hidden_dim,
+            output_dim=n_heads * head_dim,
+            key_dim=hidden_dim,
+            value_dim=symbol_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation=activation,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.symbol_dim = symbol_dim
+        self.use_relative_positional_symbols = use_relative_positional_symbols
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        if self.use_relative_positional_symbols:
+            batch_size, seq_len, _ = x.shape
+            expected_shape = (seq_len, seq_len, self.symbol_dim)
+            if tuple(symbols.shape) != expected_shape:
+                raise ValueError(
+                    f"Relative symbols must have shape {expected_shape}, got {tuple(symbols.shape)}"
+                )
+
+            q = self._reshape_for_multihead(self.q_proj(x), batch_size, seq_len)
+            k = self._reshape_for_multihead(self.k_proj(x), batch_size, seq_len)
+            values = self.v_proj(symbols).view(seq_len, seq_len, self.n_heads, self.head_dim)
+            scores = self._compute_attn_scores(q, k, pos_info)
+            weights = self._apply_activation_and_mask(scores, mask)
+            weights = self.attn_dropout(weights)
+            output = torch.einsum("bhij,ijhd->bihd", weights, values)
+            output = output.contiguous().view(batch_size, seq_len, self.output_dim)
+            output = self.o_proj(output)
+            output = self.dropout(output)
+            self.last_attn_weights = weights.detach()
+            return output, {"attention": weights}
+
+        output, weights = super().forward(
+            query=x,
+            key=x,
+            value=symbols,
+            mask=mask,
+            pos_info=pos_info,
+        )
+        return output, {"attention": weights}
+
+
+class DisentangledRelationalCrossAttention(RelationalAttentionBase):
+    def __init__(
+        self,
+        hidden_dim: int,
+        symbol_dim: int,
+        n_heads: int,
+        total_n_heads: int,
+        dropout: float = 0.0,
+        rel_activation: str = "identity",
+        use_relative_positional_symbols: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__(
+            hidden_dim=hidden_dim,
+            symbol_dim=symbol_dim,
+            n_heads=n_heads,
+            total_n_heads=total_n_heads,
+            dropout=dropout,
+            n_relations=None,
+            rel_activation=rel_activation,
+            symmetric_rels=False,
+            use_relative_positional_symbols=use_relative_positional_symbols,
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbols: torch.Tensor,
+        mask: Optional[torch.Tensor],
+        pos_info: Optional[PositionalInfo],
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
+        self._validate_symbols(x, symbols)
+        attn_weights = self._compute_base_attention(x, mask, pos_info)
+        relation_scores = self.compute_relational_scores(x, mask, return_as="scores")
+        combined_weights = self.attn_dropout(
+            self.combine_attention_and_relations(attn_weights, relation_scores)
+        )
+        output = self._process_symbols_with_attention(symbols, combined_weights)
+        output = self._apply_output_projection(output)
+        self.last_attn_weights = attn_weights.detach()
+        return output, {
+            "attention": attn_weights,
+            "relations": relation_scores,
+            "combined": combined_weights,
+        }

dat_lm.py

@@ -0,0 +1,475 @@
+"""Decoder-only LM built from dual-attention blocks."""
+
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.utils.rnn import pad_sequence
+
+from .dat_config import DatLMConfig
+from .dat_core import (
+    DisentangledRelationalCrossAttention,
+    RelationalAttention,
+    RelationalCrossAttention,
+)
+from .transformer_core import SelfAttention
+from .dat_symbols import (
+    PositionalSymbolRetriever,
+    RelationalSymbolicAttentionRetriever,
+    RelativePositionalSymbolRetriever,
+    SymbolicAttentionRetriever,
+)
+from .masks import build_decoder_attention_mask
+from .mlm import MaskClassifier
+from .transformer_components import FeedForward, PositionalEncoding, PositionalInfo, RMSNorm
+
+
+class DatDecoderBlock(nn.Module):
+    def __init__(self, config: DatLMConfig):
+        super().__init__()
+        self.norm_first = config.norm_first
+        # input_dim is the block state width; hidden_dim is the total DAT width
+        # used with total_n_heads to derive the shared SA/RA head_dim.
+        self.sensory_attention = SelfAttention(
+            input_dim=config.hidden_dim,
+            n_heads=config.n_heads_sa,
+            hidden_dim=config.hidden_dim,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            supports_relative=config.pe_type == "relative",
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+        self.relational_attention = _build_relational_attention(config)
+        if config.share_attn_params:
+            self.sensory_attention.q_proj = self.relational_attention.q_proj
+            self.sensory_attention.k_proj = self.relational_attention.k_proj
+        self.feed_forward = FeedForward(
+            input_dim=config.hidden_dim,
+            hidden_dim=config.resolved_ffn_hidden_dim,
+            dropout=config.dropout,
+            activation=config.ffn_activation,
+            use_bias=config.use_bias_ffn,
+        )
+        self.norm1 = _create_norm(config.norm_type, config.hidden_dim)
+        self.norm2 = _create_norm(config.norm_type, config.hidden_dim)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        symbol_retriever: nn.Module,
+        mask: torch.Tensor,
+        pos_info: Optional[PositionalInfo],
+    ) -> torch.Tensor:
+        if self.norm_first:
+            normed_x = self.norm1(x)
+            # Retrieve symbols at the block boundary so pre-norm attention and
+            # input-dependent symbol retrieval use the same representation.
+            symbols = symbol_retriever(normed_x)
+            sensory_output, _ = self.sensory_attention(normed_x, mask=mask, pos_info=pos_info)
+            relational_output, _ = self.relational_attention(
+                normed_x,
+                symbols,
+                mask=mask,
+                pos_info=pos_info,
+            )
+            x = x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1))
+            x = x + self.dropout(self.feed_forward(self.norm2(x)))
+            return x
+
+        symbols = symbol_retriever(x)
+        sensory_output, _ = self.sensory_attention(x, mask=mask, pos_info=pos_info)
+        relational_output, _ = self.relational_attention(
+            x,
+            symbols,
+            mask=mask,
+            pos_info=pos_info,
+        )
+        x = self.norm1(x + self.dropout(torch.cat((sensory_output, relational_output), dim=-1)))
+        x = self.norm2(x + self.dropout(self.feed_forward(x)))
+        return x
+
+
+class DatDecoderLM(nn.Module):
+    def __init__(self, config: DatLMConfig):
+        super().__init__()
+        self.config = config
+        self.token_embeddings = nn.Embedding(config.vocab_size, config.hidden_dim)
+        self.embedding_dropout = nn.Dropout(config.dropout)
+        position_dim = (
+            config.hidden_dim if config.pe_type in {"sinusoidal", "learned", "none"}
+            else config.head_dim
+        )
+        self.position_encoder = PositionalEncoding(
+            embedding_dim=position_dim,
+            pe_type=config.pe_type,
+            max_len=config.max_seq_len,
+            theta=config.rope_theta,
+            max_rel_pos=config.max_rel_pos,
+            init_range=config.init_range,
+        )
+        self.symbol_retrievers = _build_symbol_retrievers(config)
+        self.layers = nn.ModuleList(DatDecoderBlock(config) for _ in range(config.n_layers))
+        self.final_norm = _create_norm(config.norm_type, config.hidden_dim)
+        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
+        if config.mlm_head_enabled:
+            self.mlm_head = MaskClassifier(
+                hidden_dim=config.hidden_dim,
+                vocab_size=config.vocab_size,
+                norm_type=config.norm_type,
+                ffn_activation=config.ffn_activation,
+                use_bias_ffn=config.use_bias_ffn,
+                dropout=config.dropout,
+                word_embedding=self.lm_head.weight if not config.tie_lm_head else self.token_embeddings.weight,
+            )
+        self._init_weights()
+
+    @property
+    def device(self) -> torch.device:
+        return self.token_embeddings.weight.device
+
+    def _init_weights(self) -> None:
+        if self.config.init_scheme == "xavier_uniform":
+            nn.init.xavier_uniform_(
+                self.token_embeddings.weight,
+                gain=nn.init.calculate_gain("linear"),
+            )
+            if self.config.tie_lm_head:
+                self.lm_head.weight = self.token_embeddings.weight
+            else:
+                nn.init.xavier_uniform_(
+                    self.lm_head.weight,
+                    gain=nn.init.calculate_gain("linear"),
+                )
+            return
+
+        if self.config.init_scheme == "normal_0_02_scaled_projection":
+            self._init_normal_0_02_scaled_projection()
+            if self.config.tie_lm_head:
+                self.lm_head.weight = self.token_embeddings.weight
+            return
+
+        raise ValueError(f"Unsupported init_scheme: {self.config.init_scheme}")
+
+    def _init_normal_0_02_scaled_projection(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+            elif isinstance(module, nn.Embedding):
+                nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+        for module in self.modules():
+            if isinstance(module, RelationalAttention):
+                nn.init.normal_(module.wr_proj, mean=0.0, std=0.02)
+            elif isinstance(module, SymbolicAttentionRetriever):
+                nn.init.normal_(module.template_features, mean=0.0, std=1.0)
+                nn.init.normal_(module.symbol_library, mean=0.0, std=1.0)
+            elif isinstance(module, RelativePositionalSymbolRetriever) and not module.rope:
+                nn.init.xavier_uniform_(module.position_encoder.rel_pos_embeddings_table.weight)
+
+        scaled_std = 0.02 / math.sqrt(2 * self.config.n_layers)
+        ffn_scaled_suffix = (
+            "feed_forward.w_up.weight"
+            if self.config.ffn_activation == "swiglu"
+            else "feed_forward.linear2.weight"
+        )
+        for name, parameter in self.named_parameters():
+            # Exclude the MLM head from depth-scaled init: its feed_forward
+            # down-projection shares the same suffix as decoder layers but
+            # should use the regular 0.02 initialization.
+            if name.startswith("mlm_head"):
+                continue
+            if name.endswith("o_proj.weight") or name.endswith(ffn_scaled_suffix):
+                nn.init.normal_(parameter, mean=0.0, std=scaled_std)
+
+    def _build_attention_mask(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional: bool = False,
+    ) -> torch.Tensor:
+        return build_decoder_attention_mask(
+            input_ids=input_ids,
+            pad_token_id=self.config.pad_token_id,
+            eos_token_id=self.config.eos_token_id,
+            sequence_boundary_policy=self.config.sequence_boundary_policy,
+            attention_mask=attention_mask,
+            segment_boundary_token_id=self.config.segment_boundary_token_id,
+            bidirectional=bidirectional,
+        )
+
+    def encode_for_objective(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        bidirectional: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if input_ids.dim() != 2:
+            raise ValueError(f"input_ids must be rank-2 [batch, seq], got shape {tuple(input_ids.shape)}")
+        if input_ids.shape[1] > self.config.max_seq_len:
+            raise ValueError(
+                f"Sequence length {input_ids.shape[1]} exceeds max_seq_len {self.config.max_seq_len}"
+            )
+
+
+        mask = self._build_attention_mask(input_ids, attention_mask, bidirectional=bidirectional)
+        token_embeddings = self.token_embeddings(input_ids)
+        hidden_states = token_embeddings
+        pos_info = self.position_encoder.get_positional_info(input_ids.shape[1], input_ids.device)
+        if pos_info.apply_to_embeddings:
+            if pos_info.embeddings is None:
+                raise ValueError(f"Embedding-level pe_type {pos_info.pe_type} did not provide embeddings.")
+            hidden_states = hidden_states + pos_info.embeddings.unsqueeze(0)
+        hidden_states = self.embedding_dropout(hidden_states)
+
+        for symbol_retriever, layer in zip(self.symbol_retrievers, self.layers):
+            hidden_states = layer(
+                hidden_states,
+                symbol_retriever=symbol_retriever,
+                mask=mask,
+                pos_info=pos_info,
+            )
+
+        hidden_states = self.final_norm(hidden_states)
+        return token_embeddings, hidden_states
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        targets: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        _, hidden_states = self.encode_for_objective(input_ids, attention_mask=attention_mask)
+        logits = self.lm_head(hidden_states)
+
+        loss = None
+        if targets is not None:
+            if targets.shape != input_ids.shape:
+                raise ValueError(
+                    f"targets shape must match input_ids shape, got {tuple(targets.shape)} "
+                    f"vs {tuple(input_ids.shape)}"
+                )
+            loss = F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                targets.reshape(-1),
+                ignore_index=-1,
+            )
+
+        return logits, loss
+
+    def forward_mlm(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode bidirectionally and apply the MLM head.
+
+        Args:
+            input_ids: Masked input token ids [batch, seq].
+            attention_mask: Validity mask [batch, seq].
+
+        Returns:
+            MLM logits at every position [batch, seq, vocab].
+        """
+        if not self.config.mlm_head_enabled:
+            raise ValueError(
+                "forward_mlm requires mlm_head_enabled=True; "
+                "the MLM head is not instantiated."
+            )
+        _, hidden_states = self.encode_for_objective(
+            input_ids, attention_mask=attention_mask, bidirectional=True
+        )
+        return self.mlm_head(hidden_states)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        max_new_tokens: int = 100,
+        temperature: float = 1.0,
+        top_k: Optional[int] = None,
+        do_sample: bool = True,
+        eos_token_id: Optional[int] = None,
+    ) -> torch.Tensor:
+        if temperature <= 0.0:
+            raise ValueError(f"temperature must be positive, got {temperature}")
+        if max_new_tokens < 0:
+            raise ValueError(f"max_new_tokens must be non-negative, got {max_new_tokens}")
+        if top_k is not None and top_k <= 0:
+            raise ValueError(f"top_k must be positive when provided, got {top_k}")
+        if eos_token_id is None:
+            eos_token_id = self.config.eos_token_id
+
+        if attention_mask is None:
+            generated_sequences = [row.clone() for row in input_ids]
+        else:
+            current_attention_mask = attention_mask.bool()
+            generated_sequences = []
+            for row, row_mask in zip(input_ids, current_attention_mask):
+                generated_sequence = row[row_mask]
+                if generated_sequence.numel() == 0:
+                    raise ValueError("Each input row must contain at least one unmasked token for generation.")
+                generated_sequences.append(generated_sequence)
+
+        finished = torch.tensor(
+            [sequence[-1].item() == eos_token_id for sequence in generated_sequences],
+            dtype=torch.bool,
+            device=input_ids.device,
+        )
+
+        for _ in range(max_new_tokens):
+            context_sequences = [sequence[-self.config.max_seq_len :] for sequence in generated_sequences]
+            context_ids = pad_sequence(
+                context_sequences,
+                batch_first=True,
+                padding_value=self.config.pad_token_id,
+            )
+            context_mask = pad_sequence(
+                [
+                    torch.ones(sequence.shape[0], dtype=torch.bool, device=input_ids.device)
+                    for sequence in context_sequences
+                ],
+                batch_first=True,
+                padding_value=False,
+            )
+            logits, _ = self(context_ids, attention_mask=context_mask)
+            last_positions = context_mask.long().sum(dim=1) - 1
+            batch_indices = torch.arange(logits.shape[0], device=logits.device)
+            next_token_logits = logits[batch_indices, last_positions, :] / temperature
+
+            if top_k is not None:
+                k = min(top_k, next_token_logits.size(-1))
+                top_values, _ = torch.topk(next_token_logits, k=k)
+                cutoff = top_values[:, -1].unsqueeze(-1)
+                next_token_logits = next_token_logits.masked_fill(next_token_logits < cutoff, float("-inf"))
+
+            if do_sample:
+                probs = F.softmax(next_token_logits, dim=-1)
+                next_token = torch.multinomial(probs, num_samples=1)
+            else:
+                next_token = torch.argmax(next_token_logits, dim=-1, keepdim=True)
+
+            for row_index in range(next_token.shape[0]):
+                if finished[row_index]:
+                    continue
+                generated_sequences[row_index] = torch.cat(
+                    [generated_sequences[row_index], next_token[row_index]],
+                )
+                if next_token[row_index, 0].item() == eos_token_id:
+                    finished[row_index] = True
+
+            if torch.all(finished):
+                break
+
+        return pad_sequence(
+            generated_sequences,
+            batch_first=True,
+            padding_value=self.config.pad_token_id,
+        )
+
+
+def _build_symbol_retrievers(config: DatLMConfig) -> nn.ModuleList:
+    retriever = _build_symbol_retriever(config)
+    if config.shared_symbol_retriever:
+        return nn.ModuleList([retriever] * config.n_layers)
+    return nn.ModuleList(_build_symbol_retriever(config) for _ in range(config.n_layers))
+
+
+def _build_symbol_retriever(config: DatLMConfig) -> nn.Module:
+    if config.symbol_retrieval == "symbolic":
+        return SymbolicAttentionRetriever(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_symbols=config.resolved_n_symbols,
+            n_heads=config.resolved_symbolic_attn_n_heads,
+            dropout=config.dropout,
+            use_bias=config.symbolic_use_bias,
+        )
+    if config.symbol_retrieval == "positional":
+        return PositionalSymbolRetriever(
+            symbol_dim=config.resolved_symbol_dim,
+            max_len=config.max_seq_len,
+            sinusoidal=config.positional_symbols_sinusoidal,
+        )
+    if config.symbol_retrieval == "relative":
+        max_rel_distance = config.max_rel_pos if config.max_rel_pos is not None else max(1, config.max_seq_len // 2)
+        return RelativePositionalSymbolRetriever(
+            symbol_dim=config.resolved_symbol_dim,
+            max_rel_distance=max_rel_distance,
+            rope=config.relative_symbols_rope,
+            theta=config.rope_theta,
+        )
+    if config.symbol_retrieval == "relsymbolic":
+        return RelationalSymbolicAttentionRetriever(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            rel_n_heads=config.relsymbolic_rel_n_heads,
+            symbolic_attn_n_heads=config.relsymbolic_symbolic_attn_n_heads,
+            n_symbols=config.resolved_n_symbols,
+            neighborhood_size=config.relsymbolic_neighborhood_size,
+            include_self=config.relsymbolic_include_self,
+            normalize_rels=config.relsymbolic_normalize_rels,
+            dropout=config.relsymbolic_dropout,
+            trainable_symbols=config.relsymbolic_trainable_symbols,
+            rel_scale=config.relsymbolic_rel_scale,
+            symbolic_attn_scale=config.relsymbolic_symbolic_attn_scale,
+            use_bias=config.relsymbolic_use_bias,
+        )
+    raise ValueError(f"Unsupported symbol_retrieval: {config.symbol_retrieval}")
+
+
+def _build_relational_attention(config: DatLMConfig) -> nn.Module:
+    use_relative_symbols = config.symbol_retrieval == "relative"
+    if config.ra_type == "ra":
+        return RelationalAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            n_relations=config.resolved_ra_n_relations,
+            dropout=config.dropout,
+            rel_activation=config.ra_rel_activation,
+            symmetric_rels=config.ra_symmetric_rels,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    if config.ra_type == "rca":
+        return RelationalCrossAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            activation=config.ra_rel_activation,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    if config.ra_type == "disrca":
+        return DisentangledRelationalCrossAttention(
+            hidden_dim=config.hidden_dim,
+            symbol_dim=config.resolved_symbol_dim,
+            n_heads=config.n_heads_ra,
+            total_n_heads=config.total_n_heads,
+            dropout=config.dropout,
+            rel_activation=config.ra_rel_activation,
+            use_relative_positional_symbols=use_relative_symbols,
+            use_bias_qkv=config.use_bias_qkv,
+            use_bias_out=config.use_bias_out,
+        )
+    raise ValueError(f"Unsupported ra_type: {config.ra_type}")
+
+
+def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(hidden_dim)
+    if norm_type == "rmsnorm":
+        return RMSNorm(hidden_dim)
+    raise ValueError(f"Unsupported norm_type: {norm_type}")

dat_symbols.py

@@ -0,0 +1,263 @@
+"""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)

masks.py

@@ -0,0 +1,61 @@
+"""Attention-mask helpers shared by repo-owned decoder LMs."""
+
+from typing import Optional
+
+import torch
+
+
+def build_decoder_attention_mask(
+    input_ids: torch.Tensor,
+    pad_token_id: int,
+    eos_token_id: int,
+    sequence_boundary_policy: str,
+    attention_mask: Optional[torch.Tensor] = None,
+    segment_boundary_token_id: Optional[int] = None,
+    bidirectional: bool = False,
+) -> torch.Tensor:
+    if input_ids.ndim != 2:
+        raise ValueError(
+            "input_ids must be rank-2 [batch, seq], "
+            f"got shape {tuple(input_ids.shape)}"
+        )
+    _, seq_len = input_ids.shape
+    if attention_mask is None:
+        valid_tokens = input_ids != pad_token_id
+    else:
+        valid_tokens = attention_mask.bool()
+
+    query_mask = valid_tokens.unsqueeze(2)
+    key_mask = valid_tokens.unsqueeze(1)
+    if bidirectional:
+        mask = query_mask & key_mask
+    else:
+        directionality = torch.tril(
+            torch.ones(seq_len, seq_len, dtype=torch.bool, device=input_ids.device)
+        ).unsqueeze(0)
+        mask = directionality & query_mask & key_mask
+
+    if sequence_boundary_policy == "none":
+        return mask
+    if sequence_boundary_policy == "segment_document":
+        if segment_boundary_token_id is None:
+            raise ValueError(
+                "segment_boundary_token_id is required when "
+                "sequence_boundary_policy='segment_document'"
+            )
+        # The boundary token starts the next segment: cumsum increments on the
+        # boundary position, so the marker attends with the following tokens.
+        segment_ids = torch.cumsum(
+            input_ids == segment_boundary_token_id, dim=1
+        )
+        same_segment = segment_ids.unsqueeze(1) == segment_ids.unsqueeze(2)
+        return mask & same_segment
+    if sequence_boundary_policy != "eos_document":
+        raise ValueError(f"Unsupported sequence_boundary_policy: {sequence_boundary_policy}")
+
+    # Next-token training predicts EOS from the preceding document token.
+    # Once EOS is present as an input token, it starts the next segment so the
+    # following document is not predicted with prior-document context.
+    document_ids = torch.cumsum(input_ids == eos_token_id, dim=1)
+    same_document = document_ids.unsqueeze(1) == document_ids.unsqueeze(2)
+    return mask & same_document

mlm.py

@@ -0,0 +1,95 @@
+"""MLM head for the hybrid MLM/CLM objective.
+
+Ported from temp/gpt-bert-main/pretraining/model.py MaskClassifier, but
+configurable to match the trunk's customizability:
+- norm_type: "layernorm" or "rmsnorm" (reuses the existing field).
+- ffn_activation: "gelu", "swiglu", "relu", or "identity" (reuses the
+  existing field). The head's nonlinearity follows this choice.
+- use_bias_ffn: whether the linear layers in the head use bias.
+- weight tying: the final linear is tied to the word embedding, matching
+  gpt-bert's MaskClassifier and our lm_head.
+
+The head is controlled by model config (mlm_head_enabled) so checkpoint
+loading and HF export remain strict and honest.
+"""
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import RMSNorm, FeedForward
+
+
+def _create_norm(norm_type: str, hidden_dim: int) -> nn.Module:
+    if norm_type == "layernorm":
+        return nn.LayerNorm(hidden_dim)
+    if norm_type == "rmsnorm":
+        return RMSNorm(hidden_dim)
+    raise ValueError(f"Unsupported norm_type: {norm_type}")
+
+
+class MaskClassifier(nn.Module):
+    """MLM prediction head with configurable norm and activation.
+
+    Architecture (matching gpt-bert's MaskClassifier):
+        norm -> FeedForward (activation, bias configurable) -> norm -> dropout -> linear (tied)
+
+    The final linear is tied to the word embedding weight, matching
+    gpt-bert's MaskClassifier and our lm_head.
+
+    Args:
+        hidden_dim: Model hidden dimension.
+        vocab_size: Vocabulary size for the output projection.
+        norm_type: "layernorm" or "rmsnorm".
+        ffn_activation: "gelu", "swiglu", "relu", or "identity".
+        use_bias_ffn: Whether the intermediate linear layers use bias.
+        dropout: Dropout probability.
+        word_embedding: The word embedding weight to tie the final linear to.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        vocab_size: int,
+        norm_type: str,
+        ffn_activation: str,
+        use_bias_ffn: bool,
+        dropout: float,
+        word_embedding: nn.Parameter,
+    ):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.vocab_size = vocab_size
+        self.norm_type = norm_type
+        self.ffn_activation = ffn_activation
+        self.use_bias_ffn = use_bias_ffn
+        self.dropout_p = dropout
+
+        self.norm1 = _create_norm(norm_type, hidden_dim)
+        self.feed_forward = FeedForward(
+            input_dim=hidden_dim,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            activation=ffn_activation,
+            use_bias=use_bias_ffn,
+        )
+        self.norm2 = _create_norm(norm_type, hidden_dim)
+        self.dropout = nn.Dropout(dropout)
+        self.linear_out = nn.Linear(hidden_dim, vocab_size, bias=False)
+        self.linear_out.weight = word_embedding
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the MLM head to contextualized embeddings.
+
+        Args:
+            x: Contextualized embeddings [batch, seq_len, hidden_dim] or
+               flattened [num_masked, hidden_dim].
+
+        Returns:
+            Logits [batch, seq_len, vocab_size] or [num_masked, vocab_size].
+        """
+        x = self.norm1(x)
+        x = self.feed_forward(x)
+        x = self.norm2(x)
+        x = self.dropout(x)
+        x = self.linear_out(x)
+        return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8af548664053b171d726ec05dd80f7bf2f9cb46933fcf2e27071d3cc29fa4531
+size 1187885624

modeling_dat.py

@@ -0,0 +1,290 @@
+"""HuggingFace PreTrainedModel wrapper for the dual-attention (DAT) decoder LM.
+
+Source-of-truth copy. scripts/convert_dat_to_hf.py copies it into a generated HF
+repository together with configuration_dat.py and the flattened model source
+(dat_config.py, dat_core.py, dat_symbols.py, dat_lm.py, transformer_core.py,
+transformer_components.py).
+
+Why every model module is imported directly here: with trust_remote_code on a
+local directory, transformers 4.46.3 only copies the entry module's *direct*
+relative imports into its dynamic-module cache (it does not recurse). Importing
+all flattened modules here forces every one of them to be copied, after which
+their own relative imports resolve because the files are co-located.
+"""
+
+import collections
+
+import torch
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput, MaskedLMOutput, SequenceClassifierOutput
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+# Force every flattened module to be copied into the dynamic-module cache.
+from .attention_config import SUPPORTED_SEQUENCE_BOUNDARY_POLICIES  # noqa: F401
+from .masks import build_decoder_attention_mask  # noqa: F401
+from .mlm import MaskClassifier  # noqa: F401
+from .transformer_components import FeedForward, PositionalEncoding  # noqa: F401
+from .transformer_core import SelfAttention  # noqa: F401
+from .dat_symbols import SymbolicAttentionRetriever  # noqa: F401
+from .dat_core import RelationalAttention  # noqa: F401
+from .dat_config import DatLMConfig
+from .dat_lm import DatDecoderLM
+from .configuration_dat import DatConfig, DAT_LM_FIELDS
+
+
+def _refresh_rope_buffers(model: DatDecoderLM, rope_theta: float) -> None:
+    pe = model.position_encoder
+    if hasattr(pe, "pe_type") and pe.pe_type == "rope":
+        cos, sin = pe._precompute_rope_freqs(pe.embedding_dim, pe.max_len, rope_theta)
+        pe.register_buffer("rope_cos", cos, persistent=False)
+        pe.register_buffer("rope_sin", sin, persistent=False)
+        pe._rope_uninitialized = True
+
+
+def _collect_shared_weight_keys(module: torch.nn.Module) -> dict[str, str]:
+    # remove_duplicate=False is essential: by default named_parameters yields
+    # each shared tensor only once, which would hide exactly the duplicates we
+    # need to declare (and which state_dict still emits under every name).
+    names_by_storage: dict[int, list[str]] = collections.defaultdict(list)
+    for name, parameter in module.named_parameters(remove_duplicate=False):
+        names_by_storage[id(parameter)].append(name)
+    shared_keys: dict[str, str] = {}
+    for names in names_by_storage.values():
+        if len(names) > 1:
+            names = sorted(names)
+            # Prefer token_embeddings.weight as the canonical source
+            # (HF convention: output weights tied to input embeddings).
+            source = names[0]
+            for n in names:
+                if n.endswith("token_embeddings.weight"):
+                    source = n
+                    break
+            for target in names:
+                if target != source:
+                    shared_keys[target] = source
+    return shared_keys
+
+
+class _ExpandedTiedWeightsMixin:
+    def get_expanded_tied_weights_keys(self, all_submodels: bool = False) -> dict:
+        # HF's default returns {} when tie_word_embeddings=False, which would
+        # discard the mlm_head.linear_out <-> lm_head tie for NextLat models
+        # (tie_lm_head=False). Return the dynamically computed mapping directly
+        # so all shared weights are properly tied on load.
+        return dict(self._tied_weights_keys)
+
+
+class DatModel(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = ""
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        # Work around accelerate/transformers leaving persistent=False buffers uninitialized
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> BaseModelOutput:
+        _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask)
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class DatForCausalLM(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+    _tied_weights_keys = {"model.lm_head.weight": "model.token_embeddings.weight"}
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        # Work around accelerate/transformers leaving persistent=False buffers uninitialized
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def get_output_embeddings(self) -> torch.nn.Module:
+        return self.model.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None:
+        self.model.lm_head = new_embeddings
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> CausalLMOutput:
+        logits, _ = self.model(input_ids, attention_mask=attention_mask)
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = torch.nn.functional.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        return CausalLMOutput(loss=loss, logits=logits)
+
+
+class DatForMaskedLM(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def get_output_embeddings(self) -> torch.nn.Module:
+        return self.model.mlm_head.linear_out
+
+    def set_output_embeddings(self, new_embeddings: torch.nn.Module) -> None:
+        self.model.mlm_head.linear_out = new_embeddings
+        # Keep lm_head in sync with mlm_head.linear_out so that
+        # resize_token_embeddings updates both tied heads.
+        self.model.lm_head.weight = new_embeddings.weight
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> MaskedLMOutput | tuple:
+        return_dict = kwargs.pop("return_dict", None)
+        if return_dict is None:
+            return_dict = self.config.return_dict
+        if kwargs:
+            raise TypeError(f"Unexpected keyword argument(s): {list(kwargs)}")
+        logits = self.model.forward_mlm(input_ids, attention_mask=attention_mask)
+
+        loss = None
+        if labels is not None:
+            loss = CrossEntropyLoss(ignore_index=-100)(
+                logits.view(-1, logits.size(-1)),
+                labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return MaskedLMOutput(loss=loss, logits=logits)
+
+class DatForSequenceClassification(_ExpandedTiedWeightsMixin, PreTrainedModel):
+    config_class = DatConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config: DatConfig) -> None:
+        super().__init__(config)
+        self.num_labels = getattr(config, "num_labels", 2)
+        dat_config = DatLMConfig(**{field: getattr(config, field) for field in DAT_LM_FIELDS})
+        self.model = DatDecoderLM(dat_config)
+        self.score = torch.nn.Linear(config.hidden_dim, self.num_labels, bias=False)
+        self._tied_weights_keys = _collect_shared_weight_keys(self)
+        self.post_init()
+
+    def post_init(self) -> None:
+        super().post_init()
+        _refresh_rope_buffers(self.model, self.config.rope_theta)
+
+    def get_input_embeddings(self) -> torch.nn.Module:
+        return self.model.token_embeddings
+
+    def set_input_embeddings(self, value: torch.nn.Module) -> None:
+        self.model.token_embeddings = value
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        **kwargs,
+    ) -> SequenceClassifierOutput:
+        _, hidden_states = self.model.encode_for_objective(input_ids, attention_mask=attention_mask)
+        logits = self.score(hidden_states)
+
+        batch_size = input_ids.shape[0]
+
+        if self.config.pad_token_id is None:
+            last_non_pad_token = -1
+        else:
+            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+            if (non_pad_mask.sum(-1) == 0).any().item():
+                raise ValueError("Cannot pool sequence-classification logits for all-padding sequences")
+            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+        loss = None
+        if labels is not None:
+            if getattr(self.config, "problem_type", None) is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=pooled_logits,
+            hidden_states=hidden_states,
+        )

special_tokens_map.json

@@ -0,0 +1 @@
+{"bos_token": "<s>", "eos_token": "</s>", "unk_token": "<unk>", "sep_token": "</s>", "pad_token": "<pad>", "cls_token": "<s>", "mask_token": "<mask>"}

tokenizer_config.json

@@ -0,0 +1,10 @@
+{
+    "tokenizer_class": "PreTrainedTokenizerFast",
+    "bos_token": "<s>",
+    "eos_token": "</s>",
+    "unk_token": "<unk>",
+    "sep_token": "</s>",
+    "pad_token": "<pad>",
+    "cls_token": "<s>",
+    "mask_token": "<mask>"
+}

transformer_components.py

@@ -0,0 +1,282 @@
+"""Minimal transformer components for decoder LM training."""
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+
+@dataclass(frozen=True)
+class PositionalInfo:
+    pe_type: str
+    embeddings: Optional[Tensor] = None
+    rope_freqs: Optional[Tuple[Tensor, Tensor]] = None
+    rel_embeddings: Optional[Tensor] = None
+    apply_to_embeddings: bool = False
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def _norm(self, x: Tensor) -> Tensor:
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if x.dtype in (torch.float16, torch.bfloat16):
+            output = self._norm(x.float()).type_as(x)
+        else:
+            output = self._norm(x)
+        return output * self.weight
+
+
+class FeedForward(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        hidden_dim: int,
+        dropout: float = 0.0,
+        activation: str = "gelu",
+        use_bias: bool = True,
+    ):
+        super().__init__()
+        self.activation = activation
+        self.use_swiglu = activation == "swiglu"
+        self.dropout = nn.Dropout(dropout)
+        if self.use_swiglu:
+            self.w_gate = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.w_up = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.w_down = nn.Linear(hidden_dim, input_dim, bias=use_bias)
+            self.activation_fn = nn.SiLU()
+        else:
+            self.linear1 = nn.Linear(input_dim, hidden_dim, bias=use_bias)
+            self.linear2 = nn.Linear(hidden_dim, input_dim, bias=use_bias)
+            if activation == "gelu":
+                self.activation_fn = nn.GELU()
+            elif activation == "relu":
+                self.activation_fn = nn.ReLU()
+            elif activation == "identity":
+                self.activation_fn = nn.Identity()
+            else:
+                raise ValueError(f"Unsupported ffn activation: {activation}")
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        for module in self.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight, gain=1.0)
+                if module.bias is not None:
+                    nn.init.zeros_(module.bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.use_swiglu:
+            gate = self.activation_fn(self.w_gate(x))
+            up = self.w_up(x)
+            x = gate * up
+            x = self.dropout(x)
+            return self.w_down(x)
+
+        x = self.linear1(x)
+        x = self.activation_fn(x)
+        x = self.dropout(x)
+        x = self.linear2(x)
+        return x
+
+
+class PositionalEncoding(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        pe_type: str,
+        max_len: int,
+        theta: float = 10000.0,
+        max_rel_pos: Optional[int] = None,
+        init_range: float = 0.15,
+    ):
+        super().__init__()
+        if embedding_dim <= 0:
+            raise ValueError(f"embedding_dim must be positive, got {embedding_dim}")
+        if max_len <= 0:
+            raise ValueError(f"max_len must be positive, got {max_len}")
+        if pe_type not in {"sinusoidal", "learned", "relative", "rope", "none"}:
+            raise ValueError(f"Unsupported pe_type: {pe_type}")
+        if theta <= 0.0:
+            raise ValueError(f"theta must be positive, got {theta}")
+        if init_range <= 0.0:
+            raise ValueError(f"init_range must be positive, got {init_range}")
+
+        self.embedding_dim = embedding_dim
+        self.pe_type = pe_type
+        self.max_len = max_len
+
+        if pe_type == "sinusoidal":
+            if embedding_dim % 2 != 0:
+                raise ValueError(f"Sinusoidal encoding requires even embedding_dim, got {embedding_dim}")
+            pe = self._precompute_sinusoidal_embeddings(
+                embedding_dim,
+                max_len,
+                device=torch.empty(0).device,
+            )
+            self.register_buffer("sinusoidal_embeddings", pe, persistent=False)
+            self._sinusoidal_uninitialized = True
+        elif pe_type == "learned":
+            self.position_embeddings = nn.Embedding(max_len, embedding_dim)
+            nn.init.uniform_(self.position_embeddings.weight, -init_range, init_range)
+        elif pe_type == "relative":
+            self.max_relative_position = max_rel_pos if max_rel_pos is not None else max_len // 2
+            if self.max_relative_position <= 0:
+                raise ValueError(
+                    f"max_relative_position must be positive, got {self.max_relative_position}"
+                )
+            num_embeddings = 2 * self.max_relative_position + 1
+            self.rel_pos_embeddings_table = nn.Embedding(num_embeddings, embedding_dim)
+            nn.init.uniform_(self.rel_pos_embeddings_table.weight, -init_range, init_range)
+        elif pe_type == "rope":
+            if embedding_dim % 2 != 0:
+                raise ValueError(f"RoPE requires even embedding_dim, got {embedding_dim}")
+            self.theta = theta
+            cos, sin = self._precompute_rope_freqs(embedding_dim, max_len, theta)
+            self.register_buffer("rope_cos", cos, persistent=False)
+            self.register_buffer("rope_sin", sin, persistent=False)
+            self._rope_uninitialized = True
+
+    def get_positional_info(self, seq_len: int, device: torch.device) -> PositionalInfo:
+        if self.pe_type == "none":
+            return PositionalInfo(pe_type="none")
+        if self.pe_type == "sinusoidal":
+            if seq_len > self.sinusoidal_embeddings.size(0):
+                raise ValueError(
+                    f"seq_len {seq_len} exceeds precomputed sinusoidal length {self.sinusoidal_embeddings.size(0)}"
+                )
+            if (
+                getattr(self, "_sinusoidal_uninitialized", False)
+                or self.sinusoidal_embeddings.device.type == "meta"
+            ):
+                embeddings = self._precompute_sinusoidal_embeddings(
+                    self.embedding_dim,
+                    self.max_len,
+                    device=device,
+                )
+                self.register_buffer("sinusoidal_embeddings", embeddings, persistent=False)
+                self._sinusoidal_uninitialized = False
+            return PositionalInfo(
+                pe_type="sinusoidal",
+                embeddings=self.sinusoidal_embeddings[:seq_len].to(device),
+                apply_to_embeddings=True,
+            )
+        if self.pe_type == "learned":
+            if seq_len > self.max_len:
+                raise ValueError(
+                    f"seq_len {seq_len} exceeds maximum learned position length {self.max_len}"
+                )
+            positions = torch.arange(seq_len, dtype=torch.long, device=device)
+            return PositionalInfo(
+                pe_type="learned",
+                embeddings=self.position_embeddings(positions),
+                apply_to_embeddings=True,
+            )
+        if self.pe_type == "relative":
+            return PositionalInfo(
+                pe_type="relative",
+                rel_embeddings=self._generate_relative_embeddings(seq_len, device),
+                apply_to_embeddings=False,
+            )
+        if seq_len > self.rope_cos.size(0):
+            # Dynamically extend RoPE buffer (standard practice for extrapolation)
+            cos, sin = self._precompute_rope_freqs(self.embedding_dim, seq_len, self.theta)
+            self.register_buffer("rope_cos", cos, persistent=False)
+            self.register_buffer("rope_sin", sin, persistent=False)
+            self._rope_uninitialized = False
+        if getattr(self, "_rope_uninitialized", True) or torch.isnan(self.rope_cos[0, 0]):
+            cos, sin = self._precompute_rope_freqs(self.embedding_dim, self.max_len, self.theta)
+            self.rope_cos.copy_(cos)
+            self.rope_sin.copy_(sin)
+            self._rope_uninitialized = False
+        return PositionalInfo(
+            pe_type="rope",
+            rope_freqs=(
+                self.rope_cos[:seq_len].to(device),
+                self.rope_sin[:seq_len].to(device),
+            ),
+            apply_to_embeddings=False,
+        )
+
+    @staticmethod
+    def _rope_inv_freq(dim: int, theta: float, device: torch.device | str) -> Tensor:
+        return 1.0 / (theta ** (torch.arange(0, dim, 2, device=device).float() / dim))
+
+    @staticmethod
+    def _precompute_sinusoidal_embeddings(
+        dim: int,
+        end: int,
+        device: torch.device | str,
+    ) -> Tensor:
+        pe = torch.zeros(end, dim, device=device)
+        position = torch.arange(0, end, dtype=torch.float, device=device).unsqueeze(1)
+        div_term = PositionalEncoding._rope_inv_freq(dim, 10000.0, device=device)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe
+
+    @staticmethod
+    def _precompute_rope_freqs(dim: int, end: int, theta: float) -> Tuple[Tensor, Tensor]:
+        freqs = PositionalEncoding._rope_inv_freq(dim, theta, device="cpu")
+        positions = torch.arange(end, device="cpu")
+        freqs = torch.outer(positions, freqs).float()
+        return torch.cos(freqs), torch.sin(freqs)
+
+    def _generate_relative_embeddings(self, seq_len: int, device: torch.device) -> Tensor:
+        range_q = torch.arange(seq_len, device=device)
+        range_k = torch.arange(seq_len, device=device)
+        distance_mat = range_k[None, :] - range_q[:, None]
+        distance_mat_clipped = torch.clamp(
+            distance_mat,
+            -self.max_relative_position,
+            self.max_relative_position,
+        )
+        final_indices = distance_mat_clipped + self.max_relative_position
+        return self.rel_pos_embeddings_table(final_indices.long())
+
+
+def reshape_for_broadcast(freqs_cis: Tensor, x: Tensor) -> Tensor:
+    ndim = x.ndim
+    if ndim < 2:
+        raise ValueError(f"Input tensor x must have at least 2 dimensions, got {ndim}")
+    if x.shape[-2] != freqs_cis.shape[0] or x.shape[-1] != freqs_cis.shape[-1]:
+        raise ValueError(
+            f"Shape mismatch for RoPE broadcasting: freqs_cis {freqs_cis.shape}, x {x.shape}"
+        )
+    shape = [1] * ndim
+    shape[-2] = x.shape[-2]
+    shape[-1] = x.shape[-1]
+    return freqs_cis.view(shape)
+
+
+def apply_rotary_emb(xq: Tensor, xk: Tensor, freqs_cos: Tensor, freqs_sin: Tensor) -> Tuple[Tensor, Tensor]:
+    if xq.shape[-1] % 2 != 0:
+        raise ValueError(f"Query feature dimension must be even for RoPE, got {xq.shape[-1]}")
+    if xk.shape[-1] % 2 != 0:
+        raise ValueError(f"Key feature dimension must be even for RoPE, got {xk.shape[-1]}")
+
+    xq_r, xq_i = xq.float().reshape(xq.shape[:-1] + (-1, 2)).unbind(-1)
+    xk_r, xk_i = xk.float().reshape(xk.shape[:-1] + (-1, 2)).unbind(-1)
+
+    freqs_cos_q = reshape_for_broadcast(freqs_cos[: xq.shape[-2]], xq_r)
+    freqs_sin_q = reshape_for_broadcast(freqs_sin[: xq.shape[-2]], xq_r)
+    freqs_cos_k = reshape_for_broadcast(freqs_cos[: xk.shape[-2]], xk_r)
+    freqs_sin_k = reshape_for_broadcast(freqs_sin[: xk.shape[-2]], xk_r)
+
+    xq_out_r = xq_r * freqs_cos_q - xq_i * freqs_sin_q
+    xq_out_i = xq_r * freqs_sin_q + xq_i * freqs_cos_q
+    xk_out_r = xk_r * freqs_cos_k - xk_i * freqs_sin_k
+    xk_out_i = xk_r * freqs_sin_k + xk_i * freqs_cos_k
+
+    xq_out = torch.stack([xq_out_r, xq_out_i], dim=-1).flatten(-2)
+    xk_out = torch.stack([xk_out_r, xk_out_i], dim=-1).flatten(-2)
+
+    return xq_out.type_as(xq), xk_out.type_as(xk)

transformer_core.py

@@ -0,0 +1,327 @@
+"""Minimal self-attention core for decoder language models."""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from .transformer_components import PositionalInfo, apply_rotary_emb
+
+
+def _activate_scores(scores: torch.Tensor, activation: str) -> torch.Tensor:
+    if activation == "softmax":
+        return torch.softmax(scores, dim=-1)
+    if activation == "identity":
+        return scores
+    if activation == "relu":
+        return torch.relu(scores)
+    if activation == "tanh":
+        return torch.tanh(scores)
+    if activation == "sigmoid":
+        return torch.sigmoid(scores)
+    if activation == "gelu":
+        return torch.nn.functional.gelu(scores)
+    raise ValueError(f"Unsupported attention activation: {activation}")
+
+
+class MultiHeadAttentionBase(nn.Module):
+    def __init__(
+        self,
+        query_dim: int,
+        output_dim: int,
+        key_dim: Optional[int] = None,
+        value_dim: Optional[int] = None,
+        n_heads: int = 4,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        total_n_heads: Optional[int] = None,
+        activation: str = "softmax",
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        super().__init__()
+        self.query_dim = query_dim
+        self.key_dim = query_dim if key_dim is None else key_dim
+        self.value_dim = self.key_dim if value_dim is None else value_dim
+        self.output_dim = output_dim
+        self.n_heads = n_heads
+        self.hidden_dim = hidden_dim
+        self.total_n_heads = n_heads if total_n_heads is None else total_n_heads
+        self.activation = activation
+        self.use_bias_qkv = use_bias_qkv
+        self.use_bias_out = use_bias_out
+
+        if self.total_n_heads <= 0:
+            raise ValueError(f"total_n_heads must be positive, got {self.total_n_heads}")
+        if hidden_dim % self.total_n_heads != 0:
+            raise ValueError(
+                f"hidden_dim ({hidden_dim}) must be divisible by total_n_heads ({self.total_n_heads})"
+            )
+
+        self.head_dim = hidden_dim // self.total_n_heads
+        self.scale = 1.0 / math.sqrt(self.head_dim)
+        projected_dim = n_heads * self.head_dim
+        self.q_proj = nn.Linear(self.query_dim, projected_dim, bias=use_bias_qkv)
+        self.k_proj = nn.Linear(self.key_dim, projected_dim, bias=use_bias_qkv)
+        self.v_proj = nn.Linear(self.value_dim, projected_dim, bias=use_bias_qkv)
+        self.o_proj = nn.Linear(projected_dim, output_dim, bias=use_bias_out)
+        self.dropout = nn.Dropout(dropout)
+        self.attn_dropout = nn.Dropout(dropout)
+        self.last_attn_weights = None
+
+        self._init_weights()
+
+    def _init_weights(self) -> None:
+        gain = 1.0 / math.sqrt(2.0)
+        for projection in (self.q_proj, self.k_proj, self.v_proj, self.o_proj):
+            nn.init.xavier_uniform_(projection.weight, gain=gain)
+            if projection.bias is not None:
+                nn.init.zeros_(projection.bias)
+
+    def _reshape_for_multihead(self, x: torch.Tensor, batch_size: int, seq_len: int) -> torch.Tensor:
+        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        return x.transpose(1, 2)
+
+    def _process_mask(self, mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+        if mask is None:
+            return None
+        if mask.dim() == 2:
+            if mask.shape[0] != mask.shape[1]:
+                raise ValueError(
+                    "2D decoder attention masks must be square [seq, seq]; "
+                    f"got {tuple(mask.shape)}. Expand padding masks at the LM boundary."
+                )
+            return mask.bool().unsqueeze(0).unsqueeze(0)
+        if mask.dim() == 3:
+            return mask.bool().unsqueeze(1)
+        raise ValueError(f"Mask must be 2D or 3D, got {mask.dim()}D")
+
+    def _compute_attn_scores(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> torch.Tensor:
+        if pos_info is not None and pos_info.rope_freqs is not None:
+            freqs_cos, freqs_sin = pos_info.rope_freqs
+            q, k = apply_rotary_emb(q, k, freqs_cos, freqs_sin)
+        return torch.matmul(q, k.transpose(-2, -1)) * self.scale
+
+    def _apply_activation_and_mask(self, scores: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
+        processed_mask = self._process_mask(mask)
+        if processed_mask is not None and self.activation in {"softmax", "sigmoid", "tanh"}:
+            scores.masked_fill_(~processed_mask, torch.finfo(scores.dtype).min)
+        weights = _activate_scores(scores, self.activation)
+        if processed_mask is not None:
+            weights = weights.masked_fill(~processed_mask, 0.0)
+        return weights
+
+    def _apply_mask(self, scores: torch.Tensor, mask: Optional[torch.Tensor]) -> torch.Tensor:
+        return self._apply_activation_and_mask(scores, mask)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: Optional[torch.Tensor] = None,
+        value: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, _ = query.shape
+        key = query if key is None else key
+        value = key if value is None else value
+        key_len = key.shape[1]
+
+        if key.shape[0] != batch_size or value.shape[0] != batch_size:
+            raise ValueError(
+                f"Batch size mismatch: query={batch_size}, key={key.shape[0]}, value={value.shape[0]}"
+            )
+        if key.shape[1] != value.shape[1]:
+            raise ValueError(
+                f"Key and value sequence length mismatch: {key.shape[1]} vs {value.shape[1]}"
+            )
+
+        q = self.q_proj(query)
+        k = self.k_proj(key)
+        v = self.v_proj(value)
+
+        q = self._reshape_for_multihead(q, batch_size, seq_len)
+        k = self._reshape_for_multihead(k, batch_size, key_len)
+        v = self._reshape_for_multihead(v, batch_size, key_len)
+
+        attn_scores = self._compute_attn_scores(q, k, pos_info)
+        attn_weights = self._apply_activation_and_mask(attn_scores, mask)
+        attn_weights = self.attn_dropout(attn_weights)
+
+        attn_output = torch.matmul(attn_weights, v)
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.view(batch_size, seq_len, self.n_heads * self.head_dim)
+        attn_output = self.o_proj(attn_output)
+        attn_output = self.dropout(attn_output)
+
+        self.last_attn_weights = attn_weights.detach()
+        return attn_output, attn_weights
+
+
+class SelfAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        input_dim: int,
+        n_heads: int = 4,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        supports_relative: bool = False,
+        total_n_heads: Optional[int] = None,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        head_dim = hidden_dim // (n_heads if total_n_heads is None else total_n_heads)
+        super().__init__(
+            query_dim=input_dim,
+            output_dim=n_heads * head_dim,
+            key_dim=input_dim,
+            value_dim=input_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            total_n_heads=total_n_heads,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.supports_relative = supports_relative
+        if self.supports_relative:
+            self.rel_k_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            self.rel_v_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            nn.init.xavier_uniform_(self.rel_k_proj.weight)
+            nn.init.xavier_uniform_(self.rel_v_proj.weight)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.supports_relative and pos_info is not None and pos_info.rel_embeddings is not None:
+            batch_size, seq_len, _ = x.shape
+            rel_embeddings = pos_info.rel_embeddings
+            if rel_embeddings.shape[:2] != (seq_len, seq_len):
+                raise ValueError(
+                    f"Relative embeddings shape {rel_embeddings.shape} does not match sequence length {seq_len}"
+                )
+            if rel_embeddings.shape[2] != self.head_dim:
+                raise ValueError(
+                    f"Relative embeddings dim {rel_embeddings.shape[2]} does not match head_dim {self.head_dim}"
+                )
+
+            rel_k = self.rel_k_proj(rel_embeddings)
+            rel_v = self.rel_v_proj(rel_embeddings)
+
+            q = self.q_proj(x)
+            k = self.k_proj(x)
+            v = self.v_proj(x)
+
+            q = self._reshape_for_multihead(q, batch_size, seq_len)
+            k = self._reshape_for_multihead(k, batch_size, seq_len)
+            v = self._reshape_for_multihead(v, batch_size, seq_len)
+
+            attn_scores = self._compute_attn_scores(q, k, pos_info)
+            rel_scores = torch.einsum("bhid,ijd->bhij", q, rel_k) * self.scale
+            attn_scores = attn_scores + rel_scores
+
+            attn_weights = self._apply_mask(attn_scores, mask)
+            attn_weights = self.attn_dropout(attn_weights)
+
+            attn_output = torch.matmul(attn_weights, v)
+            rel_output = torch.einsum("bhij,ijd->bhid", attn_weights, rel_v)
+            attn_output = attn_output + rel_output
+
+            attn_output = attn_output.transpose(1, 2).contiguous()
+            attn_output = attn_output.view(batch_size, seq_len, self.n_heads * self.head_dim)
+            attn_output = self.o_proj(attn_output)
+            attn_output = self.dropout(attn_output)
+            self.last_attn_weights = attn_weights.detach()
+            return attn_output, attn_weights
+
+        return super().forward(query=x, mask=mask, pos_info=pos_info)
+
+
+class CrossAttention(MultiHeadAttentionBase):
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        context_dim: Optional[int] = None,
+        n_heads: int = 8,
+        hidden_dim: int = 64,
+        dropout: float = 0.0,
+        supports_relative: bool = False,
+        use_bias_qkv: bool = False,
+        use_bias_out: bool = True,
+    ):
+        resolved_context_dim = input_dim if context_dim is None else context_dim
+        super().__init__(
+            query_dim=input_dim,
+            key_dim=resolved_context_dim,
+            value_dim=resolved_context_dim,
+            output_dim=output_dim,
+            n_heads=n_heads,
+            hidden_dim=hidden_dim,
+            dropout=dropout,
+            activation="softmax",
+            use_bias_qkv=use_bias_qkv,
+            use_bias_out=use_bias_out,
+        )
+        self.supports_relative = supports_relative
+        if supports_relative:
+            self.rel_k_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            self.rel_v_proj = nn.Linear(self.head_dim, self.head_dim, bias=False)
+            nn.init.xavier_uniform_(self.rel_k_proj.weight)
+            nn.init.xavier_uniform_(self.rel_v_proj.weight)
+
+    def forward(
+        self,
+        inputs: torch.Tensor,
+        context: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        pos_info: Optional[PositionalInfo] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        if self.supports_relative and pos_info is not None and pos_info.rel_embeddings is not None:
+            batch_size, seq_len_q, _ = inputs.shape
+            _, seq_len_k, _ = context.shape
+            rel_embeddings = pos_info.rel_embeddings
+            if rel_embeddings.shape[:2] != (seq_len_q, seq_len_k):
+                raise ValueError(
+                    f"Relative embeddings shape {rel_embeddings.shape} does not match "
+                    f"sequence lengths {(seq_len_q, seq_len_k)}"
+                )
+            if rel_embeddings.shape[2] != self.head_dim:
+                raise ValueError(
+                    f"Relative embeddings dim {rel_embeddings.shape[2]} does not match head_dim {self.head_dim}"
+                )
+
+            rel_k = self.rel_k_proj(rel_embeddings)
+            rel_v = self.rel_v_proj(rel_embeddings)
+            q = self._reshape_for_multihead(self.q_proj(inputs), batch_size, seq_len_q)
+            k = self._reshape_for_multihead(self.k_proj(context), batch_size, seq_len_k)
+            v = self._reshape_for_multihead(self.v_proj(context), batch_size, seq_len_k)
+            attn_scores = self._compute_attn_scores(q, k, pos_info)
+            attn_scores = attn_scores + torch.einsum("bhid,ijd->bhij", q, rel_k) * self.scale
+            attn_weights = self._apply_activation_and_mask(attn_scores, mask)
+            attn_weights = self.attn_dropout(attn_weights)
+            self.last_attn_weights = attn_weights.detach()
+            attn_output = torch.matmul(attn_weights, v)
+            attn_output = attn_output + torch.einsum("bhij,ijd->bhid", attn_weights, rel_v)
+            attn_output = attn_output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, -1)
+            attn_output = self.o_proj(attn_output)
+            attn_output = self.dropout(attn_output)
+            return attn_output, attn_weights
+
+        return super().forward(
+            query=inputs,
+            key=context,
+            value=context,
+            mask=mask,
+            pos_info=pos_info,
+        )