Upload final NeoLLM checkpoint and inference configuration

config.json

@@ -93,7 +93,7 @@
   "stack_memory_cache_size": 2048,
   "stack_slots": 16,
   "tie_word_embeddings": false,
-  "transformers_version": "5.13.1",
+  "transformers_version": "5.14.1",
   "tweo_eps": 1e-06,
   "tweo_loss_weight": 0.01,
   "tweo_power": 4.0,

configuration_neollm.py

@@ -575,35 +575,52 @@ class NeoLLMConfig(PretrainedConfig):
             Numerical stability epsilon added to the TWEO denominator.
 
         use_nitp_temporal (:obj:`bool`, *optional*, defaults to ``False``):
-            Instantiate the NITP temporal-transition extension.  The module is
-            a separate bias-free, context-conditioned multiplicative residual
-            updater and does not reuse the existing ``NextLatDynamicsModel``.
-            The final state proposes an update while the teacher-forced token
-            embedding only gates that contextual proposal, so ``F(0, e)=0``.
-            It predicts :math:`\hat h_{t+1}^L`, may be composed for
-            ``nitp_temporal_horizon`` local rollout steps, and never materializes
-            an additional vocabulary projection.
+            Instantiate the causal NITP temporal-transition extension.  The
+            module is a separate bias-free multiplicative residual updater and
+            does not reuse ``NextLatDynamicsModel``.  At every rollout step it
+            conditions on the detached NITP prediction of the next shallow
+            state, :math:`\operatorname{sg}(P_{\mathrm{NITP}}(\hat h_t))`,
+            never on the real future token embedding.  The condition uses the
+            centered signed gate :math:`\tanh(W_z\bar z)`, which can
+            attenuate, activate, or reverse contextual update features.  With
+            ``W_z=0`` the update vanishes, so the transition cannot learn a
+            non-trivial condition-independent path.  It also preserves
+            ``F(0, z)=0``.  The branch predicts
+            :math:`\hat h_{t+1}^L`, may be composed for
+            ``nitp_temporal_horizon`` causal rollout steps, and never
+            materializes vocabulary logits.
         nitp_temporal_apply_loss (:obj:`bool`, *optional*, defaults to ``True``):
             When ``True``, add the weighted temporal state and identification
             losses to the model loss.  When ``False``, keep the transition
             parameters instantiated and compute all temporal metrics under
             ``torch.no_grad()``, allowing a monitoring-only NTP+NITP baseline.
         nitp_temporal_horizon (:obj:`int`, *optional*, defaults to ``4``):
-            Number :math:`W` of teacher-forced latent rollout steps.
+            Number :math:`W` of autonomous NITP-conditioned latent rollout steps.
         nitp_temporal_dynamics_weight (:obj:`float`, *optional*, defaults to ``1.0``):
             Weight applied to the average Smooth-L1 alignment between each
             predicted final state and the stop-gradient real final state at
             the same temporal offset.
         nitp_temporal_identification_weight (:obj:`float`, *optional*, defaults to ``0.1``):
             Weight applied to the symmetric listwise temporal-identification
-            objective over the local :math:`W\times W` cosine matrix.
+            objective over a local :math:`W\times W` joint-signature matrix.
+            The score keeps two conceptual halves, geometry and frozen-weight
+            NITP content.  Geometry is an equal-norm concatenation of direct
+            accumulated-displacement cosine :math:`A` and a relational angular
+            profile cosine :math:`R` computed against every real future in the
+            local window.  The resulting fixed score is
+            :math:`S=\tfrac14 A+\tfrac14 R+\tfrac12 Z`.  Thus a rollout that
+            makes every step parallel is penalized only when the target
+            trajectory itself has distinguishable angular structure.  No
+            direct diversity penalty, extra module, batch negative, vocabulary
+            projection, permutation in the trainable loss, or tunable mixing
+            coefficient is introduced.
         nitp_temporal_intermediate_size (:obj:`int`, *optional*):
-            Intermediate width of the context proposal and token modulation
+            Intermediate width of the context proposal and detached NITP-condition
             paths.  ``None`` resolves to ``2 * hidden_size``.  At
             ``hidden_size=512`` this branch adds ``2,098,176`` parameters, in
             addition to the NITP projector.
         nitp_temporal_similarity_temperature (:obj:`float`, *optional*, defaults to ``0.1``):
-            Temperature applied to predicted cosine similarities before the
+            Temperature applied to joint-signature similarities before the
             listwise temporal-identification softmax.
         nitp_temporal_target_temperature (:obj:`float`, *optional*, defaults to ``0.5``):
             Temperature controlling temporal label smoothing.  Smaller values
@@ -827,9 +844,10 @@ class NeoLLMConfig(PretrainedConfig):
         nitp_projector_intermediate_size=None,
         # ── NITP temporal transition / identification extension ──────────
         # This is deliberately a NITP-side module, not the existing NextLat
-        # implementation. It learns a teacher-forced latent update
-        #     (h_t^L, E[x_{t+1}]) -> h_{t+1}^L
-        # and a local temporal-identification objective over the rollout.
+        # implementation. It learns the causal latent update
+        #     z_{t+1} = stopgrad(P_NITP(h_t^L))
+        #     (h_t^L, z_{t+1}) -> h_{t+1}^L
+        # and recursively recomputes z from every predicted state.
         # `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.

generation_config.json

@@ -7,5 +7,5 @@
   "output_attentions": false,
   "output_hidden_states": false,
   "pad_token_id": 0,
-  "transformers_version": "5.13.1"
+  "transformers_version": "5.14.1"
 }

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d316c201f8bab02846009bb30491378d9f582846f2aaf97442982a5986b90740
-size 244921040
+oid sha256:e2d79a82ec12f9a950502bacacb5343dc6a3f16209e741fa6213d3195d9a0c81
+size 244921048

modeling_neollm.py

@@ -5230,6 +5230,15 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
         - W_res:         normal(0, initializer_range)
         - router:        parent handles (normal init)
         - scaler:        ones (identity at init)
+    NITPTemporalTransition:
+        - context_gate_proj/context_up_proj/condition_modulation_proj:
+                         Xavier-uniform so both the contextual proposal and the
+                         detached NITP condition have non-degenerate variance at
+                         step 0.
+        - down_proj:     normal(0, 0.01) — keeps the residual transition close
+                         to identity initially without blocking gradients into
+                         the centered condition gate.
+        - context_norm/condition_norm: weight=1.
     NeoLLMAttention (Affine-Scaled Attention):
         - alpha_proj:    normal(0, 0.02)  — near-zero so linear_clipping(≈0) ≈ 0.5
                          at init, giving a mild ~0.5× scaling of softmax weights
@@ -5258,6 +5267,19 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
     def _init_weights(self, module):
         super()._init_weights(module)
 
+        if isinstance(module, NITPTemporalTransition):
+            # A centered tanh gate has no constant 0.5 path.  Xavier
+            # initialization gives the detached NITP condition measurable
+            # variation from step 0, while a small non-zero down projection
+            # keeps the full residual transition close to identity without
+            # suppressing gradients into the condition path.
+            nn.init.xavier_uniform_(module.context_gate_proj.weight)
+            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)
+            module.context_norm.weight.data.fill_(1.0)
+            module.condition_norm.weight.data.fill_(1.0)
+
         if isinstance(module, NeoLLMAttention):
             if getattr(module, "use_fan_residual", False):
                 if hasattr(module, "lambda_1") and module.lambda_1 is not None:
@@ -6032,30 +6054,34 @@ def compute_nitp_loss(
 
 class NITPTemporalTransition(nn.Module):
     """
-    Context-conditioned multiplicative transition for temporal NITP.
+    Causal NITP-conditioned multiplicative transition.
 
-    This module is intentionally independent from ``NextLatDynamicsModel``.
-    The current final hidden state proposes the latent update, while the
-    teacher-forced next-token embedding can only modulate that contextual
-    proposal:
+    This module remains independent from ``NextLatDynamicsModel``.  The current
+    final hidden state proposes a latent update, while the detached NITP
+    prediction of the next shallow state modulates that proposal:
 
         h_bar = RMSNorm_h(h_t)
-        e_bar = RMSNorm_e(E(x_{t+1}))
+        z_bar = RMSNorm_z(stopgrad(P_NITP(h_t)))
         u_h = SiLU(W_gate h_bar) * (W_up h_bar)
-        m_e = sigmoid(W_token e_bar)
-        delta_t = W_down(u_h * m_e)
+        s_z = tanh(W_condition z_bar)
+        delta_t = W_down(u_h * s_z)
         h_hat_{t+1} = h_t + delta_t
 
-    All projections are bias-free.  Consequently, removing the context forces
-    the complete update to zero for every token embedding:
+    The centered signed gate satisfies ``s_z 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.
 
-        F(0, e) = 0
-        dF(0, e) / de = 0
+    All projections are bias-free.  Therefore the NITP condition has no
+    additive route to the output and the transition still satisfies:
 
-    The token therefore cannot reconstruct a latent state by itself; it only
-    selects and scales features made available by the current contextual state.
-    The module remains a training-only latent auxiliary and never materializes
-    vocabulary logits.
+        F(0, z) = 0
+        dF(0, z) / dz = 0
+
+    The condition cannot construct a deep state by itself; it only controls
+    features proposed by the current causal state.  The module never
+    materializes vocabulary logits.
 
     Parameter count, excluding the already-existing NITP projector:
 
@@ -6071,12 +6097,13 @@ class NITPTemporalTransition(nn.Module):
         hidden_size = int(config.hidden_size)
         intermediate_size = int(config.nitp_temporal_intermediate_size)
 
-        # Separate normalizations preserve the semantic distinction between
-        # the contextual state and the teacher-forced token observation.
+        # The two inputs have different semantics and are normalized
+        # independently: a final-layer causal state and a predicted shallow
+        # next-state condition produced by the NITP projector.
         self.context_norm = nn.RMSNorm(hidden_size, eps=config.rms_norm_eps)
-        self.token_norm = nn.RMSNorm(hidden_size, eps=config.rms_norm_eps)
+        self.condition_norm = nn.RMSNorm(hidden_size, eps=config.rms_norm_eps)
 
-        # Context path: proposes the update features.
+        # Context path: proposes the features of the deep-state update.
         self.context_gate_proj = nn.Linear(
             hidden_size, intermediate_size, bias=False
         )
@@ -6084,10 +6111,10 @@ class NITPTemporalTransition(nn.Module):
             hidden_size, intermediate_size, bias=False
         )
 
-        # Token path: only modulates contextual features.  It has no additive
-        # route to the output, so a token embedding cannot create a latent state
-        # when the contextual proposal is zero.
-        self.token_modulation_proj = nn.Linear(
+        # 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.
+        self.condition_modulation_proj = nn.Linear(
             hidden_size, intermediate_size, bias=False
         )
         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
@@ -6095,17 +6122,19 @@ class NITPTemporalTransition(nn.Module):
     def forward(
         self,
         current_states: torch.Tensor,
-        next_token_embeds: torch.Tensor,
+        next_shallow_conditions: torch.Tensor,
     ) -> torch.Tensor:
         context = self.context_norm(current_states)
-        token = self.token_norm(next_token_embeds)
+        condition = self.condition_norm(next_shallow_conditions)
 
         context_update = (
             F.silu(self.context_gate_proj(context))
             * self.context_up_proj(context)
         )
-        token_modulation = torch.sigmoid(self.token_modulation_proj(token))
-        delta = self.down_proj(context_update * token_modulation)
+        signed_modulation = torch.tanh(
+            self.condition_modulation_proj(condition)
+        )
+        delta = self.down_proj(context_update * signed_modulation)
         return current_states + delta
 
 
@@ -6172,10 +6201,10 @@ def _valid_temporal_span_mask(
     """
     Return a mask for contiguous temporal spans that stay inside one document.
 
-    A span includes the source state and every teacher-forced token consumed by
-    the rollout.  Padding / ignored labels invalidate the full span.  EOS is
-    permitted only at the final position, so no transition is trained across a
-    packed-document boundary.
+    A span includes the source state and every real future target used to
+    evaluate the causal rollout. Padding / ignored labels invalidate the full
+    span. EOS is permitted only at the final position, so no transition is
+    trained across a packed-document boundary.
     """
     token_windows = valid_tokens.unfold(dimension=1, size=span_length, step=1)
     span_mask = token_windows.all(dim=-1)
@@ -6187,45 +6216,324 @@ def _valid_temporal_span_mask(
     return span_mask
 
 
-def _build_local_temporal_similarity(
+def _project_nitp_with_frozen_weights(
+    hidden_states: torch.Tensor,
+    projector: NITPProjector,
+) -> torch.Tensor:
+    """
+    Evaluate the current NITP projector while freezing only its parameters.
+
+    ``torch.no_grad()`` would also remove the derivative with respect to
+    ``hidden_states``.  The temporal identification objective instead needs:
+
+        d P_{sg(theta_P)}(h) / d h != 0
+        d P_{sg(theta_P)}(h) / d theta_P = 0.
+
+    The explicit bias-free SwiGLU computation below therefore uses detached
+    projector weights.  Temporal gradients can move a rollout state into a
+    region where the already-defined NITP representation remains valid, but
+    they cannot turn the projector into a private code for the transition.
+    """
+    gate = F.linear(
+        hidden_states,
+        projector.gate_proj.weight.detach(),
+        bias=None,
+    )
+    up = F.linear(
+        hidden_states,
+        projector.up_proj.weight.detach(),
+        bias=None,
+    )
+    return F.linear(
+        F.silu(gate) * up,
+        projector.down_proj.weight.detach(),
+        bias=None,
+    )
+
+
+def _build_relational_temporal_similarity(
     predicted_by_step: list[torch.Tensor],
+    condition_by_step: list[torch.Tensor],
     target_hidden_states: torch.Tensor,
+    shallow_target_states: torch.Tensor,
+    initial_hidden_states: torch.Tensor,
     common_length: int,
-) -> torch.Tensor:
+) -> Tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+]:
     """
-    Build the local ``W x W`` cosine matrix used by temporal diagnostics.
-
-    Row ``k`` contains the cosine similarity between rollout state
-    ``h_hat_{t,k}`` and every real future state ``h_{t+j}``. Targets are
-    detached because all temporal targets are stop-gradient observations.
-    This helper is used only by the trained rollout diagnostics.  The
-    context-specificity counterfactual intentionally needs only Smooth-L1 and
-    therefore avoids building a second W x W similarity matrix.
+    Build the local ``W x W`` temporal-identification matrices.
+
+    The trainable identification score keeps two conceptual halves:
+
+        S_{k,j} = 0.5 * [Gamma_{k,j} + Z_{k,j}],
+
+    where the geometric half is itself an equal-norm concatenation:
+
+        Gamma_{k,j} = 0.5 * [A_{k,j} + R_{k,j}].
+
+    The expanded score is therefore
+
+        S_{k,j} = 0.25 A_{k,j} + 0.25 R_{k,j} + 0.5 Z_{k,j}.
+
+    ``A`` is the direct accumulated-displacement cosine:
+
+        u_hat_k = normalize(h_hat_k - h_0)
+        u_j^*   = normalize(stopgrad(h_j^* - h_0))
+        A_{k,j} = u_hat_k^T u_j^*.
+
+    ``R`` compares relational angular profiles instead of forcing the rollout
+    to copy the teacher's local tangent.  First compute
+
+        G^*_{j,r} = (u_j^*)^T u_r^*,
+        rho_hat_k = normalize(A_{k,:}),
+        rho_j^*   = normalize(G^*_{j,:}),
+        R_{k,j}   = rho_hat_k^T rho_j^*.
+
+    If all rollout displacements collapse to one direction while the target
+    trajectory has distinct directions, the rows ``A_{k,:}`` become nearly
+    identical and cannot match the distinct target profiles ``G^*_{j,:}``.
+    Conversely, if the target trajectory is genuinely parallel, its profiles
+    are also similar, so this term does not invent artificial diversity.
+
+    ``Z`` is the frozen-weight NITP shallow-signature cosine.  Projector
+    parameters remain stop-gradient while gradients through the rollout state
+    are preserved by ``_project_nitp_with_frozen_weights``.
+
+    No tunable mixing coefficient is introduced: 1/2 and 1/4 follow exactly
+    from concatenating independently normalized blocks.  The implementation
+    deliberately evaluates the W-by-W cosines row-by-row instead of creating
+    a ``[B,T,W,D]`` FP32 tensor, keeping peak memory small for W=4.
     """
-    similarity_rows = []
+    horizon = len(predicted_by_step)
+    initial_common = initial_hidden_states[:, :common_length, :]
+    detached_initial = initial_common.detach()
+
+    displacement_rows = []
     for predicted_states in predicted_by_step:
-        predicted_common = predicted_states[:, :common_length, :]
+        predicted_displacement = (
+            predicted_states[:, :common_length, :] - initial_common
+        )
         row = []
-        for offset in range(1, len(predicted_by_step) + 1):
-            target_common = target_hidden_states[
+        for offset in range(1, horizon + 1):
+            target_displacement = (
+                target_hidden_states[
+                    :, offset : offset + common_length, :
+                ].detach()
+                - detached_initial
+            )
+            row.append(
+                F.cosine_similarity(
+                    predicted_displacement.float(),
+                    target_displacement.float(),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+            )
+        displacement_rows.append(torch.stack(row, dim=-1))
+    displacement_similarity = torch.stack(displacement_rows, dim=-2)
+
+    target_gram_rows = []
+    for left_offset in range(1, horizon + 1):
+        left_displacement = (
+            target_hidden_states[
+                :, left_offset : left_offset + common_length, :
+            ].detach()
+            - detached_initial
+        )
+        row = []
+        for right_offset in range(1, horizon + 1):
+            right_displacement = (
+                target_hidden_states[
+                    :, right_offset : right_offset + common_length, :
+                ].detach()
+                - detached_initial
+            )
+            row.append(
+                F.cosine_similarity(
+                    left_displacement.float(),
+                    right_displacement.float(),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+            )
+        target_gram_rows.append(torch.stack(row, dim=-1))
+    target_gram = torch.stack(target_gram_rows, dim=-2)
+
+    predicted_profiles = F.normalize(
+        displacement_similarity,
+        p=2.0,
+        dim=-1,
+        eps=1.0e-8,
+    )
+    target_profiles = F.normalize(
+        target_gram,
+        p=2.0,
+        dim=-1,
+        eps=1.0e-8,
+    )
+    relational_similarity = torch.einsum(
+        "btkr,btjr->btkj",
+        predicted_profiles,
+        target_profiles,
+    )
+
+    nitp_rows = []
+    for predicted_condition in condition_by_step:
+        predicted_condition = predicted_condition[:, :common_length, :]
+        row = []
+        for offset in range(1, horizon + 1):
+            target_condition = shallow_target_states[
                 :, offset : offset + common_length, :
             ].detach()
             row.append(
                 F.cosine_similarity(
-                    predicted_common.float(),
-                    target_common.float(),
+                    predicted_condition.float(),
+                    target_condition.float(),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+            )
+        nitp_rows.append(torch.stack(row, dim=-1))
+    nitp_signature_similarity = torch.stack(nitp_rows, dim=-2)
+
+    geometric_similarity = 0.5 * (
+        displacement_similarity + relational_similarity
+    )
+    joint_similarity = 0.5 * (
+        geometric_similarity + nitp_signature_similarity
+    )
+    return (
+        joint_similarity,
+        displacement_similarity,
+        relational_similarity,
+        nitp_signature_similarity,
+    )
+
+
+def _temporal_matrix_statistics(
+    matrix: torch.Tensor,
+    window_mask: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Return ``[diagonal, off_diagonal, margin, positive_margin_fraction]``.
+
+    The margin is the diagonal score minus the hardest incorrect candidate in
+    each row.  ``positive_margin_fraction`` prevents a positive mean margin
+    from hiding a large fraction of individually misidentified rows.
+    """
+    horizon = int(matrix.shape[-1])
+    zero = matrix.new_zeros((), dtype=torch.float32)
+    diagonal = matrix.diagonal(dim1=-2, dim2=-1)
+    diagonal_mean = _masked_mean(diagonal.mean(dim=-1), window_mask)
+    if horizon <= 1:
+        return torch.stack((diagonal_mean, zero, zero, zero))
+
+    identity = torch.eye(horizon, dtype=torch.bool, device=matrix.device)
+    identity = identity.view(1, 1, horizon, horizon)
+    off_diagonal = matrix.masked_select(~identity).view(
+        matrix.shape[0],
+        matrix.shape[1],
+        horizon * (horizon - 1),
+    )
+    off_diagonal_mean = _masked_mean(
+        off_diagonal.mean(dim=-1),
+        window_mask,
+    )
+    hard_negative = matrix.masked_fill(
+        identity,
+        torch.finfo(matrix.dtype).min,
+    ).max(dim=-1).values
+    row_margin = diagonal - hard_negative
+    margin = _masked_mean(row_margin.mean(dim=-1), window_mask)
+    positive_fraction = _masked_mean(
+        (row_margin > 0.0).float().mean(dim=-1),
+        window_mask,
+    )
+    return torch.stack(
+        (diagonal_mean, off_diagonal_mean, margin, positive_fraction)
+    )
+
+
+def _temporal_neighbor_cosines(
+    matrix: torch.Tensor,
+    window_mask: torch.Tensor,
+) -> torch.Tensor:
+    """Average matrix entries grouped by temporal distance ``|k-j|``."""
+    horizon = int(matrix.shape[-1])
+    if horizon <= 1:
+        return matrix.new_zeros((0,), dtype=torch.float32)
+
+    indices = torch.arange(horizon, device=matrix.device)
+    temporal_distance = (indices[:, None] - indices[None, :]).abs()
+    values = []
+    for offset in range(1, horizon):
+        offset_mask = temporal_distance.eq(offset).view(
+            1, 1, horizon, horizon
+        )
+        selected = matrix.masked_select(offset_mask).view(
+            matrix.shape[0],
+            matrix.shape[1],
+            -1,
+        )
+        values.append(_masked_mean(selected.mean(dim=-1), window_mask))
+    return torch.stack(values)
+
+
+def _pairwise_cosine_matrix(states: torch.Tensor) -> torch.Tensor:
+    """Return the within-trajectory cosine Gram matrix for ``[B,T,W,D]``."""
+    unit = F.normalize(states.float(), p=2.0, dim=-1, eps=1.0e-8)
+    return torch.einsum("btkd,btjd->btkj", unit, unit)
+
+
+def _pairwise_cosine_from_steps(
+    states_by_step: list[torch.Tensor],
+    origin: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """Low-peak-memory pair matrix from W tensors shaped ``[B,T,D]``."""
+    rows = []
+    for left in states_by_step:
+        left_value = left if origin is None else left - origin
+        row = []
+        for right in states_by_step:
+            right_value = right if origin is None else right - origin
+            row.append(
+                F.cosine_similarity(
+                    left_value.float(),
+                    right_value.float(),
                     dim=-1,
-                    eps=1e-8,
+                    eps=1.0e-8,
                 )
             )
-        similarity_rows.append(torch.stack(row, dim=-1))
-    return torch.stack(similarity_rows, dim=-2)
+        rows.append(torch.stack(row, dim=-1))
+    return torch.stack(rows, dim=-2)
+
+
+def _off_diagonal_pair_mean(
+    matrix: torch.Tensor,
+    window_mask: torch.Tensor,
+) -> torch.Tensor:
+    horizon = int(matrix.shape[-1])
+    zero = matrix.new_zeros((), dtype=torch.float32)
+    if horizon <= 1:
+        return zero
+    identity = torch.eye(horizon, dtype=torch.bool, device=matrix.device)
+    selected = matrix.masked_select(
+        ~identity.view(1, 1, horizon, horizon)
+    ).view(
+        matrix.shape[0],
+        matrix.shape[1],
+        horizon * (horizon - 1),
+    )
+    return _masked_mean(selected.mean(dim=-1), window_mask)
 
 
 def compute_nitp_temporal_objective(
     hidden_states: torch.Tensor,
     shallow_target_states: torch.Tensor,
-    token_embeds: torch.Tensor,
     labels: torch.LongTensor,
     attention_mask: Optional[torch.Tensor],
     transition_model: NITPTemporalTransition,
@@ -6237,79 +6545,63 @@ def compute_nitp_temporal_objective(
     target_temperature: float,
     pad_token_id: Optional[int] = None,
     eos_token_id: Optional[int] = None,
-) -> Tuple[
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-    torch.Tensor,
-]:
+) -> Tuple[torch.Tensor, ...]:
     """
-    NITP temporal transition and local temporal-identification objective.
+    Causal NITP temporal rollout with relational joint identification.
 
-    The branch learns a teacher-forced rollout in final-hidden-state space:
+    Autonomous rollout:
 
         h_hat_{t,0} = h_t
-        h_hat_{t,k} = F_psi(h_hat_{t,k-1}, E[x_{t+k}]).
-
-    It contains two trainable losses:
-
-    1. ``state_loss``: Smooth-L1 alignment of each rollout state with the
-       stop-gradient real final state at the same offset.
-    2. ``identification_loss``: a symmetric listwise classification objective
-       over the W x W cosine-similarity matrix.  Each predicted step k must be
-       identifiable as the real state at offset k.  The target is temporally
-       softened, so adjacent states may remain legitimately similar.
-
-    The objective never projects rollouts to vocabulary space.  The existing
-    NITP projector is used only for ``rollout_nitp_cosine``, a detached metric
-    that checks whether predicted final states retain NITP's ability to predict
-    the next shallow implicit token.
-
-    The shortcut diagnostic is an evaluation-only, on-manifold context
-    counterfactual.  For example ``i`` it keeps the exact teacher-forced future
-    tokens but replaces the initial state with a different example's state:
-
-        pi(i) = (i + 1) mod B
-        h_tilde_{i,t,0} = stopgrad(h^L_{pi(i),t})
-        h_tilde_{i,t,k} = F_psi(
-            h_tilde_{i,t,k-1}, E[x_{i,t+k}]
-        ).
-
-    Let L_full be ``state_loss`` and L_wrong the Smooth-L1 loss of that
-    wrong-context rollout against the original targets.  The single logged
-    diagnostic is the bounded context-specificity score
-
-        C_ctx = (L_wrong - L_full) / (L_wrong + L_full + eps).
-
-    ``C_ctx`` near zero means that an arbitrary plausible context works almost
-    as well as the correct one; positive values mean the exact contextual state
-    matters; negative values mean the wrong context unexpectedly performs
-    better.  The counterfactual runs under ``torch.no_grad()`` and does not
-    modify the training objective.  A deterministic cyclic batch shift avoids
-    RNG, graph recompilation, and cross-rank communication.
-
-    Returned scalars, in order:
-        weighted_total, state_loss, identification_loss,
-        temporal_top1_accuracy, mean_absolute_offset,
-        diagonal_cosine, off_diagonal_cosine, identification_margin,
-        last_step_state_loss, rollout_nitp_cosine,
-        rollout_pair_cosine, valid_window_fraction, context_specificity.
+        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.
     """
     horizon = int(horizon)
     seq_len = int(hidden_states.shape[1])
     max_horizon = max(min(horizon, seq_len - 1), 0)
     zero = hidden_states.new_zeros((), dtype=torch.float32)
     if max_horizon == 0:
-        return (zero,) * 13
+        return (
+            *((zero,) * 14),
+            zero.new_zeros((4, 4)),
+            zero.new_zeros((11, 0)),
+            zero.new_zeros((4, 0)),
+            zero.new_zeros((6,)),
+        )
 
     device = hidden_states.device
     labels_on_device = labels.to(device=device)
@@ -6322,69 +6614,167 @@ def compute_nitp_temporal_objective(
         valid_tokens = valid_tokens & (labels_on_device != pad_token_id)
 
     current_states = hidden_states
-    shifted_token_embeds = token_embeds
     shifted_targets = hidden_states
-    predicted_by_step = []
-    step_masks = []
+    predicted_by_step: list[torch.Tensor] = []
+    condition_by_step: list[torch.Tensor] = []
+    step_masks: list[torch.Tensor] = []
+    state_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] = []
+    predicted_turn_by_step: list[torch.Tensor] = []
+    target_turn_by_step: list[torch.Tensor] = []
+
     state_loss_total = zero
-    rollout_nitp_cosine_total = zero
-    last_step_state_loss = 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):
-        # At step k, current_states contains either the real h_t (k=1) or the
-        # recursively predicted h_hat_{t,k-1}.  Every next token is the real
-        # teacher-forced token, which isolates latent-transition quality from
-        # token-sampling error.
+        # 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_token_embeds = shifted_token_embeds[:, 1:, :]
         shifted_targets = shifted_targets[:, 1:, :]
-        predicted_states = transition_model(source_states, shifted_token_embeds)
+        next_shallow_signature = _project_nitp_with_frozen_weights(
+            source_states,
+            nitp_projector,
+        )
+        predicted_states = transition_model(
+            source_states,
+            next_shallow_signature.detach(),
+        )
 
-        step_span = step + 2  # source + (step + 1) consumed future tokens
         step_mask = _valid_temporal_span_mask(
             valid_tokens=valid_tokens,
             labels=labels_on_device,
-            span_length=step_span,
+            span_length=step + 2,
             eos_token_id=eos_token_id,
         )
         step_masks.append(step_mask)
+
         step_state_loss = _masked_smooth_l1(
             predicted=predicted_states,
             target=shifted_targets,
             mask=step_mask,
         )
         state_loss_total = state_loss_total + step_state_loss
-        last_step_state_loss = step_state_loss
+        state_loss_by_step.append(step_state_loss.detach().float())
 
-        # NITP interaction metric only.  It does not contribute gradients or
-        # alter the objective.  For k=1 this reproduces the ordinary NITP
-        # alignment; later k measure whether recursively predicted states still
-        # point toward the correct shallow implicit token.
         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
+
+            step_delta_cosine = F.cosine_similarity(
+                predicted_delta.float(),
+                real_delta.float(),
+                dim=-1,
+                eps=1.0e-8,
+            )
+            delta_cosine_by_step.append(
+                _masked_mean(step_delta_cosine, step_mask)
+            )
+
+            predicted_norm = predicted_delta.float().norm(dim=-1)
+            real_norm = real_delta.float().norm(dim=-1)
+            predicted_norm_mean = _masked_mean(predicted_norm, step_mask)
+            real_norm_mean = _masked_mean(real_norm, step_mask)
+            delta_norm_ratio_by_step.append(
+                predicted_norm_mean / real_norm_mean.clamp_min(1.0e-8)
+            )
+            step_weights = step_mask.to(dtype=torch.float32)
+            path_predicted_norm_sum = path_predicted_norm_sum + (
+                predicted_norm * step_weights
+            ).sum()
+            path_target_norm_sum = path_target_norm_sum + (
+                real_norm * step_weights
+            ).sum()
+
             shallow_target = shallow_target_states[:, step + 1 :, :]
-            projected_shallow = nitp_projector(source_states)
-            step_cosine = F.cosine_similarity(
-                projected_shallow.float(),
+            step_nitp_cosine = F.cosine_similarity(
+                next_shallow_signature.detach().float(),
                 shallow_target.float(),
                 dim=-1,
-                eps=1e-8,
+                eps=1.0e-8,
             )
-            rollout_nitp_cosine_total = rollout_nitp_cosine_total + _masked_mean(
-                step_cosine,
-                step_mask,
+            rollout_nitp_cosine_by_step.append(
+                _masked_mean(step_nitp_cosine, step_mask)
             )
 
+            if previous_predicted_delta is None:
+                predicted_turn_by_step.append(zero)
+                target_turn_by_step.append(zero)
+            else:
+                previous_predicted_aligned = previous_predicted_delta[
+                    :, : predicted_delta.shape[1], :
+                ]
+                previous_real_aligned = previous_real_delta[
+                    :, : real_delta.shape[1], :
+                ]
+                predicted_turn = F.cosine_similarity(
+                    previous_predicted_aligned.float(),
+                    predicted_delta.float(),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+                target_turn = F.cosine_similarity(
+                    previous_real_aligned.float(),
+                    real_delta.float(),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+                predicted_turn_by_step.append(
+                    _masked_mean(predicted_turn, step_mask)
+                )
+                target_turn_by_step.append(
+                    _masked_mean(target_turn, step_mask)
+                )
+            previous_predicted_delta = predicted_delta.detach()
+            previous_real_delta = real_delta.detach()
+
         predicted_by_step.append(predicted_states)
+        condition_by_step.append(next_shallow_signature)
         current_states = predicted_states
 
     inv_horizon = 1.0 / float(max_horizon)
     state_loss = state_loss_total * inv_horizon
-    rollout_nitp_cosine = rollout_nitp_cosine_total * inv_horizon
+    state_step_tensor = torch.stack(state_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)
+    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)
+    )
 
-    # All W rollout steps must share the same source positions to form the
-    # local W x W identification matrix.  The common number of starts is T-W.
     common_length = seq_len - max_horizon
     if common_length <= 0:
+        zero_steps = zero.new_zeros((max_horizon,))
+        step_metrics = torch.stack(
+            (
+                state_step_tensor,
+                delta_step_tensor,
+                delta_norm_ratio_tensor,
+                rollout_nitp_step_tensor,
+                zero_steps,
+                zero_steps,
+                zero_steps,
+                zero_steps,
+                predicted_turn_tensor,
+                target_turn_tensor,
+                predicted_turn_tensor - target_turn_tensor,
+            )
+        )
         weighted_total = float(dynamics_weight) * state_loss
         return (
             weighted_total,
@@ -6396,21 +6786,30 @@ def compute_nitp_temporal_objective(
             zero,
             zero,
             last_step_state_loss,
+            delta_cosine,
             rollout_nitp_cosine,
             zero,
             zero,
             zero,
+            zero.new_zeros((4, 4)),
+            step_metrics,
+            zero.new_zeros((4, max(max_horizon - 1, 0))),
+            torch.stack((path_length_ratio, zero, zero, zero, zero, zero)),
         )
 
-    # Build only the small [B, T-W, W, W] matrix.  Do not materialize a
-    # [B, T-W, W, D] FP32 tensor: at the production shape B=64, T=512,
-    # W=4, D=512 that unnecessary representation would consume hundreds of
-    # MiB.  Pairwise cosine calls keep only one [B, T-W, D] FP32 temporary.
-    similarity = _build_local_temporal_similarity(
+    (
+        similarity,
+        displacement_similarity,
+        relational_similarity,
+        nitp_signature_similarity,
+    ) = _build_relational_temporal_similarity(
         predicted_by_step=predicted_by_step,
+        condition_by_step=condition_by_step,
         target_hidden_states=hidden_states,
+        shallow_target_states=shallow_target_states,
+        initial_hidden_states=hidden_states,
         common_length=common_length,
-    )  # [B, T-W, W_pred, W_real]
+    )
 
     window_mask = _valid_temporal_span_mask(
         valid_tokens=valid_tokens,
@@ -6438,137 +6837,191 @@ def compute_nitp_temporal_objective(
 
     row_prediction = similarity.argmax(dim=-1)
     expected_offsets = offsets.view(1, 1, -1)
-    top1_per_window = (row_prediction == expected_offsets).float().mean(dim=-1)
-    temporal_top1_accuracy = _masked_mean(top1_per_window, window_mask)
+    temporal_top1_accuracy = _masked_mean(
+        (row_prediction == expected_offsets).float().mean(dim=-1),
+        window_mask,
+    )
     mean_absolute_offset = _masked_mean(
         (row_prediction - expected_offsets).abs().float().mean(dim=-1),
         window_mask,
     )
 
-    diagonal = similarity.diagonal(dim1=-2, dim2=-1)
-    diagonal_cosine = _masked_mean(diagonal.mean(dim=-1), window_mask)
+    # 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.
+    component_matrices = (
+        similarity.detach(),
+        displacement_similarity.detach(),
+        relational_similarity.detach(),
+        nitp_signature_similarity.detach(),
+    )
+    component_metrics = torch.stack(
+        [
+            _temporal_matrix_statistics(component, window_mask)
+            for component in component_matrices
+        ]
+    )
+    offset_metrics = torch.stack(
+        [
+            _temporal_neighbor_cosines(component, window_mask)
+            for component in component_matrices
+        ]
+    )
+    diagonal_cosine = component_metrics[0, 0]
+    off_diagonal_cosine = component_metrics[0, 1]
+    identification_margin = component_metrics[0, 2]
 
-    if max_horizon > 1:
-        identity = torch.eye(max_horizon, dtype=torch.bool, device=device)
-        off_diagonal = similarity.masked_select(
-            ~identity.view(1, 1, max_horizon, max_horizon)
-        )
-        off_diagonal = off_diagonal.view(
-            similarity.shape[0],
-            similarity.shape[1],
-            max_horizon * (max_horizon - 1),
+    with torch.no_grad():
+        initial_common = hidden_states[:, :common_length, :].detach()
+        predicted_state_steps = [
+            state[:, :common_length, :].detach()
+            for state in predicted_by_step
+        ]
+        target_state_steps = [
+            hidden_states[
+                :, offset : offset + common_length, :
+            ].detach()
+            for offset in range(1, max_horizon + 1)
+        ]
+        rollout_pair_cosine = _off_diagonal_pair_mean(
+            _pairwise_cosine_from_steps(predicted_state_steps),
+            window_mask,
         )
-        off_diagonal_cosine = _masked_mean(
-            off_diagonal.mean(dim=-1),
+        target_pair_cosine = _off_diagonal_pair_mean(
+            _pairwise_cosine_from_steps(target_state_steps),
             window_mask,
         )
-
-        hard_negative = similarity.masked_fill(
-            identity.view(1, 1, max_horizon, max_horizon),
-            torch.finfo(similarity.dtype).min,
-        ).max(dim=-1).values
-        identification_margin = _masked_mean(
-            (diagonal - hard_negative).mean(dim=-1),
+        rollout_displacement_pair_cosine = _off_diagonal_pair_mean(
+            _pairwise_cosine_from_steps(
+                predicted_state_steps,
+                origin=initial_common,
+            ),
             window_mask,
         )
-
-        predicted_pair_rows = []
-        for predicted_left in predicted_by_step:
-            predicted_left = predicted_left[:, :common_length, :]
-            pair_row = []
-            for predicted_right in predicted_by_step:
-                predicted_right = predicted_right[:, :common_length, :]
-                pair_row.append(
-                    F.cosine_similarity(
-                        predicted_left.float(),
-                        predicted_right.float(),
-                        dim=-1,
-                        eps=1e-8,
-                    )
-                )
-            predicted_pair_rows.append(torch.stack(pair_row, dim=-1))
-        predicted_pair_similarity = torch.stack(predicted_pair_rows, dim=-2)
-        rollout_off_diagonal = predicted_pair_similarity.masked_select(
-            ~identity.view(1, 1, max_horizon, max_horizon)
-        ).view(
-            similarity.shape[0],
-            similarity.shape[1],
-            max_horizon * (max_horizon - 1),
-        )
-        rollout_pair_cosine = _masked_mean(
-            rollout_off_diagonal.mean(dim=-1),
+        target_displacement_pair_cosine = _off_diagonal_pair_mean(
+            _pairwise_cosine_from_steps(
+                target_state_steps,
+                origin=initial_common,
+            ),
             window_mask,
         )
-    else:
-        off_diagonal_cosine = zero
-        identification_margin = zero
-        rollout_pair_cosine = zero
-
-    # ------------------------------------------------------------------
-    # Context-specificity diagnostic (monitoring only, no gradients).
-    # ------------------------------------------------------------------
-    # The multiplicative transition already guarantees F(0, e) = 0, so a
-    # zero-context/token-only rollout is now true by construction and no longer
-    # diagnoses learned behavior.  Instead, use an on-manifold but incorrect
-    # context from the next example in the local batch while retaining each
-    # example's exact teacher-forced future tokens:
-    #
-    #   pi(i) = (i + 1) mod B
-    #   h_tilde_{i,t,0} = stopgrad(h^L_{pi(i),t})
-    #   h_tilde_{i,t,k} = F_psi(h_tilde_{i,t,k-1}, E[x_{i,t+k}]).
-    #
-    # With L_full = state_loss and L_wrong the counterfactual Smooth-L1 loss,
-    # report the bounded, scale-normalized score
-    #
-    #   C_ctx = (L_wrong - L_full) / (L_wrong + L_full + eps).
-    #
-    # C_ctx ~= 0: any plausible context works almost equally well.
-    # C_ctx > 0:  the exact context improves the transition.
-    # C_ctx < 0:  the wrong context unexpectedly performs better.
-    #
-    # This replaces the three old token-only metrics at lower cost: it reuses
-    # one no-grad rollout but does not build a second W x W similarity matrix.
-    with torch.no_grad():
-        if int(hidden_states.shape[0]) > 1:
-            wrong_context_current = torch.roll(
-                hidden_states.detach(), shifts=1, dims=0
-            )
-            wrong_context_shifted_embeds = token_embeds.detach()
-            wrong_context_targets = hidden_states.detach()
-            wrong_context_state_loss_total = zero
 
+        cross_predicted_by_step = []
+        cross_target_by_step = []
+        if int(hidden_states.shape[0]) > 1:
             for step in range(max_horizon):
-                wrong_context_source = wrong_context_current[:, :-1, :]
-                wrong_context_shifted_embeds = (
-                    wrong_context_shifted_embeds[:, 1:, :]
+                step_mask_common = step_masks[step][:, :common_length]
+                paired_mask = step_mask_common & torch.roll(
+                    step_mask_common,
+                    shifts=1,
+                    dims=0,
                 )
-                wrong_context_targets = wrong_context_targets[:, 1:, :]
-                wrong_context_predicted = transition_model(
-                    wrong_context_source,
-                    wrong_context_shifted_embeds,
+                predicted_displacement = (
+                    predicted_by_step[step][:, :common_length, :].detach()
+                    - initial_common
                 )
-                wrong_context_state_loss_total = (
-                    wrong_context_state_loss_total
-                    + _masked_smooth_l1(
-                        predicted=wrong_context_predicted,
-                        target=wrong_context_targets,
-                        mask=step_masks[step],
-                    )
+                target_displacement = (
+                    hidden_states[
+                        :, step + 1 : step + 1 + common_length, :
+                    ].detach()
+                    - initial_common
+                )
+                cross_predicted = F.cosine_similarity(
+                    predicted_displacement.float(),
+                    torch.roll(
+                        predicted_displacement.float(),
+                        shifts=1,
+                        dims=0,
+                    ),
+                    dim=-1,
+                    eps=1.0e-8,
+                )
+                cross_target = F.cosine_similarity(
+                    target_displacement.float(),
+                    torch.roll(
+                        target_displacement.float(),
+                        shifts=1,
+                        dims=0,
+                    ),
+                    dim=-1,
+                    eps=1.0e-8,
                 )
-                wrong_context_current = wrong_context_predicted
+                cross_predicted_by_step.append(
+                    _masked_mean(cross_predicted, paired_mask)
+                )
+                cross_target_by_step.append(
+                    _masked_mean(cross_target, paired_mask)
+                )
+        else:
+            cross_predicted_by_step = [zero] * max_horizon
+            cross_target_by_step = [zero] * max_horizon
+        cross_predicted_tensor = torch.stack(cross_predicted_by_step)
+        cross_target_tensor = torch.stack(cross_target_by_step)
 
-            wrong_context_state_loss = (
-                wrong_context_state_loss_total * inv_horizon
-            )
-            full_state_loss = state_loss.detach().float()
-            wrong_state_loss = wrong_context_state_loss.float()
-            context_specificity = (
+        # 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:
-            # A distinct on-manifold context does not exist for batch size one.
-            context_specificity = zero
+            trajectory_specificity_by_step = zero.new_zeros((max_horizon,))
+            trajectory_specificity = zero
+
+    step_metrics = torch.stack(
+        (
+            state_step_tensor,
+            delta_step_tensor,
+            delta_norm_ratio_tensor,
+            rollout_nitp_step_tensor,
+            trajectory_specificity_by_step,
+            cross_predicted_tensor,
+            cross_target_tensor,
+            cross_predicted_tensor - cross_target_tensor,
+            predicted_turn_tensor,
+            target_turn_tensor,
+            predicted_turn_tensor - target_turn_tensor,
+        )
+    )
+    scalar_metrics = torch.stack(
+        (
+            path_length_ratio,
+            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
@@ -6584,10 +7037,15 @@ def compute_nitp_temporal_objective(
         off_diagonal_cosine,
         identification_margin,
         last_step_state_loss,
+        delta_cosine,
         rollout_nitp_cosine,
         rollout_pair_cosine,
         valid_window_fraction,
-        context_specificity,
+        trajectory_specificity,
+        component_metrics,
+        step_metrics,
+        offset_metrics,
+        scalar_metrics,
     )
 
 
@@ -6761,9 +7219,14 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         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_context_specificity = 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
+        self._last_nitp_temporal_scalar_metrics = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
@@ -6913,7 +7376,12 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         nitp_temporal_rollout_nitp_cosine = None
         nitp_temporal_rollout_pair_cosine = None
         nitp_temporal_valid_window_fraction = None
-        nitp_temporal_context_specificity = 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
@@ -6934,9 +7402,14 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         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_context_specificity = 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
+        self._last_nitp_temporal_scalar_metrics = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
@@ -7018,7 +7491,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                         )
                     nitp_target_hidden_states = hidden_states_tuple[target_layer]
 
-                token_embeds = hidden_states_tuple[0]
                 temporal_apply_loss = bool(
                     getattr(self.config, "nitp_temporal_apply_loss", True)
                 )
@@ -7032,7 +7504,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     temporal_results = compute_nitp_temporal_objective(
                         hidden_states=hidden_states,
                         shallow_target_states=nitp_target_hidden_states,
-                        token_embeds=token_embeds,
                         labels=labels,
                         attention_mask=attention_mask,
                         transition_model=self.nitp_temporal_transition,
@@ -7059,7 +7530,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                         temporal_results = compute_nitp_temporal_objective(
                             hidden_states=hidden_states,
                             shallow_target_states=nitp_target_hidden_states,
-                            token_embeds=token_embeds,
                             labels=labels,
                             attention_mask=attention_mask,
                             transition_model=self.nitp_temporal_transition,
@@ -7091,10 +7561,15 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     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_context_specificity,
+                    nitp_temporal_trajectory_specificity,
+                    nitp_temporal_component_metrics,
+                    nitp_temporal_step_metrics,
+                    nitp_temporal_offset_metrics,
+                    nitp_temporal_scalar_metrics,
                 ) = temporal_results
 
                 if temporal_apply_loss:
@@ -7113,10 +7588,10 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                         "NeoLLMForCausalLM.forward enables them automatically "
                         "while labels are present."
                     )
+                # NextLat keeps its original teacher-forced token-embedding
+                # path.  This assignment is local to NextLat; causal Temporal
+                # NITP no longer reads ``hidden_states_tuple[0]``.
                 token_embeds = hidden_states_tuple[0]
-                # NextLat uses final hidden states h_t and input embeddings
-                # x_{t+1}; p_psi is training-only and inference remains the
-                # standard autoregressive LM path.
                 nextlat_loss, nextlat_mse_loss, nextlat_kl_loss = compute_nextlat_loss(
                     hidden_states=hidden_states,
                     token_embeds=token_embeds,
@@ -7203,9 +7678,34 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 if nitp_temporal_valid_window_fraction is not None
                 else None
             )
-            self._last_nitp_temporal_context_specificity = (
-                nitp_temporal_context_specificity.detach()
-                if nitp_temporal_context_specificity is not 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
+                else None
+            )
+            self._last_nitp_temporal_step_metrics = (
+                nitp_temporal_step_metrics.detach()
+                if nitp_temporal_step_metrics is not None
+                else None
+            )
+            self._last_nitp_temporal_offset_metrics = (
+                nitp_temporal_offset_metrics.detach()
+                if nitp_temporal_offset_metrics is not None
+                else None
+            )
+            self._last_nitp_temporal_scalar_metrics = (
+                nitp_temporal_scalar_metrics.detach()
+                if nitp_temporal_scalar_metrics is not None
                 else None
             )
             self._last_nitp_temporal_to_nitp_ratio = (