Update configuration_neollm.py

configuration_neollm.py

@@ -207,6 +207,12 @@ class NeoLLMConfig(PretrainedConfig):
         stack_memory_cache_size (:obj:`int`, *optional*, defaults to ``2048``):
             Cache length kept for parity with the released StackTrans source.
             The default training path keeps this cache disabled.
+        stack_entropy_loss_weight (:obj:`float`, *optional*, defaults to ``0.001``):
+            Weight ``lambda_St`` for the StackTrans action-entropy regularizer.
+            When StackMemory is active, training adds
+            ``lambda_St * mean(H(push, pop, noop))`` across valid tokens,
+            memory heads, and decoder layers. Set to ``0.0`` to monitor the
+            entropy without applying it to the objective.
         fan_ratio (:obj:`float`, *optional*, defaults to 0.125):
             Ratio controlling the periodic-dimension size in FANformer attention.
             The transformed representation has dimension
@@ -707,11 +713,12 @@ class NeoLLMConfig(PretrainedConfig):
         use_attn_res=False,
         attn_res_num_blocks=4,
         # ── StackMemory / STACKTRANS (Zhang et al., NeurIPS 2025) ───────
-            use_stack_memory=False,
-        stack_d_model=24,  # H * d_s = 4 * 16
+        use_stack_memory=False,
+        stack_d_model=24,  # H * d_s = 4 * 6
         num_mem_heads=4,  # H = 4
-        stack_slots=16,  # S = 24
+        stack_slots=16,  # S = 16
         stack_memory_cache_size=2048,
+        stack_entropy_loss_weight=1e-3,
         # ── ResFormer cross-layer FAN residual (He et al., 2023) ─────────
         use_fan_residual=False,
         fan_ratio=0.125,
@@ -1190,6 +1197,14 @@ class NeoLLMConfig(PretrainedConfig):
                     f"`stack_d_model` must be divisible by `num_mem_heads`, "
                     f"got stack_d_model={stack_d_model}, num_mem_heads={num_mem_heads}."
                 )
+            if (
+                not math.isfinite(float(stack_entropy_loss_weight))
+                or float(stack_entropy_loss_weight) < 0.0
+            ):
+                raise ValueError(
+                    "`stack_entropy_loss_weight` must be finite and >= 0, "
+                    f"got {stack_entropy_loss_weight}."
+                )
 
         super().__init__(tie_word_embeddings=tie_word_embeddings, **kwargs)
 
@@ -1277,6 +1292,7 @@ class NeoLLMConfig(PretrainedConfig):
         self.num_mem_heads = num_mem_heads
         self.stack_slots = stack_slots
         self.stack_memory_cache_size = stack_memory_cache_size
+        self.stack_entropy_loss_weight = float(stack_entropy_loss_weight)
         rope_config_validation(self)
 
         # ── FANformer periodicity ─────────────────────────────────────────