SURE-tools 2.4.13__py3-none-any.whl → 2.4.22__py3-none-any.whl

SURE/EfficientTranscriptomeDecoder.py

@@ -12,7 +12,7 @@ warnings.filterwarnings('ignore')
 class EfficientTranscriptomeDecoder:
     """
     High-performance, memory-efficient transcriptome decoder
-    Combines latest research techniques for optimal performance
+    Fixed version with corrected RMSNorm implementation
     """
     
     def __init__(self, 
@@ -20,7 +20,7 @@ class EfficientTranscriptomeDecoder:
                  gene_dim: int = 60000,
                  hidden_dims: List[int] = [512, 1024, 2048],
                  bottleneck_dim: int = 256,
-                 num_experts: int = 8,
+                 num_experts: int = 4,
                  dropout_rate: float = 0.1,
                  device: str = None):
         """
@@ -43,8 +43,8 @@ class EfficientTranscriptomeDecoder:
         self.dropout_rate = dropout_rate
         self.device = device or ('cuda' if torch.cuda.is_available() else 'cpu')
         
-        # Initialize model with advanced architecture
-        self.model = self._build_advanced_model()
+        # Initialize model with corrected architecture
+        self.model = self._build_corrected_model()
         self.model.to(self.device)
         
         # Training state
@@ -58,74 +58,50 @@ class EfficientTranscriptomeDecoder:
         print(f"   - Hidden Dimensions: {hidden_dims}")
         print(f"   - Bottleneck Dimension: {bottleneck_dim}")
         print(f"   - Number of Experts: {num_experts}")
-        print(f"   - Estimated GPU Memory: ~6-8GB")
         print(f"   - Parameters: {sum(p.numel() for p in self.model.parameters()):,}")
     
-    class SwiGLU(nn.Module):
-        """SwiGLU activation - better than GELU (PaLM, LLaMA)"""
-        def forward(self, x):
-            x, gate = x.chunk(2, dim=-1)
-            return x * F.silu(gate)
-    
-    class RMSNorm(nn.Module):
-        """RMS Normalization - more stable than LayerNorm (GPT-3)"""
+    class CorrectedRMSNorm(nn.Module):
+        """Corrected RMS Normalization with proper dimension handling"""
         def __init__(self, dim: int, eps: float = 1e-8):
             super().__init__()
             self.eps = eps
-            self.weight = nn.Parameter(torch.ones(dim))
+            self.dim = dim
+            self.weight = nn.Parameter(torch.ones(dim))  # Correct: weight has same dim as input
         
         def forward(self, x):
-            norm_x = x.norm(2, dim=-1, keepdim=True)
-            rms_x = norm_x * (x.shape[-1] ** -0.5)
-            return x / (rms_x + self.eps) * self.weight
+            # Ensure input has the right dimension
+            if x.size(-1) != self.dim:
+                raise ValueError(f"Input dimension {x.size(-1)} doesn't match RMSNorm dimension {self.dim}")
+            
+            # Calculate RMS
+            rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + self.eps)
+            # Normalize and apply weight
+            return x / rms * self.weight
     
-    class MixtureOfExperts(nn.Module):
-        """Mixture of Experts for conditional computation"""
-        def __init__(self, input_dim: int, expert_dim: int, num_experts: int):
-            super().__init__()
-            self.num_experts = num_experts
-            self.experts = nn.ModuleList([
-                nn.Sequential(
-                    nn.Linear(input_dim, expert_dim),
-                    nn.Dropout(0.1),
-                    nn.Linear(expert_dim, input_dim)
-                ) for _ in range(num_experts)
-            ])
-            self.gate = nn.Linear(input_dim, num_experts)
-            self.expert_dim = expert_dim
-        
+    class SimplifiedSwiGLU(nn.Module):
+        """Simplified SwiGLU activation"""
         def forward(self, x):
-            # Gate network
-            gate_logits = self.gate(x)
-            gate_weights = F.softmax(gate_logits, dim=-1)
-            
-            # Expert outputs
-            expert_outputs = []
-            for i, expert in enumerate(self.experts):
-                expert_out = expert(x)
-                expert_outputs.append(expert_out.unsqueeze(-1))
-            
-            # Combine expert outputs
-            expert_outputs = torch.cat(expert_outputs, dim=-1)
-            output = torch.einsum('bd, bde -> be', gate_weights, expert_outputs)
-            
-            return output + x  # Residual connection
+            # Split into two parts
+            x, gate = x.chunk(2, dim=-1)
+            return x * F.silu(gate)
     
-    class AdaptiveBottleneck(nn.Module):
-        """Adaptive bottleneck for memory efficiency"""
+    class MemoryEfficientBottleneck(nn.Module):
+        """Memory-efficient bottleneck with corrected dimensions"""
         def __init__(self, input_dim: int, bottleneck_dim: int, output_dim: int):
             super().__init__()
+            # Ensure proper dimension matching
             self.compress = nn.Linear(input_dim, bottleneck_dim)
-            self.norm1 = EfficientTranscriptomeDecoder.RMSNorm(bottleneck_dim)
+            self.norm1 = EfficientTranscriptomeDecoder.CorrectedRMSNorm(bottleneck_dim)
             self.expand = nn.Linear(bottleneck_dim, output_dim)
-            self.norm2 = EfficientTranscriptomeDecoder.RMSNorm(output_dim)
+            self.norm2 = EfficientTranscriptomeDecoder.CorrectedRMSNorm(output_dim)
+            self.activation = nn.SiLU()
             self.dropout = nn.Dropout(0.1)
             
         def forward(self, x):
             # Compress
             compressed = self.compress(x)
             compressed = self.norm1(compressed)
-            compressed = F.silu(compressed)
+            compressed = self.activation(compressed)
             compressed = self.dropout(compressed)
             
             # Expand
@@ -134,37 +110,57 @@ class EfficientTranscriptomeDecoder:
             
             return expanded
     
-    class GeneSpecificProjection(nn.Module):
-        """Gene-specific projection with weight sharing"""
-        def __init__(self, latent_dim: int, gene_dim: int, proj_dim: int = 64):
+    class StableMixtureOfExperts(nn.Module):
+        """Stable mixture of experts without dimension issues"""
+        def __init__(self, input_dim: int, num_experts: int = 4):
             super().__init__()
-            self.proj_dim = proj_dim
-            self.gene_embeddings = nn.Parameter(torch.randn(gene_dim, proj_dim) * 0.02)
-            self.latent_projection = nn.Linear(latent_dim, proj_dim)
-            self.output_layer = nn.Linear(proj_dim, 1)
-            
-        def forward(self, latent):
-            batch_size = latent.shape[0]
+            self.num_experts = num_experts
+            self.input_dim = input_dim
             
-            # Project latent to gene space
-            latent_proj = self.latent_projection(latent)  # [batch_size, proj_dim]
+            # Shared expert with different scaling factors
+            self.shared_expert = nn.Sequential(
+                nn.Linear(input_dim, input_dim * 2),
+                nn.SiLU(),
+                nn.Dropout(0.1),
+                nn.Linear(input_dim * 2, input_dim)
+            )
             
-            # Efficient matrix multiplication
-            gene_output = torch.matmul(latent_proj, self.gene_embeddings.T)  # [batch_size, gene_dim]
+            # Gating network
+            self.gate = nn.Sequential(
+                nn.Linear(input_dim, num_experts * 4),
+                nn.SiLU(),
+                nn.Linear(num_experts * 4, num_experts)
+            )
             
-            return gene_output
+        def forward(self, x):
+            # Get gate weights
+            gate_weights = F.softmax(self.gate(x), dim=-1)  # [batch, num_experts]
+            
+            # Process through shared expert
+            expert_output = self.shared_expert(x)  # [batch, input_dim]
+            
+            # Apply expert-specific scaling
+            weighted_output = torch.zeros_like(expert_output)
+            for i in range(self.num_experts):
+                expert_scale = 0.5 + 0.5 * i  # Different scaling for each expert
+                expert_contribution = expert_output * expert_scale
+                expert_weight = gate_weights[:, i].unsqueeze(-1)  # [batch, 1]
+                weighted_output += expert_weight * expert_contribution
+            
+            # Residual connection
+            return x + weighted_output
     
-    class AdvancedDecoder(nn.Module):
-        """Advanced decoder combining multiple techniques"""
+    class CorrectedDecoder(nn.Module):
+        """Corrected decoder with proper dimension handling"""
         
         def __init__(self, latent_dim: int, gene_dim: int, hidden_dims: List[int], 
-                    bottleneck_dim: int, num_experts: int, dropout_rate: float):
+                     bottleneck_dim: int, num_experts: int, dropout_rate: float):
             super().__init__()
             
-            # Initial projection
+            # Input projection
             self.input_projection = nn.Sequential(
                 nn.Linear(latent_dim, hidden_dims[0]),
-                EfficientTranscriptomeDecoder.RMSNorm(hidden_dims[0]),
+                EfficientTranscriptomeDecoder.CorrectedRMSNorm(hidden_dims[0]),
                 nn.SiLU(),
                 nn.Dropout(dropout_rate)
             )
@@ -173,71 +169,74 @@ class EfficientTranscriptomeDecoder:
             self.blocks = nn.ModuleList()
             current_dim = hidden_dims[0]
             
-            for i, hidden_dim in enumerate(hidden_dims[1:], 1):
-                block = nn.ModuleList([
-                    # Mixture of Experts
-                    EfficientTranscriptomeDecoder.MixtureOfExperts(current_dim, hidden_dim, num_experts),
-                    
-                    # Adaptive Bottleneck
-                    EfficientTranscriptomeDecoder.AdaptiveBottleneck(current_dim, bottleneck_dim, hidden_dim),
-                    
-                    # SwiGLU activation
-                    nn.Sequential(
-                        nn.Linear(hidden_dim, hidden_dim * 2),
-                        EfficientTranscriptomeDecoder.SwiGLU(),
-                        nn.Dropout(dropout_rate)
+            for i, next_dim in enumerate(hidden_dims[1:], 1):
+                block = nn.ModuleDict({
+                    'swiglu': nn.Sequential(
+                        nn.Linear(current_dim, current_dim * 2),
+                        EfficientTranscriptomeDecoder.SimplifiedSwiGLU(),
+                        nn.Dropout(dropout_rate),
+                        nn.Linear(current_dim, current_dim)  # Project back to same dimension
+                    ),
+                    'bottleneck': EfficientTranscriptomeDecoder.MemoryEfficientBottleneck(
+                        current_dim, bottleneck_dim, next_dim
+                    ),
+                    'experts': EfficientTranscriptomeDecoder.StableMixtureOfExperts(
+                        next_dim, num_experts
                     )
-                ])
+                })
                 self.blocks.append(block)
-                current_dim = hidden_dim
+                current_dim = next_dim
             
-            # Gene-specific projection
-            self.gene_projection = EfficientTranscriptomeDecoder.GeneSpecificProjection(
-                current_dim, gene_dim, proj_dim=128
+            # Final projection to gene dimension
+            self.final_projection = nn.Sequential(
+                nn.Linear(current_dim, current_dim * 2),
+                nn.SiLU(),
+                nn.Dropout(dropout_rate),
+                nn.Linear(current_dim * 2, gene_dim)
             )
             
-            # Output scaling
+            # Output parameters
             self.output_scale = nn.Parameter(torch.ones(1))
             self.output_bias = nn.Parameter(torch.zeros(1))
             
             self._init_weights()
         
         def _init_weights(self):
-            """Advanced weight initialization"""
+            """Proper weight initialization"""
             for module in self.modules():
                 if isinstance(module, nn.Linear):
-                    # Kaiming init for SiLU/SwiGLU
-                    nn.init.kaiming_normal_(module.weight, nonlinearity='linear')
+                    nn.init.xavier_uniform_(module.weight)
                     if module.bias is not None:
                         nn.init.zeros_(module.bias)
         
         def forward(self, x):
-            # Initial projection
+            # Input projection
             x = self.input_projection(x)
             
             # Process through blocks
             for block in self.blocks:
-                # Mixture of Experts
-                expert_out = block[0](x)
+                # SwiGLU with residual
+                residual = x
+                x_swiglu = block['swiglu'](x)
+                x = x + x_swiglu  # Residual connection
                 
-                # Adaptive Bottleneck
-                bottleneck_out = block[1](expert_out)
+                # Bottleneck
+                x = block['bottleneck'](x)
                 
-                # SwiGLU activation with residual
-                swiglu_out = block[2](bottleneck_out)
-                x = x + swiglu_out  # Residual connection
+                # Mixture of Experts with residual
+                x = block['experts'](x)
             
-            # Final gene projection
-            output = self.gene_projection(x)
+            # Final projection
+            x = self.final_projection(x)
             
             # Ensure non-negative output
-            output = F.softplus(output * self.output_scale + self.output_bias)
+            x = F.softplus(x * self.output_scale + self.output_bias)
             
-            return output
+            return x
     
-    def _build_advanced_model(self):
-        """Build the advanced decoder model"""
-        return self.AdvancedDecoder(
+    def _build_corrected_model(self):
+        """Build the corrected model"""
+        return self.CorrectedDecoder(
             self.latent_dim, self.gene_dim, self.hidden_dims,
             self.bottleneck_dim, self.num_experts, self.dropout_rate
         )
@@ -247,24 +246,14 @@ class EfficientTranscriptomeDecoder:
               train_expression: np.ndarray,
               val_latent: np.ndarray = None,
               val_expression: np.ndarray = None,
-              batch_size: int = 16,  # Smaller batches for memory efficiency
-              num_epochs: int = 200,
+              batch_size: int = 32,
+              num_epochs: int = 100,
               learning_rate: float = 1e-4,
-              checkpoint_path: str = 'efficient_decoder.pth') -> Dict:
+              checkpoint_path: str = 'transcriptome_decoder.pth') -> Dict:
         """
-        Train with advanced optimization techniques
-        
-        Args:
-            train_latent: Training latent variables
-            train_expression: Training expression data
-            val_latent: Validation latent variables
-            val_expression: Validation expression data
-            batch_size: Batch size (optimized for memory)
-            num_epochs: Number of epochs
-            learning_rate: Learning rate
-            checkpoint_path: Model save path
+        Train the corrected decoder
         """
-        print("🚀 Starting Advanced Training...")
+        print("🚀 Starting Training...")
         
         # Data preparation
         train_dataset = self._create_dataset(train_latent, train_expression)
@@ -273,114 +262,83 @@ class EfficientTranscriptomeDecoder:
             val_dataset = self._create_dataset(val_latent, val_expression)
             train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, pin_memory=True)
             val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, pin_memory=True)
+            print(f"📈 Using provided validation data: {len(val_dataset)} samples")
         else:
+            # Auto split
             train_size = int(0.9 * len(train_dataset))
             val_size = len(train_dataset) - train_size
             train_subset, val_subset = torch.utils.data.random_split(train_dataset, [train_size, val_size])
             train_loader = DataLoader(train_subset, batch_size=batch_size, shuffle=True, pin_memory=True)
             val_loader = DataLoader(val_subset, batch_size=batch_size, shuffle=False, pin_memory=True)
+            print(f"📈 Auto-split validation: {val_size} samples")
         
         print(f"📊 Training samples: {len(train_loader.dataset)}")
         print(f"📊 Validation samples: {len(val_loader.dataset)}")
         print(f"📊 Batch size: {batch_size}")
         
-        # Advanced optimizer configuration
+        # Optimizer
         optimizer = optim.AdamW(
             self.model.parameters(),
             lr=learning_rate,
-            weight_decay=0.1,  # Stronger regularization
-            betas=(0.9, 0.95),  # Tuned betas
-            eps=1e-8
+            weight_decay=0.01,
+            betas=(0.9, 0.999)
         )
         
-        # Cosine annealing with warmup
-        scheduler = optim.lr_scheduler.OneCycleLR(
-            optimizer,
-            max_lr=learning_rate * 5,
-            epochs=num_epochs,
-            steps_per_epoch=len(train_loader),
-            pct_start=0.1,
-            div_factor=10.0,
-            final_div_factor=100.0
-        )
+        # Scheduler
+        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)
         
-        # Advanced loss function
-        def advanced_loss(pred, target):
-            # 1. MSE loss for overall accuracy
+        # Loss function
+        def combined_loss(pred, target):
             mse_loss = F.mse_loss(pred, target)
-            
-            # 2. Poisson loss for count data
             poisson_loss = (pred - target * torch.log(pred + 1e-8)).mean()
-            
-            # 3. Correlation loss for pattern matching
-            correlation_loss = 1 - self._pearson_correlation(pred, target)
-            
-            # 4. Sparsity loss for realistic distribution
-            sparsity_loss = F.mse_loss(
-                (pred < 1e-3).float().mean(),
-                torch.tensor(0.85, device=pred.device)  # Target sparsity
-            )
-            
-            # 5. Spectral loss for smoothness
-            spectral_loss = self._spectral_loss(pred, target)
-            
-            # Weighted combination
-            total_loss = (mse_loss + 0.3 * poisson_loss + 0.2 * correlation_loss + 
-                         0.1 * sparsity_loss + 0.05 * spectral_loss)
-            
-            return total_loss, {
-                'mse': mse_loss.item(),
-                'poisson': poisson_loss.item(),
-                'correlation': correlation_loss.item(),
-                'sparsity': sparsity_loss.item(),
-                'spectral': spectral_loss.item()
-            }
+            correlation = self._pearson_correlation(pred, target)
+            correlation_loss = 1 - correlation
+            return mse_loss + 0.3 * poisson_loss + 0.1 * correlation_loss
         
         # Training history
         history = {
             'train_loss': [], 'val_loss': [],
             'train_mse': [], 'val_mse': [],
             'train_correlation': [], 'val_correlation': [],
-            'learning_rates': [], 'grad_norms': []
+            'learning_rates': []
         }
         
         best_val_loss = float('inf')
-        patience = 25
+        patience = 20
         patience_counter = 0
         
-        print("\n📈 Starting training with advanced techniques...")
+        print("\n📈 Starting training loop...")
         for epoch in range(1, num_epochs + 1):
-            # Training phase
-            train_loss, train_components, grad_norm = self._train_epoch_advanced(
-                train_loader, optimizer, scheduler, advanced_loss
-            )
+            # Training
+            train_metrics = self._train_epoch(train_loader, optimizer, combined_loss)
             
-            # Validation phase
-            val_loss, val_components = self._validate_epoch_advanced(val_loader, advanced_loss)
+            # Validation
+            val_metrics = self._validate_epoch(val_loader, combined_loss)
+            
+            # Update scheduler
+            scheduler.step()
+            current_lr = optimizer.param_groups[0]['lr']
             
             # Record history
-            history['train_loss'].append(train_loss)
-            history['val_loss'].append(val_loss)
-            history['train_mse'].append(train_components['mse'])
-            history['val_mse'].append(val_components['mse'])
-            history['train_correlation'].append(train_components['correlation'])
-            history['val_correlation'].append(val_components['correlation'])
-            history['learning_rates'].append(optimizer.param_groups[0]['lr'])
-            history['grad_norms'].append(grad_norm)
-            
-            # Print detailed progress
+            history['train_loss'].append(train_metrics['loss'])
+            history['val_loss'].append(val_metrics['loss'])
+            history['train_mse'].append(train_metrics['mse'])
+            history['val_mse'].append(val_metrics['mse'])
+            history['train_correlation'].append(train_metrics['correlation'])
+            history['val_correlation'].append(val_metrics['correlation'])
+            history['learning_rates'].append(current_lr)
+            
+            # Print progress
             if epoch % 10 == 0 or epoch == 1:
-                lr = optimizer.param_groups[0]['lr']
                 print(f"📍 Epoch {epoch:3d}/{num_epochs} | "
-                      f"Train: {train_loss:.4f} | "
-                      f"Val: {val_loss:.4f} | "
-                      f"Corr: {val_components['correlation']:.4f} | "
-                      f"LR: {lr:.2e} | "
-                      f"Grad: {grad_norm:.4f}")
-            
-            # Early stopping with patience
-            if val_loss < best_val_loss:
-                best_val_loss = val_loss
+                      f"Train Loss: {train_metrics['loss']:.4f} | "
+                      f"Val Loss: {val_metrics['loss']:.4f} | "
+                      f"Correlation: {val_metrics['correlation']:.4f} | "
+                      f"LR: {current_lr:.2e}")
+            
+            # Early stopping
+            if val_metrics['loss'] < best_val_loss:
+                best_val_loss = val_metrics['loss']
                 patience_counter = 0
                 self._save_checkpoint(epoch, optimizer, scheduler, best_val_loss, history, checkpoint_path)
                 if epoch % 20 == 0:
@@ -402,8 +360,8 @@ class EfficientTranscriptomeDecoder:
         return history
     
     def _create_dataset(self, latent_data, expression_data):
-        """Create memory-efficient dataset"""
-        class EfficientDataset(Dataset):
+        """Create dataset"""
+        class SimpleDataset(Dataset):
             def __init__(self, latent, expression):
                 self.latent = torch.FloatTensor(latent)
                 self.expression = torch.FloatTensor(expression)
@@ -414,7 +372,7 @@ class EfficientTranscriptomeDecoder:
             def __getitem__(self, idx):
                 return self.latent[idx], self.expression[idx]
         
-        return EfficientDataset(latent_data, expression_data)
+        return SimpleDataset(latent_data, expression_data)
     
     def _pearson_correlation(self, pred, target):
         """Calculate Pearson correlation"""
@@ -426,68 +384,48 @@ class EfficientTranscriptomeDecoder:
         
         return (numerator / (denominator + 1e-8)).mean()
     
-    def _spectral_loss(self, pred, target):
-        """Spectral loss for frequency domain matching"""
-        pred_fft = torch.fft.fft(pred, dim=1)
-        target_fft = torch.fft.fft(target, dim=1)
-        
-        magnitude_loss = F.mse_loss(torch.abs(pred_fft), torch.abs(target_fft))
-        phase_loss = F.mse_loss(torch.angle(pred_fft), torch.angle(target_fft))
-        
-        return magnitude_loss + 0.5 * phase_loss
-    
-    def _train_epoch_advanced(self, train_loader, optimizer, scheduler, loss_fn):
-        """Advanced training with gradient accumulation"""
+    def _train_epoch(self, train_loader, optimizer, loss_fn):
+        """Train one epoch"""
         self.model.train()
         total_loss = 0
-        total_components = {'mse': 0, 'poisson': 0, 'correlation': 0, 'sparsity': 0, 'spectral': 0}
-        grad_norms = []
-        
-        # Gradient accumulation for effective larger batch size
-        accumulation_steps = 4
-        optimizer.zero_grad()
+        total_mse = 0
+        total_correlation = 0
         
-        for i, (latent, target) in enumerate(train_loader):
+        for latent, target in train_loader:
             latent = latent.to(self.device, non_blocking=True)
             target = target.to(self.device, non_blocking=True)
             
-            # Forward pass with mixed precision
-            with torch.cuda.amp.autocast():  # Mixed precision for memory efficiency
-                pred = self.model(latent)
-                loss, components = loss_fn(pred, target)
+            optimizer.zero_grad()
+            pred = self.model(latent)
             
-            # Scale loss for gradient accumulation
-            loss = loss / accumulation_steps
+            loss = loss_fn(pred, target)
             loss.backward()
             
-            # Gradient accumulation
-            if (i + 1) % accumulation_steps == 0:
-                # Gradient clipping
-                grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
-                optimizer.step()
-                scheduler.step()
-                optimizer.zero_grad()
-                
-                grad_norms.append(grad_norm.item())
+            # Gradient clipping
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            optimizer.step()
+            
+            # Calculate metrics
+            mse_loss = F.mse_loss(pred, target).item()
+            correlation = self._pearson_correlation(pred, target).item()
             
-            # Accumulate losses
-            total_loss += loss.item() * accumulation_steps
-            for key in total_components:
-                total_components[key] += components[key]
+            total_loss += loss.item()
+            total_mse += mse_loss
+            total_correlation += correlation
         
-        # Average metrics
         num_batches = len(train_loader)
-        avg_loss = total_loss / num_batches
-        avg_components = {key: value / num_batches for key, value in total_components.items()}
-        avg_grad_norm = np.mean(grad_norms) if grad_norms else 0.0
-        
-        return avg_loss, avg_components, avg_grad_norm
+        return {
+            'loss': total_loss / num_batches,
+            'mse': total_mse / num_batches,
+            'correlation': total_correlation / num_batches
+        }
     
-    def _validate_epoch_advanced(self, val_loader, loss_fn):
-        """Advanced validation"""
+    def _validate_epoch(self, val_loader, loss_fn):
+        """Validate one epoch"""
         self.model.eval()
         total_loss = 0
-        total_components = {'mse': 0, 'poisson': 0, 'correlation': 0, 'sparsity': 0, 'spectral': 0}
+        total_mse = 0
+        total_correlation = 0
         
         with torch.no_grad():
             for latent, target in val_loader:
@@ -495,17 +433,20 @@ class EfficientTranscriptomeDecoder:
                 target = target.to(self.device, non_blocking=True)
                 
                 pred = self.model(latent)
-                loss, components = loss_fn(pred, target)
+                loss = loss_fn(pred, target)
+                mse_loss = F.mse_loss(pred, target).item()
+                correlation = self._pearson_correlation(pred, target).item()
                 
                 total_loss += loss.item()
-                for key in total_components:
-                    total_components[key] += components[key]
+                total_mse += mse_loss
+                total_correlation += correlation
         
         num_batches = len(val_loader)
-        avg_loss = total_loss / num_batches
-        avg_components = {key: value / num_batches for key, value in total_components.items()}
-        
-        return avg_loss, avg_components
+        return {
+            'loss': total_loss / num_batches,
+            'mse': total_mse / num_batches,
+            'correlation': total_correlation / num_batches
+        }
     
     def _save_checkpoint(self, epoch, optimizer, scheduler, best_loss, history, path):
         """Save checkpoint"""
@@ -525,8 +466,8 @@ class EfficientTranscriptomeDecoder:
             }
         }, path)
     
-    def predict(self, latent_data: np.ndarray, batch_size: int = 16) -> np.ndarray:
-        """Memory-efficient prediction"""
+    def predict(self, latent_data: np.ndarray, batch_size: int = 32) -> np.ndarray:
+        """Predict gene expression"""
         if not self.is_trained:
             warnings.warn("⚠️ Model not trained. Predictions may be inaccurate.")
         
@@ -539,15 +480,8 @@ class EfficientTranscriptomeDecoder:
         with torch.no_grad():
             for i in range(0, len(latent_data), batch_size):
                 batch_latent = latent_data[i:i+batch_size].to(self.device)
-                
-                with torch.cuda.amp.autocast():  # Mixed precision for memory
-                    batch_pred = self.model(batch_latent)
-                
+                batch_pred = self.model(batch_latent)
                 predictions.append(batch_pred.cpu())
-                
-                # Clear memory
-                if torch.cuda.is_available():
-                    torch.cuda.empty_cache()
         
         return torch.cat(predictions).numpy()
     
@@ -559,11 +493,23 @@ class EfficientTranscriptomeDecoder:
         self.training_history = checkpoint.get('training_history')
         self.best_val_loss = checkpoint.get('best_val_loss', float('inf'))
         print(f"✅ Model loaded! Best val loss: {self.best_val_loss:.4f}")
+    
+    def get_model_info(self) -> Dict:
+        """Get model information"""
+        return {
+            'is_trained': self.is_trained,
+            'best_val_loss': self.best_val_loss,
+            'parameters': sum(p.numel() for p in self.model.parameters()),
+            'latent_dim': self.latent_dim,
+            'gene_dim': self.gene_dim,
+            'hidden_dims': self.hidden_dims,
+            'device': str(self.device)
+        }
 
 '''
 # Example usage
 def example_usage():
-    """Demonstrate the advanced decoder"""
+    """Example demonstration"""
     
     # Initialize decoder
     decoder = EfficientTranscriptomeDecoder(
@@ -590,7 +536,7 @@ def example_usage():
     history = decoder.train(
         train_latent=latent_data,
         train_expression=expression_data,
-        batch_size=16,
+        batch_size=32,
         num_epochs=50
     )
     
@@ -603,5 +549,4 @@ def example_usage():
 
 if __name__ == "__main__":
     example_usage()
-    
 '''