Update core/trainer.py

core/trainer.py

@@ -17,12 +17,12 @@ class GraphMambaTrainer:
         self.config = config
         self.device = device
         
-        # Conservative learning rate
+        # Balanced learning rate
         self.lr = config['training']['learning_rate']
         self.epochs = config['training']['epochs']
-        self.patience = config['training'].get('patience', 8)
+        self.patience = config['training'].get('patience', 12)
         self.min_lr = config['training'].get('min_lr', 1e-6)
-        self.max_gap = config['training'].get('max_gap', 0.25)  # New gap threshold
+        self.max_gap = config['training'].get('max_gap', 0.35)  # More lenient threshold
         
         # Heavily regularized optimizer
         self.optimizer = optim.AdamW(
@@ -35,15 +35,15 @@ class GraphMambaTrainer:
         
         # Proper loss function with label smoothing
         self.criterion = nn.CrossEntropyLoss(
-            label_smoothing=config['training'].get('label_smoothing', 0.15)
+            label_smoothing=config['training'].get('label_smoothing', 0.1)
         )
         
-        # Conservative scheduler
+        # Balanced scheduler
         self.scheduler = ReduceLROnPlateau(
             self.optimizer,
             mode='max',
             factor=0.5,
-            patience=4,  # Reduced from 5
+            patience=5,
             min_lr=self.min_lr
         )
         
@@ -58,7 +58,7 @@ class GraphMambaTrainer:
         
         # Track overfitting
         self.best_gap = float('inf')
-        self.overfitting_threshold = 0.2  # Reduced from 0.3
+        self.overfitting_threshold = 0.25  # Balanced threshold
     
     def train_node_classification(self, data, verbose=True):
         """Anti-overfitting training with gap monitoring"""