Update core/trainer.py

core/trainer.py

@@ -20,8 +20,9 @@ class GraphMambaTrainer:
         # Conservative learning rate
         self.lr = config['training']['learning_rate']
         self.epochs = config['training']['epochs']
-        self.patience = config['training'].get('patience', 10)
+        self.patience = config['training'].get('patience', 8)
         self.min_lr = config['training'].get('min_lr', 1e-6)
+        self.max_gap = config['training'].get('max_gap', 0.25)  # New gap threshold
         
         # Heavily regularized optimizer
         self.optimizer = optim.AdamW(
@@ -33,14 +34,16 @@ class GraphMambaTrainer:
         )
         
         # Proper loss function with label smoothing
-        self.criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
+        self.criterion = nn.CrossEntropyLoss(
+            label_smoothing=config['training'].get('label_smoothing', 0.15)
+        )
         
-        # Conservative scheduler - FIXED: removed verbose parameter
+        # Conservative scheduler
         self.scheduler = ReduceLROnPlateau(
             self.optimizer,
             mode='max',
             factor=0.5,
-            patience=5,
+            patience=4,  # Reduced from 5
             min_lr=self.min_lr
         )
         
@@ -55,10 +58,10 @@ class GraphMambaTrainer:
         
         # Track overfitting
         self.best_gap = float('inf')
-        self.overfitting_threshold = 0.3
+        self.overfitting_threshold = 0.2  # Reduced from 0.3
     
     def train_node_classification(self, data, verbose=True):
-        """Anti-overfitting training"""
+        """Anti-overfitting training with gap monitoring"""
         
         if verbose:
             total_params = sum(p.numel() for p in self.model.parameters())
@@ -72,8 +75,9 @@ class GraphMambaTrainer:
             print(f"⚙️ Parameters: {total_params:,}")
             print(f"📚 Training samples: {train_samples}")
             print(f"⚠️ Params per sample: {params_per_sample:.1f}")
+            print(f"🚨 Max allowed gap: {self.max_gap:.3f}")
             
-            if params_per_sample > 1000:
+            if params_per_sample > 500:
                 print(f"🚨 WARNING: High params per sample ratio - overfitting risk!")
         
         # Initialize classifier
@@ -114,16 +118,16 @@ class GraphMambaTrainer:
             else:
                 self.patience_counter += 1
             
-            # Overfitting detection
+            # Aggressive overfitting detection
             if acc_gap > self.overfitting_threshold:
                 if verbose:
                     print(f"🚨 OVERFITTING detected: {acc_gap:.3f} gap")
                     print(f"   Train: {train_metrics['acc']:.3f}, Val: {val_metrics['acc']:.3f}")
             
             # Progress logging
-            if verbose and (epoch == 0 or (epoch + 1) % 10 == 0 or epoch == self.epochs - 1):
+            if verbose and (epoch == 0 or (epoch + 1) % 5 == 0 or epoch == self.epochs - 1):
                 elapsed = time.time() - start_time
-                gap_indicator = "🚨" if acc_gap > 0.2 else "⚠️" if acc_gap > 0.1 else "✅"
+                gap_indicator = "🚨" if acc_gap > 0.25 else "⚠️" if acc_gap > 0.15 else "✅"
                 
                 print(f"Epoch {epoch:3d} | "
                       f"Train: {train_metrics['loss']:.4f} ({train_metrics['acc']:.4f}) | "
@@ -131,16 +135,22 @@ class GraphMambaTrainer:
                       f"Gap: {acc_gap:.3f} {gap_indicator} | "
                       f"LR: {self.optimizer.param_groups[0]['lr']:.6f}")
             
-            # Early stopping conditions
+            # Enhanced early stopping conditions
             if self.patience_counter >= self.patience:
                 if verbose:
                     print(f"🛑 Early stopping at epoch {epoch} (patience)")
                 break
                 
-            # Stop if severe overfitting
-            if acc_gap > 0.5:
+            # Stop if gap exceeds threshold
+            if acc_gap > self.max_gap:
+                if verbose:
+                    print(f"🛑 Stopping due to overfitting gap: {acc_gap:.3f} > {self.max_gap:.3f}")
+                break
+                
+            # Stop if severe overfitting (backup check)
+            if acc_gap > 0.6:
                 if verbose:
-                    print(f"🛑 Stopping due to severe overfitting (gap: {acc_gap:.3f})")
+                    print(f"🛑 Emergency stop - severe overfitting (gap: {acc_gap:.3f})")
                 break
         
         if verbose:
@@ -170,15 +180,15 @@ class GraphMambaTrainer:
         # Compute loss on training nodes only
         train_loss = self.criterion(logits[data.train_mask], data.y[data.train_mask])
         
-        # Add L2 regularization manually
+        # Add stronger L2 regularization
         l2_reg = 0.0
         for param in self.model.parameters():
             l2_reg += torch.norm(param, p=2)
-        train_loss += 1e-5 * l2_reg
+        train_loss += 5e-5 * l2_reg  # Increased from 1e-5
         
         # Backward pass with gradient clipping
         train_loss.backward()
-        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+        torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=0.5)  # Reduced from 1.0
         self.optimizer.step()
         
         # Compute accuracy