Update configuration_neollm.py

configuration_neollm.py

@@ -1,4 +1,6 @@
 # ==================== configuration_neollm.py ====================
+import math
+
 from transformers.configuration_utils import PretrainedConfig
 from transformers.modeling_rope_utils import rope_config_validation
 from transformers.utils import logging
@@ -62,6 +64,33 @@ class NeoLLMConfig(PretrainedConfig):
             CCE implementation used when ``ntp_loss_backend="cce"``. The default
             keeps the conservative Kahan/full-classifier-gradient path used for
             pretraining from scratch.
+        use_mile_loss (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable MiLe Loss (Su et al., 2026, arXiv:2310.19531) in the extended
+            CCE backend. MiLe reweights token losses with detached,
+            mean-normalized predictive-entropy weights.
+        mile_loss_gamma (:obj:`float`, *optional*, defaults to ``1.0``):
+            Non-negative exponent applied to MiLe's ``1 + entropy`` weight.
+        use_mu_loss (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable the output-embedding centering regularizer from Stollenwerk
+            et al. (2026, arXiv:2601.02031) in the extended CCE backend.
+        mu_loss_lambda (:obj:`float`, *optional*, defaults to ``1e-4``):
+            Non-negative coefficient in ``L_mu = lambda * ||mean(C)||^2``.
+        use_meap (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable Mask-Enhanced Autoregressive Prediction during training only
+            (Zhuang et al., 2026, arXiv:2502.07490). Labels and attention masks
+            remain clean; only selected input token IDs are replaced.
+        meap_mask_ratio (:obj:`float`, *optional*, defaults to ``0.15``):
+            Fraction of eligible input positions replaced in each sequence.
+        meap_mask_token_id (:obj:`int`, *optional*):
+            Token ID used as the MEAP replacement. Required when MEAP is enabled.
+        meap_seed (:obj:`int`, *optional*, defaults to ``0``):
+            Base 32-bit seed. The Trainer derives a distinct deterministic seed
+            for each microstep and distributed process.
+        meap_implementation (:obj:`str`, *optional*, defaults to ``"triton"``):
+            MEAP implementation, either ``"triton"`` or the reference
+            ``"torch"`` path.
+        meap_exclude_last (:obj:`bool`, *optional*, defaults to ``True``):
+            Exclude the final eligible input because it has no shifted target.
         liger_loss_accum_dtype (:obj:`str`, *optional*, defaults to ``"float32"``):
             Accumulator dtype passed to Liger FLCE for weight/bias gradient
             accumulation. ``"float32"`` is the stability-first setting for BF16
@@ -742,6 +771,17 @@ class NeoLLMConfig(PretrainedConfig):
         ntp_loss_backend="cce",
         use_liger_kernel=False,
         cce_loss_impl="cce_kahan_full_c",
+        # ── Optional extended-CCE objectives and MEAP input corruption ────
+        use_mile_loss=True,
+        mile_loss_gamma=1.0,
+        use_mu_loss=True,
+        mu_loss_lambda=1e-4,
+        use_meap=False,
+        meap_mask_ratio=0.15,
+        meap_mask_token_id=None,
+        meap_seed=0,
+        meap_implementation="triton",
+        meap_exclude_last=True,
         liger_loss_accum_dtype="float32",
         rope_theta=10000.0,
         rope_scaling=None,
@@ -851,7 +891,7 @@ class NeoLLMConfig(PretrainedConfig):
         # `nitp_temporal_apply_loss=False` keeps the module instantiated and
         # computes every diagnostic metric, but does not add its objective to
         # the model loss. This is the requested monitoring-only ablation.
-        use_nitp_temporal=True,
+        use_nitp_temporal=False,
         nitp_temporal_apply_loss=False,
         nitp_temporal_horizon=4,
         nitp_temporal_dynamics_weight=1.0,
@@ -888,6 +928,33 @@ class NeoLLMConfig(PretrainedConfig):
         if not str(cce_loss_impl).strip():
             raise ValueError("`cce_loss_impl` must be a non-empty string.")
 
+        if (use_mile_loss or use_mu_loss) and ntp_loss_backend != "cce":
+            raise ValueError(
+                "MiLe and mu-loss require `ntp_loss_backend='cce'`; they are "
+                "not implemented by the Liger loss backend."
+            )
+        if not math.isfinite(float(mile_loss_gamma)) or mile_loss_gamma < 0.0:
+            raise ValueError(
+                f"`mile_loss_gamma` must be finite and >= 0, got {mile_loss_gamma}."
+            )
+        if not math.isfinite(float(mu_loss_lambda)) or mu_loss_lambda < 0.0:
+            raise ValueError(
+                f"`mu_loss_lambda` must be finite and >= 0, got {mu_loss_lambda}."
+            )
+        if not math.isfinite(float(meap_mask_ratio)) or not 0.0 <= meap_mask_ratio <= 1.0:
+            raise ValueError(
+                f"`meap_mask_ratio` must be finite and in [0, 1], got {meap_mask_ratio}."
+            )
+        if use_meap and meap_mask_token_id is None:
+            raise ValueError("`meap_mask_token_id` is required when `use_meap=True`.")
+        if meap_mask_token_id is not None and int(meap_mask_token_id) < 0:
+            raise ValueError("`meap_mask_token_id` must be a non-negative integer.")
+        if not isinstance(meap_seed, int) or not 0 <= meap_seed <= 0xFFFFFFFF:
+            raise ValueError("`meap_seed` must be an integer in [0, 2**32 - 1].")
+        meap_implementation = str(meap_implementation).strip().lower()
+        if meap_implementation not in {"triton", "torch"}:
+            raise ValueError("`meap_implementation` must be 'triton' or 'torch'.")
+
         _valid_liger_accum_dtypes = {
             None,
             "",
@@ -1217,6 +1284,18 @@ class NeoLLMConfig(PretrainedConfig):
         self.ntp_loss_backend = ntp_loss_backend
         self.use_liger_kernel = bool(use_liger_kernel)
         self.cce_loss_impl = str(cce_loss_impl).strip()
+        self.use_mile_loss = bool(use_mile_loss)
+        self.mile_loss_gamma = float(mile_loss_gamma)
+        self.use_mu_loss = bool(use_mu_loss)
+        self.mu_loss_lambda = float(mu_loss_lambda)
+        self.use_meap = bool(use_meap)
+        self.meap_mask_ratio = float(meap_mask_ratio)
+        self.meap_mask_token_id = (
+            None if meap_mask_token_id is None else int(meap_mask_token_id)
+        )
+        self.meap_seed = int(meap_seed)
+        self.meap_implementation = meap_implementation
+        self.meap_exclude_last = bool(meap_exclude_last)
         self.liger_loss_accum_dtype = liger_loss_accum_dtype
 
         # ── Positional encoding ───────────────────────────────────────────