Update configuration_neollm.py

configuration_neollm.py

@@ -576,12 +576,13 @@ class NeoLLMConfig(PretrainedConfig):
 
         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.
+            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.
         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
@@ -597,9 +598,10 @@ class NeoLLMConfig(PretrainedConfig):
             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.
+            Intermediate width of the context proposal and token modulation
+            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
             listwise temporal-identification softmax.
@@ -1121,9 +1123,9 @@ class NeoLLMConfig(PretrainedConfig):
                 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).
+            # For d=512 this resolves to 1024.  The conditioned
+            # multiplicative transition adds 4dm + 2d = 2,098,176 trainable
+            # parameters, including the two affine RMSNorm scales.
             nitp_temporal_intermediate_size = 2 * hidden_size
         if int(nitp_temporal_intermediate_size) <= 0:
             raise ValueError(