Update modeling_neollm.py

modeling_neollm.py

@@ -72,6 +72,7 @@ References:
         FP8 Training And Quantization For Dummies." arXiv:2511.23225.
 """
 
+import inspect
 import math
 from typing import Optional, Union, Tuple
 
@@ -87,6 +88,11 @@ except ImportError:
     linear_cross_entropy = None
     _CCE_AVAILABLE = False
 
+try:
+    from cut_cross_entropy import meap_mask_inputs  # type: ignore[import-not-found]
+except ImportError:
+    meap_mask_inputs = None
+
 try:
     from liger_kernel.transformers import (  # type: ignore[import-not-found]
         LigerFusedLinearCrossEntropyLoss,
@@ -109,7 +115,7 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_u
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, logging
-from .configuration_neollm import NeoLLMConfig
+from configuration_neollm import NeoLLMConfig
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
@@ -5238,6 +5244,12 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
         - down_proj:     normal(0, 0.01) — keeps the residual transition close
                          to identity initially without blocking gradients into
                          the centered condition gate.
+        - temporal_step_bias: zeros — all horizons start with the same phase.
+                         During the forward pass the bias table is projected to
+                         zero mean over the horizon, so it can represent only
+                         relative phase differences, not one shared phase offset.
+        - temporal_step_gain_logits: zeros — 2*sigmoid(0)=1, so every horizon
+                         starts with exact unit residual gain.
         - context_norm/condition_norm: weight=1.
     NeoLLMAttention (Affine-Scaled Attention):
         - alpha_proj:    normal(0, 0.02)  — near-zero so linear_clipping(≈0) ≈ 0.5
@@ -5277,6 +5289,8 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
             nn.init.xavier_uniform_(module.context_up_proj.weight)
             nn.init.xavier_uniform_(module.condition_modulation_proj.weight)
             nn.init.normal_(module.down_proj.weight, mean=0.0, std=0.01)
+            nn.init.zeros_(module.temporal_step_bias)
+            nn.init.zeros_(module.temporal_step_gain_logits)
             module.context_norm.weight.data.fill_(1.0)
             module.condition_norm.weight.data.fill_(1.0)
 
@@ -5885,6 +5899,27 @@ def _resolve_liger_accum_dtype(accum_dtype):
     )
 
 
+_EXTENDED_CCE_INSTALL = (
+    "cut-cross-entropy @ "
+    "git+https://github.com/Kitsunp/ml-cross-entropy.git@main"
+)
+
+
+def _require_extended_cce_options(*option_names: str) -> None:
+    """Fail clearly only when an explicitly enabled CCE extension is unavailable."""
+    if linear_cross_entropy is None:
+        return
+    parameters = inspect.signature(linear_cross_entropy).parameters
+    missing = [name for name in option_names if name not in parameters]
+    if missing:
+        raise ImportError(
+            "The installed cut-cross-entropy package does not provide the enabled "
+            f"extension options {missing}. Install the extended backend with "
+            f"`pip install --force-reinstall --no-deps \"{_EXTENDED_CCE_INSTALL}\"`, "
+            "or disable the corresponding NeoLLM flags."
+        )
+
+
 @torch.compiler.disable
 def compute_cce_loss(
     hidden_states,
@@ -5893,6 +5928,11 @@ def compute_cce_loss(
     lm_head_bias=None,
     pad_token_id=None,
     cce_impl="cce_kahan_full_c",
+    use_mile_loss=False,
+    mile_loss_gamma=1.0,
+    use_mu_loss=False,
+    mu_loss_lambda=1e-4,
+    return_loss_metrics=False,
 ):
     """CCE loss excluded from torch.compile, preserving the configured CCE impl."""
     if linear_cross_entropy is None:
@@ -5905,6 +5945,22 @@ def compute_cce_loss(
     processed_labels = _prepare_lm_labels(
         labels, device=hidden_states.device, pad_token_id=pad_token_id
     )
+    extension_kwargs = {}
+    if use_mile_loss:
+        _require_extended_cce_options("mile_enabled", "mile_gamma")
+        extension_kwargs.update(
+            mile_enabled=True,
+            mile_gamma=float(mile_loss_gamma),
+        )
+    if use_mu_loss:
+        _require_extended_cce_options("mu_loss_enabled", "mu_loss_lambda")
+        extension_kwargs.update(
+            mu_loss_enabled=True,
+            mu_loss_lambda=float(mu_loss_lambda),
+        )
+    if return_loss_metrics:
+        _require_extended_cce_options("return_loss_metrics")
+        extension_kwargs["return_loss_metrics"] = True
     return linear_cross_entropy(
         hidden_states,
         lm_head_weight,
@@ -5913,6 +5969,7 @@ def compute_cce_loss(
         shift=1,
         impl=cce_impl,
         reduction="mean",
+        **extension_kwargs,
     )
 
 
@@ -6063,15 +6120,38 @@ class NITPTemporalTransition(nn.Module):
         h_bar = RMSNorm_h(h_t)
         z_bar = RMSNorm_z(stopgrad(P_NITP(h_t)))
         u_h = SiLU(W_gate h_bar) * (W_up h_bar)
-        s_z = tanh(W_condition z_bar)
-        delta_t = W_down(u_h * s_z)
-        h_hat_{t+1} = h_t + delta_t
+        b_tilde_k = b_k - mean_j(b_j)
+        s_{z,k} = tanh(W_condition z_bar + b_tilde_k)
+        g_k = 2 * sigmoid(a_k)
+        delta_{t,k} = g_k * W_down(u_h * s_{z,k})
+        h_hat_{t+1} = h_t + delta_{t,k}
 
-    The centered signed gate satisfies ``s_z in (-1, 1)``.  Unlike a
+    The centered signed gate satisfies ``s_{z,k} in (-1, 1)``.  Unlike a
     positive sigmoid gate, it can attenuate, activate, or reverse an
-    individual contextual feature.  More importantly, ``W_condition=0``
-    now implies ``s_z=0`` and therefore ``T(h, z)=h``: the transition can
-    no longer learn a non-trivial update while ignoring the NITP condition.
+    individual contextual feature.  A zero-initialized learned vector ``b_k``
+    modulates only the recurrent phase.  Before indexing a horizon, the table is
+    reparameterized as
+
+        b_tilde_k = b_k - mean_j(b_j),
+
+    so ``sum_k b_tilde_k = 0``.  All horizons still share the same transition
+    weights, but the bias table can encode only relative differences between
+    rollout phases; a single offset shared by every horizon is removed exactly.
+    A second zero-initialized scalar logit ``a_k`` supplies a positive bounded
+    residual gain ``g_k = 2 sigmoid(a_k)`` for each horizon.  Thus
+    ``b_tilde_k`` controls the relative phase/channel pattern, while ``g_k``
+    controls only its global amplitude.  At initialization ``b_k=0`` and
+    ``a_k=0`` for every horizon, hence ``b_tilde_k=0``, ``g_k=1``, and the model
+    starts exactly from the previous transition.  The step index is
+    deterministic and contains no future-token information.
+
+    At initialization, ``W_condition=0`` implies ``s_{z,k}=0`` and therefore
+    ``T_k(h, z)=h``.  After training, a non-zero relative phase bias can provide
+    a horizon-specific gate even when the condition projection is weak, but it
+    cannot create the same bias offset at every horizon.  The gain can rescale
+    the resulting residual, yet neither phase nor gain creates an additive route
+    through ``down_proj``: every output update remains multiplicatively tied to
+    contextual features proposed by the current causal state.
 
     All projections are bias-free.  Therefore the NITP condition has no
     additive route to the output and the transition still satisfies:
@@ -6085,17 +6165,21 @@ class NITPTemporalTransition(nn.Module):
 
     Parameter count, excluding the already-existing NITP projector:
 
-        4 * d * m + 2d
+        4 * d * m + 2d + W * m + W
 
-    where ``m = nitp_temporal_intermediate_size`` and ``2d`` comes from the two
-    affine RMSNorm scales.  With d=512 and m=1024 this is 2,098,176 trainable
-    parameters.
+    where ``m = nitp_temporal_intermediate_size``, ``W`` is the temporal
+    horizon, and ``2d`` comes from the two affine RMSNorm scales.  With d=512,
+    m=1024, and W=4 this is 2,102,276 trainable parameters: 4,100 more than
+    the shared transition, of which only four are residual-gain parameters.
     """
 
     def __init__(self, config: NeoLLMConfig):
         super().__init__()
         hidden_size = int(config.hidden_size)
         intermediate_size = int(config.nitp_temporal_intermediate_size)
+        horizon = int(config.nitp_temporal_horizon)
+        if horizon < 1:
+            raise ValueError("nitp_temporal_horizon must be at least 1")
 
         # The two inputs have different semantics and are normalized
         # independently: a final-layer causal state and a predicted shallow
@@ -6112,17 +6196,33 @@ class NITPTemporalTransition(nn.Module):
         )
 
         # NITP-condition path: controls contextual features but has no additive
-        # route to the output.  The centered tanh gate forces every non-zero
-        # update to depend on the detached NITP condition.
+        # route to the output.  The signed tanh gate is centered at zero; the
+        # condition and the zero-mean relative phase bias modulate features
+        # proposed by the current causal state.
         self.condition_modulation_proj = nn.Linear(
             hidden_size, intermediate_size, bias=False
         )
         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
 
+        # Minimal phase conditioning.  Each horizon receives one learned raw
+        # bias in the shared intermediate gate.  Forward reparameterizes this
+        # table to zero mean over horizons, so only relative phase differences
+        # are expressible through this branch.  Zero initialization preserves
+        # the exact previous transition at step 0 and avoids disturbing warmup.
+        self.temporal_step_bias = nn.Parameter(
+            torch.zeros(horizon, intermediate_size)
+        )
+
+        # Minimal amplitude conditioning.  One scalar logit per horizon is
+        # mapped to a positive bounded gain in (0, 2).  Zero initialization
+        # gives exactly unit gain: 2 * sigmoid(0) = 1.
+        self.temporal_step_gain_logits = nn.Parameter(torch.zeros(horizon))
+
     def forward(
         self,
         current_states: torch.Tensor,
         next_shallow_conditions: torch.Tensor,
+        step_index: int,
     ) -> torch.Tensor:
         context = self.context_norm(current_states)
         condition = self.condition_norm(next_shallow_conditions)
@@ -6131,11 +6231,22 @@ class NITPTemporalTransition(nn.Module):
             F.silu(self.context_gate_proj(context))
             * self.context_up_proj(context)
         )
+        # Remove the horizon-common bias mode without materializing a full
+        # broadcast tensor.  For horizon=1 this term is exactly zero, as there
+        # is no relative phase to represent.
+        relative_step_bias = (
+            self.temporal_step_bias[step_index]
+            - self.temporal_step_bias.mean(dim=0)
+        )
         signed_modulation = torch.tanh(
             self.condition_modulation_proj(condition)
+            + relative_step_bias
         )
         delta = self.down_proj(context_update * signed_modulation)
-        return current_states + delta
+        step_gain = 2.0 * torch.sigmoid(
+            self.temporal_step_gain_logits[step_index]
+        )
+        return current_states + step_gain * delta
 
 
 class NextLatDynamicsModel(nn.Module):
@@ -6531,6 +6642,170 @@ def _off_diagonal_pair_mean(
     return _masked_mean(selected.mean(dim=-1), window_mask)
 
 
+def _temporal_two_token_readout(
+    student_states: torch.Tensor,
+    teacher_states: torch.Tensor,
+    candidate_labels: torch.LongTensor,
+    candidate_valid: torch.Tensor,
+    target_index: int,
+    base_mask: torch.Tensor,
+    lm_head_weight: torch.Tensor,
+    lm_head_bias: Optional[torch.Tensor],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Cheap functional anchor over the local temporal candidate set.
+
+    For the real state at horizon ``k`` the candidate set contains only the
+    tokens attached to the W local future horizons.  The strongest *different*
+    temporal token under the real state is selected as the competitor.  The
+    rollout state then matches the teacher's binary distribution over
+
+        [token at the correct horizon, strongest incorrect temporal token].
+
+    This avoids a [B,T,W,V] or [B,T,V] vocabulary projection.  It evaluates W
+    teacher dot-products and two student dot-products, using detached LM-head
+    rows so the auxiliary loss trains the temporal trajectory rather than the
+    vocabulary head.  No new temperature, loss weight, or trainable parameter
+    is introduced.
+    """
+    horizon = int(candidate_labels.shape[-1])
+    zero = student_states.new_zeros((), dtype=torch.float32)
+    if horizon <= 1 or int(student_states.shape[1]) == 0:
+        return zero, zero, zero, zero
+
+    detached_weight = lm_head_weight.detach()
+    detached_bias = lm_head_bias.detach() if lm_head_bias is not None else None
+    teacher = teacher_states.detach()
+
+    teacher_scores_by_candidate = []
+    for candidate_index in range(horizon):
+        token_ids = candidate_labels[..., candidate_index].clamp_min(0)
+        token_rows = F.embedding(token_ids, detached_weight)
+        score = (teacher * token_rows).sum(dim=-1, dtype=torch.float32)
+        if detached_bias is not None:
+            score = score + detached_bias[token_ids].float()
+        teacher_scores_by_candidate.append(score)
+    teacher_scores = torch.stack(teacher_scores_by_candidate, dim=-1)
+
+    target_ids = candidate_labels[..., target_index]
+    safe_target_ids = target_ids.clamp_min(0)
+    target_valid = candidate_valid[..., target_index]
+    target_teacher_score = teacher_scores[..., target_index]
+
+    candidate_indices = torch.arange(
+        horizon,
+        device=candidate_labels.device,
+    ).view(1, 1, horizon)
+    competitor_mask = (
+        candidate_valid
+        & candidate_indices.ne(int(target_index))
+        & candidate_labels.ne(target_ids.unsqueeze(-1))
+    )
+    has_competitor = competitor_mask.any(dim=-1)
+    masked_teacher_scores = teacher_scores.masked_fill(
+        ~competitor_mask,
+        torch.finfo(teacher_scores.dtype).min,
+    )
+    competitor_index = masked_teacher_scores.argmax(dim=-1)
+    competitor_ids = candidate_labels.gather(
+        dim=-1,
+        index=competitor_index.unsqueeze(-1),
+    ).squeeze(-1)
+    safe_competitor_ids = competitor_ids.clamp_min(0)
+    competitor_teacher_score = teacher_scores.gather(
+        dim=-1,
+        index=competitor_index.unsqueeze(-1),
+    ).squeeze(-1)
+
+    # Compute the two student logits sequentially so only one gathered
+    # [B,T,H] LM-head block is live at a time.
+    target_rows = F.embedding(safe_target_ids, detached_weight)
+    target_student_score = (
+        student_states * target_rows
+    ).sum(dim=-1, dtype=torch.float32)
+    competitor_rows = F.embedding(safe_competitor_ids, detached_weight)
+    competitor_student_score = (
+        student_states * competitor_rows
+    ).sum(dim=-1, dtype=torch.float32)
+    if detached_bias is not None:
+        target_student_score = (
+            target_student_score + detached_bias[safe_target_ids].float()
+        )
+        competitor_student_score = (
+            competitor_student_score
+            + detached_bias[safe_competitor_ids].float()
+        )
+
+    teacher_pair_logits = torch.stack(
+        (target_teacher_score, competitor_teacher_score),
+        dim=-1,
+    )
+    student_pair_logits = torch.stack(
+        (target_student_score, competitor_student_score),
+        dim=-1,
+    )
+    readout_mask = base_mask & target_valid & has_competitor
+
+    teacher_log_probs = F.log_softmax(teacher_pair_logits.detach(), dim=-1)
+    student_log_probs = F.log_softmax(student_pair_logits, dim=-1)
+    per_token_kl = F.kl_div(
+        student_log_probs,
+        teacher_log_probs,
+        log_target=True,
+        reduction="none",
+    ).sum(dim=-1)
+    readout_loss = _masked_mean(per_token_kl, readout_mask)
+
+    with torch.no_grad():
+        readout_agreement = _masked_mean(
+            (
+                teacher_pair_logits.argmax(dim=-1)
+                == student_pair_logits.argmax(dim=-1)
+            ).float(),
+            readout_mask,
+        )
+        teacher_margin = (
+            target_teacher_score - competitor_teacher_score
+        )
+        student_margin = (
+            target_student_score - competitor_student_score
+        )
+        readout_margin_error = _masked_mean(
+            (student_margin - teacher_margin).abs(),
+            readout_mask,
+        )
+        readout_valid_fraction = readout_mask.float().mean()
+
+    return (
+        readout_loss,
+        readout_agreement,
+        readout_margin_error,
+        readout_valid_fraction,
+    )
+
+
+NITP_TEMPORAL_STEP_METRIC_NAMES = (
+    # Rows 0..12 are stable for backward-compatible trainer logging.
+    "state_loss",
+    "composite_step_loss",
+    "delta_cosine",
+    "delta_norm_ratio",
+    "rollout_nitp_cosine",
+    "readout_loss",
+    "readout_agreement",
+    "readout_margin_error",
+    "readout_valid_fraction",
+    "cross_example_pred_displacement_cosine",
+    "cross_example_target_displacement_cosine",
+    "predicted_delta_turn_cosine",
+    "target_delta_turn_cosine",
+    # Append-only component diagnostics for the new cheap moment objective.
+    "raw_step_loss",
+    "mean_error_loss",
+    "std_match_loss",
+)
+
+
 def compute_nitp_temporal_objective(
     hidden_states: torch.Tensor,
     shallow_target_states: torch.Tensor,
@@ -6538,6 +6813,8 @@ def compute_nitp_temporal_objective(
     attention_mask: Optional[torch.Tensor],
     transition_model: NITPTemporalTransition,
     nitp_projector: NITPProjector,
+    lm_head_weight: torch.Tensor,
+    lm_head_bias: Optional[torch.Tensor],
     horizon: int,
     dynamics_weight: float,
     identification_weight: float,
@@ -6547,48 +6824,77 @@ def compute_nitp_temporal_objective(
     eos_token_id: Optional[int] = None,
 ) -> Tuple[torch.Tensor, ...]:
     """
-    Causal NITP temporal rollout with relational joint identification.
-
-    Autonomous rollout:
-
-        h_hat_{t,0} = h_t
-        z_hat_{t,k} = P_{stopgrad(theta_P)}(h_hat_{t,k-1})
-        h_hat_{t,k} = T_psi(h_hat_{t,k-1}, stopgrad(z_hat_{t,k})).
-
-    The state objective remains the average Smooth-L1 alignment with real final
-    states.  The identification objective now uses
-
-        S = 0.25 A + 0.25 R + 0.5 Z,
-
-    where ``A`` is accumulated-displacement alignment, ``R`` is the relational
-    angular-profile alignment across the full local future window, and ``Z``
-    is frozen-weight NITP shallow-content alignment.  This discourages an
-    artificially parallel rollout only when the target trajectory itself has
-    non-parallel angular structure.  It does not add a direct diversity
-    penalty, a new parameter, a new rollout, vocabulary logits, batch negatives
-    or a tunable mixing coefficient.
-
-    Packed monitoring outputs are appended after the original fourteen scalar
-    returns so old dashboards can remain intact:
-
-    ``component_metrics`` [4,4]
-        rows: joint, displacement A, relational R, NITP signature Z
-        cols: diagonal, off-diagonal, margin, positive-margin fraction.
-
-    ``step_metrics`` [11,W]
-        rows: state loss, delta cosine, delta norm ratio, rollout-NITP cosine,
-        trajectory specificity, cross-example predicted displacement cosine,
-        cross-example target displacement cosine, their gap, predicted
-        consecutive-delta cosine, target consecutive-delta cosine, their gap.
-
-    ``offset_metrics`` [4,W-1]
-        rows use the same component order and columns group incorrect pairs by
-        temporal distance 1..W-1.
-
-    ``scalar_metrics`` [6]
-        path-length ratio, target absolute-state pair cosine, absolute pair gap,
-        predicted centered-displacement pair cosine, target centered-
-        displacement pair cosine, centered pair gap.
+    Causal Temporal-NITP with a cheap first-order trajectory objective.
+
+    The autonomous rollout and relational identification are unchanged:
+
+        h_hat_0 = h_0,
+        z_hat_k = P_{sg(theta_P)}(h_hat_{k-1}),
+        h_hat_k = T_k(h_hat_{k-1}, sg(z_hat_k)),
+
+        S = 0.25 A + 0.25 R + 0.5 Z.
+
+    Only the *dynamics objective* changes.  The old objective supervised final
+    states alone.  The new objective is the fixed arithmetic mean
+
+        L_dyn = (L_position + L_step + L_readout) / 3.
+
+    ``L_position`` is the original Smooth-L1 state alignment. ``L_step`` uses
+    target-RMS-normalized local increments and is now
+
+        L_step = L_raw + L_mean_error + 0.5 * L_std.
+
+    ``L_raw`` keeps the full per-window correspondence. ``L_mean_error`` applies
+    Smooth-L1 only to the valid-window mean prediction error, directly penalizing
+    a horizon-common direction that is absent from the targets. ``L_std`` matches
+    the per-channel standard deviation around those means, retaining the useful
+    lesson from the removed explicit centered loss: the rollout must preserve
+    context-dependent dispersion instead of shrinking every example toward its
+    batch mean.  Means and second moments are reductions to [1,1,H]; no centered
+    [B,S,H] tensors or second full-size Smooth-L1 branch are materialized.
+    ``L_readout`` matches a two-token temporal readout selected from the existing
+    W future labels; it protects functionally sensitive LM-head directions
+    without projecting to the full vocabulary.
+
+    The constants 1/3 and 0.5 are not configurable hyperparameters.  The first
+    defines an equal mean of the three trajectory views; the second keeps the
+    dispersion match auxiliary to the exact per-window and common-error terms.
+    Existing ``dynamics_weight`` and ``identification_weight`` remain the only
+    outer weights.
+
+    Packed diagnostics are intentionally compact and non-redundant:
+
+    ``component_metrics`` [3,4]
+        rows: displacement A, relational R, NITP signature Z.
+        cols: diagonal, off-diagonal, hardest-negative margin,
+        positive-margin fraction.  The joint diagonal/off-diagonal are omitted
+        because they are exactly reconstructed as 0.25*A + 0.25*R + 0.5*Z;
+        joint positive-margin fraction is identical to joint top-1.
+
+    ``step_metrics`` [16,W]
+        Rows 0..12 preserve the previous order for logger compatibility: state
+        loss, composite normalized step loss, delta cosine, delta norm ratio,
+        rollout-NITP cosine, two-token readout KL, readout rank agreement,
+        readout margin error, readout valid fraction, cross-example predicted
+        displacement cosine, cross-example target displacement cosine,
+        predicted consecutive-delta cosine, target consecutive-delta cosine.
+        Rows 13..15 append the three step-loss components: raw local loss,
+        common mean-error loss, and standard-deviation matching loss.  The
+        composite row is exactly raw + mean_error + 0.5*std.  Aggregate component
+        means are reconstructible by averaging the appended rows over W.
+
+    ``offset_metrics`` [3,W-1]
+        neighbor similarities for A, R, and Z.  The joint neighbor profile is
+        omitted because it is exactly reconstructed from these three rows.
+
+    ``scalar_metrics`` [5]
+        path-length ratio, predicted and target absolute-state pair cosine,
+        predicted and target centered-displacement pair cosine.  Pair gaps are
+        omitted because they are simple differences.
+
+    The previous counterfactual trajectory-specificity rollout is removed.  It
+    required an additional W-step transition rollout and is replaced by the
+    much cheaper per-step readout agreement plus cross-example diversity.
     """
     horizon = int(horizon)
     seq_len = int(hidden_states.shape[1])
@@ -6596,11 +6902,11 @@ def compute_nitp_temporal_objective(
     zero = hidden_states.new_zeros((), dtype=torch.float32)
     if max_horizon == 0:
         return (
-            *((zero,) * 14),
-            zero.new_zeros((4, 4)),
-            zero.new_zeros((11, 0)),
-            zero.new_zeros((4, 0)),
-            zero.new_zeros((6,)),
+            *((zero,) * 9),
+            zero.new_zeros((3, 4)),
+            zero.new_zeros((16, 0)),
+            zero.new_zeros((3, 0)),
+            zero.new_zeros((5,)),
         )
 
     device = hidden_states.device
@@ -6613,29 +6919,71 @@ def compute_nitp_temporal_objective(
     if pad_token_id is not None:
         valid_tokens = valid_tokens & (labels_on_device != pad_token_id)
 
+    # The functional readout needs one token beyond the farthest hidden-state
+    # target because h_{t+k} predicts label_{t+k+1}.  Candidate labels are the W
+    # true tokens attached to the local temporal horizons, so selecting the
+    # strongest incorrect one costs O(W*H), not O(V*H).
+    functional_length = max(seq_len - max_horizon - 1, 0)
+    if functional_length > 0:
+        candidate_labels = torch.stack(
+            [
+                labels_on_device[
+                    :, offset + 1 : offset + 1 + functional_length
+                ]
+                for offset in range(1, max_horizon + 1)
+            ],
+            dim=-1,
+        )
+        candidate_valid = torch.stack(
+            [
+                valid_tokens[
+                    :, offset + 1 : offset + 1 + functional_length
+                ]
+                for offset in range(1, max_horizon + 1)
+            ],
+            dim=-1,
+        )
+        functional_window_mask = _valid_temporal_span_mask(
+            valid_tokens=valid_tokens,
+            labels=labels_on_device,
+            span_length=max_horizon + 2,
+            eos_token_id=eos_token_id,
+        )
+    else:
+        candidate_labels = None
+        candidate_valid = None
+        functional_window_mask = None
+
     current_states = hidden_states
     shifted_targets = hidden_states
     predicted_by_step: list[torch.Tensor] = []
     condition_by_step: list[torch.Tensor] = []
     step_masks: list[torch.Tensor] = []
+
     state_loss_by_step: list[torch.Tensor] = []
+    normalized_step_loss_by_step: list[torch.Tensor] = []
+    raw_step_loss_by_step: list[torch.Tensor] = []
+    mean_error_loss_by_step: list[torch.Tensor] = []
+    std_match_loss_by_step: list[torch.Tensor] = []
     delta_cosine_by_step: list[torch.Tensor] = []
     delta_norm_ratio_by_step: list[torch.Tensor] = []
     rollout_nitp_cosine_by_step: list[torch.Tensor] = []
+    readout_loss_by_step: list[torch.Tensor] = []
+    readout_agreement_by_step: list[torch.Tensor] = []
+    readout_margin_error_by_step: list[torch.Tensor] = []
+    readout_valid_fraction_by_step: list[torch.Tensor] = []
     predicted_turn_by_step: list[torch.Tensor] = []
     target_turn_by_step: list[torch.Tensor] = []
 
     state_loss_total = zero
+    normalized_step_loss_total = zero
+    readout_loss_total = zero
     path_predicted_norm_sum = zero
     path_target_norm_sum = zero
     previous_predicted_delta: Optional[torch.Tensor] = None
     previous_real_delta: Optional[torch.Tensor] = None
 
     for step in range(max_horizon):
-        # The projector weights are detached but the derivative with respect to
-        # the rollout state is preserved for the Z component.  The transition
-        # consumes a fully detached copy, preventing a private projector/
-        # transition code.
         source_states = current_states[:, :-1, :]
         shifted_targets = shifted_targets[:, 1:, :]
         next_shallow_signature = _project_nitp_with_frozen_weights(
@@ -6645,6 +6993,7 @@ def compute_nitp_temporal_objective(
         predicted_states = transition_model(
             source_states,
             next_shallow_signature.detach(),
+            step_index=step,
         )
 
         step_mask = _valid_temporal_span_mask(
@@ -6663,13 +7012,125 @@ def compute_nitp_temporal_objective(
         state_loss_total = state_loss_total + step_state_loss
         state_loss_by_step.append(step_state_loss.detach().float())
 
-        with torch.no_grad():
-            real_source_states = hidden_states[
-                :, step : step + predicted_states.shape[1], :
-            ]
-            predicted_delta = predicted_states - source_states
-            real_delta = shifted_targets - real_source_states
+        real_source_states = hidden_states[
+            :, step : step + predicted_states.shape[1], :
+        ].detach()
+        predicted_delta = predicted_states - source_states
+        real_delta = shifted_targets.detach() - real_source_states
+
+        # Scale each target delta to unit RMS.  This is data-derived and detached,
+        # so there is no new tuned scale and no radial shortcut.
+        target_delta_rms = real_delta.float().pow(2).mean(
+            dim=-1,
+            keepdim=True,
+        ).sqrt().clamp_min(1.0e-4)
+        normalized_predicted_delta = predicted_delta.float() / target_delta_rms
+        normalized_real_delta = real_delta.float() / target_delta_rms
+
+        # Absolute view: preserve the legitimate common temporal component as
+        # well as each local direction and magnitude.
+        raw_step_error = F.smooth_l1_loss(
+            normalized_predicted_delta,
+            normalized_real_delta,
+            reduction="none",
+        ).mean(dim=-1)
+        raw_step_loss = _masked_mean(raw_step_error, step_mask)
+
+        # Cheap moment supervision over valid windows.  Unlike the removed
+        # explicit centered branch, this directly penalizes the common *error*
+        # while preserving the full raw loss.  The standard-deviation term keeps
+        # the context-specific spread from collapsing.  Only [1,1,H] reductions
+        # survive; no centered [B,S,H] tensors are constructed.
+        step_weights = step_mask.to(
+            device=normalized_predicted_delta.device,
+            dtype=normalized_predicted_delta.dtype,
+        ).unsqueeze(-1)
+        valid_step_count = step_weights.sum(
+            dim=(0, 1),
+            keepdim=True,
+        ).clamp_min(1.0)
+
+        predicted_step_mean = (
+            normalized_predicted_delta * step_weights
+        ).sum(dim=(0, 1), keepdim=True) / valid_step_count
+        target_step_mean = (
+            normalized_real_delta * step_weights
+        ).sum(dim=(0, 1), keepdim=True) / valid_step_count
+
+        mean_error_loss = F.smooth_l1_loss(
+            predicted_step_mean,
+            target_step_mean.detach(),
+            reduction="mean",
+        )
+
+        predicted_second_moment = (
+            normalized_predicted_delta.square() * step_weights
+        ).sum(dim=(0, 1), keepdim=True) / valid_step_count
+        target_second_moment = (
+            normalized_real_delta.square() * step_weights
+        ).sum(dim=(0, 1), keepdim=True) / valid_step_count
+
+        # E[x^2] - E[x]^2 is clamped only for roundoff.  The same epsilon on
+        # both paths makes exactly constant prediction/target channels match.
+        predicted_step_std = (
+            predicted_second_moment - predicted_step_mean.square()
+        ).clamp_min(0.0).add(1.0e-6).sqrt()
+        target_step_std = (
+            target_second_moment - target_step_mean.square()
+        ).clamp_min(0.0).add(1.0e-6).sqrt()
+        std_match_loss = F.smooth_l1_loss(
+            predicted_step_std,
+            target_step_std.detach(),
+            reduction="mean",
+        )
+
+        # Preserve the full local supervision, add direct pressure on the
+        # horizon-common error, and retain a smaller dispersion constraint.
+        step_delta_loss = (
+            raw_step_loss + mean_error_loss + 0.5 * std_match_loss
+        )
+        normalized_step_loss_total = (
+            normalized_step_loss_total + step_delta_loss
+        )
+        normalized_step_loss_by_step.append(step_delta_loss.detach().float())
+        raw_step_loss_by_step.append(raw_step_loss.detach().float())
+        mean_error_loss_by_step.append(mean_error_loss.detach().float())
+        std_match_loss_by_step.append(std_match_loss.detach().float())
+
+        if functional_length > 0:
+            (
+                step_readout_loss,
+                step_readout_agreement,
+                step_readout_margin_error,
+                step_readout_valid_fraction,
+            ) = _temporal_two_token_readout(
+                student_states=predicted_states[:, :functional_length, :],
+                teacher_states=shifted_targets[:, :functional_length, :],
+                candidate_labels=candidate_labels,
+                candidate_valid=candidate_valid,
+                target_index=step,
+                base_mask=functional_window_mask,
+                lm_head_weight=lm_head_weight,
+                lm_head_bias=lm_head_bias,
+            )
+        else:
+            step_readout_loss = zero
+            step_readout_agreement = zero
+            step_readout_margin_error = zero
+            step_readout_valid_fraction = zero
+        readout_loss_total = readout_loss_total + step_readout_loss
+        readout_loss_by_step.append(step_readout_loss.detach().float())
+        readout_agreement_by_step.append(
+            step_readout_agreement.detach().float()
+        )
+        readout_margin_error_by_step.append(
+            step_readout_margin_error.detach().float()
+        )
+        readout_valid_fraction_by_step.append(
+            step_readout_valid_fraction.detach().float()
+        )
 
+        with torch.no_grad():
             step_delta_cosine = F.cosine_similarity(
                 predicted_delta.float(),
                 real_delta.float(),
@@ -6743,15 +7204,32 @@ def compute_nitp_temporal_objective(
 
     inv_horizon = 1.0 / float(max_horizon)
     state_loss = state_loss_total * inv_horizon
+    normalized_step_loss = normalized_step_loss_total * inv_horizon
+    readout_loss = readout_loss_total * inv_horizon
+    trajectory_loss = (
+        state_loss + normalized_step_loss + readout_loss
+    ) / 3.0
+
     state_step_tensor = torch.stack(state_loss_by_step)
+    normalized_step_loss_tensor = torch.stack(
+        normalized_step_loss_by_step
+    )
+    raw_step_loss_tensor = torch.stack(raw_step_loss_by_step)
+    mean_error_loss_tensor = torch.stack(mean_error_loss_by_step)
+    std_match_loss_tensor = torch.stack(std_match_loss_by_step)
     delta_step_tensor = torch.stack(delta_cosine_by_step)
     delta_norm_ratio_tensor = torch.stack(delta_norm_ratio_by_step)
     rollout_nitp_step_tensor = torch.stack(rollout_nitp_cosine_by_step)
+    readout_loss_tensor = torch.stack(readout_loss_by_step)
+    readout_agreement_tensor = torch.stack(readout_agreement_by_step)
+    readout_margin_error_tensor = torch.stack(
+        readout_margin_error_by_step
+    )
+    readout_valid_fraction_tensor = torch.stack(
+        readout_valid_fraction_by_step
+    )
     predicted_turn_tensor = torch.stack(predicted_turn_by_step)
     target_turn_tensor = torch.stack(target_turn_by_step)
-    delta_cosine = delta_step_tensor.mean()
-    rollout_nitp_cosine = rollout_nitp_step_tensor.mean()
-    last_step_state_loss = state_step_tensor[-1]
     path_length_ratio = (
         path_predicted_norm_sum
         / path_target_norm_sum.clamp_min(1.0e-8)
@@ -6763,38 +7241,38 @@ def compute_nitp_temporal_objective(
         step_metrics = torch.stack(
             (
                 state_step_tensor,
+                normalized_step_loss_tensor,
                 delta_step_tensor,
                 delta_norm_ratio_tensor,
                 rollout_nitp_step_tensor,
-                zero_steps,
-                zero_steps,
+                readout_loss_tensor,
+                readout_agreement_tensor,
+                readout_margin_error_tensor,
+                readout_valid_fraction_tensor,
                 zero_steps,
                 zero_steps,
                 predicted_turn_tensor,
                 target_turn_tensor,
-                predicted_turn_tensor - target_turn_tensor,
+                raw_step_loss_tensor,
+                mean_error_loss_tensor,
+                std_match_loss_tensor,
             )
         )
-        weighted_total = float(dynamics_weight) * state_loss
+        weighted_total = float(dynamics_weight) * trajectory_loss
         return (
             weighted_total,
             state_loss,
+            normalized_step_loss,
+            readout_loss,
             zero,
             zero,
             zero,
             zero,
             zero,
-            zero,
-            last_step_state_loss,
-            delta_cosine,
-            rollout_nitp_cosine,
-            zero,
-            zero,
-            zero,
-            zero.new_zeros((4, 4)),
+            zero.new_zeros((3, 4)),
             step_metrics,
-            zero.new_zeros((4, max(max_horizon - 1, 0))),
-            torch.stack((path_length_ratio, zero, zero, zero, zero, zero)),
+            zero.new_zeros((3, max(max_horizon - 1, 0))),
+            torch.stack((path_length_ratio, zero, zero, zero, zero)),
         )
 
     (
@@ -6845,12 +7323,14 @@ def compute_nitp_temporal_objective(
         (row_prediction - expected_offsets).abs().float().mean(dim=-1),
         window_mask,
     )
+    identification_margin = _temporal_matrix_statistics(
+        similarity.detach(),
+        window_mask,
+    )[2]
 
-    # All decomposed diagnostics use detached matrices so they cannot retain
-    # additional backward graphs.  Their values still describe the exact
-    # matrices used by the live identification objective.
+    # A/R/Z component rows only.  Joint diagonal/off-diagonal and neighbor
+    # profiles are exact fixed-weight reconstructions and are not logged.
     component_matrices = (
-        similarity.detach(),
         displacement_similarity.detach(),
         relational_similarity.detach(),
         nitp_signature_similarity.detach(),
@@ -6867,9 +7347,6 @@ def compute_nitp_temporal_objective(
             for component in component_matrices
         ]
     )
-    diagonal_cosine = component_metrics[0, 0]
-    off_diagonal_cosine = component_metrics[0, 1]
-    identification_margin = component_metrics[0, 2]
 
     with torch.no_grad():
         initial_common = hidden_states[:, :common_length, :].detach()
@@ -6958,97 +7435,56 @@ def compute_nitp_temporal_objective(
         cross_predicted_tensor = torch.stack(cross_predicted_by_step)
         cross_target_tensor = torch.stack(cross_target_by_step)
 
-        # Reuse the existing counterfactual rollout and retain each horizon's
-        # specificity instead of discarding it into one final average.
-        if int(hidden_states.shape[0]) > 1:
-            wrong_current = torch.roll(hidden_states.detach(), shifts=1, dims=0)
-            wrong_targets = hidden_states.detach()
-            wrong_state_loss_by_step = []
-            for step in range(max_horizon):
-                wrong_source = wrong_current[:, :-1, :]
-                wrong_targets = wrong_targets[:, 1:, :]
-                wrong_condition = nitp_projector(wrong_source)
-                wrong_predicted = transition_model(
-                    wrong_source,
-                    wrong_condition,
-                )
-                wrong_step_loss = _masked_smooth_l1(
-                    predicted=wrong_predicted,
-                    target=wrong_targets,
-                    mask=step_masks[step],
-                )
-                wrong_state_loss_by_step.append(wrong_step_loss.float())
-                wrong_current = wrong_predicted
-
-            wrong_step_tensor = torch.stack(wrong_state_loss_by_step)
-            correct_step_tensor = state_step_tensor.detach().float()
-            trajectory_specificity_by_step = (
-                (wrong_step_tensor - correct_step_tensor)
-                / (wrong_step_tensor + correct_step_tensor + 1.0e-8)
-            ).clamp(min=-1.0, max=1.0)
-            wrong_state_loss = wrong_step_tensor.mean()
-            full_state_loss = correct_step_tensor.mean()
-            trajectory_specificity = (
-                (wrong_state_loss - full_state_loss)
-                / (wrong_state_loss + full_state_loss + 1.0e-8)
-            ).clamp(min=-1.0, max=1.0)
-        else:
-            trajectory_specificity_by_step = zero.new_zeros((max_horizon,))
-            trajectory_specificity = zero
-
     step_metrics = torch.stack(
         (
             state_step_tensor,
+            normalized_step_loss_tensor,
             delta_step_tensor,
             delta_norm_ratio_tensor,
             rollout_nitp_step_tensor,
-            trajectory_specificity_by_step,
+            readout_loss_tensor,
+            readout_agreement_tensor,
+            readout_margin_error_tensor,
+            readout_valid_fraction_tensor,
             cross_predicted_tensor,
             cross_target_tensor,
-            cross_predicted_tensor - cross_target_tensor,
             predicted_turn_tensor,
             target_turn_tensor,
-            predicted_turn_tensor - target_turn_tensor,
+            raw_step_loss_tensor,
+            mean_error_loss_tensor,
+            std_match_loss_tensor,
         )
     )
     scalar_metrics = torch.stack(
         (
             path_length_ratio,
+            rollout_pair_cosine,
             target_pair_cosine,
-            rollout_pair_cosine - target_pair_cosine,
             rollout_displacement_pair_cosine,
             target_displacement_pair_cosine,
-            rollout_displacement_pair_cosine
-            - target_displacement_pair_cosine,
         )
     )
 
     weighted_total = (
-        float(dynamics_weight) * state_loss
+        float(dynamics_weight) * trajectory_loss
         + float(identification_weight) * identification_loss
     )
     return (
         weighted_total,
         state_loss,
+        normalized_step_loss,
+        readout_loss,
         identification_loss,
         temporal_top1_accuracy,
         mean_absolute_offset,
-        diagonal_cosine,
-        off_diagonal_cosine,
         identification_margin,
-        last_step_state_loss,
-        delta_cosine,
-        rollout_nitp_cosine,
-        rollout_pair_cosine,
         valid_window_fraction,
-        trajectory_specificity,
         component_metrics,
         step_metrics,
         offset_metrics,
         scalar_metrics,
     )
 
-
 def _categorical_kl_loss(
     teacher_logits: torch.Tensor,
     student_logits: torch.Tensor,
@@ -7204,25 +7640,44 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 "`ntp_loss_backend` must be 'cce' or 'liger'. "
                 f"Got {self.ntp_loss_backend!r}."
             )
+        if bool(getattr(config, "use_mile_loss", False)):
+            _require_extended_cce_options("mile_enabled", "mile_gamma")
+        if bool(getattr(config, "use_mu_loss", False)):
+            _require_extended_cce_options("mu_loss_enabled", "mu_loss_lambda")
+        if bool(getattr(config, "use_meap", False)) and meap_mask_inputs is None:
+            raise ImportError(
+                "`use_meap=True` requires the extended cut-cross-entropy package. "
+                f"Install it with `pip install --force-reinstall --no-deps "
+                f"\"{_EXTENDED_CCE_INSTALL}\"`, or set `use_meap=False`."
+            )
+        if bool(getattr(config, "use_meap", False)) and (
+            "return_metrics" not in inspect.signature(meap_mask_inputs).parameters
+        ):
+            raise ImportError(
+                "The installed cut-cross-entropy package does not provide kernel-side "
+                "MEAP metrics. Reinstall the extended backend or set `use_meap=False`."
+            )
         self._last_ntp_loss = None
+        self._last_ntp_ce_unweighted = None
+        self._last_mile_reweighting_delta = None
+        self._last_mu_loss = None
+        self._last_meap_mask_fraction = None
+        self._last_meap_masked_tokens = None
+        self._last_meap_seed = None
         self._last_tweo_loss = None
         self._last_nitp_loss = None
         self._last_nitp_temporal_loss = None
         self._last_nitp_temporal_state_loss = None
+        self._last_nitp_temporal_step_loss = None
+        self._last_nitp_temporal_raw_step_loss = None
+        self._last_nitp_temporal_mean_error_loss = None
+        self._last_nitp_temporal_std_match_loss = None
+        self._last_nitp_temporal_readout_loss = None
         self._last_nitp_temporal_identification_loss = None
         self._last_nitp_temporal_top1_accuracy = None
         self._last_nitp_temporal_mean_absolute_offset = None
-        self._last_nitp_temporal_diagonal_cosine = None
-        self._last_nitp_temporal_off_diagonal_cosine = None
         self._last_nitp_temporal_identification_margin = None
-        self._last_nitp_temporal_last_step_state_loss = None
-        self._last_nitp_temporal_rollout_nitp_cosine = None
-        self._last_nitp_temporal_rollout_pair_cosine = None
         self._last_nitp_temporal_valid_window_fraction = None
-        self._last_nitp_temporal_delta_cosine = None
-        self._last_nitp_temporal_trajectory_specificity = None
-        self._last_nitp_temporal_to_nitp_ratio = None
-        self._last_nitp_temporal_loss_applied = None
         self._last_nitp_temporal_component_metrics = None
         self._last_nitp_temporal_step_metrics = None
         self._last_nitp_temporal_offset_metrics = None
@@ -7281,11 +7736,71 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         position_ids: Optional[torch.LongTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
         labels: Optional[torch.LongTensor] = None,
+        meap_seed: Optional[int] = None,
         logits_to_keep: Union[int, torch.Tensor] = 0,
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
+        self._last_meap_mask_fraction = None
+        self._last_meap_masked_tokens = None
+        self._last_meap_seed = None
+
+        meap_enabled = (
+            bool(getattr(self.config, "use_meap", False))
+            and self.training
+            and labels is not None
+        )
+        if meap_enabled:
+            if input_ids is None or inputs_embeds is not None:
+                raise ValueError(
+                    "MEAP requires `input_ids` and does not support precomputed "
+                    "`inputs_embeds`."
+                )
+            if meap_mask_inputs is None:
+                raise ImportError(
+                    "MEAP is enabled but is unavailable in the installed "
+                    "cut-cross-entropy package."
+                )
+            mask_token_id = getattr(self.config, "meap_mask_token_id", None)
+            if mask_token_id is None:
+                raise ValueError("`meap_mask_token_id` is required when MEAP is enabled.")
+            effective_meap_seed = int(
+                getattr(self.config, "meap_seed", 0)
+                if meap_seed is None
+                else meap_seed
+            )
+            eligible_mask = (
+                attention_mask.to(device=input_ids.device, dtype=torch.bool)
+                if attention_mask is not None
+                else None
+            )
+            exclude_last = bool(getattr(self.config, "meap_exclude_last", True))
+            mask_ratio = float(getattr(self.config, "meap_mask_ratio", 0.15))
+            input_ids, meap_metrics = meap_mask_inputs(
+                input_ids,
+                int(mask_token_id),
+                enabled=True,
+                mask_ratio=mask_ratio,
+                eligible_mask=eligible_mask,
+                seed=effective_meap_seed,
+                exclude_last=exclude_last,
+                return_mask=False,
+                return_metrics=True,
+                implementation=str(
+                    getattr(self.config, "meap_implementation", "triton")
+                ),
+            )
+            eligible_count = meap_metrics[0]
+            masked_count = meap_metrics[1]
+            self._last_meap_mask_fraction = (
+                masked_count.float() / eligible_count.clamp_min(1).float()
+            ).detach()
+            self._last_meap_masked_tokens = masked_count.detach()
+            self._last_meap_seed = input_ids.new_tensor(
+                effective_meap_seed, dtype=torch.long
+            )
+
         tweo_enabled = (
             bool(getattr(self.config, "use_tweo", False))
             and labels is not None
@@ -7366,46 +7881,39 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         nitp_loss = None
         nitp_temporal_loss = None
         nitp_temporal_state_loss = None
+        nitp_temporal_step_loss = None
+        nitp_temporal_readout_loss = None
         nitp_temporal_identification_loss = None
         nitp_temporal_top1_accuracy = None
         nitp_temporal_mean_absolute_offset = None
-        nitp_temporal_diagonal_cosine = None
-        nitp_temporal_off_diagonal_cosine = None
         nitp_temporal_identification_margin = None
-        nitp_temporal_last_step_state_loss = None
-        nitp_temporal_rollout_nitp_cosine = None
-        nitp_temporal_rollout_pair_cosine = None
         nitp_temporal_valid_window_fraction = None
-        nitp_temporal_delta_cosine = None
-        nitp_temporal_trajectory_specificity = None
         nitp_temporal_component_metrics = None
         nitp_temporal_step_metrics = None
         nitp_temporal_offset_metrics = None
         nitp_temporal_scalar_metrics = None
-        nitp_temporal_to_nitp_ratio = None
         nitp_target_hidden_states = None
         nextlat_loss = None
         nextlat_mse_loss = None
         nextlat_kl_loss = None
         self._last_ntp_loss = None
+        self._last_ntp_ce_unweighted = None
+        self._last_mile_reweighting_delta = None
+        self._last_mu_loss = None
         self._last_tweo_loss = None
         self._last_nitp_loss = None
         self._last_nitp_temporal_loss = None
         self._last_nitp_temporal_state_loss = None
+        self._last_nitp_temporal_step_loss = None
+        self._last_nitp_temporal_raw_step_loss = None
+        self._last_nitp_temporal_mean_error_loss = None
+        self._last_nitp_temporal_std_match_loss = None
+        self._last_nitp_temporal_readout_loss = None
         self._last_nitp_temporal_identification_loss = None
         self._last_nitp_temporal_top1_accuracy = None
         self._last_nitp_temporal_mean_absolute_offset = None
-        self._last_nitp_temporal_diagonal_cosine = None
-        self._last_nitp_temporal_off_diagonal_cosine = None
         self._last_nitp_temporal_identification_margin = None
-        self._last_nitp_temporal_last_step_state_loss = None
-        self._last_nitp_temporal_rollout_nitp_cosine = None
-        self._last_nitp_temporal_rollout_pair_cosine = None
         self._last_nitp_temporal_valid_window_fraction = None
-        self._last_nitp_temporal_delta_cosine = None
-        self._last_nitp_temporal_trajectory_specificity = None
-        self._last_nitp_temporal_to_nitp_ratio = None
-        self._last_nitp_temporal_loss_applied = None
         self._last_nitp_temporal_component_metrics = None
         self._last_nitp_temporal_step_metrics = None
         self._last_nitp_temporal_offset_metrics = None
@@ -7425,14 +7933,41 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     self.liger_ntp_loss,
                 )
             else:
-                ntp_loss = compute_cce_loss(
+                return_loss_metrics = bool(
+                    getattr(self.config, "use_mile_loss", False)
+                    or getattr(self.config, "use_mu_loss", False)
+                    or meap_enabled
+                )
+                ntp_result = compute_cce_loss(
                     hidden_states,
                     labels,
                     self.lm_head.weight,
                     getattr(self.lm_head, "bias", None),
                     self.config.pad_token_id,
                     getattr(self.config, "cce_loss_impl", "cce_kahan_full_c"),
+                    use_mile_loss=bool(
+                        getattr(self.config, "use_mile_loss", False)
+                    ),
+                    mile_loss_gamma=float(
+                        getattr(self.config, "mile_loss_gamma", 1.0)
+                    ),
+                    use_mu_loss=bool(getattr(self.config, "use_mu_loss", False)),
+                    mu_loss_lambda=float(
+                        getattr(self.config, "mu_loss_lambda", 1e-4)
+                    ),
+                    return_loss_metrics=return_loss_metrics,
                 )
+                if return_loss_metrics:
+                    ntp_loss, ntp_metrics = ntp_result
+                    self._last_ntp_ce_unweighted = ntp_metrics[
+                        "ntp_ce_unweighted"
+                    ].detach()
+                    self._last_mile_reweighting_delta = ntp_metrics[
+                        "mile_reweighting_delta"
+                    ].detach()
+                    self._last_mu_loss = ntp_metrics["mu_loss"].detach()
+                else:
+                    ntp_loss = ntp_result
             loss = ntp_loss
 
             if tweo_enabled:
@@ -7508,6 +8043,8 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                         attention_mask=attention_mask,
                         transition_model=self.nitp_temporal_transition,
                         nitp_projector=self.nitp_projector,
+                        lm_head_weight=self.lm_head.weight,
+                        lm_head_bias=getattr(self.lm_head, "bias", None),
                         horizon=self.config.nitp_temporal_horizon,
                         dynamics_weight=self.config.nitp_temporal_dynamics_weight,
                         identification_weight=(
@@ -7534,6 +8071,8 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                             attention_mask=attention_mask,
                             transition_model=self.nitp_temporal_transition,
                             nitp_projector=self.nitp_projector,
+                            lm_head_weight=self.lm_head.weight,
+                            lm_head_bias=getattr(self.lm_head, "bias", None),
                             horizon=self.config.nitp_temporal_horizon,
                             dynamics_weight=(
                                 self.config.nitp_temporal_dynamics_weight
@@ -7554,18 +8093,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 (
                     nitp_temporal_loss,
                     nitp_temporal_state_loss,
+                    nitp_temporal_step_loss,
+                    nitp_temporal_readout_loss,
                     nitp_temporal_identification_loss,
                     nitp_temporal_top1_accuracy,
                     nitp_temporal_mean_absolute_offset,
-                    nitp_temporal_diagonal_cosine,
-                    nitp_temporal_off_diagonal_cosine,
                     nitp_temporal_identification_margin,
-                    nitp_temporal_last_step_state_loss,
-                    nitp_temporal_delta_cosine,
-                    nitp_temporal_rollout_nitp_cosine,
-                    nitp_temporal_rollout_pair_cosine,
                     nitp_temporal_valid_window_fraction,
-                    nitp_temporal_trajectory_specificity,
                     nitp_temporal_component_metrics,
                     nitp_temporal_step_metrics,
                     nitp_temporal_offset_metrics,
@@ -7575,12 +8109,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if temporal_apply_loss:
                     loss = loss + nitp_temporal_loss
 
-                if nitp_loss is not None:
-                    nitp_temporal_to_nitp_ratio = (
-                        nitp_temporal_loss.detach().float()
-                        / nitp_loss.detach().float().clamp_min(1e-8)
-                    )
-
             if nextlat_enabled:
                 if hidden_states_tuple is None or len(hidden_states_tuple) < 2:
                     raise ValueError(
@@ -7628,6 +8156,16 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if nitp_temporal_state_loss is not None
                 else None
             )
+            self._last_nitp_temporal_step_loss = (
+                nitp_temporal_step_loss.detach()
+                if nitp_temporal_step_loss is not None
+                else None
+            )
+            self._last_nitp_temporal_readout_loss = (
+                nitp_temporal_readout_loss.detach()
+                if nitp_temporal_readout_loss is not None
+                else None
+            )
             self._last_nitp_temporal_identification_loss = (
                 nitp_temporal_identification_loss.detach()
                 if nitp_temporal_identification_loss is not None
@@ -7643,51 +8181,16 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if nitp_temporal_mean_absolute_offset is not None
                 else None
             )
-            self._last_nitp_temporal_diagonal_cosine = (
-                nitp_temporal_diagonal_cosine.detach()
-                if nitp_temporal_diagonal_cosine is not None
-                else None
-            )
-            self._last_nitp_temporal_off_diagonal_cosine = (
-                nitp_temporal_off_diagonal_cosine.detach()
-                if nitp_temporal_off_diagonal_cosine is not None
-                else None
-            )
             self._last_nitp_temporal_identification_margin = (
                 nitp_temporal_identification_margin.detach()
                 if nitp_temporal_identification_margin is not None
                 else None
             )
-            self._last_nitp_temporal_last_step_state_loss = (
-                nitp_temporal_last_step_state_loss.detach()
-                if nitp_temporal_last_step_state_loss is not None
-                else None
-            )
-            self._last_nitp_temporal_rollout_nitp_cosine = (
-                nitp_temporal_rollout_nitp_cosine.detach()
-                if nitp_temporal_rollout_nitp_cosine is not None
-                else None
-            )
-            self._last_nitp_temporal_rollout_pair_cosine = (
-                nitp_temporal_rollout_pair_cosine.detach()
-                if nitp_temporal_rollout_pair_cosine is not None
-                else None
-            )
             self._last_nitp_temporal_valid_window_fraction = (
                 nitp_temporal_valid_window_fraction.detach()
                 if nitp_temporal_valid_window_fraction is not None
                 else None
             )
-            self._last_nitp_temporal_delta_cosine = (
-                nitp_temporal_delta_cosine.detach()
-                if nitp_temporal_delta_cosine is not None
-                else None
-            )
-            self._last_nitp_temporal_trajectory_specificity = (
-                nitp_temporal_trajectory_specificity.detach()
-                if nitp_temporal_trajectory_specificity is not None
-                else None
-            )
             self._last_nitp_temporal_component_metrics = (
                 nitp_temporal_component_metrics.detach()
                 if nitp_temporal_component_metrics is not None
@@ -7698,6 +8201,24 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if nitp_temporal_step_metrics is not None
                 else None
             )
+            # Aggregate component metrics are derived from the append-only rows
+            # 13..15.  The existing `_last_nitp_temporal_step_loss` remains the
+            # composite raw + mean_error + 0.5*std value for compatibility.
+            if (
+                nitp_temporal_step_metrics is not None
+                and nitp_temporal_step_metrics.shape[0] >= 16
+                and nitp_temporal_step_metrics.shape[1] > 0
+            ):
+                detached_step_metrics = nitp_temporal_step_metrics.detach()
+                self._last_nitp_temporal_raw_step_loss = (
+                    detached_step_metrics[13].mean()
+                )
+                self._last_nitp_temporal_mean_error_loss = (
+                    detached_step_metrics[14].mean()
+                )
+                self._last_nitp_temporal_std_match_loss = (
+                    detached_step_metrics[15].mean()
+                )
             self._last_nitp_temporal_offset_metrics = (
                 nitp_temporal_offset_metrics.detach()
                 if nitp_temporal_offset_metrics is not None
@@ -7708,21 +8229,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if nitp_temporal_scalar_metrics is not None
                 else None
             )
-            self._last_nitp_temporal_to_nitp_ratio = (
-                nitp_temporal_to_nitp_ratio.detach()
-                if nitp_temporal_to_nitp_ratio is not None
-                else None
-            )
-            self._last_nitp_temporal_loss_applied = (
-                hidden_states.new_tensor(
-                    1.0
-                    if bool(getattr(self.config, "nitp_temporal_apply_loss", True))
-                    else 0.0,
-                    dtype=torch.float32,
-                )
-                if nitp_temporal_enabled
-                else None
-            )
             self._last_nextlat_loss = (
                 nextlat_loss.detach() if nextlat_loss is not None else None
             )