Update modeling_neollm.py

modeling_neollm.py

@@ -65,6 +65,8 @@ References:
         arXiv:2601.15380.
     StackMemory / STACKTRANS: Zhang et al. (2025). "Recursive Transformer:
         Boosting Reasoning Ability with State Stack." NeurIPS 2025.
+    TWEO: Liang et al. (2026). "Transformers Without Extreme Outliers Enables
+        FP8 Training And Quantization For Dummies." arXiv:2511.23225.
 """
 
 import math
@@ -5268,6 +5270,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         inputs_embeds: Optional[torch.FloatTensor] = None,
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
+        output_tweo_activations: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Tuple:
@@ -5281,6 +5284,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             if output_attentions is not None
             else self.config.output_attentions
         )
+        output_tweo_activations = bool(output_tweo_activations)
         if return_dict is None:
             cfg_dict = vars(self.config)
             return_dict = cfg_dict.get(
@@ -5341,6 +5345,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         siamese_y_states = inputs_embeds if use_siamesenorm else None
         all_hidden_states = () if output_hidden_states else None
         all_attentions = () if output_attentions else None
+        tweo_activations = () if output_tweo_activations else None
         all_aux_stats = []
 
         # ── StackMemory state ─────────────────────────────────────────────
@@ -5492,6 +5497,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                         all_aux_stats.append(item)
                         break
 
+            if output_tweo_activations:
+                tweo_activations = tweo_activations + (hidden_states,)
+
             if (
                 getattr(self.config, "use_fan_residual", True)
                 and self.first_layer_fan is None
@@ -5511,18 +5519,24 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             all_hidden_states = all_hidden_states + (hidden_states,)
 
         if not return_dict:
-            return tuple(
+            base_outputs = tuple(
                 v
                 for v in [hidden_states, None, all_hidden_states, all_attentions]
                 if v is not None
-            ) + (all_aux_stats,)
+            )
+            if output_tweo_activations:
+                return base_outputs + (all_aux_stats, tweo_activations)
+            return base_outputs + (all_aux_stats,)
 
-        return BaseModelOutputWithPast(
+        outputs = BaseModelOutputWithPast(
             last_hidden_state=hidden_states,
             past_key_values=None,
             hidden_states=all_hidden_states,
             attentions=all_attentions,
-        ), all_aux_stats
+        )
+        if output_tweo_activations:
+            return outputs, all_aux_stats, tweo_activations
+        return outputs, all_aux_stats
 
 
 def _prepare_lm_labels(labels, device, pad_token_id=None):
@@ -5626,6 +5640,31 @@ def compute_liger_loss(
     )
 
 
+def compute_tweo_loss(
+    block_activations: Tuple[torch.Tensor, ...],
+    tau: float,
+    power: float,
+    eps: float,
+) -> torch.Tensor:
+    """
+    Transformers Without Extreme Outliers activation regularizer.
+
+    Liang et al. (2026) apply a scaled Lp penalty to the final output
+    activations of every Transformer block. The penalty is intentionally
+    smooth rather than hinge-thresholded, so normal activations remain cheap
+    while extreme outliers become expensive.
+    """
+    if not block_activations:
+        raise ValueError("TWEO requires at least one Transformer block activation.")
+
+    scale = float(tau) + float(eps)
+    total = block_activations[0].new_zeros((), dtype=torch.float32)
+    for activation in block_activations:
+        scaled = activation.float().abs() / scale
+        total = total + scaled.pow(float(power)).mean()
+    return total / len(block_activations)
+
+
 class NITPProjector(nn.Module):
     """
     SwiGLU projection head for Next Implicit Token Prediction.
@@ -5895,6 +5934,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 f"Got {self.ntp_loss_backend!r}."
             )
         self._last_ntp_loss = None
+        self._last_tweo_loss = None
         self._last_nitp_loss = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
@@ -5955,6 +5995,11 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
+        tweo_enabled = (
+            bool(getattr(self.config, "use_tweo", False))
+            and labels is not None
+            and float(getattr(self.config, "tweo_loss_weight", 0.0)) > 0.0
+        )
         nitp_enabled = (
             bool(getattr(self.config, "use_nitp", False))
             and self.nitp_projector is not None
@@ -5980,30 +6025,48 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             position_ids=position_ids,
             inputs_embeds=inputs_embeds,
             output_hidden_states=model_output_hidden_states,
+            output_tweo_activations=tweo_enabled,
             return_dict=return_dict,
             **kwargs,
         )
 
         # Unpack: model returns (BaseModelOutputWithPast, aux_stats_list)
         # or a tuple when return_dict=False
+        tweo_activations = None
+        hidden_states_tuple = None
         if isinstance(model_out, tuple):
-            outputs, all_aux_stats = model_out[0], model_out[-1]
-            if isinstance(outputs, tuple):
+            outputs = model_out[0]
+            if tweo_enabled and len(model_out) >= 3:
+                all_aux_stats = model_out[-2]
+                tweo_activations = model_out[-1]
+            else:
+                all_aux_stats = model_out[-1]
+            if isinstance(outputs, torch.Tensor):
+                hidden_states = outputs
+                if len(model_out) > 1 and isinstance(model_out[1], tuple):
+                    hidden_states_tuple = model_out[1]
+            elif isinstance(outputs, tuple):
                 hidden_states = outputs[0]
+                if len(outputs) > 2:
+                    hidden_states_tuple = outputs[2]
             else:
                 hidden_states = outputs.last_hidden_state
+                hidden_states_tuple = outputs.hidden_states
         else:
             outputs = model_out
             all_aux_stats = []
             hidden_states = outputs.last_hidden_state
+            hidden_states_tuple = outputs.hidden_states
 
         loss = None
         ntp_loss = None
+        tweo_loss = None
         nitp_loss = None
         nextlat_loss = None
         nextlat_mse_loss = None
         nextlat_kl_loss = None
         self._last_ntp_loss = None
+        self._last_tweo_loss = None
         self._last_nitp_loss = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
@@ -6030,10 +6093,22 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 )
             loss = ntp_loss
 
-            if nitp_enabled:
-                hidden_states_tuple = (
-                    outputs.hidden_states if hasattr(outputs, "hidden_states") else None
+            if tweo_enabled:
+                if not tweo_activations:
+                    raise ValueError(
+                        "`use_tweo=True` requires Transformer block activations. "
+                        "NeoLLMForCausalLM.forward enables them automatically "
+                        "while labels are present."
+                    )
+                tweo_loss = compute_tweo_loss(
+                    block_activations=tweo_activations,
+                    tau=self.config.tweo_tau,
+                    power=self.config.tweo_power,
+                    eps=self.config.tweo_eps,
                 )
+                loss = loss + self.config.tweo_loss_weight * tweo_loss
+
+            if nitp_enabled:
                 target_layer = int(getattr(self.config, "nitp_target_layer", 1))
                 if (
                     hidden_states_tuple is None
@@ -6056,9 +6131,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 loss = loss + self.config.nitp_loss_weight * nitp_loss
 
             if nextlat_enabled:
-                hidden_states_tuple = (
-                    outputs.hidden_states if hasattr(outputs, "hidden_states") else None
-                )
                 if hidden_states_tuple is None or len(hidden_states_tuple) < 2:
                     raise ValueError(
                         "`use_nextlat=True` requires backbone hidden states. "
@@ -6093,6 +6165,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 )
                 loss = loss + aux_loss
             self._last_ntp_loss = ntp_loss.detach() if ntp_loss is not None else None
+            self._last_tweo_loss = tweo_loss.detach() if tweo_loss is not None else None
             self._last_nitp_loss = nitp_loss.detach() if nitp_loss is not None else None
             self._last_nextlat_loss = (
                 nextlat_loss.detach() if nextlat_loss is not None else None
@@ -6117,9 +6190,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             loss=loss,
             logits=logits,
             past_key_values=None,
-            hidden_states=outputs.hidden_states
-            if hasattr(outputs, "hidden_states")
-            else None,
+            hidden_states=hidden_states_tuple,
             attentions=outputs.attentions if hasattr(outputs, "attentions") else None,
         )
 
@@ -6148,4 +6219,4 @@ __all__ = [
 
 AutoConfig.register("neollm", NeoLLMConfig)
 AutoModel.register(NeoLLMConfig, NeoLLMModel)
-AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)
+AutoModelForCausalLM.register(NeoLLMConfig, NeoLLMForCausalLM)