Update modeling_neollm.py

modeling_neollm.py

@@ -4946,9 +4946,13 @@ class SpellingBeeEmbedding(nn.Module):
 
         e_bee(t) = 0.5 * (e_tok(t) + e_chars(t))
 
-        e_chars(t) = inv_sqrt_len(t) * Σ_{i=0}^{15} RoPE(e_byte[b_i], i)
+        e_chars(t) = (1 / MAX_BYTES) * Σ_{i=0}^{15}
+                     RoPE(e_byte[b_i], i)
 
-    where inv_sqrt_len = 1/√|t| is precomputed per token type at setup time.
+    The active path uses the fixed 16-position mean from the reference
+    implementation.  A theoretical alternative is to mask padded positions
+    and normalize only the valid bytes with ``1 / sqrt(|t|)``; that variant is
+    documented in the configuration but is not active here.
 
     Key design decisions vs. a naïve per-occurrence implementation:
 
@@ -4966,10 +4970,11 @@ class SpellingBeeEmbedding(nn.Module):
        dims) into a single kernel.  Token-level e_chars is then a gather +
        sum over this table, also fully static.
 
-    3. **Precomputed inv_sqrt_lens** — 1/√byte_len is computed once in
-       set_byte_table and stored as a persistent buffer.  The per-forward
-       normalisation becomes a single elementwise multiply, with no sqrt or
-       division in the hot path.
+    3. **Fixed-length mean** — every token contributes exactly 16
+       position-aware byte slots, including the learned null-byte padding,
+       and the sum is divided by ``MAX_BYTES``.  This keeps the active path
+       faithful to the reference experiment and avoids a per-token scale
+       factor.
 
     Compatible with both the standard embed_tokens path and the
     LeviathanGenerator path.
@@ -4981,7 +4986,7 @@ class SpellingBeeEmbedding(nn.Module):
 
     **Setup: call ``set_byte_table(tokenizer)`` once after model init** (and
     before any .to(device) / FP8 conversion) before training.  The byte table
-    and inv_sqrt_lens are persistent buffers saved in checkpoints.
+    token_bytes is a persistent buffer saved in checkpoints.
 
     References:
         Rabe, Clymo & Dong (2026). "Spelling Bee Embeddings for Language
@@ -5009,19 +5014,6 @@ class SpellingBeeEmbedding(nn.Module):
             torch.zeros(config.vocab_size, self.MAX_BYTES, dtype=torch.long),
             persistent=True,
         )
-        # inv_sqrt_lens [vocab_size]: precomputed 1/sqrt(byte_len) per token.
-        # Replaces the runtime sqrt+division of the naïve implementation.
-        self.register_buffer(
-            "inv_sqrt_lens",
-            torch.ones(config.vocab_size, dtype=torch.float),
-            persistent=True,
-        )
-
-        # LayerNorm over character embeddings — mirrors the reference impl
-        # (character_norm=True by default in littletrainingloop).  Runs in
-        # float32 for stability, applied at vocab level before the batch gather.
-        self.char_norm = nn.LayerNorm(d)
-
     def _load_from_state_dict(
         self,
         state_dict,
@@ -5056,7 +5048,7 @@ class SpellingBeeEmbedding(nn.Module):
 
     def set_byte_table(self, tokenizer) -> None:
         """
-        Precompute the UTF-8 byte table and inv_sqrt_lens from a tokenizer.
+        Precompute the UTF-8 byte table from a tokenizer.
 
         Must be called **once** after model instantiation and **before**
         ``.to(device)`` / FP8 conversion so the buffers land on the correct
@@ -5064,30 +5056,36 @@ class SpellingBeeEmbedding(nn.Module):
         be saved/restored from checkpoints automatically.
 
         Args:
-            tokenizer: Any HuggingFace tokenizer with
-                ``convert_ids_to_tokens(int) -> str | None``.
+            tokenizer: Any tokenizer with ``decode([token_id]) -> str``.
         """
         vocab_size = self.token_bytes.shape[0]
         byte_ids = torch.zeros(vocab_size, self.MAX_BYTES, dtype=torch.long)
-        inv_sqrt = torch.ones(vocab_size, dtype=torch.float)  # default 1/√1
 
         for token_id in range(vocab_size):
-            token_str = tokenizer.convert_ids_to_tokens(token_id)
-            if token_str is None:
-                continue
-            # Some tokenizers use a special space character (Ġ / ▁); encode
-            # directly to UTF-8 so byte values match raw text bytes.
+            # Decode through the active tokenizer.  convert_ids_to_tokens()
+            # returns an internal BPE/SentencePiece symbol (for example
+            # ``Ġapple`` or ``Ċ``), not the text surface (`` apple`` or ``\n``)
+            # whose UTF-8 spelling Spelling Bee is meant to represent.
             try:
+                try:
+                    token_str = tokenizer.decode(
+                        [token_id],
+                        skip_special_tokens=False,
+                        clean_up_tokenization_spaces=False,
+                    )
+                except TypeError:
+                    # Keep compatibility with tokenizer implementations that
+                    # do not expose clean_up_tokenization_spaces.
+                    token_str = tokenizer.decode(
+                        [token_id], skip_special_tokens=False
+                    )
                 raw = token_str.encode("utf-8")
             except Exception:
                 raw = b"\x00"
             n = min(len(raw), self.MAX_BYTES)
             for i in range(n):
                 byte_ids[token_id, i] = raw[i]
-            inv_sqrt[token_id] = 1.0 / math.sqrt(max(n, 1))
-
         self.token_bytes.copy_(byte_ids.to(self.token_bytes.device))
-        self.inv_sqrt_lens.copy_(inv_sqrt.to(self.inv_sqrt_lens.device))
 
     # ── Core helpers ──────────────────────────────────────────────────────────
 
@@ -5160,18 +5158,7 @@ class SpellingBeeEmbedding(nn.Module):
             torch.arange(
                 self.MAX_BYTES, device=rope_bytes.device
             ),  # [MAX_BYTES] — col index
-        ].sum(1)  # [V, d]
-
-        # ── Step 3: apply precomputed 1/√byte_len per vocab type ────────────
-        # No sqrt or division in the hot path — pure multiply.
-        e_chars_vocab = e_chars_vocab * self.inv_sqrt_lens.unsqueeze(-1)  # [V, d]
-
-        # ── Step 3b: LayerNorm over character embeddings (float32 for stability)
-        # Mirrors reference impl (character_norm=True).  Ensures E[‖e_chars‖²]
-        # matches E[‖e_tok‖²] regardless of token byte-length distribution.
-        e_chars_vocab = self.char_norm(
-            e_chars_vocab.to(self.char_norm.weight.dtype)
-        ).to(token_embeds.dtype)
+        ].mean(1)  # [V, d], fixed 16-position mean
 
         # ── Step 4: gather only the tokens present in this batch ────────────
         # This is the only B×S operation — a single embedding lookup.
@@ -5212,10 +5199,7 @@ class SpellingBeeEmbedding(nn.Module):
         e_chars_vocab = rope_bytes[
             self.token_bytes,
             torch.arange(self.MAX_BYTES, device=rope_bytes.device),
-        ].sum(1) * self.inv_sqrt_lens.unsqueeze(-1)  # [V, d]
-        e_chars_vocab = self.char_norm(
-            e_chars_vocab.to(self.char_norm.weight.dtype)
-        ).to(token_emb_weight.dtype)
+        ].mean(1)  # [V, d], fixed 16-position mean
         return (token_emb_weight + e_chars_vocab) * 0.5
 
 
@@ -5413,9 +5397,10 @@ class NeoLLMPreTrainedModel(PreTrainedModel):
                 module.attn_res_query_mlp.data.zero_()
 
         elif isinstance(module, SpellingBeeEmbedding):
-            # byte_emb initialised identically to token embeddings: std=1/√d.
-            # Ensures E[‖e_byte‖²] ≈ 1 at init, matching etok, so the
-            # normalisation factor α = sqrt(byte_len) is calibrated from step 0.
+            # Keep the byte lookup on the same base scale as the model's
+            # embedding initialization.  The active SBE path uses the fixed
+            # 16-position mean; the masked 1/sqrt(byte_len) alternative is
+            # theoretical only and is documented separately.
             d = module.byte_emb.embedding_dim
             nn.init.normal_(module.byte_emb.weight, mean=0.0, std=1.0 / math.sqrt(d))
 
@@ -6045,6 +6030,56 @@ def compute_nitp_loss(
     return (per_token_loss * valid_weights).sum() / valid_count
 
 
+class NITPTemporalTransition(nn.Module):
+    """
+    NITP-native latent transition used by the temporal extension.
+
+    This module is intentionally independent from ``NextLatDynamicsModel``.
+    It follows the existing NITP design language by using a bias-free SwiGLU
+    update over the concatenation of the current final hidden state and the
+    teacher-forced next-token embedding:
+
+        delta_t = W_down(SiLU(W_gate([h_t; e_{t+1}])) * W_up([h_t; e_{t+1}]))
+        h_hat_{t+1} = h_t + delta_t
+
+    The residual update makes the learned object an update rule rather than a
+    free reconstruction.  It is a training-only auxiliary module and never
+    materializes vocabulary logits.
+
+    Parameter count, excluding the already-existing NITP projector:
+
+        2 * (2d * m) + (m * d) + 2d
+
+    where ``m = nitp_temporal_intermediate_size`` and the final ``2d`` term is
+    the affine RMSNorm scale.  With d=512 and m=1024 this is 2,622,464
+    trainable parameters.
+    """
+
+    def __init__(self, config: NeoLLMConfig):
+        super().__init__()
+        hidden_size = int(config.hidden_size)
+        input_size = 2 * hidden_size
+        intermediate_size = int(config.nitp_temporal_intermediate_size)
+
+        self.input_norm = nn.RMSNorm(input_size, eps=config.rms_norm_eps)
+        self.gate_proj = nn.Linear(input_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(input_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(
+        self,
+        current_states: torch.Tensor,
+        next_token_embeds: torch.Tensor,
+    ) -> torch.Tensor:
+        transition_input = torch.cat([current_states, next_token_embeds], dim=-1)
+        transition_input = self.input_norm(transition_input)
+        delta = self.down_proj(
+            F.silu(self.gate_proj(transition_input))
+            * self.up_proj(transition_input)
+        )
+        return current_states + delta
+
+
 class NextLatDynamicsModel(nn.Module):
     """
     Latent dynamics model p_psi(h_t, x_{t+1}) from NextLat.
@@ -6099,6 +6134,321 @@ def _masked_smooth_l1(
     return (loss * weights).sum() / weights.expand_as(loss).sum().clamp_min(1.0)
 
 
+def _valid_temporal_span_mask(
+    valid_tokens: torch.Tensor,
+    labels: torch.Tensor,
+    span_length: int,
+    eos_token_id: Optional[int],
+) -> torch.Tensor:
+    """
+    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.
+    """
+    token_windows = valid_tokens.unfold(dimension=1, size=span_length, step=1)
+    span_mask = token_windows.all(dim=-1)
+    if eos_token_id is not None and span_length > 1:
+        label_windows = labels.unfold(dimension=1, size=span_length, step=1)
+        # The last token may itself be EOS; all earlier positions must remain
+        # inside the same document.
+        span_mask = span_mask & (label_windows[..., :-1] != eos_token_id).all(dim=-1)
+    return span_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,
+    nitp_projector: NITPProjector,
+    horizon: int,
+    dynamics_weight: float,
+    identification_weight: float,
+    similarity_temperature: float,
+    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,
+]:
+    """
+    NITP temporal transition and local temporal-identification objective.
+
+    The branch learns a teacher-forced rollout in final-hidden-state space:
+
+        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.
+
+    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.
+    """
+    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,) * 12
+
+    device = hidden_states.device
+    labels_on_device = labels.to(device=device)
+    if attention_mask is None:
+        valid_tokens = torch.ones_like(labels_on_device, dtype=torch.bool)
+    else:
+        valid_tokens = attention_mask.to(device=device).bool()
+    valid_tokens = valid_tokens & (labels_on_device != -100)
+    if pad_token_id is not None:
+        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 = []
+    state_loss_total = zero
+    rollout_nitp_cosine_total = zero
+    last_step_state_loss = zero
+
+    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.
+        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)
+
+        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,
+            eos_token_id=eos_token_id,
+        )
+        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
+
+        # 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():
+            shallow_target = shallow_target_states[:, step + 1 :, :]
+            projected_shallow = nitp_projector(source_states)
+            step_cosine = F.cosine_similarity(
+                projected_shallow.float(),
+                shallow_target.float(),
+                dim=-1,
+                eps=1e-8,
+            )
+            rollout_nitp_cosine_total = rollout_nitp_cosine_total + _masked_mean(
+                step_cosine,
+                step_mask,
+            )
+
+        predicted_by_step.append(predicted_states)
+        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
+
+    # 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:
+        weighted_total = float(dynamics_weight) * state_loss
+        return (
+            weighted_total,
+            state_loss,
+            zero,
+            zero,
+            zero,
+            zero,
+            zero,
+            zero,
+            last_step_state_loss,
+            rollout_nitp_cosine,
+            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_rows = []
+    for predicted_states in predicted_by_step:
+        predicted_common = predicted_states[:, :common_length, :]
+        row = []
+        for offset in range(1, max_horizon + 1):
+            target_common = hidden_states[
+                :, offset : offset + common_length, :
+            ].detach()
+            row.append(
+                F.cosine_similarity(
+                    predicted_common.float(),
+                    target_common.float(),
+                    dim=-1,
+                    eps=1e-8,
+                )
+            )
+        similarity_rows.append(torch.stack(row, dim=-1))
+    similarity = torch.stack(
+        similarity_rows,
+        dim=-2,
+    )  # [B, T-W, W_pred, W_real]
+
+    window_mask = _valid_temporal_span_mask(
+        valid_tokens=valid_tokens,
+        labels=labels_on_device,
+        span_length=max_horizon + 1,
+        eos_token_id=eos_token_id,
+    )
+    valid_window_fraction = window_mask.float().mean()
+
+    offsets = torch.arange(max_horizon, device=device)
+    distance = (offsets[:, None] - offsets[None, :]).abs().float()
+    target_scores = -distance / float(target_temperature)
+    row_target = F.softmax(target_scores, dim=-1)
+    col_target = F.softmax(target_scores, dim=0)
+
+    identification_logits = similarity / float(similarity_temperature)
+    row_log_probs = F.log_softmax(identification_logits, dim=-1)
+    col_log_probs = F.log_softmax(identification_logits, dim=-2)
+    row_loss = -(row_target * row_log_probs).sum(dim=-1).mean(dim=-1)
+    col_loss = -(col_target * col_log_probs).sum(dim=-2).mean(dim=-1)
+    identification_loss = _masked_mean(
+        0.5 * (row_loss + col_loss),
+        window_mask,
+    )
+
+    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)
+    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)
+
+    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),
+        )
+        off_diagonal_cosine = _masked_mean(
+            off_diagonal.mean(dim=-1),
+            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),
+            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),
+            window_mask,
+        )
+    else:
+        off_diagonal_cosine = zero
+        identification_margin = zero
+        rollout_pair_cosine = zero
+
+    weighted_total = (
+        float(dynamics_weight) * state_loss
+        + float(identification_weight) * identification_loss
+    )
+    return (
+        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,
+    )
+
+
 def _categorical_kl_loss(
     teacher_logits: torch.Tensor,
     student_logits: torch.Tensor,
@@ -6224,6 +6574,11 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.nitp_projector = NITPProjector(config) if config.use_nitp else None
+        self.nitp_temporal_transition = (
+            NITPTemporalTransition(config)
+            if bool(getattr(config, "use_nitp_temporal", False))
+            else None
+        )
         self.nextlat_dynamics = (
             NextLatDynamicsModel(config) if config.use_nextlat else None
         )
@@ -6252,6 +6607,20 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_ntp_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_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_to_nitp_ratio = None
+        self._last_nitp_temporal_loss_applied = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
@@ -6316,11 +6685,20 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             and labels is not None
             and float(getattr(self.config, "tweo_loss_weight", 0.0)) > 0.0
         )
+        nitp_temporal_enabled = (
+            bool(getattr(self.config, "use_nitp_temporal", False))
+            and self.nitp_temporal_transition is not None
+            and self.nitp_projector is not None
+            and labels is not None
+        )
         nitp_enabled = (
             bool(getattr(self.config, "use_nitp", False))
             and self.nitp_projector is not None
             and labels is not None
-            and float(getattr(self.config, "nitp_loss_weight", 0.0)) > 0.0
+            and (
+                float(getattr(self.config, "nitp_loss_weight", 0.0)) > 0.0
+                or nitp_temporal_enabled
+            )
         )
         nextlat_enabled = (
             bool(getattr(self.config, "use_nextlat", False))
@@ -6332,7 +6710,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             )
         )
         model_output_hidden_states = (
-            True if (nitp_enabled or nextlat_enabled) else output_hidden_states
+            True
+            if (nitp_enabled or nitp_temporal_enabled or nextlat_enabled)
+            else output_hidden_states
         )
 
         model_out = self.model(
@@ -6378,12 +6758,40 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         ntp_loss = None
         tweo_loss = None
         nitp_loss = None
+        nitp_temporal_loss = None
+        nitp_temporal_state_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_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_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_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_to_nitp_ratio = None
+        self._last_nitp_temporal_loss_applied = None
         self._last_nextlat_loss = None
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
@@ -6445,7 +6853,112 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     projector=self.nitp_projector,
                     pad_token_id=self.config.pad_token_id,
                 )
-                loss = loss + self.config.nitp_loss_weight * nitp_loss
+                nitp_target_hidden_states = hidden_states_tuple[target_layer]
+                if float(self.config.nitp_loss_weight) > 0.0:
+                    loss = loss + self.config.nitp_loss_weight * nitp_loss
+
+            if nitp_temporal_enabled:
+                if hidden_states_tuple is None or len(hidden_states_tuple) < 2:
+                    raise ValueError(
+                        "`use_nitp_temporal=True` requires backbone hidden "
+                        "states. NeoLLMForCausalLM.forward enables them "
+                        "automatically while labels are present."
+                    )
+                if nitp_target_hidden_states is None:
+                    target_layer = int(getattr(self.config, "nitp_target_layer", 1))
+                    if len(hidden_states_tuple) <= target_layer:
+                        raise ValueError(
+                            "`use_nitp_temporal=True` requires the NITP shallow "
+                            f"target layer {target_layer}."
+                        )
+                    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)
+                )
+                eos_token_id = (
+                    self.config.eos_token_id[0]
+                    if isinstance(self.config.eos_token_id, (list, tuple))
+                    and self.config.eos_token_id
+                    else self.config.eos_token_id
+                )
+                if temporal_apply_loss:
+                    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,
+                        nitp_projector=self.nitp_projector,
+                        horizon=self.config.nitp_temporal_horizon,
+                        dynamics_weight=self.config.nitp_temporal_dynamics_weight,
+                        identification_weight=(
+                            self.config.nitp_temporal_identification_weight
+                        ),
+                        similarity_temperature=(
+                            self.config.nitp_temporal_similarity_temperature
+                        ),
+                        target_temperature=(
+                            self.config.nitp_temporal_target_temperature
+                        ),
+                        pad_token_id=self.config.pad_token_id,
+                        eos_token_id=eos_token_id,
+                    )
+                else:
+                    # Monitoring-only mode: the branch is instantiated (and
+                    # therefore contributes parameters) but its graph is not
+                    # retained and its objective is not added to training.
+                    with torch.no_grad():
+                        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,
+                            nitp_projector=self.nitp_projector,
+                            horizon=self.config.nitp_temporal_horizon,
+                            dynamics_weight=(
+                                self.config.nitp_temporal_dynamics_weight
+                            ),
+                            identification_weight=(
+                                self.config.nitp_temporal_identification_weight
+                            ),
+                            similarity_temperature=(
+                                self.config.nitp_temporal_similarity_temperature
+                            ),
+                            target_temperature=(
+                                self.config.nitp_temporal_target_temperature
+                            ),
+                            pad_token_id=self.config.pad_token_id,
+                            eos_token_id=eos_token_id,
+                        )
+
+                (
+                    nitp_temporal_loss,
+                    nitp_temporal_state_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_rollout_nitp_cosine,
+                    nitp_temporal_rollout_pair_cosine,
+                    nitp_temporal_valid_window_fraction,
+                ) = temporal_results
+
+                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:
@@ -6484,6 +6997,81 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             self._last_ntp_loss = ntp_loss.detach() if ntp_loss is not None else None
             self._last_tweo_loss = tweo_loss.detach() if tweo_loss is not None else None
             self._last_nitp_loss = nitp_loss.detach() if nitp_loss is not None else None
+            self._last_nitp_temporal_loss = (
+                nitp_temporal_loss.detach()
+                if nitp_temporal_loss is not None
+                else None
+            )
+            self._last_nitp_temporal_state_loss = (
+                nitp_temporal_state_loss.detach()
+                if nitp_temporal_state_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
+                else None
+            )
+            self._last_nitp_temporal_top1_accuracy = (
+                nitp_temporal_top1_accuracy.detach()
+                if nitp_temporal_top1_accuracy is not None
+                else None
+            )
+            self._last_nitp_temporal_mean_absolute_offset = (
+                nitp_temporal_mean_absolute_offset.detach()
+                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_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
             )