Update configuration_neollm.py

configuration_neollm.py

@@ -518,6 +518,34 @@ class NeoLLMConfig(PretrainedConfig):
             (paper §3.3).  Ignored when ``use_laurel=False`` or
             ``use_laurel_lr=False``.
 
+        use_tweo (:obj:`bool`, *optional*, defaults to ``False``):
+            Enable Transformers Without Extreme Outliers (TWEO), a
+            training-only activation regularizer from Liang et al. (2026).
+            TWEO adds a scaled :math:`L_p` penalty over the final output
+            activations of every Transformer block:
+
+            .. math::
+                L_{total} = L_{task} + \lambda L_{TWEO}, \quad
+                L_{TWEO} = \frac{1}{L}\sum_l
+                \mathbb{E}\left[\left(\frac{|A^{(l)}|}
+                {\tau + \epsilon}\right)^p\right]
+
+            This is non-invasive: no architecture, inference path, or
+            quantization backend is changed. It is intended for FP8/low-bit
+            training runs where extreme activation outliers are harmful.
+        tweo_loss_weight (:obj:`float`, *optional*, defaults to ``0.01``):
+            Weight :math:`\lambda` applied to the TWEO regularizer. The paper
+            uses ``0.01`` in its experiments, optionally with schedule control
+            outside the model.
+        tweo_tau (:obj:`float`, *optional*, defaults to ``3.0``):
+            Soft activation magnitude scale :math:`\tau`. The paper reports
+            robust results for :math:`\tau \in [2, 5]` and defaults to ``3``.
+        tweo_power (:obj:`float`, *optional*, defaults to ``4.0``):
+            Power :math:`p` for the non-linear TWEO penalty. ``4`` preserves
+            small activations while sharply penalizing extreme outliers.
+        tweo_eps (:obj:`float`, *optional*, defaults to ``1e-6``):
+            Numerical stability epsilon added to the TWEO denominator.
+
     Constraints:
         - ``use_jtokm=True`` requires ``use_token_generator=True``.
         - ``1 ≤ jtokm_top_k < jtokm_num_experts`` when ``use_jtokm=True``.
@@ -714,6 +742,12 @@ class NeoLLMConfig(PretrainedConfig):
         use_laurel_rw=False,
         use_laurel_lr=False,
         laurel_lr_rank=32,
+        # ── TWEO: Transformers Without Extreme Outliers (Liang et al., 2026)
+        use_tweo=False,
+        tweo_loss_weight=0.01,
+        tweo_tau=3.0,
+        tweo_power=4.0,
+        tweo_eps=1e-6,
         # ── NITP: Next Implicit Token Prediction (Zhang et al., 2026) ─────
         use_nitp=False,
         nitp_target_layer=None,
@@ -924,6 +958,22 @@ class NeoLLMConfig(PretrainedConfig):
                 "Set `use_laurel_rw=True` and/or `use_laurel_lr=True`."
             )
 
+        # ── TWEO: resolve paper-faithful defaults and validate ────────────
+        # Liang et al. (2026) regularize final output activations from all
+        # Transformer blocks with mean((abs(A) / (tau + eps)) ** p). The paper
+        # uses tau=3, p=4, lambda=0.01 and leaves scheduling of lambda(t) to
+        # the training loop.
+        if tweo_loss_weight < 0.0:
+            raise ValueError(
+                f"`tweo_loss_weight` must be >= 0, got {tweo_loss_weight}."
+            )
+        if tweo_tau <= 0.0:
+            raise ValueError(f"`tweo_tau` must be > 0, got {tweo_tau}.")
+        if tweo_power <= 0.0:
+            raise ValueError(f"`tweo_power` must be > 0, got {tweo_power}.")
+        if tweo_eps < 0.0:
+            raise ValueError(f"`tweo_eps` must be >= 0, got {tweo_eps}.")
+
         # ── NITP: resolve paper-faithful defaults and validate ────────────
         # Zhang et al. (2026) select a shallow contextualized target layer
         # around 20% of total depth, use cosine loss, stop-gradient targets,
@@ -961,9 +1011,7 @@ class NeoLLMConfig(PretrainedConfig):
         # graph-output error seen with NITP because both auxiliary branches
         # share final hidden states with the CCE path.
         if int(nextlat_horizon) < 1:
-            raise ValueError(
-                f"`nextlat_horizon` must be >= 1, got {nextlat_horizon}."
-            )
+            raise ValueError(f"`nextlat_horizon` must be >= 1, got {nextlat_horizon}.")
         if nextlat_mse_weight < 0.0:
             raise ValueError(
                 f"`nextlat_mse_weight` must be >= 0, got {nextlat_mse_weight}."
@@ -1131,6 +1179,13 @@ class NeoLLMConfig(PretrainedConfig):
         self.use_laurel_lr = use_laurel_lr
         self.laurel_lr_rank = laurel_lr_rank
 
+        # ── TWEO: Transformers Without Extreme Outliers ──────────────────
+        self.use_tweo = bool(use_tweo)
+        self.tweo_loss_weight = float(tweo_loss_weight)
+        self.tweo_tau = float(tweo_tau)
+        self.tweo_power = float(tweo_power)
+        self.tweo_eps = float(tweo_eps)
+
         # ── NITP: Next Implicit Token Prediction ─────────────────────────
         self.use_nitp = bool(use_nitp)
         self.nitp_target_layer = int(nitp_target_layer)
@@ -1184,4 +1239,4 @@ class NeoLLMConfig(PretrainedConfig):
         }
 
 
-__all__ = ["NeoLLMConfig"]
+__all__ = ["NeoLLMConfig"]