Upload fractus/continuous_engine.py with huggingface_hub

fractus/continuous_engine.py

@@ -120,11 +120,11 @@ class CTEBlock(nn.Module):
         attn_out = attn_out @ attn.w_out + attn.b_out
         h = h + attn_out
 
-        # Kuramoto: advance phases by one Euler step.
+        # Kuramoto: ALIGNED with tick_chunk_core (RK4), not single Euler step.
+        # Train/gen mismatch was: gen used Euler 0.1, train used full RK4 integrate.
         h_kur = self.norm_kur(h)
         theta = self.kuramoto._encode_from_hidden(h_kur)
-        theta = theta + 0.1 * self.kuramoto._derivative(theta)
-        theta = torch.remainder(theta, self.kuramoto.TWO_PI)
+        theta = self.kuramoto._rk4_integrate(theta)
         self.kuramoto_phases = theta.detach()
 
         # MoE: transform the thought, routed by Kuramoto phases.
@@ -475,7 +475,11 @@ class ContinuousThoughtEngine(nn.Module):
         return output_logits
 
     def tick_chunk_train(self, observations: torch.Tensor) -> torch.Tensor:
-        """Fast training: head on LAST position only. Returns logits (B, vocab)."""
+        """Train over the full chunk: logits (B, C, vocab) for dense next-token CE.
+
+        Stage-2 surgery: supervise every position so the model must chain tokens,
+        not only collapse to a single last-position attractor.
+        """
         B, C = observations.shape
 
         obs_vecs = self.observe(observations)
@@ -485,12 +489,12 @@ class ContinuousThoughtEngine(nn.Module):
         total_lb = torch.tensor(0.0, device=h.device)
         for blk in self.blocks:
             h, lb = blk.tick_chunk_core(h)
-            total_lb = total_lb + lb  # keep graph — load-balance must train
+            total_lb = total_lb + lb
         self.last_lb_loss = total_lb
 
         self.thought_state = h[:, -1:, :].detach()
-        last_logits = self.output_head(h[:, -1, :])
-        return last_logits, total_lb
+        logits = self.output_head(h)  # (B, C, vocab)
+        return logits, total_lb
 
     def think(self, observations: torch.Tensor, max_ticks: int = 10,
               confidence_threshold: float = 0.7) -> torch.Tensor: