Update modeling_neollm.py

modeling_neollm.py

@@ -676,7 +676,9 @@ class VectorMultiplier(nn.Module):
 class LinearWithMultipliers(nn.Module):
     """
     Linear layer with optional row and/or column learnable multipliers.
-    Implements: y = (r ⊙ (W @ (c ⊙ x))) + b
+    Implements: y = (r ⊙ (W @ (c ⊙ x))) + b when enabled.
+    With ``enable_multipliers=False`` no multiplier parameters are instantiated
+    and the module reduces to its wrapped ``nn.Linear``.
     """
     def __init__(
         self,
@@ -684,16 +686,18 @@ class LinearWithMultipliers(nn.Module):
         out_features: int,
         bias: bool = True,
         use_row_multiplier: bool = False,
-        use_column_multiplier: bool = False
+        use_column_multiplier: bool = False,
+        enable_multipliers: bool = True,
     ):
         super().__init__()
         self.linear = nn.Linear(in_features, out_features, bias=bias)
-        self.use_row_multiplier    = use_row_multiplier
-        self.use_column_multiplier = use_column_multiplier
+        self.enable_multipliers    = bool(enable_multipliers)
+        self.use_row_multiplier    = bool(use_row_multiplier and self.enable_multipliers)
+        self.use_column_multiplier = bool(use_column_multiplier and self.enable_multipliers)
 
-        if use_row_multiplier:
+        if self.use_row_multiplier:
             self.row_multiplier = VectorMultiplier(out_features, multiplier_type="row")
-        if use_column_multiplier:
+        if self.use_column_multiplier:
             self.column_multiplier = VectorMultiplier(in_features, multiplier_type="column")
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -705,6 +709,53 @@ class LinearWithMultipliers(nn.Module):
         return x
 
 
+
+class EmbeddingWithMultipliers(nn.Module):
+    """
+    Token embedding matrix with optional vocabulary-row and hidden-channel
+    learnable multipliers.
+
+    Effective embedding:
+        E_eff[token, channel] = r_token * E[token, channel] * c_channel
+
+    This mirrors the Learnable Multipliers paper's embedding recommendation,
+    while keeping the actual weight available as ``.weight`` for tooling that
+    expects a standard ``nn.Embedding``-like object.  It is only intended for
+    the untied, non-generator embedding path.
+    """
+    def __init__(
+        self,
+        num_embeddings: int,
+        embedding_dim: int,
+        padding_idx: Optional[int] = None,
+        enable_multipliers: bool = True,
+    ):
+        super().__init__()
+        self.embedding = nn.Embedding(num_embeddings, embedding_dim, padding_idx=padding_idx)
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.padding_idx = padding_idx
+        self.enable_multipliers = bool(enable_multipliers)
+        self.use_row_multiplier = self.enable_multipliers
+        self.use_column_multiplier = self.enable_multipliers
+
+        if self.enable_multipliers:
+            self.row_multiplier = nn.Parameter(torch.ones(num_embeddings))
+            self.column_multiplier = VectorMultiplier(embedding_dim, multiplier_type="column")
+
+    @property
+    def weight(self) -> torch.nn.Parameter:
+        return self.embedding.weight
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        x = self.embedding(input_ids)
+        if self.use_column_multiplier:
+            x = self.column_multiplier(x)
+        if self.use_row_multiplier:
+            row_scale = F.embedding(input_ids, self.row_multiplier.unsqueeze(-1)).to(dtype=x.dtype)
+            x = x * row_scale
+        return x
+
 # ==================== LEVIATHAN CONTINUOUS TOKEN GENERATOR ====================
 
 class LeviathanGenerator(nn.Module):
@@ -3204,7 +3255,9 @@ class NeoLLMAttention(nn.Module):
 
     o_proj variants (controlled by config.use_hadamard_o_proj):
         False (default): dense LinearWithMultipliers — full expressivity,
-                         develops high κ during training (FP8 risk).
+                         with Learnable Multipliers controlled by
+                         config.use_learnable_multipliers; develops high κ
+                         during training (FP8 risk).
         True:            HadamardOProj — fixed WHT + learnable α/β,
                          κ = 1 by construction, 25% fewer attention params,
                          FP8-friendly (Aggarwal & Kumar, 2026, arXiv:2603.08343).
@@ -3242,6 +3295,7 @@ class NeoLLMAttention(nn.Module):
             getattr(config, "lucid_attention_eps", config.rms_norm_eps)
         )
         self.use_hadamard_o_proj           = getattr(config, "use_hadamard_o_proj", False)
+        self.use_learnable_multipliers     = getattr(config, "use_learnable_multipliers", True)
 
         self.fan_layer = FANLayer(
             hidden_size=config.hidden_size,
@@ -3254,6 +3308,7 @@ class NeoLLMAttention(nn.Module):
         self.q_proj = LinearWithMultipliers(
             fan_output_dim, config.num_attention_heads * self.head_dim * 2,
             bias=config.attention_bias, use_row_multiplier=True, use_column_multiplier=False,
+            enable_multipliers=self.use_learnable_multipliers,
         )
         self.num_mea_component_heads = (
             self.mea_component_key_value_heads
@@ -3283,6 +3338,7 @@ class NeoLLMAttention(nn.Module):
             self.o_proj = LinearWithMultipliers(
                 _o_in, config.hidden_size,
                 bias=config.attention_bias, use_row_multiplier=True, use_column_multiplier=True,
+                enable_multipliers=self.use_learnable_multipliers,
             )
 
         self.q_norm = SeeDNorm(self.head_dim, eps=config.rms_norm_eps)
@@ -4800,11 +4856,13 @@ class VersatileFFN(nn.Module):
         self.gate_proj = LinearWithMultipliers(
             fan_dim, config.intermediate_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=False,
+            enable_multipliers=getattr(config, "use_learnable_multipliers", True),
         )
         self.up_proj   = nn.Linear(fan_dim, config.intermediate_size, bias=False)
         self.down_proj = LinearWithMultipliers(
             config.intermediate_size, config.hidden_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=True,
+            enable_multipliers=getattr(config, "use_learnable_multipliers", True),
         )
         self.act_fn  = PolyNorm(exclusive=False)
         self.dropout = nn.Dropout(config.dropout_rate)
@@ -4868,10 +4926,11 @@ class VersatileFFN(nn.Module):
 
         Returns x_in + expert_out  (residual included, matching Eq. 8).
         """
-        # Gate projection: sliced rows + row multiplier
-        g_w   = self.gate_proj.linear.weight[idx]                      # [seg, fan_dim]
-        g_mul = self.gate_proj.row_multiplier.multiplier[idx]          # [seg]
-        gate  = F.linear(x_fan, g_w) * g_mul                          # [N, seg]
+        # Gate projection: sliced rows + optional row multiplier
+        g_w  = self.gate_proj.linear.weight[idx]                       # [seg, fan_dim]
+        gate = F.linear(x_fan, g_w)                                    # [N, seg]
+        if self.gate_proj.use_row_multiplier:
+            gate = gate * self.gate_proj.row_multiplier.multiplier[idx]
 
         # Up projection: sliced rows, no multiplier on up_proj
         u_w = self.up_proj.weight[idx]                                 # [seg, fan_dim]
@@ -4881,11 +4940,13 @@ class VersatileFFN(nn.Module):
         act = self.act_fn(gate) * up                                   # [N, seg]
         act = self.dropout(act)
 
-        # Down projection: column multiplier on input, then row on output
-        col_mul = self.down_proj.column_multiplier.multiplier[idx]     # [seg]
-        d_w     = self.down_proj.linear.weight[:, idx]                 # [hidden, seg]
-        row_mul = self.down_proj.row_multiplier.multiplier             # [hidden]
-        out     = F.linear(act * col_mul, d_w) * row_mul              # [N, hidden]
+        # Down projection: optional column multiplier on input, optional row on output
+        d_w = self.down_proj.linear.weight[:, idx]                     # [hidden, seg]
+        if self.down_proj.use_column_multiplier:
+            act = act * self.down_proj.column_multiplier.multiplier[idx]
+        out = F.linear(act, d_w)                                      # [N, hidden]
+        if self.down_proj.use_row_multiplier:
+            out = out * self.down_proj.row_multiplier.multiplier
 
         return x_in + out
 
@@ -5186,11 +5247,13 @@ class NeoLLMMLP(nn.Module):
         self.gate_proj = LinearWithMultipliers(
             fan_dim, config.intermediate_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=False,
+            enable_multipliers=getattr(config, "use_learnable_multipliers", True),
         )
         self.up_proj   = nn.Linear(fan_dim, config.intermediate_size, bias=False)
         self.down_proj = LinearWithMultipliers(
             config.intermediate_size, config.hidden_size,
             bias=False, use_row_multiplier=True, use_column_multiplier=True,
+            enable_multipliers=getattr(config, "use_learnable_multipliers", True),
         )
         self.act_fn  = PolyNorm(exclusive_init=0.00, exclusive=getattr(config, "polynorm_exclusive", True))
         self.dropout = nn.Dropout(config.dropout_rate)
@@ -6179,6 +6242,10 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             if hasattr(module, "multiplier"):
                 module.multiplier.data.fill_(1.0)
 
+        elif isinstance(module, EmbeddingWithMultipliers):
+            if hasattr(module, "row_multiplier"):
+                module.row_multiplier.data.fill_(1.0)
+
         elif isinstance(module, LeviathanGenerator):
             nn.init.normal_(module.codebooks,
                             mean=0.0, std=self.config.initializer_range)
@@ -6297,9 +6364,22 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         if config.use_token_generator:
             self.token_generator = LeviathanGenerator(config)
         else:
-            self.embed_tokens = nn.Embedding(
-                config.vocab_size, config.hidden_size, config.pad_token_id
+            use_embedding_multipliers = (
+                bool(getattr(config, "use_learnable_multipliers", True))
+                and bool(getattr(config, "use_embedding_multipliers", True))
+                and not bool(getattr(config, "tie_word_embeddings", False))
             )
+            if use_embedding_multipliers:
+                self.embed_tokens = EmbeddingWithMultipliers(
+                    config.vocab_size,
+                    config.hidden_size,
+                    padding_idx=config.pad_token_id,
+                    enable_multipliers=True,
+                )
+            else:
+                self.embed_tokens = nn.Embedding(
+                    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