Update configuration_neollm.py

configuration_neollm.py

@@ -12,8 +12,9 @@ class NeoLLMConfig(PretrainedConfig):
 
     Instantiates a NeoLLM model according to the specified arguments, defining the
     full architecture including attention mechanisms, normalization, periodicity
-    modeling, an optional Leviathan continuous token embedding generator, and an
-    optional Leviathan-JTok-M token-indexed modulation module.
+    modeling, an optional Leviathan continuous token embedding generator, an
+    optional Leviathan-JTok-M token-indexed modulation module, and optional
+    Spelling Bee character-level embedding augmentation.
 
     Configuration objects inherit from :class:`~transformers.PretrainedConfig` and
     can be used to control the model outputs. Read the documentation from
@@ -247,6 +248,36 @@ class NeoLLMConfig(PretrainedConfig):
             Coefficient ``λ`` for the load-balancing auxiliary loss.
         jtokm_norm_eps (:obj:`float`, *optional*, defaults to 1e-6):
             Epsilon for L2 normalisation of modulation vectors.
+        use_spelling_bee_embeddings (:obj:`bool`, *optional*, defaults to ``False``):
+            Augment token embeddings with character-level byte information
+            (Rabe, Clymo & Dong, 2026).
+
+            Each token's UTF-8 encoding (up to 16 bytes) is embedded through a
+            shared ``nn.Embedding(256, d)`` table. Byte embeddings are
+            position-encoded with RoPE using intra-token byte positions (not
+            sequence positions), summed and normalised by ``√byte_len``, then
+            averaged with the standard token embedding:
+
+            .. math::
+                e_{\\text{bee}}(t) = \\frac{1}{2}\\left(e_{\\text{tok}}(t) +
+                \\frac{1}{\\sqrt{|t|}}\\sum_{i=1}^{16}
+                \\text{RoPE}(e_{\\text{byte}}[b_i], i)\\right)
+
+            Adds ``256 × hidden_size`` parameters (≈0.13M for d=512).
+            Zero inference overhead when ``bake_inference_table()`` is called
+            after training.
+
+            Compatible with all four combinations of ``use_token_generator``
+            and ``use_spelling_bee_embeddings``.
+
+            **Setup required**: call
+            ``model.model.spelling_bee.set_byte_table(tokenizer)`` once after
+            model instantiation (handled automatically by ``setup_model`` in
+            ``train.py``).
+
+            Reference: Rabe, Clymo & Dong (2026). *Spelling Bee Embeddings for
+            Language Modeling.* arXiv:2601.18030.
+
         use_hadamard_o_proj (:obj:`bool`, *optional*, defaults to ``False``):
             Replace the dense ``W_O ∈ R^{d×d}`` output projection in every
             multi-head attention block with a fixed Walsh–Hadamard Transform
@@ -269,10 +300,64 @@ class NeoLLMConfig(PretrainedConfig):
             Projection: Structured Hadamard Transforms for Efficient
             Transformers.* arXiv:2603.08343.
 
+        use_repo (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable Context Re-Positioning (REPO) in attention layers at or
+            above ``repo_start_layer`` (Li et al., 2026).
+
+            REPO replaces the fixed linear position indices ``0…L-1`` fed to
+            RoPE with continuous, data-dependent positions ``z_i = f_ϕ(h_i)``
+            learned end-to-end. The attention score between tokens ``i`` and
+            ``j`` becomes:
+
+            .. math::
+                A^{\\text{REPO}}_{i,j} = q_i^\\top\\, g_\\theta(z_j - z_i)\\, k_j
+
+            where ``g_θ`` is the standard RoPE rotation and ``z_i`` is
+            predicted from the hidden state ``h_i`` by a lightweight SwiGLU
+            sub-layer ``f_ϕ``:
+
+            .. math::
+                r_i = \\text{Swish}(h_i W_g) \\odot (h_i W_c), \\quad
+                z_i^{(h)} = r_i w_z^{(h)}
+
+            ``W_g, W_c \\in \\mathbb{R}^{d \\times d_p}`` are shared across
+            all query heads within a layer; ``w_z^{(h)} \\in \\mathbb{R}^{d_p}``
+            is learned independently per head. The assigned positions are
+            real-valued and unconstrained — the model may learn constant
+            (NoPE-like), monotonic (RoPE-like), or hybrid patterns as needed.
+
+            Lower layers (``layer_idx < repo_start_layer``) retain the
+            standard integer RoPE positions because they primarily capture
+            surface-level, locally-dependent features that benefit less from
+            re-positioning (Li et al., 2026, §3).
+
+            Overhead: +0.9% parameters; inference latency negligible.
+
+        repo_start_layer (:obj:`int`, *optional*, defaults to
+            ``num_hidden_layers // 3``):
+            Index of the first decoder layer to which REPO is applied.
+            Layers ``[0, repo_start_layer)`` continue to use standard
+            integer RoPE positions. Must satisfy
+            ``0 <= repo_start_layer < num_hidden_layers``.
+            Ignored when ``use_repo=False``.
+
+        repo_d_p (:obj:`int`, *optional*, defaults to
+            ``hidden_size // 8``):
+            Dimensionality of the intermediate position representation
+            ``r_i \\in \\mathbb{R}^{d_p}`` inside ``f_ϕ``. The paper sets
+            ``d_p = d/8`` on the assumption that positional information
+            is less rich than the full hidden representation. Ignored
+            when ``use_repo=False``.
 
     Constraints:
         - ``use_jtokm=True`` requires ``use_token_generator=True``.
         - ``1 ≤ jtokm_top_k < jtokm_num_experts`` when ``use_jtokm=True``.
+        - ``use_spelling_bee_embeddings=True`` requires calling
+          ``model.model.spelling_bee.set_byte_table(tokenizer)`` before
+          training (handled automatically by ``setup_model``).
+        - ``repo_start_layer`` must satisfy
+          ``0 <= repo_start_layer < num_hidden_layers`` when
+          ``use_repo=True``.
 
     Examples::
 
@@ -298,6 +383,13 @@ class NeoLLMConfig(PretrainedConfig):
         ...     use_jtokm=True,
         ... )
 
+        >>> # REPO: context re-positioning from layer 4 onward (default for 12 layers)
+        >>> config_repo = NeoLLMConfig(
+        ...     use_repo=True,
+        ...     # repo_start_layer defaults to num_hidden_layers // 3 = 4
+        ...     # repo_d_p          defaults to hidden_size // 8       = 64
+        ... )
+
     References:
         Bae, J. et al. (2026). *Affine-Scaled Attention: Towards Flexible and
         Stable Transformer Attention.* arXiv:2602.23057.
@@ -315,6 +407,12 @@ class NeoLLMConfig(PretrainedConfig):
 
         Robinson, M. et al. (2025). *Token Embeddings Violate the Manifold
         Hypothesis.* arXiv:2504.01002.
+
+        Rabe, M. N., Clymo, J. & Dong, Z. (2026). *Spelling Bee Embeddings for
+        Language Modeling.* arXiv:2601.18030.
+
+        Li, H., Zhao, T., Cai, D. & Sproat, R. (2026). *REPO: Language Models
+        with Context Re-Positioning.* arXiv:2512.14391.
     """
 
     model_type = "neollm"
@@ -377,6 +475,12 @@ class NeoLLMConfig(PretrainedConfig):
         jtokm_norm_eps=1e-6,
         # ── Hadamard output projection (Aggarwal & Kumar, 2026) ───────────
         use_hadamard_o_proj=True,
+        # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
+        use_spelling_bee_embeddings=False,
+        # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
+        use_repo=True,
+        repo_start_layer=None,
+        repo_d_p=None,
         **kwargs,
     ):
         # ── Generator / tying consistency ─────────────────────────────────
@@ -407,6 +511,24 @@ class NeoLLMConfig(PretrainedConfig):
                 f"got jtokm_top_k={jtokm_top_k}, jtokm_num_experts={jtokm_num_experts}."
             )
 
+        # ── REPO: resolve defaults and validate ───────────────────────────
+        # repo_start_layer defaults to num_hidden_layers // 3, matching the
+        # paper's 1/3-of-depth heuristic (Li et al., 2026, §3).
+        # repo_d_p defaults to hidden_size // 8, matching the paper's
+        # assumption that positional information is less rich than the full
+        # hidden representation (Li et al., 2026, §3.2).
+        if repo_start_layer is None:
+            repo_start_layer = num_hidden_layers // 3
+        if repo_d_p is None:
+            repo_d_p = hidden_size // 8
+        if use_repo and not (0 <= repo_start_layer < num_hidden_layers):
+            raise ValueError(
+                f"`repo_start_layer` must satisfy "
+                f"0 <= repo_start_layer < num_hidden_layers, "
+                f"got repo_start_layer={repo_start_layer}, "
+                f"num_hidden_layers={num_hidden_layers}."
+            )
+
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
 
         # ── Core Transformer ──────────────────────────────────────────────
@@ -490,6 +612,14 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Hadamard output projection (Aggarwal & Kumar, 2026) ───────────
         self.use_hadamard_o_proj           = use_hadamard_o_proj
 
+        # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
+        self.use_spelling_bee_embeddings   = use_spelling_bee_embeddings
+
+        # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
+        self.use_repo                      = use_repo
+        self.repo_start_layer              = repo_start_layer
+        self.repo_d_p                      = repo_d_p
+
         self.auto_map = {
             "AutoConfig":           "configuration_neollm.NeoLLMConfig",
             "AutoModel":            "modeling_neollm.NeoLLMModel",