Phase 7: Curriculum Learning (20K steps, BPC 1.78)

src/models/agiformer.py

@@ -1,103 +1,63 @@
 ## Developer: inkbytefo
-## Modified: 2025-11-22
+## Modified: 2025-11-23
 
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from typing import Optional
 from .encoder import ByteLatentEncoder
 from .layers import HybridBlock
 from .reasoning import RecurrentReasoningBlock
 
 class LocalAutoregressiveHead(nn.Module):
-    """
-    Latent vector -> Bytes (Autoregressive).
-    Global Model -> Latent -> Local Model -> Bytes
-    """
     def __init__(self, d_model, patch_size, hidden_dim=256):
         super().__init__()
         self.patch_size = patch_size
-        
-        # Project latent to be the initial state or context
         self.proj_latent = nn.Linear(d_model, hidden_dim)
-        
-        # Byte embedding for the local decoder
         self.byte_emb = nn.Embedding(256, hidden_dim)
-        
-        # Small, fast RNN (GRU) for local decoding
-        # Input size is now hidden_dim (embedding) + hidden_dim (latent context)
         self.rnn = nn.GRU(hidden_dim * 2, hidden_dim, batch_first=True)
-        
         self.head = nn.Linear(hidden_dim, 256)
 
     def forward(self, latents, target_bytes=None, temperature=0.0):
         B, N, D = latents.shape
-        # (B * N, 1, Hidden)
         latent_context = self.proj_latent(latents).view(B * N, 1, -1)
         
         if target_bytes is not None:
-            # --- TRAINING MODE ---
-            # Reshape targets to (B, N, Patch_Size)
             targets = target_bytes.view(B, N, self.patch_size)
-            
-            # Flatten: (B*N, Patch_Size)
             flat_targets = targets.contiguous().view(B * N, self.patch_size)
-            
-            # Shift targets right to get inputs
             sos = torch.zeros(B * N, 1, dtype=torch.long, device=latents.device)
-            rnn_inputs_bytes = torch.cat([sos, flat_targets[:, :-1]], dim=1) # (B*N, P)
-            
-            emb = self.byte_emb(rnn_inputs_bytes) # (B*N, P, Hidden)
-            
-            # Concatenate latent context to every step
+            rnn_inputs_bytes = torch.cat([sos, flat_targets[:, :-1]], dim=1)
+            emb = self.byte_emb(rnn_inputs_bytes)
             latent_expanded = latent_context.expand(-1, self.patch_size, -1)
-            
-            # Concatenation instead of addition to preserve signal
-            rnn_input = torch.cat([emb, latent_expanded], dim=-1) # (B*N, P, Hidden * 2)
-            
+            rnn_input = torch.cat([emb, latent_expanded], dim=-1)
             out, _ = self.rnn(rnn_input)
-            logits = self.head(out) # (B*N, P, 256)
-            
+            logits = self.head(out)
             return logits.view(B, N, self.patch_size, 256)
-            
         else:
-            # INFERENCE MODE
-            pred_bytes = []
-            # Start with SOS (0)
-            current_input = torch.zeros(B * N, 1, dtype=torch.long, device=latents.device)
-            
-            # Initialize hidden state
-            hidden = None 
+            # Inference logic (omitted for brevity, same as before)
+            # ...
+            return self._inference(latents, latent_context, temperature)
 
-            for i in range(self.patch_size):
-                emb = self.byte_emb(current_input) # (B*N, 1, H)
-                
-                # Concatenate latent
-                rnn_in = torch.cat([emb, latent_context], dim=-1) # (B*N, 1, H*2)
-                
-                # GRU State Preservation
-                out, hidden = self.rnn(rnn_in, hidden)
-                logit = self.head(out) # (B*N, 1, 256)
-                
-                # SAMPLING LOGIC
-                if temperature > 0:
-                    # Apply temperature
-                    probs = F.softmax(logit / temperature, dim=-1)
-                    # Sample from distribution
-                    next_byte = torch.multinomial(probs.squeeze(1), 1)
-                else:
-                    # Greedy
-                    next_byte = torch.argmax(logit, dim=-1)
-                
-                pred_bytes.append(next_byte)
-                current_input = next_byte
-            
-            return torch.cat(pred_bytes, dim=1).view(B, N, self.patch_size)
+    def _inference(self, latents, latent_context, temperature):
+        # Helper for inference to keep code clean
+        B, N, _ = latents.shape
+        pred_bytes = []
+        current_input = torch.zeros(B * N, 1, dtype=torch.long, device=latents.device)
+        hidden = None 
+        for i in range(self.patch_size):
+            emb = self.byte_emb(current_input)
+            rnn_in = torch.cat([emb, latent_context], dim=-1)
+            out, hidden = self.rnn(rnn_in, hidden)
+            logit = self.head(out)
+            if temperature > 0:
+                probs = torch.nn.functional.softmax(logit / temperature, dim=-1)
+                next_byte = torch.multinomial(probs.squeeze(1), 1)
+            else:
+                next_byte = torch.argmax(logit, dim=-1)
+            pred_bytes.append(next_byte)
+            current_input = next_byte
+        return torch.cat(pred_bytes, dim=1).view(B, N, self.patch_size)
 
 class AGIFORMER(nn.Module):
-    """
-    AGIFORMER Phase 3: System 2 Enabled
-    """
     def __init__(
         self,
         d_model: int = 512,
@@ -111,54 +71,23 @@ class AGIFORMER(nn.Module):
     ):
         super().__init__()
         
-        self.encoder = ByteLatentEncoder(
-            d_model=d_model,
-            patch_size=patch_size,
-            dropout=dropout
-        )
+        self.encoder = ByteLatentEncoder(d_model, patch_size, dropout)
         
+        # Hybrid Blocks now use Hebbian Memory
         self.layers = nn.ModuleList([
-            HybridBlock(
-                d_model=d_model,
-                num_heads=num_heads,
-                window_size=window_size,
-                dropout=dropout
-            )
+            HybridBlock(d_model, num_heads, window_size, dropout)
             for _ in range(n_layers)
         ])
         
         self.norm_f = nn.LayerNorm(d_model)
-        
-        # SYSTEM 2 MODULE
         self.reasoning = RecurrentReasoningBlock(d_model, thinking_steps, dropout)
-        
-        # Local Autoregressive Head
         self.head = LocalAutoregressiveHead(d_model, patch_size)
 
-    def forward(self, x: torch.Tensor, target_bytes: Optional[torch.Tensor] = None, temperature: float = 0.0) -> torch.Tensor:
-        """
-        Args:
-            x: (Batch, Seq_Len) uint8 - Input Context
-            target_bytes: (Batch, Seq_Len_Target) - Required for training the local head
-            temperature: float - Sampling temperature (0.0 = Greedy)
-            
-        Returns:
-            logits: (Batch, Num_Patches, Patch_Size, 256)
-        """
-        # 1. System 1 (Intuition / Perception)
-        x = self.encoder(x) # (B, N_Patches, D)
-        
-        # 2. Backbone
+    def forward(self, x, target_bytes=None, temperature=0.0):
+        x = self.encoder(x)
         for layer in self.layers:
             x = layer(x)
-            
         x = self.norm_f(x)
-        
-        # 3. System 2 (Reasoning / Thinking Loop)
-        # Refine the latent state before speaking
         x = self.reasoning(x)
-        
-        # 4. Output (Articulation)
         logits = self.head(x, target_bytes, temperature=temperature)
-        
         return logits