Update configuration_neollm.py

configuration_neollm.py

@@ -126,6 +126,23 @@ class NeoLLMConfig(PretrainedConfig):
             Epsilon for the denominator of the XSA projection to prevent
             division by zero: ``‖v_ref‖² + xsa_eps``. Ignored when
             ``use_xsa=False``.
+        use_dca (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable **DeepCrossAttention** (DCA; Heddes et al., 2025). DCA
+            replaces the single residual-stream attention input with three
+            independent GRN-v3 compositions over previous decoder outputs,
+            producing separate q/k/v input streams for the attention module.
+            The attention operator itself is left unchanged.
+
+            Mutually exclusive with ``use_laurel=True``, ``use_attn_res=True``
+            and ``use_stack_memory=True`` because all four mechanisms modify
+            cross-layer residual flow during training.
+
+        dca_past_layers_k (:obj:`int`, *optional*, defaults to ``2``):
+            Efficient k-DCA stack size. Each layer keeps the model input, a
+            summed intermediate residual block, and the last ``k`` decoder
+            outputs as explicit GRN sources, following the paper's first/input
+            + intermediate-sum + last-k memory reduction.
+
         use_stack_memory (:obj:`bool`, *optional*, defaults to ``False``):
             Enable the StackMemory / STACKTRANS hidden-state stack between
             decoder layers (Zhang et al., NeurIPS 2025). The module keeps the
@@ -480,7 +497,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``.
-        - ``use_laurel=True`` is incompatible with ``use_attn_res=True``.
+        - At most one of ``use_dca=True``, ``use_laurel=True``, ``use_attn_res=True`` or
+          ``use_stack_memory=True`` may be active.
         - 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
@@ -559,6 +577,9 @@ class NeoLLMConfig(PretrainedConfig):
         Litman, E. & Guo, G. (2026). *You Need Better Attention Priors.*
         arXiv:2601.15380.
 
+        Heddes, M. et al. (2025). *DeepCrossAttention: Supercharging
+        Transformer Residual Connections.* ICML 2025. arXiv:2502.06785.
+
         Menghani, G., Kumar, R. & Kumar, S. (2025). *LAUREL: Learned Augmented
         Residual Layer.* ICML 2025. arXiv:2411.07501.
     """
@@ -587,7 +608,7 @@ class NeoLLMConfig(PretrainedConfig):
         head_dim=64,
         use_momentum_attention=True,
         momentum_gamma=0.10,
-        use_mea_attention=False,
+        use_mea_attention=True,
         mea_component_key_value_heads=None,
         mea_groupnorm_eps=1e-6,
         use_lucid_attention=False,
@@ -603,6 +624,9 @@ class NeoLLMConfig(PretrainedConfig):
         # ── Attention Residuals (Kimi Team, 2026) ─────────────────────────
         use_attn_res=False,
         attn_res_num_blocks=4,
+        # ── DeepCrossAttention (Heddes et al., ICML 2025) ─────────────────
+        use_dca=False,
+        dca_past_layers_k=1,
         # ── StackMemory / STACKTRANS (Zhang et al., NeurIPS 2025) ───────
         use_stack_memory=False,
         stack_d_model=32,          # H * d_s = 4 * 16
@@ -615,7 +639,7 @@ class NeoLLMConfig(PretrainedConfig):
         fan_ratio_ffn=0.0625,
         dropout_rate=0.1,
         # ── Embedding input normalization ─────────────────────────────────
-        use_embedding_input_norm=False,
+        use_embedding_input_norm=True,
         # ── Leviathan continuous token generator ──────────────────────────
         use_token_generator=False,
         generator_d_seed=128,
@@ -695,6 +719,13 @@ class NeoLLMConfig(PretrainedConfig):
                 f"got jtokm_top_k={jtokm_top_k}, jtokm_num_experts={jtokm_num_experts}."
             )
 
+        # ── DCA: validate k-DCA source selection ──────────────────────────
+        if use_dca and dca_past_layers_k < 1:
+            raise ValueError(
+                f"`dca_past_layers_k` must be >= 1 when `use_dca=True`, "
+                f"got {dca_past_layers_k}."
+            )
+
         # ── REPO: resolve defaults and validate ───────────────────────────
         # repo_start_layer defaults to num_hidden_layers // 3, matching the
         # paper's 1/3-of-depth heuristic (Li et al., 2026, §3).
@@ -761,18 +792,25 @@ class NeoLLMConfig(PretrainedConfig):
                 f"`iha_num_pseudo_heads` must be >= 1, got {iha_num_pseudo_heads}."
             )
 
-        # ── LAuReL: mutual exclusion and sub-flag consistency ─────────────
-        # use_laurel and use_attn_res both modify the residual stream and are
-        # structurally incompatible: AttnRes replaces the accumulation entirely
-        # with learned depth-wise attention, while LAuReL scales/augments the
-        # additive residual in-place.
-        if use_laurel and use_attn_res:
+        # ── Residual-flow mechanism mutex ─────────────────────────────────
+        # DCA, AttnRes, LAuReL and StackMemory all alter cross-layer residual
+        # routing during training.  Keep exactly one experimental residual-flow
+        # mechanism active at a time so ablations remain interpretable.
+        _residual_flow_flags = {
+            "use_dca": bool(use_dca),
+            "use_laurel": bool(use_laurel),
+            "use_attn_res": bool(use_attn_res),
+            "use_stack_memory": bool(use_stack_memory),
+        }
+        _active_residual_flow = [name for name, enabled in _residual_flow_flags.items() if enabled]
+        if len(_active_residual_flow) > 1:
             raise ValueError(
-                "`use_laurel=True` is incompatible with `use_attn_res=True`. "
-                "Both methods modify the residual stream: AttnRes replaces it "
-                "with depth-wise softmax attention, while LAuReL applies learned "
-                "scalar/low-rank augmentation in-place. Enable at most one."
+                "Enable at most one residual-flow mechanism among "
+                "`use_dca`, `use_laurel`, `use_attn_res`, and `use_stack_memory`. "
+                f"Active flags: {_active_residual_flow}."
             )
+
+        # ── LAuReL: sub-flag consistency ──────────────────────────────────
         if use_laurel and not use_laurel_rw and not use_laurel_lr:
             raise ValueError(
                 "`use_laurel=True` requires at least one sub-variant to be active. "
@@ -848,6 +886,10 @@ class NeoLLMConfig(PretrainedConfig):
         self.use_attn_res                  = use_attn_res
         self.attn_res_num_blocks           = attn_res_num_blocks
 
+        # ── DeepCrossAttention ───────────────────────────────────────────
+        self.use_dca                       = use_dca
+        self.dca_past_layers_k             = int(dca_past_layers_k)
+
         # ── StackMemory / STACKTRANS ──────────────────────────────────────
         self.use_stack_memory              = use_stack_memory
         self.stack_d_model                 = stack_d_model