Update configuration_neollm.py

configuration_neollm.py

@@ -194,13 +194,36 @@ class NeoLLMConfig(PretrainedConfig):
             tied input/output embeddings and for the Leviathan generator path so
             it does not break parameter sharing or wrap a non-matrix generator.
         use_seednorm (:obj:`bool`, *optional*, defaults to ``True``):
-            Use SeeDNorm as the model's main Transformer normalisation operator.
-            When ``True``, the pre-attention, pre-MLP, final and Q/K
-            normalisations use ``SeeDNorm``. When ``False``, those same backbone
-            normalisation sites instantiate standard ``nn.RMSNorm`` instead,
-            removing SeeDNorm's ``alpha`` and ``beta`` parameters and dynamic
-            rescaling path. MEA output normalisation is controlled separately by
+            Backward-compatible boolean selector for the model's main
+            Transformer normalisation operator.  When ``norm_type`` is left as
+            ``None``, ``use_seednorm=True`` resolves to ``norm_type="seednorm"``
+            and ``use_seednorm=False`` resolves to ``norm_type="rmsnorm"``.
+            MEA output normalisation is still controlled separately by
             ``mea_norm_type``.
+        norm_type (:obj:`str`, *optional*):
+            Main normalisation topology. Valid values are ``"seednorm"``,
+            ``"rmsnorm"`` and ``"siamesenorm"``.  ``"seednorm"`` and
+            ``"rmsnorm"`` keep the original single-stream Pre-Norm topology.
+            ``"siamesenorm"`` activates the two-stream SiameseNorm topology
+            from Li et al. (2026): an X stream with Hybrid/Post-like scaled
+            residual updates and a Y stream with Pre-Norm-like identity
+            accumulation.  The heavy Attention/MLP modules are shared.
+        use_siamesenorm (:obj:`bool`, *optional*, defaults to ``False``):
+            Convenience flag. When ``True``, it forces
+            ``norm_type="siamesenorm"``. This flag exists so training scripts can
+            turn SiameseNorm on without also changing the legacy
+            ``use_seednorm`` field.
+        siamese_norm_base_type (:obj:`str`, *optional*, defaults to ``"rmsnorm"``):
+            Local normaliser used inside the SiameseNorm streams. The paper
+            treats LN generically and the released OLMo configs use RMS-style
+            normalization, so ``"rmsnorm"`` is the default. ``"seednorm"`` is
+            allowed for NeoLLM ablations, but it is not the paper-faithful
+            baseline.
+        siamese_depth_scaling (:obj:`bool`, *optional*, defaults to ``True``):
+            Apply the paper's depth-wise residual scaling to the X stream:
+            residual updates are divided by
+            ``sqrt(2 * (layer_idx + 1))``.  The Y stream keeps the clean
+            unscaled Pre-Norm residual update.
         use_embedding_input_norm (:obj:`bool`, *optional*, defaults to ``True``):
             Controls only the first pre-attention normalisation applied to the
             raw embedding stream. When ``True``, layer 0 instantiates and applies
@@ -363,50 +386,6 @@ class NeoLLMConfig(PretrainedConfig):
             Reference: Rabe, Clymo & Dong (2026). *Spelling Bee Embeddings for
             Language Modeling.* arXiv:2601.18030.
 
-        use_yurii_spelling_velocity (:obj:`bool`, *optional*, defaults to ``False``):
-            Enable the price-quality YuriiFormer × Spelling Bee variant used in
-            ``modeling_neollm.py``. The original Spelling Bee embedding is kept
-            unchanged:
-
-            .. math::
-                X_0(t)=\frac{1}{2}\left(E_{\text{tok}}[t]+S_x(t)\right)
-
-            and an additional byte table builds a spelling-only velocity seed:
-
-            .. math::
-                V_0(t)=\frac{\lambda_v}{2}\left(S_x(t)+S_v(t)\right)
-
-            where ``S_x`` is the existing Spelling Bee character embedding and
-            ``S_v`` is built from a new ``256 × hidden_size`` byte table. The
-            decoder then uses the Nesterov + Lie--Trotter update from
-            YuriiFormer, the article's best loss/price variant because it keeps
-            one attention call and one MLP call per block while adding only a
-            velocity stream, learned scalar coefficients and velocity LayerNorms.
-            This avoids YuriiFormer's full vocabulary velocity table
-            ``vocab_size × hidden_size``.
-
-            Requires ``use_spelling_bee_embeddings=True`` because the velocity
-            branch intentionally reuses the original Spelling Bee spelling table
-            ``S_x`` instead of replacing it.
-
-        yurii_spelling_velocity_scale (:obj:`float`, *optional*, defaults to ``0.10``):
-            Multiplicative scale ``λ_v`` applied to the initial spelling
-            velocity. Keep this small when combining with the original SBE path
-            because ``S_x`` contributes to both ``X_0`` and ``V_0``.
-
-        yurii_mu_init (:obj:`float`, *optional*, defaults to ``0.50``):
-            Initial value for the learned Nesterov lookahead coefficients
-            ``μ`` in the attention and MLP substeps. Stored as logits and
-            constrained to ``(0, 1)`` during forward.
-
-        yurii_beta_init (:obj:`float`, *optional*, defaults to ``0.90``):
-            Initial value for the learned velocity carry coefficients ``β``.
-            Stored as logits and constrained to ``(0, 1)`` during forward.
-
-        yurii_gamma_init (:obj:`float`, *optional*, defaults to ``1.00``):
-            Initial value for the positive learned oracle scale coefficients
-            ``γ``. Stored through an inverse-softplus parameterization.
-
         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
@@ -567,17 +546,13 @@ class NeoLLMConfig(PretrainedConfig):
         - ``use_spelling_bee_embeddings=True`` requires calling
           ``model.model.spelling_bee.set_byte_table(tokenizer)`` before
           training (handled automatically by ``setup_model``).
-        - ``use_yurii_spelling_velocity=True`` requires
-          ``use_spelling_bee_embeddings=True`` and uses the same tokenizer byte
-          table through ``model.model.spelling_bee``. It is mutually exclusive
-          with DCA, AttnRes, LAuReL and StackMemory because it changes the
-          residual update rule into a two-stream YuriiFormer update.
         - ``repo_start_layer`` must satisfy
           ``0 <= repo_start_layer < num_hidden_layers`` when
           ``use_repo=True``, ``use_repo_grape=True`` or
           ``use_repo_goat_prior=True``.
-        - At most one of ``use_dca=True``, ``use_laurel=True``, ``use_attn_res=True`` or
-          ``use_stack_memory=True`` may be active.
+        - At most one of ``use_dca=True``, ``use_laurel=True``,
+          ``use_attn_res=True``, ``use_stack_memory=True`` or
+          ``use_siamesenorm=True`` may be active.
         - When ``use_laurel=True``, at least one of ``use_laurel_rw`` or
           ``use_laurel_lr`` must be ``True``.
         - ``stack_d_model`` must be divisible by ``num_mem_heads`` when
@@ -627,6 +602,9 @@ class NeoLLMConfig(PretrainedConfig):
         ... )
 
     References:
+        Li, T. et al. (2026). *SiameseNorm: Breaking the Barrier to
+        Reconciling Pre/Post-Norm.* arXiv:2602.08064.
+
         Bae, J. et al. (2026). *Affine-Scaled Attention: Towards Flexible and
         Stable Transformer Attention.* arXiv:2602.23057.
 
@@ -723,6 +701,10 @@ class NeoLLMConfig(PretrainedConfig):
         use_embedding_multipliers=False,
         # ── Normalization controls ────────────────────────────────────────
         use_seednorm=False,
+        norm_type=None,
+        use_siamesenorm=True,
+        siamese_norm_base_type="rmsnorm",
+        siamese_depth_scaling=True,
         # ── Embedding input normalization ─────────────────────────────────
         use_embedding_input_norm=True,
         # ── Leviathan continuous token generator ──────────────────────────
@@ -747,12 +729,6 @@ class NeoLLMConfig(PretrainedConfig):
         polynorm_exclusive=False,
         # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
         use_spelling_bee_embeddings=True,
-        # ── YuriiFormer × Spelling Bee velocity seed (Zimin et al., 2026) ──
-        use_yurii_spelling_velocity=True,
-        yurii_spelling_velocity_scale=0.03,
-        yurii_mu_init=0.25,
-        yurii_beta_init=0.70,
-        yurii_gamma_init=0.5,
         # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
         use_repo=True,
         repo_start_layer=None,
@@ -782,6 +758,50 @@ class NeoLLMConfig(PretrainedConfig):
         iha_sliding_window=None,          # auto = N // (2*P^2) usando la longitud real del batch
         **kwargs,
     ):
+        # ── Normalization topology resolution ─────────────────────────────
+        def _canon_norm_name(value, *, allow_siamese: bool) -> str:
+            raw = str(value).lower().replace("_", "").replace("-", "")
+            aliases = {
+                "rms": "rmsnorm",
+                "rmsnorm": "rmsnorm",
+                "seed": "seednorm",
+                "seednorm": "seednorm",
+                "seeednorm": "seednorm",
+                "siamese": "siamesenorm",
+                "siamesenorm": "siamesenorm",
+            }
+            out = aliases.get(raw)
+            valid = {"seednorm", "rmsnorm"} | ({"siamesenorm"} if allow_siamese else set())
+            if out not in valid:
+                raise ValueError(
+                    f"Invalid normalization type {value!r}. Valid values are "
+                    f"{sorted(valid)}."
+                )
+            return out
+
+        if norm_type is None:
+            norm_type = (
+                "siamesenorm"
+                if bool(use_siamesenorm)
+                else ("seednorm" if bool(use_seednorm) else "rmsnorm")
+            )
+        norm_type = _canon_norm_name(norm_type, allow_siamese=True)
+        if bool(use_siamesenorm) and norm_type != "siamesenorm":
+            logger.warning(
+                "`use_siamesenorm=True` overrides `norm_type=%s`; "
+                "using `norm_type='siamesenorm'`.",
+                norm_type,
+            )
+            norm_type = "siamesenorm"
+        use_siamesenorm = norm_type == "siamesenorm"
+        siamese_norm_base_type = _canon_norm_name(siamese_norm_base_type, allow_siamese=False)
+        if use_siamesenorm and not use_embedding_input_norm:
+            logger.warning(
+                "`norm_type='siamesenorm'` is enabled with "
+                "`use_embedding_input_norm=False`. The best SiameseNorm ablation "
+                "uses normalized fused inputs; keep this False only for ablation."
+            )
+
         # ── Generator / tying consistency ─────────────────────────────────
         if use_token_generator and tie_word_embeddings:
             logger.warning(
@@ -831,30 +851,6 @@ class NeoLLMConfig(PretrainedConfig):
                 f"got jtokm_top_k={jtokm_top_k}, jtokm_num_experts={jtokm_num_experts}."
             )
 
-        # ── YuriiFormer × Spelling Bee velocity consistency ───────────────
-        if use_yurii_spelling_velocity and not use_spelling_bee_embeddings:
-            raise ValueError(
-                "`use_yurii_spelling_velocity=True` requires "
-                "`use_spelling_bee_embeddings=True`. The velocity seed is "
-                "defined as V0 = λ/2 · (S_x + S_v), so the original Spelling "
-                "Bee character table S_x must remain active."
-            )
-        if yurii_spelling_velocity_scale < 0.0:
-            raise ValueError(
-                f"`yurii_spelling_velocity_scale` must be >= 0, "
-                f"got {yurii_spelling_velocity_scale}."
-            )
-        for _name, _value in {
-            "yurii_mu_init": yurii_mu_init,
-            "yurii_beta_init": yurii_beta_init,
-        }.items():
-            if not (0.0 < float(_value) < 1.0):
-                raise ValueError(f"`{_name}` must satisfy 0 < value < 1, got {_value}.")
-        if float(yurii_gamma_init) <= 0.0:
-            raise ValueError(
-                f"`yurii_gamma_init` must be > 0, got {yurii_gamma_init}."
-            )
-
         # ── DCA: validate k-DCA source selection ──────────────────────────
         if use_dca and dca_past_layers_k < 1:
             raise ValueError(
@@ -937,22 +933,23 @@ class NeoLLMConfig(PretrainedConfig):
             )
 
         # ── Residual-flow mechanism mutex ─────────────────────────────────
-        # DCA, AttnRes, LAuReL and StackMemory all alter cross-layer residual
-        # routing during training.  Keep exactly one experimental residual-flow
-        # mechanism active at a time so ablations remain interpretable.
+        # DCA, AttnRes, LAuReL, StackMemory and SiameseNorm all alter
+        # cross-layer residual routing during training. Keep exactly one
+        # residual-flow mechanism active at a time so ablations remain
+        # interpretable and the Siamese two-stream invariant stays intact.
         _residual_flow_flags = {
             "use_dca": bool(use_dca),
             "use_laurel": bool(use_laurel),
             "use_attn_res": bool(use_attn_res),
             "use_stack_memory": bool(use_stack_memory),
-            "use_yurii_spelling_velocity": bool(use_yurii_spelling_velocity),
+            "use_siamesenorm": bool(use_siamesenorm),
         }
         _active_residual_flow = [name for name, enabled in _residual_flow_flags.items() if enabled]
         if len(_active_residual_flow) > 1:
             raise ValueError(
                 "Enable at most one residual-flow mechanism among "
                 "`use_dca`, `use_laurel`, `use_attn_res`, `use_stack_memory`, "
-                "and `use_yurii_spelling_velocity`. "
+                "and `use_siamesenorm`. "
                 f"Active flags: {_active_residual_flow}."
             )
 
@@ -1058,7 +1055,14 @@ class NeoLLMConfig(PretrainedConfig):
         self.use_embedding_multipliers     = use_embedding_multipliers
 
         # ── Normalization controls ────────────────────────────────────────
-        self.use_seednorm                  = bool(use_seednorm)
+        self.norm_type                     = norm_type
+        self.use_siamesenorm               = bool(use_siamesenorm)
+        self.siamese_norm_base_type        = siamese_norm_base_type
+        self.siamese_depth_scaling         = bool(siamese_depth_scaling)
+        # Backward-compatible legacy field. It stays true only for the
+        # single-stream SeeDNorm topology; SiameseNorm reads
+        # ``siamese_norm_base_type`` instead.
+        self.use_seednorm                  = norm_type == "seednorm"
 
         # ── Embedding input normalization ─────────────────────────────────
         # True preserves the previous layer-0 pre-norm path, using the active
@@ -1096,13 +1100,6 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
         self.use_spelling_bee_embeddings   = use_spelling_bee_embeddings
 
-        # ── YuriiFormer × Spelling Bee velocity seed (Zimin et al., 2026) ──
-        self.use_yurii_spelling_velocity   = bool(use_yurii_spelling_velocity)
-        self.yurii_spelling_velocity_scale = float(yurii_spelling_velocity_scale)
-        self.yurii_mu_init                 = float(yurii_mu_init)
-        self.yurii_beta_init               = float(yurii_beta_init)
-        self.yurii_gamma_init              = float(yurii_gamma_init)
-
         # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
         self.use_repo                      = use_repo
         self.repo_start_layer              = repo_start_layer