Update configuration_neollm.py

configuration_neollm.py

@@ -171,6 +171,23 @@ class NeoLLMConfig(PretrainedConfig):
             periodicities in the feature space.
         dropout_rate (:obj:`float`, *optional*, defaults to 0.1):
             General dropout probability for attention outputs and MLP states.
+        use_learnable_multipliers (:obj:`bool`, *optional*, defaults to ``True``):
+            Enable Learnable Multipliers on the matrix layers where the
+            Velikanov et al. placement recommends them. When ``True``, the
+            model keeps row multipliers on ``q_proj`` and MLP ``gate_proj``,
+            row+column multipliers on dense ``o_proj`` and MLP ``down_proj``,
+            and no extra multipliers on ``k_proj``, ``v_proj`` or ``up_proj``.
+            When ``False``, those multiplier parameters are not instantiated
+            and the wrapped linear layers behave as ordinary ``nn.Linear``
+            modules.
+        use_embedding_multipliers (:obj:`bool`, *optional*, defaults to ``True``):
+            Add vocabulary-row and hidden-channel multipliers to the standard
+            ``nn.Embedding`` matrix, matching the Learnable Multipliers paper's
+            recommendation for embeddings. This flag is effective only when
+            ``use_learnable_multipliers=True``, ``use_token_generator=False``
+            and ``tie_word_embeddings=False``. It is intentionally disabled for
+            tied input/output embeddings and for the Leviathan generator path so
+            it does not break parameter sharing or wrap a non-matrix generator.
         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
@@ -638,6 +655,9 @@ class NeoLLMConfig(PretrainedConfig):
         fan_ratio=0.125,
         fan_ratio_ffn=0.0625,
         dropout_rate=0.1,
+        # ── Learnable Multipliers (Velikanov et al., 2026) ─────────────────
+        use_learnable_multipliers=False,
+        use_embedding_multipliers=False,
         # ── Embedding input normalization ─────────────────────────────────
         use_embedding_input_norm=True,
         # ── Leviathan continuous token generator ──────────────────────────
@@ -703,6 +723,27 @@ class NeoLLMConfig(PretrainedConfig):
             )
             tie_word_embeddings = False
 
+        # ── Embedding multiplier applicability ─────────────────────────────
+        if use_embedding_multipliers:
+            if not use_learnable_multipliers:
+                logger.warning(
+                    "`use_embedding_multipliers=True` is inactive because "
+                    "`use_learnable_multipliers=False`. No embedding multiplier "
+                    "parameters will be instantiated."
+                )
+            elif use_token_generator:
+                logger.warning(
+                    "`use_embedding_multipliers=True` is inactive because "
+                    "`use_token_generator=True` replaces the standard embedding "
+                    "matrix with LeviathanGenerator."
+                )
+            elif tie_word_embeddings:
+                logger.warning(
+                    "`use_embedding_multipliers=True` is inactive because "
+                    "`tie_word_embeddings=True`; embedding multipliers would "
+                    "break clean input/output parameter sharing."
+                )
+
         # ── JTok-M / generator dependency ─────────────────────────────────
         if use_jtokm and not use_token_generator:
             raise ValueError(
@@ -906,6 +947,10 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Regularization ────────────────────────────────────────────────
         self.dropout_rate                  = dropout_rate
 
+        # ── Learnable Multipliers ─────────────────────────────────────────
+        self.use_learnable_multipliers     = use_learnable_multipliers
+        self.use_embedding_multipliers     = use_embedding_multipliers
+
         # ── Embedding input normalization ─────────────────────────────────
         # True preserves the previous layer-0 pre-norm path:
         #     embeddings -> SeeDNorm -> LNS -> attention.