Update configuration_neollm.py

configuration_neollm.py

@@ -475,12 +475,23 @@ class NeoLLMConfig(PretrainedConfig):
         jtokm_norm_eps=1e-6,
         # ── Hadamard output projection (Aggarwal & Kumar, 2026) ───────────
         use_hadamard_o_proj=True,
+        # ── PolyNorm exclusivity ──────────────────────────────────────────
+        polynorm_exclusive=False,
         # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
         use_spelling_bee_embeddings=False,
         # ── Context Re-Positioning (Li et al., 2026) ──────────────────────
-        use_repo=True,
+        use_repo=False,
         repo_start_layer=None,
         repo_d_p=None,
+        # ── VersatileFFN (Nie et al., 2026) ───────────────────────────────
+        use_versatile_ffn=False,
+        versatile_total_experts=8,
+        versatile_active_experts=2,
+        versatile_max_depth=2,
+        versatile_gumbel_temp_start=5.0,
+        versatile_gumbel_temp_end=0.1,
+        versatile_gumbel_temp_decay=0.99984,
+        versatile_aux_loss_weight=1e-5,
         **kwargs,
     ):
         # ── Generator / tying consistency ─────────────────────────────────
@@ -529,6 +540,20 @@ class NeoLLMConfig(PretrainedConfig):
                 f"num_hidden_layers={num_hidden_layers}."
             )
 
+        # ── VersatileFFN: validate expert configuration ────────────────────
+        if use_versatile_ffn:
+            if not (1 <= versatile_active_experts < versatile_total_experts):
+                raise ValueError(
+                    f"`versatile_active_experts` must satisfy "
+                    f"1 <= versatile_active_experts < versatile_total_experts, "
+                    f"got {versatile_active_experts} vs {versatile_total_experts}."
+                )
+            if intermediate_size % versatile_total_experts != 0:
+                raise ValueError(
+                    f"`intermediate_size` ({intermediate_size}) must be divisible by "
+                    f"`versatile_total_experts` ({versatile_total_experts})."
+                )
+
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
 
         # ── Core Transformer ──────────────────────────────────────────────
@@ -612,6 +637,9 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Hadamard output projection (Aggarwal & Kumar, 2026) ───────────
         self.use_hadamard_o_proj           = use_hadamard_o_proj
 
+        # ── PolyNorm exclusivity ──────────────────────────────────────────
+        self.polynorm_exclusive            = polynorm_exclusive
+
         # ── Spelling Bee Embeddings (Rabe et al., 2026) ───────────────────
         self.use_spelling_bee_embeddings   = use_spelling_bee_embeddings
 
@@ -620,6 +648,16 @@ class NeoLLMConfig(PretrainedConfig):
         self.repo_start_layer              = repo_start_layer
         self.repo_d_p                      = repo_d_p
 
+        # ── VersatileFFN (Nie et al., 2026) ───────────────────────────────
+        self.use_versatile_ffn             = use_versatile_ffn
+        self.versatile_total_experts       = versatile_total_experts
+        self.versatile_active_experts      = versatile_active_experts
+        self.versatile_max_depth           = versatile_max_depth
+        self.versatile_gumbel_temp_start   = versatile_gumbel_temp_start
+        self.versatile_gumbel_temp_end     = versatile_gumbel_temp_end
+        self.versatile_gumbel_temp_decay   = versatile_gumbel_temp_decay
+        self.versatile_aux_loss_weight     = versatile_aux_loss_weight
+
         self.auto_map = {
             "AutoConfig":           "configuration_neollm.NeoLLMConfig",
             "AutoModel":            "modeling_neollm.NeoLLMModel",