Update configuration_neollm.py

configuration_neollm.py

@@ -364,14 +364,19 @@ class NeoLLMConfig(PretrainedConfig):
             Each token's UTF-8 encoding (up to 16 bytes) is embedded through a
             shared ``nn.Embedding(256, d)`` table. Byte embeddings are
             position-encoded with RoPE using intra-token byte positions (not
-            sequence positions), summed and normalised by ``√byte_len``, then
+            sequence positions), averaged over the fixed 16 byte slots, then
             averaged with the standard token embedding:
 
             .. math::
                 e_{\\text{bee}}(t) = \\frac{1}{2}\\left(e_{\\text{tok}}(t) +
-                \\frac{1}{\\sqrt{|t|}}\\sum_{i=1}^{16}
+                \\frac{1}{16}\\sum_{i=1}^{16}
                 \\text{RoPE}(e_{\\text{byte}}[b_i], i)\\right)
 
+            The active configuration follows the reference fixed-length mean.
+            A separate theoretical variant would mask padded positions and use
+            ``1 / sqrt(|t|)`` over valid bytes; it is intentionally not active
+            in this training configuration.
+
             Adds ``256 × hidden_size`` parameters (≈0.13M for d=512).
             Zero inference overhead when ``bake_inference_table()`` is called
             after training.
@@ -569,6 +574,40 @@ class NeoLLMConfig(PretrainedConfig):
         tweo_eps (:obj:`float`, *optional*, defaults to ``1e-6``):
             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 SwiGLU residual updater and does not reuse the
+            existing ``NextLatDynamicsModel``.  Given final state
+            :math:`h_t^L` and teacher-forced next-token embedding
+            :math:`E(x_{t+1})`, it predicts :math:`\hat h_{t+1}^L` and may be
+            composed for ``nitp_temporal_horizon`` local rollout steps.  It
+            never materializes an additional vocabulary projection.
+        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.
+        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.
+        nitp_temporal_intermediate_size (:obj:`int`, *optional*):
+            SwiGLU width of the transition module.  ``None`` resolves to
+            ``2 * hidden_size``.  At ``hidden_size=512`` this branch adds
+            ``2,622,464`` 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
+            listwise temporal-identification softmax.
+        nitp_temporal_target_temperature (:obj:`float`, *optional*, defaults to ``0.5``):
+            Temperature controlling temporal label smoothing.  Smaller values
+            make the correct diagonal step more nearly one-hot, while retaining
+            limited tolerance for genuinely similar adjacent states.
+
     Constraints:
         - ``use_jtokm=True`` requires ``use_token_generator=True``.
         - ``1 ≤ jtokm_top_k < jtokm_num_experts`` when ``use_jtokm=True``.
@@ -585,6 +624,7 @@ class NeoLLMConfig(PretrainedConfig):
           ``use_laurel_lr`` must be ``True``.
         - ``stack_d_model`` must be divisible by ``num_mem_heads`` when
           ``use_stack_memory=True``.
+        - ``use_nitp_temporal=True`` requires ``use_nitp=True``.
 
     Examples::
 
@@ -783,6 +823,22 @@ class NeoLLMConfig(PretrainedConfig):
         nitp_target_layer=None,
         nitp_loss_weight=1.0,
         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.
+        # `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.
+        use_nitp_temporal=True,
+        nitp_temporal_apply_loss=False,
+        nitp_temporal_horizon=4,
+        nitp_temporal_dynamics_weight=1.0,
+        nitp_temporal_identification_weight=0.1,
+        nitp_temporal_intermediate_size=None,
+        nitp_temporal_similarity_temperature=0.1,
+        nitp_temporal_target_temperature=0.5,
         # ── NextLat: Next-Latent Prediction (Teoh et al., 2026) ───────────
         use_nextlat=False,
         nextlat_horizon=1,
@@ -1039,6 +1095,52 @@ class NeoLLMConfig(PretrainedConfig):
                 f"got {nitp_projector_intermediate_size}."
             )
 
+        # ── NITP temporal transition / identification ────────────────────
+        # The extension is intentionally anchored to the active NITP branch:
+        # it reuses NITP's final hidden states, shallow target states and
+        # projector for diagnostics, but owns a separate SwiGLU transition
+        # module so enabling it genuinely increases the parameter count.
+        if use_nitp_temporal and not use_nitp:
+            raise ValueError(
+                "`use_nitp_temporal=True` requires `use_nitp=True` because "
+                "the temporal branch is defined as an extension of NITP."
+            )
+        if int(nitp_temporal_horizon) < 1:
+            raise ValueError(
+                "`nitp_temporal_horizon` must be >= 1, "
+                f"got {nitp_temporal_horizon}."
+            )
+        if nitp_temporal_dynamics_weight < 0.0:
+            raise ValueError(
+                "`nitp_temporal_dynamics_weight` must be >= 0, "
+                f"got {nitp_temporal_dynamics_weight}."
+            )
+        if nitp_temporal_identification_weight < 0.0:
+            raise ValueError(
+                "`nitp_temporal_identification_weight` must be >= 0, "
+                f"got {nitp_temporal_identification_weight}."
+            )
+        if nitp_temporal_intermediate_size is None:
+            # For d=512 this resolves to 1024. With the NITP-style SwiGLU
+            # transition below, the branch adds 10d^2 + 2d = 2,622,464
+            # trainable parameters (including the affine RMSNorm scale).
+            nitp_temporal_intermediate_size = 2 * hidden_size
+        if int(nitp_temporal_intermediate_size) <= 0:
+            raise ValueError(
+                "`nitp_temporal_intermediate_size` must be positive, "
+                f"got {nitp_temporal_intermediate_size}."
+            )
+        if nitp_temporal_similarity_temperature <= 0.0:
+            raise ValueError(
+                "`nitp_temporal_similarity_temperature` must be > 0, "
+                f"got {nitp_temporal_similarity_temperature}."
+            )
+        if nitp_temporal_target_temperature <= 0.0:
+            raise ValueError(
+                "`nitp_temporal_target_temperature` must be > 0, "
+                f"got {nitp_temporal_target_temperature}."
+            )
+
         # ── NextLat: resolve paper-faithful defaults and validate ─────────
         # Teoh et al. train a latent dynamics model p_psi(h_t, x_{t+1}) to
         # recursively predict final-layer hidden states with Smooth L1 and KL
@@ -1236,6 +1338,26 @@ class NeoLLMConfig(PretrainedConfig):
         self.nitp_loss_weight = float(nitp_loss_weight)
         self.nitp_projector_intermediate_size = int(nitp_projector_intermediate_size)
 
+        # ── NITP temporal transition / identification ────────────────────
+        self.use_nitp_temporal = bool(use_nitp_temporal)
+        self.nitp_temporal_apply_loss = bool(nitp_temporal_apply_loss)
+        self.nitp_temporal_horizon = int(nitp_temporal_horizon)
+        self.nitp_temporal_dynamics_weight = float(
+            nitp_temporal_dynamics_weight
+        )
+        self.nitp_temporal_identification_weight = float(
+            nitp_temporal_identification_weight
+        )
+        self.nitp_temporal_intermediate_size = int(
+            nitp_temporal_intermediate_size
+        )
+        self.nitp_temporal_similarity_temperature = float(
+            nitp_temporal_similarity_temperature
+        )
+        self.nitp_temporal_target_temperature = float(
+            nitp_temporal_target_temperature
+        )
+
         # ── NextLat: Next-Latent Prediction ──────────────────────────────
         self.use_nextlat = bool(use_nextlat)
         self.nextlat_horizon = int(nextlat_horizon)