Update configuration_neollm.py

configuration_neollm.py

@@ -556,8 +556,8 @@ class NeoLLMConfig(PretrainedConfig):
           ``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_laurel=True``, ``use_attn_res=True``,
-          ``use_stack_memory=True`` or ``use_siamesenorm=True`` may be active.
+        - At most one residual-flow mechanism may be active, except the
+          explicit ``use_stack_memory=True`` + ``use_siamesenorm=True`` pair.
         - 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
@@ -686,8 +686,8 @@ 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=32,  # H * d_s = 4 * 16
+            use_stack_memory=False,
+        stack_d_model=24,  # H * d_s = 4 * 16
         num_mem_heads=4,  # H = 4
         stack_slots=16,  # S = 24
         stack_memory_cache_size=2048,
@@ -931,9 +931,9 @@ class NeoLLMConfig(PretrainedConfig):
                 )
 
         # ── Residual-flow mechanism mutex ─────────────────────────────────
-        # AttnRes, LAuReL, StackMemory and SiameseNorm all alter the residual
-        # topology.  Keep exactly one experimental residual-flow mechanism
-        # active at a time so ablations remain interpretable.
+        # AttnRes and LAuReL compete directly with residual topology.
+        # StackMemory may be paired with SiameseNorm because it acts on the
+        # shared block input, while SiameseNorm owns the X/Y stream updates.
         _residual_flow_flags = {
             "use_laurel": bool(use_laurel),
             "use_attn_res": bool(use_attn_res),
@@ -943,11 +943,18 @@ class NeoLLMConfig(PretrainedConfig):
         _active_residual_flow = [
             name for name, enabled in _residual_flow_flags.items() if enabled
         ]
-        if len(_active_residual_flow) > 1:
+        allow_stack_siamese = (
+            bool(use_stack_memory)
+            and bool(use_siamesenorm)
+            and not bool(use_laurel)
+            and not bool(use_attn_res)
+        )
+        if len(_active_residual_flow) > 1 and not allow_stack_siamese:
             raise ValueError(
                 "Enable at most one residual-flow mechanism among "
                 "`use_laurel`, `use_attn_res`, `use_stack_memory`, "
-                "and `use_siamesenorm`. "
+                "and `use_siamesenorm`, except the explicit "
+                "`use_stack_memory` + `use_siamesenorm` pair. "
                 f"Active flags: {_active_residual_flow}."
             )