Update modeling_neollm.py

modeling_neollm.py

@@ -4,7 +4,9 @@ NeoLLM model with FANformer, SeeDNorm, ResFormer, Learnable Multipliers,
 full attention augmented with optional Momentum, MEA, and LUCID operators,
 Gated Attention (Qiu et al., 2025) combined with Affine-Scaled Attention
 (Bae et al., 2026), an optional Leviathan continuous token embedding
-generator, and an optional Leviathan-JTok-M token-indexed modulation module.
+generator, an optional Leviathan-JTok-M token-indexed modulation module,
+optional Spelling Bee Embeddings (Rabe et al., 2026), and optional Context
+Re-Positioning (Li et al., 2026).
 
 Attention stack (orthogonal, all active simultaneously when enabled):
     1. Gated Attention (use_gated_attention implicit via q_proj gate chunk):
@@ -49,6 +51,10 @@ References:
         for Rapid, Resource-Efficient Scientific Computation." arXiv:2505.13315.
     JTok / JTok-M: Yang et al. (2026). "JTok: On Token Embedding as Another
         Axis of Scaling Law via Joint Token Self-Modulation." arXiv:2602.00800.
+    Spelling Bee Embeddings: Rabe, Clymo & Dong (2026). "Spelling Bee
+        Embeddings for Language Modeling." arXiv:2601.18030.
+    Context Re-Positioning: Li, Zhao, Cai & Sproat (2026). "REPO: Language
+        Models with Context Re-Positioning." arXiv:2512.14391.
 """
 
 import math
@@ -165,6 +171,27 @@ class HadamardAnalysis:
     alpha_snapshot: Optional[torch.Tensor] = None  # self.alpha [D] — learned scale
 
 
+@dataclass
+class REPOAnalysis:
+    """
+    Internals of a REPOModule forward pass.
+    Only populated when use_repo=True and layer_idx >= repo_start_layer.
+
+    Reference: Li, Zhao, Cai & Sproat (2026). arXiv:2512.14391.
+
+    positions: continuous per-head positions z [B, H, S] produced by f_ϕ.
+               Captures what position pattern the model learned for each head:
+               constant (NoPE-like), monotonic (RoPE-like), or hybrid.
+               Use this field with the attention interpretability toolkit to
+               reproduce the position-pattern analysis of Li et al. (2026) §5.2.
+    r_repr:    intermediate position representation r [B, S, d_p] — output of
+               the SwiGLU sub-layer before the per-head linear W_z.
+               Shared across heads within the layer.
+    """
+    positions: Optional[torch.Tensor] = None  # z [B, H, S] — predicted positions
+    r_repr:    Optional[torch.Tensor] = None  # r [B, S, d_p] — shared repr
+
+
 @dataclass
 class AttentionAnalysis:
     """
@@ -183,9 +210,15 @@ class AttentionAnalysis:
     k_post_norm:  Optional[torch.Tensor] = None  # K after SeeDNorm k_norm [B,H,S,d]
     v_raw:        Optional[torch.Tensor] = None  # V raw (pre MEA/LUCID) [B,H,S,d]
 
-    # ── RoPE (always active) ──────────────────────────────────────────
-    q_post_rope: Optional[torch.Tensor] = None   # Q after RoPE [B,H,S,d]
-    k_post_rope: Optional[torch.Tensor] = None   # K after RoPE [B,H,S,d]
+    # ── RoPE / REPO (always active) ──────────────────────────────────
+    # When use_repo=False or layer_idx < repo_start_layer: standard integer
+    #   RoPE — q_post_rope/k_post_rope are Q/K after apply_rotary_pos_emb.
+    # When use_repo=True and layer_idx >= repo_start_layer: REPO path —
+    #   q_post_rope/k_post_rope are Q/K after _apply_repo_rope with
+    #   continuous per-head positions z from REPOModule.  The positions
+    #   themselves and the intermediate r_repr are in the .repo sub-object.
+    q_post_rope: Optional[torch.Tensor] = None   # Q after RoPE/REPO [B,H,S,d]
+    k_post_rope: Optional[torch.Tensor] = None   # K after RoPE/REPO [B,H,S,d]
 
     # ── Momentum (conditional on use_momentum_attention) ──────────────
     q_momentum_delta: Optional[torch.Tensor] = None  # causal_first_difference(Q)
@@ -238,6 +271,9 @@ class AttentionAnalysis:
     # ── HadamardOProj internals (conditional on use_hadamard_o_proj) ──
     hadamard: Optional["HadamardAnalysis"] = None  # None when dense o_proj active
 
+    # ── REPO position prediction (conditional on use_repo) ────────────
+    repo: Optional["REPOAnalysis"] = None  # None when layer_idx < repo_start_layer
+
 
 @dataclass
 class MLPAnalysis:
@@ -352,6 +388,20 @@ class AnalysisState:
         _ = model(input_ids)
         state = model.last_analysis
         alpha = state.layers[3].attention.alpha_per_head
+
+    REPO access (use_repo=True, layers >= repo_start_layer):
+        # Per-head predicted positions [B, H, S] for layer i:
+        z    = state.layers[i].attention.repo.positions
+
+        # Shared position representation r [B, S, d_p] for layer i:
+        r    = state.layers[i].attention.repo.r_repr
+
+        # Q/K after REPO rotation (identical field name as standard RoPE path):
+        q    = state.layers[i].attention.q_post_rope   # [B, H, S, head_dim]
+        k    = state.layers[i].attention.k_post_rope   # [B, H, S, head_dim]
+
+        # For layers below repo_start_layer: .repo is None, q/k_post_rope
+        # contain standard integer-RoPE rotated Q/K — same field, same shape.
     """
     input_ids:              Optional[torch.Tensor] = None
     embeddings:             Optional[torch.Tensor] = None
@@ -1741,12 +1791,170 @@ class HadamardOProj(nn.Module):
         return out
 
 
+class REPOModule(nn.Module):
+    """
+    Context Re-Positioning module f_ϕ (Li et al., 2026, arXiv:2512.14391).
+
+    Replaces the fixed linear integer indices ``0…L-1`` fed to RoPE with
+    continuous, data-dependent positions ``z_i`` learned end-to-end.
+
+    Architecture (Eq. 4–6 of the paper):
+
+        # Position representation — shared across all heads in this layer
+        r_i = Swish(h_i W_g) ⊙ (h_i W_c)       r_i ∈ R^{d_p}
+
+        # Position assignment — independent per head
+        z_i^(h) = r_i w_z^(h)                    z_i^(h) ∈ R  (scalar)
+
+    where ``h_i ∈ R^d`` is the hidden state of token ``i`` entering the
+    decoder layer (pre-FANLayer), and ``d_p = hidden_size // 8`` by default.
+
+    The resulting positions ``z [B, H, S]`` are real-valued and
+    unconstrained.  They are used to compute per-head ``cos/sin`` embeddings
+    inline from ``inv_freq``, replacing the standard integer-based
+    ``position_embeddings`` for this layer.
+
+    Design notes:
+    - ``W_g`` and ``W_c`` are shared across heads (parameter efficiency).
+    - ``W_z`` is a single ``[d_p, num_heads]`` matrix; each column is the
+      per-head assignment vector ``w_z^(h)``.  Vectorized as one matmul.
+    - The raw hidden state ``h_i`` (not the FAN-augmented or normed variant)
+      is used as input, matching the paper's formulation and avoiding
+      circular dependency with q/k norm.
+    - No bias on any projection — consistent with the paper's Eq. 4–5.
+
+    Reference:
+        Li, H., Zhao, T., Cai, D. & Sproat, R. (2026). "REPO: Language
+        Models with Context Re-Positioning." arXiv:2512.14391.
+    """
+
+    def __init__(self, hidden_size: int, d_p: int, num_heads: int):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.d_p         = d_p
+        self.num_heads   = num_heads
+
+        # SwiGLU position representation (shared across heads, Eq. 4)
+        self.W_g = nn.Linear(hidden_size, d_p, bias=False)
+        self.W_c = nn.Linear(hidden_size, d_p, bias=False)
+
+        # Per-head position assignment (vectorized, Eq. 5)
+        # W_z[:, h] is w_z^(h) for head h
+        self.W_z = nn.Linear(d_p, num_heads, bias=False)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        repo_analysis: Optional[REPOAnalysis] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: [B, S, hidden_size] — residual stream entering the
+                           decoder layer, before FANLayer augmentation.
+            repo_analysis: REPOAnalysis container populated when analysis mode
+                           is active. None during training (zero overhead).
+
+        Returns:
+            z: [B, H, S] — continuous per-head position scalars.
+               z[:, h, i] is the position assigned to token i by head h.
+        """
+        # Position representation (Eq. 4): Swish(h W_g) ⊙ (h W_c)
+        r = F.silu(self.W_g(hidden_states)) * self.W_c(hidden_states)  # [B, S, d_p]
+
+        # Per-head assignment (Eq. 5): z^(h) = r W_z[:, h]
+        # W_z output: [B, S, H] → transpose to [B, H, S]
+        z = self.W_z(r).transpose(1, 2).contiguous()                   # [B, H, S]
+
+        if repo_analysis is not None:
+            repo_analysis.r_repr    = r.detach()
+            repo_analysis.positions = z.detach()
+
+        return z
+
+
+def _apply_repo_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    z: torch.Tensor,
+    inv_freq: torch.Tensor,
+    attention_scaling: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply RoPE to Q and K using continuous per-head positions from REPO.
+
+    Replaces the standard ``apply_rotary_pos_emb(q, k, cos, sin)`` call for
+    layers where REPO is active.  Builds ``cos/sin`` inline from ``z`` and
+    ``inv_freq`` so that the rotation is differentiable w.r.t. ``z`` and
+    therefore w.r.t. the parameters of REPOModule.
+
+    Args:
+        q:                 [B, H, S, head_dim]
+        k:                 [B, H_kv, S, head_dim]  (GQA: H_kv ≤ H)
+        z:                 [B, H, S] — per-head positions from REPOModule
+        inv_freq:          [rotary_dim/2] — frozen RoPE frequency vector
+        attention_scaling: float — scaling factor from NeoLLMRotaryEmbedding
+
+    Returns:
+        (q_embed, k_embed) with the same shapes as (q, k).
+
+    Implementation note on GQA:
+        Q has ``num_attention_heads`` heads; K/V have ``num_key_value_heads``
+        heads (fewer under GQA).  REPO produces one position per Q head.
+        For K we average the positions of the Q heads that map to each KV
+        head (groups of size ``num_key_value_groups``).  This is the minimal
+        approach consistent with the paper's per-head independence claim:
+        each KV head receives a position that is representative of the Q
+        heads it serves.
+    """
+    B, H, S    = z.shape
+    H_kv       = k.shape[1]
+    n_groups   = H // H_kv
+    rotary_dim = inv_freq.shape[0] * 2   # inv_freq covers half the rotary dim
+
+    # inv_freq is already float32 on the correct device (registered as buffer
+    # via set_repo_inv_freq) — no .to() needed, no DeviceCopy op.
+    # No autocast barrier: explicit .float() casts on z_q/z_k are sufficient
+    # to maintain float32 precision for the trig ops. Removing the context
+    # manager lets Inductor plan all intermediate tensors as part of a single
+    # static memory graph, eliminating mid-forward allocations that cause
+    # VRAM variance under max-autotune.
+    inv_freq_f = inv_freq
+
+    # z_q: [B, H, S, 1]  ×  inv_freq: [rotary_dim/2]  →  [B, H, S, rotary_dim/2]
+    z_q     = z.float().unsqueeze(-1)                                # [B, H, S, 1]
+    freqs_q = z_q * inv_freq_f                                       # [B, H, S, r/2]
+    emb_q   = torch.cat([freqs_q, freqs_q], dim=-1)                 # [B, H, S, r]
+    cos_q   = (emb_q.cos() * attention_scaling).to(q.dtype)
+    sin_q   = (emb_q.sin() * attention_scaling).to(q.dtype)
+
+    # KV positions: mean over the Q heads in each GQA group → [B, H_kv, S]
+    z_k     = z.view(B, H_kv, n_groups, S).mean(dim=2)              # [B, H_kv, S]
+    z_k     = z_k.float().unsqueeze(-1)                              # [B, H_kv, S, 1]
+    freqs_k = z_k * inv_freq_f                                       # [B, H_kv, S, r/2]
+    emb_k   = torch.cat([freqs_k, freqs_k], dim=-1)                 # [B, H_kv, S, r]
+    cos_k   = (emb_k.cos() * attention_scaling).to(k.dtype)
+    sin_k   = (emb_k.sin() * attention_scaling).to(k.dtype)
+
+    # Rotate only the first rotary_dim channels; pass the rest through unchanged.
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    q_embed = torch.cat(
+        [(q_rot * cos_q) + (rotate_half(q_rot) * sin_q), q_pass], dim=-1
+    )
+    k_embed = torch.cat(
+        [(k_rot * cos_k) + (rotate_half(k_rot) * sin_k), k_pass], dim=-1
+    )
+    return q_embed, k_embed
+
+
 class NeoLLMAttention(nn.Module):
     """
     Full attention with FANformer, SeeDNorm, ResFormer, Learnable Multipliers,
     optional Momentum, MEA head-level composition, optional LUCID preconditioning,
-    optional Affine-Scaled Attention, optional Exclusive Self Attention, and
-    optional Directional Routing (Taylor, 2026).
+    optional Affine-Scaled Attention, optional Exclusive Self Attention,
+    optional Directional Routing (Taylor, 2026), and optional Context
+    Re-Positioning (Li et al., 2026).
 
     Directional Routing inserts at position C — post-XSA, pre-reshape — where
     the output is already normalized (MEAHeadSeeDNorm) and has auto-position
@@ -1754,11 +1962,23 @@ class NeoLLMAttention(nn.Module):
     orthogonal to the self-position already cleaned by XSA.
 
     Pipeline (all active simultaneously when enabled):
-        FANLayer → q_proj(gate) → q_norm/k_norm → RoPE → Momentum
+        FANLayer → q_proj(gate) → q_norm/k_norm → REPO/RoPE → Momentum
         → MEA(K,V) → LUCID(V) → v_ref → Affine-Scaled SDPA
         → MEAHeadSeeDNorm → XSA → Directional Routing → reshape
         → o_proj · sigmoid(gate) → dropout
 
+    RoPE variants (controlled by config.use_repo and layer_idx):
+        use_repo=False (default): standard integer RoPE via pre-computed
+                         position_embeddings — identical to prior behaviour.
+        use_repo=True, layer_idx >= repo_start_layer:
+                         REPOModule f_ϕ predicts continuous per-head positions
+                         z [B, H, S] from hidden_states.  cos/sin are built
+                         inline from z and inv_freq so the rotation is
+                         differentiable w.r.t. f_ϕ parameters.
+        use_repo=True, layer_idx < repo_start_layer:
+                         standard integer RoPE (lower layers capture surface
+                         features that benefit less from re-positioning).
+
     o_proj variants (controlled by config.use_hadamard_o_proj):
         False (default): dense LinearWithMultipliers — full expressivity,
                          develops high κ during training (FP8 risk).
@@ -1769,6 +1989,7 @@ class NeoLLMAttention(nn.Module):
     References:
         Directional Routing: Taylor (2026). arXiv:2603.14923.
         Hadamard o_proj:     Aggarwal & Kumar (2026). arXiv:2603.08343.
+        Context Re-Positioning: Li et al. (2026). arXiv:2512.14391.
     """
 
     def __init__(self, config: NeoLLMConfig, layer_idx: int):
@@ -1945,6 +2166,53 @@ class NeoLLMAttention(nn.Module):
             self.direction_vecs   = None
             self.direction_router = None
 
+        # ── Context Re-Positioning (Li et al., 2026) ──────────���──────────
+        # Active for layers at or above repo_start_layer only.
+        # Layers below repo_start_layer use standard integer RoPE positions.
+        # inv_freq is accessed from the model's rotary_emb at forward time;
+        # stored here as a non-persistent buffer reference set by NeoLLMModel.
+        self.use_repo = (
+            getattr(config, "use_repo", False)
+            and layer_idx >= getattr(config, "repo_start_layer", config.num_hidden_layers // 3)
+        )
+        if self.use_repo:
+            _d_p = getattr(config, "repo_d_p", config.hidden_size // 8)
+            self.repo_module = REPOModule(
+                hidden_size=config.hidden_size,
+                d_p=_d_p,
+                num_heads=config.num_attention_heads,
+            )
+            # _repo_inv_freq is registered as a non-persistent buffer by
+            # set_repo_inv_freq(), called from NeoLLMModel.__init__ after
+            # rotary_emb is built. Declaring it here would conflict.
+            self._repo_attn_scaling: float = 1.0
+        else:
+            self.repo_module = None
+
+    def set_repo_inv_freq(
+        self,
+        inv_freq: torch.Tensor,
+        attention_scaling: float,
+    ) -> None:
+        """
+        Inject the rotary frequency vector from NeoLLMRotaryEmbedding so that
+        REPO can build cos/sin inline from continuous positions.
+
+        Called once by NeoLLMModel.__init__ after rotary_emb is constructed.
+        Only has effect when use_repo=True for this layer.
+
+        Args:
+            inv_freq:          [rotary_dim/2] — frozen inv_freq buffer from
+                               NeoLLMRotaryEmbedding.
+            attention_scaling: float — attention_scaling from the same module.
+        """
+        if self.use_repo:
+            # Register as non-persistent buffer so .to(device) / .cuda() moves
+            # it automatically — eliminates the DeviceCopy op that splits the
+            # CUDAGraph into 2 partitions when _apply_repo_rope runs.
+            self.register_buffer("_repo_inv_freq", inv_freq.float(), persistent=False)
+            self._repo_attn_scaling = attention_scaling
+
     def _apply_momentum_attention(
         self,
         q: torch.Tensor,
@@ -2113,7 +2381,21 @@ class NeoLLMAttention(nn.Module):
             attn_analysis.v_raw       = v.detach()
 
         cos, sin = position_embeddings
-        q, k     = apply_rotary_pos_emb(q, k, cos, sin)
+        if self.use_repo:
+            # REPO path: f_ϕ predicts continuous per-head positions from the
+            # residual stream, then cos/sin are built inline from those positions
+            # so the rotation is differentiable w.r.t. REPOModule parameters.
+            # (Li et al., 2026, §3.2 — Eq. 6–7)
+            repo_a = attn_analysis.repo if attn_analysis is not None else None
+            z = self.repo_module(hidden_states, repo_analysis=repo_a)  # [B, H, S]
+            q, k = _apply_repo_rope(
+                q, k, z,
+                self._repo_inv_freq,
+                self._repo_attn_scaling,
+            )
+        else:
+            # Standard path: integer positions pre-computed by NeoLLMModel.
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
 
         if attn_analysis is not None:
             attn_analysis.q_post_rope = q.detach()
@@ -2353,7 +2635,7 @@ class NeoLLMMLP(nn.Module):
             config.intermediate_size, config.hidden_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=True,
         )
-        self.act_fn  = PolyNorm(exclusive_init=0.15)
+        self.act_fn  = PolyNorm(exclusive_init=0.05)
         self.dropout = nn.Dropout(config.dropout_rate)
 
     def forward(
@@ -2596,6 +2878,246 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         return outputs
 
 
+class SpellingBeeEmbedding(nn.Module):
+    """
+    Spelling Bee Embeddings (Rabe et al., 2026, arXiv:2601.18030).
+
+    Augments token embeddings with character-level information derived from
+    the UTF-8 byte sequence of each token.  The spelling bee embedding is the
+    mean of the standard token embedding and a character-level summary:
+
+        e_bee(t) = 0.5 * (e_tok(t) + e_chars(t))
+
+        e_chars(t) = inv_sqrt_len(t) * Σ_{i=0}^{15} RoPE(e_byte[b_i], i)
+
+    where inv_sqrt_len = 1/√|t| is precomputed per token type at setup time.
+
+    Key design decisions vs. a naïve per-occurrence implementation:
+
+    1. **Vocab-level computation** — e_chars is built over the full vocabulary
+       once per forward (shape [V, d]), then gathered by token_ids.  A naïve
+       implementation would compute [B*S, 16, d] per step, repeating identical
+       work for every occurrence of a frequent token.  This approach reduces
+       the dominant intermediate from O(B·S·16·d) to O(V·16·d), where V ≪ B·S
+       in practice for most batches.
+
+    2. **Static [256, 16, d] rope_bytes table** — RoPE is applied once over
+       all 256 possible byte values at all 16 positions, producing a table
+       with fully static shapes.  torch.compile / max_autotune can fuse the
+       construction of this table (two elementwise ops + concat over fixed
+       dims) into a single kernel.  Token-level e_chars is then a gather +
+       sum over this table, also fully static.
+
+    3. **Precomputed inv_sqrt_lens** — 1/√byte_len is computed once in
+       set_byte_table and stored as a persistent buffer.  The per-forward
+       normalisation becomes a single elementwise multiply, with no sqrt or
+       division in the hot path.
+
+    Compatible with both the standard embed_tokens path and the
+    LeviathanGenerator path.
+
+    **Inference cost: zero overhead after baking.**
+    Call ``bake_inference_table(token_embeds_weight)`` once after training to
+    collapse the SBE into a single embedding table indistinguishable from a
+    standard nn.Embedding lookup.
+
+    **Setup: call ``set_byte_table(tokenizer)`` once after model init** (and
+    before any .to(device) / FP8 conversion) before training.  The byte table
+    and inv_sqrt_lens are persistent buffers saved in checkpoints.
+
+    References:
+        Rabe, Clymo & Dong (2026). "Spelling Bee Embeddings for Language
+        Modeling." arXiv:2601.18030.
+    """
+
+    MAX_BYTES: int = 16
+
+    def __init__(self, config: "NeoLLMConfig"):
+        super().__init__()
+        d    = config.hidden_size
+        base = getattr(config, "rope_theta", 10000.0)
+
+        # 256 × d byte embedding lookup (one per UTF-8 byte value 0..255).
+        self.byte_emb = nn.Embedding(256, d)
+
+        # ── Persistent buffers (saved in checkpoints) ─────────────────────
+        # token_bytes [vocab_size, MAX_BYTES]: UTF-8 byte values per token,
+        # padded with 0x00 up to MAX_BYTES positions.
+        self.register_buffer(
+            "token_bytes",
+            torch.zeros(config.vocab_size, self.MAX_BYTES, dtype=torch.long),
+            persistent=True,
+        )
+        # inv_sqrt_lens [vocab_size]: precomputed 1/sqrt(byte_len) per token.
+        # Replaces the runtime sqrt+division of the naïve implementation.
+        self.register_buffer(
+            "inv_sqrt_lens",
+            torch.ones(config.vocab_size, dtype=torch.float),
+            persistent=True,
+        )
+
+        # ── Non-persistent buffers (recomputed from fixed formula on load) ─
+        # RoPE cos/sin for intra-token positions 0..MAX_BYTES-1.
+        # Shape [MAX_BYTES, d//2] — applied over the 256-type axis in
+        # _build_rope_bytes, not over the batch/sequence axis.
+        half  = d // 2
+        theta = 1.0 / (base ** (torch.arange(0, half, dtype=torch.float) * 2.0 / d))
+        pos   = torch.arange(self.MAX_BYTES, dtype=torch.float)
+        freqs = torch.outer(pos, theta)                  # [MAX_BYTES, half]
+        self.register_buffer("intra_cos", freqs.cos(), persistent=False)
+        self.register_buffer("intra_sin", freqs.sin(), persistent=False)
+
+        # Static position index [MAX_BYTES] used as the column index in the
+        # vocab-level gather.  Registered as buffer to avoid dynamic tensor
+        # creation inside forward (which would trigger torch.compile retracing).
+        self.register_buffer(
+            "pos_idx",
+            torch.arange(self.MAX_BYTES, dtype=torch.long),
+            persistent=False,
+        )
+
+    # ── Setup ─────────────────────────────────────────────────────────────────
+
+    def set_byte_table(self, tokenizer) -> None:
+        """
+        Precompute the UTF-8 byte table and inv_sqrt_lens from a tokenizer.
+
+        Must be called **once** after model instantiation and **before**
+        ``.to(device)`` / FP8 conversion so the buffers land on the correct
+        device after those transforms.  Both buffers are persistent and will
+        be saved/restored from checkpoints automatically.
+
+        Args:
+            tokenizer: Any HuggingFace tokenizer with
+                ``convert_ids_to_tokens(int) -> str | None``.
+        """
+        vocab_size = self.token_bytes.shape[0]
+        byte_ids   = torch.zeros(vocab_size, self.MAX_BYTES, dtype=torch.long)
+        inv_sqrt   = torch.ones(vocab_size, dtype=torch.float)  # default 1/√1
+
+        for token_id in range(vocab_size):
+            token_str = tokenizer.convert_ids_to_tokens(token_id)
+            if token_str is None:
+                continue
+            # Some tokenizers use a special space character (Ġ / ▁); encode
+            # directly to UTF-8 so byte values match raw text bytes.
+            try:
+                raw = token_str.encode("utf-8")
+            except Exception:
+                raw = b"\x00"
+            n = min(len(raw), self.MAX_BYTES)
+            for i in range(n):
+                byte_ids[token_id, i] = raw[i]
+            inv_sqrt[token_id] = 1.0 / math.sqrt(max(n, 1))
+
+        self.token_bytes.copy_(byte_ids.to(self.token_bytes.device))
+        self.inv_sqrt_lens.copy_(inv_sqrt.to(self.inv_sqrt_lens.device))
+
+    # ── Core helpers ──────────────────────────────────────────────────────────
+
+    def _build_rope_bytes(self) -> torch.Tensor:
+        """
+        Build the static [256, MAX_BYTES, d] RoPE-encoded byte table.
+
+        For each of the 256 possible byte values and each of the MAX_BYTES
+        intra-token positions, applies RoPE rotation using the current
+        byte_emb.weight.  All shapes are fully static, so torch.compile can
+        fuse this into a single kernel.
+
+        Called once per forward pass; the result is discarded afterward.
+        The cost is two broadcast elementwise ops + one cat over fixed dims.
+
+        Returns:
+            rope_bytes [256, MAX_BYTES, d]
+        """
+        w    = self.byte_emb.weight              # [256, d]
+        half = w.shape[-1] // 2
+        w1   = w[:, :half].unsqueeze(1)          # [256, 1, half]
+        w2   = w[:, half:].unsqueeze(1)          # [256, 1, half]
+        cos  = self.intra_cos.unsqueeze(0)       # [1, MAX_BYTES, half]
+        sin  = self.intra_sin.unsqueeze(0)       # [1, MAX_BYTES, half]
+        return torch.cat(
+            [w1 * cos - w2 * sin,
+             w1 * sin + w2 * cos],
+            dim=-1,
+        )                                        # [256, MAX_BYTES, d]
+
+    # ── Forward ───────────────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        token_ids:    torch.Tensor,  # [B, S] or [N]
+        token_embeds: torch.Tensor,  # [B, S, d] or [N, d]
+    ) -> torch.Tensor:
+        """
+        Args:
+            token_ids:    integer token indices to look up byte sequences.
+            token_embeds: embeddings from embed_tokens or LeviathanGenerator.
+        Returns:
+            Spelling bee embeddings — same shape as token_embeds.
+        """
+        # ── Step 1: build rope_bytes over 256 byte types × 16 positions ───
+        # Shape [256, MAX_BYTES, d] — fully static, one kernel via compile.
+        rope_bytes = self._build_rope_bytes()            # [256, MAX_BYTES, d]
+
+        # ── Step 2: build e_chars over vocab types, not occurrences ────────
+        # token_bytes [V, MAX_BYTES]: byte value at each position per token.
+        # pos_idx     [MAX_BYTES]:    column selector 0..MAX_BYTES-1.
+        # rope_bytes[token_bytes, pos_idx[None, :], :] selects, for each
+        # vocab token and each position, the RoPE-rotated embedding of that
+        # byte at that position.  Result [V, MAX_BYTES, d], then sum → [V, d].
+        e_chars_vocab = rope_bytes[
+            self.token_bytes,             # [V, MAX_BYTES] — row index
+            self.pos_idx.unsqueeze(0),    # [1, MAX_BYTES] → broadcast [V, MAX_BYTES]
+        ].sum(1)                          # [V, d]
+
+        # ── Step 3: apply precomputed 1/√byte_len per vocab type ────────────
+        # No sqrt or division in the hot path — pure multiply.
+        e_chars_vocab = e_chars_vocab * self.inv_sqrt_lens.unsqueeze(-1)  # [V, d]
+
+        # ── Step 4: gather only the tokens present in this batch ────────────
+        # This is the only B×S operation — a single embedding lookup.
+        e_chars = e_chars_vocab[token_ids]               # [B, S, d] or [N, d]
+
+        # ── Step 5: mean with token embeddings ──────────────────────────────
+        return (token_embeds + e_chars) * 0.5
+
+    # ── Inference utility ─────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def bake_inference_table(
+        self,
+        token_emb_weight: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Collapse SBE into a single [vocab_size, d] embedding table.
+
+        After baking, the SBE computation is indistinguishable from a standard
+        nn.Embedding lookup — zero additional overhead at inference time.
+
+        Args:
+            token_emb_weight: [vocab_size, d] — weight matrix of embed_tokens
+                              or the equivalent table (e.g. after Leviathan).
+        Returns:
+            [vocab_size, d] — baked spelling bee embedding table.
+
+        Usage::
+
+            baked = model.model.spelling_bee.bake_inference_table(
+                model.model.embed_tokens.weight
+            )
+            model.model.embed_tokens.weight.copy_(baked)
+            # Optionally free byte_emb parameters:
+            # del model.model.spelling_bee
+        """
+        rope_bytes = self._build_rope_bytes()            # [256, MAX_BYTES, d]
+        e_chars_vocab = rope_bytes[
+            self.token_bytes,
+            self.pos_idx.unsqueeze(0),
+        ].sum(1) * self.inv_sqrt_lens.unsqueeze(-1)      # [V, d]
+        return (token_emb_weight + e_chars_vocab) * 0.5
+
+
 class NeoLLMPreTrainedModel(PreTrainedModel):
     """
     Base class with custom weight initialization for all NeoLLM components.
@@ -2629,6 +3151,13 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
                          per head rather than collapsing to 0 or 1.
         - alpha_ma:      zeros  — running EMA starts at 0, β starts as −α/N ≈ small
                          negative offset; model quickly learns to adjust both.
+    REPOModule (Context Re-Positioning):
+        - W_g, W_c, W_z: default normal init from parent _init_weights.
+                         No special initialization required — the SwiGLU
+                         sub-layer starts near-zero, so z_i ≈ 0 for all tokens
+                         at step 0, which is equivalent to constant position
+                         assignment (NoPE-like).  The model quickly learns to
+                         differentiate positions as needed.
     """
     config: NeoLLMConfig
     base_model_prefix            = "model"
@@ -2729,6 +3258,14 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
                 module.attn_res_query_attn.data.zero_()
                 module.attn_res_query_mlp.data.zero_()
 
+        elif isinstance(module, SpellingBeeEmbedding):
+            # byte_emb initialised identically to token embeddings: std=1/√d.
+            # Ensures E[‖e_byte‖²] ≈ 1 at init, matching etok, so the
+            # normalisation factor α = sqrt(byte_len) is calibrated from step 0.
+            d = module.byte_emb.embedding_dim
+            nn.init.normal_(module.byte_emb.weight, mean=0.0, std=1.0 / math.sqrt(d))
+
+
 class NeoLLMModel(NeoLLMPreTrainedModel):
     """
     NeoLLM base decoder-only Transformer.
@@ -2741,11 +3278,30 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
     outputs (or block summaries for Block AttnRes) and passes them to each
     decoder layer, replacing fixed residual accumulation with learned
     depth-wise softmax attention (Kimi Team, 2026, arXiv:2603.15031).
+
+    Spelling Bee Embeddings (Rabe et al., 2026, arXiv:2601.18030):
+        ``use_spelling_bee_embeddings`` is independent of the Leviathan flag:
+        SBE is applied post-embedding regardless of which path produced the
+        token embeddings (embed_tokens or LeviathanGenerator).
+
+        Flag coupling:
+            - use_token_generator=False, use_spelling_bee_embeddings=False
+              → standard embed_tokens, no SBE                          [default]
+            - use_token_generator=False, use_spelling_bee_embeddings=True
+              → standard embed_tokens + SBE
+            - use_token_generator=True,  use_spelling_bee_embeddings=False
+              → LeviathanGenerator only, no SBE
+            - use_token_generator=True,  use_spelling_bee_embeddings=True
+              → LeviathanGenerator + SBE
+
+        Setup: call ``model.model.spelling_bee.set_byte_table(tokenizer)``
+        once after model init and before training (and after any .to(device)).
     """
 
     def __init__(self, config: NeoLLMConfig):
         super().__init__(config)
 
+        # ── Embedding path ────────────────────────────────────────────────────
         if config.use_token_generator:
             self.token_generator = LeviathanGenerator(config)
         else:
@@ -2753,6 +3309,15 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 config.vocab_size, config.hidden_size, config.pad_token_id
             )
 
+        # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────────
+        # Active when use_spelling_bee_embeddings=True, compatible with both
+        # the embed_tokens and LeviathanGenerator paths.
+        use_sbe = getattr(config, "use_spelling_bee_embeddings", False)
+        if use_sbe:
+            self.spelling_bee = SpellingBeeEmbedding(config)
+        else:
+            self.spelling_bee = None
+
         self.layers = nn.ModuleList(
             [NeoLLMDecoderLayer(config, layer_idx)
              for layer_idx in range(config.num_hidden_layers)]
@@ -2764,6 +3329,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
 
         self.post_init()
 
+        # ── REPO: inject inv_freq into every attention layer that uses it ─────
+        # Done after post_init so rotary_emb.inv_freq is already initialized.
+        # Layers below repo_start_layer never call set_repo_inv_freq (their
+        # use_repo flag is False) so the call is harmless for those layers.
+        if getattr(config, "use_repo", False):
+            for layer in self.layers:
+                layer.self_attn.set_repo_inv_freq(
+                    self.rotary_emb.inv_freq,
+                    self.rotary_emb.attention_scaling,
+                )
+
     def get_input_embeddings(self):
         if self.config.use_token_generator:
             return self.token_generator
@@ -2775,7 +3351,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         else:
             self.embed_tokens = value
 
-    def _build_layer_analysis(self) -> LayerAnalysis:
+    def _build_layer_analysis(self, layer_idx: int = 0) -> LayerAnalysis:
         """
         Construct a LayerAnalysis with sub-objects pre-allocated for every
         component that is active in the current config.
@@ -2785,12 +3361,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         Called once per layer per forward when analysis is active.
         """
         cfg = self.config
+        _repo_active = (
+            getattr(cfg, "use_repo", False)
+            and layer_idx >= getattr(cfg, "repo_start_layer", cfg.num_hidden_layers // 3)
+        )
         return LayerAnalysis(
             seednorm_pre_attn  = SeeDNormAnalysis(),
             seednorm_post_attn = SeeDNormAnalysis(),
             attention          = AttentionAnalysis(
                 fan      = FANAnalysis(),
                 hadamard = HadamardAnalysis() if getattr(cfg, "use_hadamard_o_proj", False) else None,
+                repo     = REPOAnalysis() if _repo_active else None,
             ),
             mlp                = MLPAnalysis(
                 fan      = FANAnalysis(),
@@ -2856,6 +3437,13 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             else:
                 inputs_embeds = self.embed_tokens(input_ids)
 
+        # ── Spelling Bee Embeddings (applied post-embedding, pre-decoder) ──────
+        # input_ids may be None when inputs_embeds was passed directly by the
+        # caller; in that case SBE cannot run (no token_ids available) and is
+        # silently skipped — consistent with the standard embedding bypass path.
+        if self.spelling_bee is not None and input_ids is not None:
+            inputs_embeds = self.spelling_bee(input_ids, inputs_embeds)
+
         if analysis_state is not None:
             analysis_state.embeddings = inputs_embeds.detach()
 
@@ -2928,7 +3516,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             # Build per-layer analysis container (only in eval + analysis mode)
             layer_analysis = None
             if analysis_state is not None:
-                layer_analysis           = self._build_layer_analysis()
+                layer_analysis           = self._build_layer_analysis(layer_idx)
                 layer_analysis.layer_idx = layer_idx
                 analysis_state.layers.append(layer_analysis)
 
@@ -3239,6 +3827,7 @@ __all__ = [
     "NeoLLMConfig",
     "LeviathanGenerator",
     "LeviathanJTokM",
+    "SpellingBeeEmbedding",
     "FANLayer",
     "SeeDNorm",
     "ScalarMultiplier",
@@ -3246,6 +3835,7 @@ __all__ = [
     "LinearWithMultipliers",
     "MEAHeadSeeDNorm",
     "HadamardOProj",
+    "REPOModule",
     # Analysis dataclasses — exported so external tools can type-hint against them
     "AnalysisState",
     "LayerAnalysis",
@@ -3256,6 +3846,7 @@ __all__ = [
     "GPASAnalysis",
     "PolyNormAnalysis",
     "HadamardAnalysis",
+    "REPOAnalysis",
     "JTokMAnalysis",
     "AttnResAnalysis",
     "GeneratorAnalysis",