Upload architectures.py with huggingface_hub

architectures.py

@@ -0,0 +1,1496 @@
+import logging
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import json
+import time
+import pytorch_lightning as pl
+from torch.optim import AdamW
+from torchmetrics import MeanSquaredError, PearsonCorrCoef, SpearmanCorrCoef, R2Score
+
+logger = logging.getLogger(__name__)
+
+# ===================== VERSION STRING FOR CLUSTER VERIFICATION =====================
+ARCH_VERSION = "2024-12-24-stability-fix"
+print(f"[ARCH] architectures.py loaded: {ARCH_VERSION}")
+# ====================================================================================
+
+class Interp1d(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, y, xnew):
+        is_flat = {}
+        vals = {'x': x, 'y': y, 'xnew': xnew}
+        for name, arr in vals.items():
+            is_flat[name] = (arr.dim() == 1)
+            if is_flat[name]:
+                vals[name] = arr.unsqueeze(0)
+        x_2d, y_2d, xnew_2d = vals['x'], vals['y'], vals['xnew']
+        B, Nx = x_2d.shape
+
+        # SAFETY: Handle edge case where sequence length is < 5
+        if Nx < 5:
+            # Return constant interpolation (repeat/average the values)
+            ynew_2d = y_2d.mean(dim=1, keepdim=True).expand(-1, xnew_2d.shape[1])
+            ctx.save_for_backward(x_2d, y_2d, xnew_2d, 
+                                  torch.zeros_like(xnew_2d, dtype=torch.long),
+                                  torch.zeros_like(xnew_2d))
+            ctx.Nx_was_small = True
+            if is_flat['x'] and is_flat['xnew']:
+                ynew_2d = ynew_2d.squeeze(0)
+            return ynew_2d
+        
+        ctx.Nx_was_small = False
+        idx = torch.searchsorted(x_2d, xnew_2d, right=False) - 1
+        idx = idx.clamp(min=0, max=Nx-2)
+
+        xL = torch.gather(x_2d, 1, idx)
+        xR = torch.gather(x_2d, 1, idx+1)
+        yL = torch.gather(y_2d, 1, idx)
+        yR = torch.gather(y_2d, 1, idx+1)
+
+        denom = (xR - xL)
+        denom[denom == 0] = 1e-12
+        t = (xnew_2d - xL)/denom
+        ynew_2d = yL + (yR - yL)*t
+
+        ctx.save_for_backward(x_2d, y_2d, xnew_2d, idx, t)
+        if is_flat['x'] and is_flat['xnew']:
+            ynew_2d = ynew_2d.squeeze(0)
+        return ynew_2d
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        x_2d, y_2d, xnew_2d, idx, t = ctx.saved_tensors
+        grad_x = grad_y = grad_xnew = None
+        
+        # Handle edge case from forward
+        if getattr(ctx, 'Nx_was_small', False):
+            if ctx.needs_input_grad[1]:
+                grad_y = grad_out.sum(dim=-1, keepdim=True).expand_as(y_2d)
+            return grad_x, grad_y, grad_xnew
+            
+        if ctx.needs_input_grad[1]:
+            grad_y_tmp = torch.zeros_like(y_2d)
+            idxp1 = (idx + 1).clamp(max=y_2d.shape[1] - 1)  # SAFETY: clamp idxp1
+            
+            # Calculate gradients
+            grad_yL = (1.0 - t) * grad_out
+            grad_yR = t * grad_out
+            
+            # Ensure consistent dtype between source and destination tensors
+            grad_yL = grad_yL.to(dtype=grad_y_tmp.dtype)
+            grad_yR = grad_yR.to(dtype=grad_y_tmp.dtype)
+            
+            grad_y_tmp.scatter_add_(1, idx, grad_yL)
+            grad_y_tmp.scatter_add_(1, idxp1, grad_yR)
+            grad_y = grad_y_tmp
+        return grad_x, grad_y, grad_xnew
+
+def interp1d(x, y, xnew):
+    return Interp1d.apply(x, y, xnew)
+
+
+class SWE_Pooling(nn.Module):
+    """
+    Sliced-Wasserstein Embedding (SWE) Pooling.
+    Maps token embeddings [B, L, d_in] => [B, num_slices].
+    """
+    def __init__(self, d_in, num_slices, num_ref_points, freeze_swe=False):
+        super().__init__()
+        self.num_slices = num_slices
+        self.num_ref_points = num_ref_points
+
+        ref = torch.linspace(-1,1,num_ref_points).unsqueeze(1).repeat(1,num_slices)
+        self.reference = nn.Parameter(ref, requires_grad=not freeze_swe)
+
+        self.theta = nn.utils.weight_norm(nn.Linear(d_in, num_slices, bias=False), dim=0)
+        self.theta.weight_g.data = torch.ones_like(self.theta.weight_g.data)
+        self.theta.weight_g.requires_grad=False
+        nn.init.normal_(self.theta.weight_v)
+
+        self.weight = nn.Linear(num_ref_points,1,bias=False)
+
+        if freeze_swe:
+            self.theta.weight_v.requires_grad=False
+            self.reference.requires_grad=False
+
+    def forward(self, X, mask=None):
+        B, N, D = X.shape
+        device = X.device
+
+        X_slices = self.theta(X)  # => [B,N,num_slices]
+        X_slices_sorted, _ = torch.sort(X_slices, dim=1)
+
+        x_coord = torch.linspace(0,1,N,device=device).unsqueeze(0).repeat(B*self.num_slices,1)
+        X_flat = X_slices_sorted.permute(0,2,1).reshape(B*self.num_slices, N)
+        xnew = torch.linspace(0,1,self.num_ref_points,device=device).unsqueeze(0).repeat(B*self.num_slices,1)
+
+        y_intp = interp1d(x_coord, X_flat, xnew)
+        X_slices_sorted_interp = y_intp.view(B,self.num_slices,self.num_ref_points).permute(0,2,1)
+
+        r_expanded = self.reference.expand_as(X_slices_sorted_interp)
+        embeddings = (r_expanded - X_slices_sorted_interp).permute(0,2,1)  # => [B,num_slices,num_ref_points]
+        weighted = self.weight(embeddings).sum(dim=-1)  # => [B, num_slices]
+        return weighted
+
+
+#############################################################################
+#             Enhanced Mutation-Aware SWE_Pooling                           #
+#############################################################################
+
+class MutationAwareSWEPooling(nn.Module):
+    """
+    Enhanced Sliced-Wasserstein Embedding Pooling with explicit mutation position handling.
+    Maps token embeddings [B, L, d_in] => [B, num_slices].
+    
+    - Preserves mutation position information through weighted aggregation
+    - Uses mutation positions to guide the pooling process
+    """
+    def __init__(self, d_in, num_slices, num_ref_points, freeze_swe=False):
+        super().__init__()
+        self.num_slices = num_slices
+        self.num_ref_points = num_ref_points
+        self.d_esm = 1152  # FIXED: Hardcode to 1152 to avoid channel indexing bugs with context window
+
+        # Standard SWE components
+        ref = torch.linspace(-1, 1, num_ref_points).unsqueeze(1).repeat(1, num_slices)
+        self.reference = nn.Parameter(ref, requires_grad=not freeze_swe)
+
+        # For ESM features (without mutation channel)
+        self.theta = nn.utils.weight_norm(nn.Linear(self.d_esm, num_slices, bias=False), dim=0)
+        self.theta.weight_g.data = torch.ones_like(self.theta.weight_g.data)
+        self.theta.weight_g.requires_grad = False
+        nn.init.normal_(self.theta.weight_v)
+
+        # Mutation-aware components
+        self.mutation_importance = nn.Sequential(
+            nn.Linear(1, 32),
+            nn.ReLU(),
+            nn.Linear(32, num_slices),
+            nn.Sigmoid()
+        )
+        
+        # Position-specific weighting for each slice
+        self.pos_weighting = nn.Linear(1, num_slices, bias=False)
+
+        # FIXED: Direct projection of mutation channel (the 1153rd dim)
+        # Previously, this channel was only used as a multiplier, meaning if ESM 
+        # features had no diff, the result was 0. Now we project it directly.
+        self.mut_projection = nn.Linear(1, num_slices, bias=False)
+        
+        # Final weighting
+        self.weight = nn.Linear(num_ref_points, 1, bias=False)
+
+        if freeze_swe:
+            self.theta.weight_v.requires_grad = False
+            self.reference.requires_grad = False
+
+    def forward(self, X, mask=None):
+        """
+        X: [B, L, d_in] where d_in = d_esm + 1 (mutation channel)
+        mask: [B, L] boolean mask
+        """
+        B, N, D = X.shape
+        device = X.device
+        
+        # Check if using context window (additional channel)
+        use_context = (D > self.d_esm + 1)
+        
+        if use_context:
+            # Split ESM features and channels
+            X_esm = X[:, :, :-2]  # [B, L, d_esm]
+            X_mut = X[:, :, -2:-1]  # [B, L, 1] - mutation indicator
+        else:
+            # Split ESM features and mutation channel
+            X_esm = X[:, :, :-1]  # [B, L, d_esm]
+            X_mut = X[:, :, -1:]  # [B, L, 1] - mutation indicator
+        
+        # Regular SWE on ESM features
+        X_slices = self.theta(X_esm)  # => [B, L, num_slices]
+        
+        # Compute mutation importance weights
+        mut_weights = self.mutation_importance(X_mut)  # [B, L, num_slices]
+        
+        # Create position encodings (0 to 1 for each sequence)
+        pos_tensor = torch.linspace(0, 1, N, device=device).view(1, N, 1).expand(B, N, 1)
+        pos_weights = self.pos_weighting(pos_tensor)  # [B, L, num_slices]
+        
+        # Apply mutation-aware weighting to slices
+        # Use both mutation indicator and position information
+        # BUGFIX: We ALSO add the projected mutation signal directly. 
+        # This ensures the model 'sees' the 1.0 signal even if ESM features are identical.
+        X_slices = X_slices * (1.0 + mut_weights * pos_weights) + self.mut_projection(X_mut)
+        
+        # Sort slices as in standard SWE
+        X_slices_sorted, _ = torch.sort(X_slices, dim=1)
+        
+        # Continue with standard SWE interpolation
+        x_coord = torch.linspace(0, 1, N, device=device).unsqueeze(0).repeat(B*self.num_slices, 1)
+        X_flat = X_slices_sorted.permute(0, 2, 1).reshape(B*self.num_slices, N)
+        xnew = torch.linspace(0, 1, self.num_ref_points, device=device).unsqueeze(0).repeat(B*self.num_slices, 1)
+        
+        y_intp = interp1d(x_coord, X_flat, xnew)
+        X_slices_sorted_interp = y_intp.view(B, self.num_slices, self.num_ref_points).permute(0, 2, 1)
+        
+        r_expanded = self.reference.expand_as(X_slices_sorted_interp)
+        embeddings = (r_expanded - X_slices_sorted_interp).permute(0, 2, 1)  # => [B, num_slices, num_ref_points]
+        weighted = self.weight(embeddings).sum(dim=-1)  # => [B, num_slices]
+        
+        return weighted
+
+
+#############################################################################
+#            Mutation-Specific Cross-Attention with Gating                  #
+#############################################################################
+
+class MutationSpecificAttention(nn.Module):
+    """
+    Enhanced cross-attention that explicitly handles mutation positions with gating.
+    - Keeps ESM embeddings (1152-dim) and mutation channel separate
+    - Uses specific mutation positions to guide attention
+    - Preserves position-specific information throughout the network
+    - Adds gating mechanism to control information flow
+    - Includes memory-efficient computation for long sequences
+    """
+    def __init__(self, d_model=1152, num_heads=4, dropout=0.1):
+        super().__init__()
+        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.head_dim = d_model // num_heads
+        
+        # Core attention for ESM embeddings only (1152-dim)
+        self.query = nn.Linear(d_model, d_model)
+        self.key = nn.Linear(d_model, d_model)
+        self.value = nn.Linear(d_model, d_model)
+        
+        # Absolute position encoding
+        self.pos_encoder = nn.Sequential(
+            nn.Linear(1, 32),
+            nn.ReLU(),
+            nn.Linear(32, d_model)
+        )
+        
+        # Mutation-position specific attention
+        self.mut_encoder = nn.Sequential(
+            nn.Linear(2, 64),  # Input: [mut_binary, position_normalized]
+            nn.ReLU(),
+            nn.Linear(64, num_heads)
+        )
+        
+        # Gating mechanism to control information flow
+        self.gate = nn.Sequential(
+            nn.Linear(d_model*2, d_model),
+            nn.Sigmoid()
+        )
+        
+        self.dropout = nn.Dropout(dropout)
+        self.out_proj = nn.Linear(d_model, d_model)
+        
+    def split_heads(self, x):
+        """Split the last dimension into (heads, head_dim)"""
+        batch_size, seq_len, _ = x.shape
+        x = x.view(batch_size, seq_len, self.num_heads, self.head_dim)
+        return x.permute(0, 2, 1, 3)  # [batch, heads, seq_len, head_dim]
+    
+    def merge_heads(self, x):
+        """Merge the (heads, head_dim) into d_model"""
+        batch_size, _, seq_len, _ = x.shape
+        x = x.permute(0, 2, 1, 3)  # [batch, seq_len, heads, head_dim]
+        return x.reshape(batch_size, seq_len, self.d_model)
+    
+    def forward(self, q_esm, k_esm, v_esm, q_mut, k_mut, mask=None):
+        """
+        Inputs:
+            q_esm, k_esm, v_esm: ESM embeddings [B, L, 1152]
+            q_mut, k_mut: Mutation information [B, L, 1]
+            mask: Optional attention mask [B, L] or [B, 1, L]
+        """
+        batch_size = q_esm.shape[0]
+        q_len, k_len = q_esm.shape[1], k_esm.shape[1]
+        
+        # Create position tensors (0-1 range for each sequence)
+        q_pos = torch.linspace(0, 1, q_len, device=q_esm.device).view(1, -1, 1).expand(batch_size, q_len, 1)
+        k_pos = torch.linspace(0, 1, k_len, device=k_esm.device).view(1, -1, 1).expand(batch_size, k_len, 1)
+        
+        # Position encoding
+        q_pos_enc = self.pos_encoder(q_pos)
+        k_pos_enc = self.pos_encoder(k_pos)
+        
+        # Add position encodings to ESM features
+        q_esm_pos = q_esm + q_pos_enc
+        k_esm_pos = k_esm + k_pos_enc
+        
+        # Process core ESM embeddings with position information
+        q = self.split_heads(self.query(q_esm_pos))  # [B, h, q_len, d_k]
+        k = self.split_heads(self.key(k_esm_pos))    # [B, h, k_len, d_k]
+        v = self.split_heads(self.value(v_esm))      # [B, h, v_len, d_v]
+        
+        # Concatenate mutation indicator with position
+        q_mut_pos = torch.cat([q_mut, q_pos], dim=-1)  # [B, q_len, 2]
+        k_mut_pos = torch.cat([k_mut, k_pos], dim=-1)  # [B, k_len, 2]
+        
+        # Encode position-aware mutation information
+        q_mut_enc = self.mut_encoder(q_mut_pos)  # [B, q_len, num_heads]
+        k_mut_enc = self.mut_encoder(k_mut_pos)  # [B, k_len, num_heads]
+        
+        # Standard scaled dot-product attention
+        d_k = q.size(-1)
+        scores = torch.matmul(q, k.transpose(-2, -1)) / (d_k ** 0.5)  # [B, h, q_len, k_len]
+        
+        # Create mutation-position attention bias
+        # This explicitly boosts attention between positions based on mutation status
+        mut_attn_bias = torch.matmul(
+            q_mut_enc.permute(0, 2, 1).unsqueeze(3),  # [B, h, q_len, 1]
+            k_mut_enc.permute(0, 2, 1).unsqueeze(2)   # [B, h, 1, k_len]
+        )  # [B, h, q_len, k_len]
+        
+        # Apply mutation bias to attention scores
+        # This makes mutations and their surrounding context attend more to each other
+        scores = scores + mut_attn_bias
+        
+        # Apply mask if provided
+        if mask is not None:
+            # Fix mask dimension to match scores
+            # mask shape should be [B, L] or [B, 1, L]
+            if mask.dim() == 2:  # [B, L]
+                # For keys mask [B, k_len] -> [B, 1, 1, k_len]
+                mask = mask.unsqueeze(1).unsqueeze(2)
+            elif mask.dim() == 3 and mask.size(1) == 1:  # [B, 1, L]
+                # For keys mask [B, 1, k_len] -> [B, 1, 1, k_len]
+                mask = mask.unsqueeze(2)
+                
+            # Expand mask to match scores dimensions
+            # [B, 1, 1, k_len] -> [B, h, q_len, k_len]
+            mask = mask.expand(-1, scores.size(1), scores.size(2), -1)
+            
+            # FIXED: Use -1e4 instead of -1e9 to avoid half-precision overflow
+            scores = scores.masked_fill(mask == 0, -1e4)
+        
+        # Apply softmax and dropout
+        attention_weights = F.softmax(scores, dim=-1)
+        attention_weights = self.dropout(attention_weights)
+        
+        # Apply attention to values
+        context = torch.matmul(attention_weights, v)  # [B, h, q_len, d_v]
+        context = self.merge_heads(context)  # [B, q_len, d_model]
+        attn_output = self.out_proj(context)
+        
+        # Apply gating mechanism (new addition)
+        # Concatenate the original query with the attention output to determine the gate
+        gate_input = torch.cat([q_esm, attn_output], dim=-1)
+        gate_value = self.gate(gate_input)
+        
+        # Memory optimization for long sequences
+        # Processing the gating operation in chunks to prevent OOM errors
+        if q_len > 1000:  # Only use chunking for very long sequences
+            chunk_size = 500
+            output_chunks = []
+            
+            for i in range(0, q_len, chunk_size):
+                end_idx = min(i + chunk_size, q_len)
+                # Process chunks
+                chunk_gate = gate_value[:, i:end_idx, :]
+                chunk_attn = attn_output[:, i:end_idx, :]
+                chunk_q = q_esm[:, i:end_idx, :]
+                
+                # Apply gating equation to this chunk
+                chunk_output = chunk_gate * chunk_attn + (1 - chunk_gate) * chunk_q
+                output_chunks.append(chunk_output)
+            
+            # Combine chunks
+            output = torch.cat(output_chunks, dim=1)
+        else:
+            # Original operation for shorter sequences
+            output = gate_value * attn_output + (1 - gate_value) * q_esm
+        
+        return output
+
+
+class MutationSpecificCrossAttentionBlock(nn.Module):
+    """
+    Cross-attention block with explicit mutation position handling.
+    Each block processes ESM embeddings and mutation channels separately,
+    with special emphasis on mutation positions.
+    """
+    def __init__(self, d_model=1152, num_heads=4, ffn_dim=2048, dropout=0.1):
+        super().__init__()
+        # Mutation-aware cross attention
+        self.attn_c12 = MutationSpecificAttention(d_model, num_heads, dropout)
+        self.attn_c21 = MutationSpecificAttention(d_model, num_heads, dropout)
+        
+        # Layer normalization for ESM embeddings
+        self.norm_c1 = nn.LayerNorm(d_model)
+        self.norm_c2 = nn.LayerNorm(d_model)
+        
+        # FFN for ESM embeddings
+        self.ffn_c1 = nn.Sequential(
+            nn.Linear(d_model, ffn_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ffn_dim, d_model)
+        )
+        self.ffn_c2 = nn.Sequential(
+            nn.Linear(d_model, ffn_dim),
+            nn.ReLU(),
+            nn.Dropout(dropout),
+            nn.Linear(ffn_dim, d_model)
+        )
+        self.norm_ffn_c1 = nn.LayerNorm(d_model)
+        self.norm_ffn_c2 = nn.LayerNorm(d_model)
+        
+        # Mutation importance update layer
+        self.mut_update = nn.Sequential(
+            nn.Linear(d_model + 1, 64),
+            nn.ReLU(),
+            nn.Linear(64, 1),
+            nn.Sigmoid()
+        )
+
+    def forward(self, c1_esm, c1_mut, c2_esm, c2_mut, mask1=None, mask2=None):
+        """
+        Inputs:
+            c1_esm, c2_esm: ESM embeddings [B, L, 1152]
+            c1_mut, c2_mut: Mutation channels [B, L, 1]
+            mask1, mask2: Optional masks
+        """
+        # c1->c2 cross-attention
+        c1_attn = self.attn_c12(c1_esm, c2_esm, c2_esm, c1_mut, c2_mut, mask2)
+        c1_out = self.norm_c1(c1_esm + c1_attn)
+        
+        # c2->c1 cross-attention
+        c2_attn = self.attn_c21(c2_esm, c1_esm, c1_esm, c2_mut, c1_mut, mask1)
+        c2_out = self.norm_c2(c2_esm + c2_attn)
+        
+        # Feed-forward
+        c1_ffn = self.ffn_c1(c1_out)
+        c1_ffn_out = self.norm_ffn_c1(c1_out + c1_ffn)
+        
+        c2_ffn = self.ffn_c2(c2_out)
+        c2_ffn_out = self.norm_ffn_c2(c2_out + c2_ffn)
+        
+        # Update mutation importance based on attention output
+        # This creates a feedback loop where mutation effect is refined
+        c1_mut_in = torch.cat([c1_ffn_out, c1_mut], dim=-1)
+        c2_mut_in = torch.cat([c2_ffn_out, c2_mut], dim=-1)
+        
+        # Stabilized update: convex combination ensures values stay in [0, 1]
+        # Avoids exponential decay (old bug) and unbounded growth (additive bug)
+        c1_mut_updated = 0.9 * c1_mut + 0.1 * self.mut_update(c1_mut_in)
+        c2_mut_updated = 0.9 * c2_mut + 0.1 * self.mut_update(c2_mut_in)
+        
+        return c1_ffn_out, c2_ffn_out, c1_mut_updated, c2_mut_updated
+
+
+class MutationSpecificCrossAttentionStack(nn.Module):
+    """
+    Stack of Mutation-Specific Cross-Attention blocks.
+    Emphasizes mutation positions throughout the network.
+    Now includes gradient checkpointing for memory efficiency.
+    """
+    def __init__(self, d_model=1152, num_heads=4, ffn_dim=2048, dropout=0.1, num_layers=2):
+        super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.use_checkpoint = True  # Enable gradient checkpointing by default
+        
+        self.blocks = nn.ModuleList([
+            MutationSpecificCrossAttentionBlock(
+                d_model=d_model,
+                num_heads=num_heads,
+                ffn_dim=ffn_dim,
+                dropout=dropout
+            ) for _ in range(num_layers)
+        ])
+
+    def forward(self, c1, c2, mask1=None, mask2=None):
+        """
+        Process protein chains with mutation-specific attention.
+        c1, c2: [B, L, D] where D can be 1153 (original) or 1154 (with context window)
+        Uses gradient checkpointing when in training mode to save memory.
+        """
+        # Check input dimension to determine if context window is used
+        d_in = c1.shape[2]
+        use_context = (d_in > 1153)
+        
+        if use_context:
+            # Split ESM embeddings from mutation+context channels
+            c1_esm, c1_channels = c1[:, :, :-2], c1[:, :, -2:]  # [B, L, 1152], [B, L, 2]
+            c2_esm, c2_channels = c2[:, :, :-2], c2[:, :, -2:]  # [B, L, 1152], [B, L, 2]
+            
+            # Extract mutation channel (first channel)
+            c1_mut = c1_channels[:, :, :1]  # [B, L, 1]
+            c2_mut = c2_channels[:, :, :1]  # [B, L, 1]
+        else:
+            # Original behavior - just split ESM and mutation
+            c1_esm, c1_mut = c1[:, :, :-1], c1[:, :, -1:]  # [B, L, 1152], [B, L, 1]
+            c2_esm, c2_mut = c2[:, :, :-1], c2[:, :, -1:]  # [B, L, 1152], [B, L, 1]
+        
+        # Process through attention blocks with optional checkpointing
+        for block in self.blocks:
+            # Use gradient checkpointing in training mode for memory efficiency
+            if self.use_checkpoint and self.training:
+                # Define helper function for checkpointing that handles None masks
+                def create_checkpoint_fn(block_fn):
+                    def checkpoint_fn(esm1, mut1, esm2, mut2, has_mask1, has_mask2, mask1_val, mask2_val):
+                        # Conditionally use the masks based on the has_mask flags
+                        m1 = mask1_val if has_mask1 else None
+                        m2 = mask2_val if has_mask2 else None
+                        return block_fn(esm1, mut1, esm2, mut2, m1, m2)
+                    return checkpoint_fn
+                
+                # Convert None masks to flags and dummy tensors for checkpointing
+                has_mask1 = mask1 is not None
+                has_mask2 = mask2 is not None
+                mask1_val = mask1 if has_mask1 else torch.zeros(1, device=c1_esm.device)
+                mask2_val = mask2 if has_mask2 else torch.zeros(1, device=c1_esm.device)
+                
+                # Apply checkpointing
+                c1_esm, c2_esm, c1_mut, c2_mut = torch.utils.checkpoint.checkpoint(
+                    create_checkpoint_fn(block),
+                    c1_esm, c1_mut, c2_esm, c2_mut,
+                    torch.tensor(has_mask1, device=c1_esm.device),
+                    torch.tensor(has_mask2, device=c1_esm.device),
+                    mask1_val, mask2_val
+                )
+            else:
+                c1_esm, c2_esm, c1_mut, c2_mut = block(c1_esm, c1_mut, c2_esm, c2_mut, mask1, mask2)
+        
+        # Recombine with appropriate channels
+        if use_context:
+            # Need to preserve the context channel
+            context_channels_c1 = c1_channels[:, :, 1:]  # [B, L, 1]
+            context_channels_c2 = c2_channels[:, :, 1:]  # [B, L, 1]
+            c1_out = torch.cat([c1_esm, c1_mut, context_channels_c1], dim=-1)  # [B, L, 1154]
+            c2_out = torch.cat([c2_esm, c2_mut, context_channels_c2], dim=-1)  # [B, L, 1154]
+        else:
+            # Original behavior
+            c1_out = torch.cat([c1_esm, c1_mut], dim=-1)  # [B, L, 1153]
+            c2_out = torch.cat([c2_esm, c2_mut], dim=-1)  # [B, L, 1153]
+        
+        return c1_out, c2_out
+
+
+#############################################################################
+#        AffinityPredictor with Improved Memory Efficiency                  #
+#############################################################################
+
+class AffinityPredictor(nn.Module):
+    """
+    Enhanced AffinityPredictor with explicit mutation position handling.
+    embedding_method => "difference", "cosine", "cross_attention", or "cross_attention_swe".
+    """
+    def __init__(
+        self,
+        input_dim=1153,  # 1152 (ESM) + 1 (mutation)
+        latent_dim=1024,
+        num_slices=1024,
+        num_ref_points=128,
+        dropout_rate=0.2,
+        freeze_swe=False,
+        embedding_method="difference",
+        normalize_difference=False,
+        num_hidden_layers=2,
+        # cross-attn
+        num_cross_attn_layers=2,
+        num_attention_heads=4,
+        cross_ffn_dim=2048,
+    ):
+        super().__init__()
+        self.embedding_method = embedding_method.lower()
+        self.normalize_difference = normalize_difference
+        self.input_dim = input_dim
+        
+        # ESM dimension (without mutation channel)
+        self.esm_dim = input_dim - 1  # 1152
+        
+        # Define cross-attention stack if needed
+        self.cross_stack = None
+        if "cross_attention" in self.embedding_method:
+            self.cross_stack = MutationSpecificCrossAttentionStack(
+                d_model=self.esm_dim,  # 1152
+                num_heads=num_attention_heads,
+                ffn_dim=cross_ffn_dim,
+                dropout=dropout_rate,
+                num_layers=num_cross_attn_layers
+            )
+        
+        # Enhanced Mutation-Aware SWE Pooling
+        self.swe_pooling = None
+        if self.embedding_method in ["difference", "cosine", "cross_attention_swe"]:
+            # For SWE, we use the full input_dim (1153)
+            self.swe_pooling = MutationAwareSWEPooling(
+                d_in=input_dim,
+                num_slices=num_slices,
+                num_ref_points=num_ref_points,
+                freeze_swe=freeze_swe
+            )
+        
+        # Define aggregator MLP in-dimensions
+        if self.embedding_method == "cosine":
+            in_features = 1
+        else:
+            in_features = num_slices  # difference or cross_attention_swe => [B, num_slices]
+        
+        # Add projection layer for cross_attention to avoid dynamic creation
+        self.cross_attn_projection = None
+        if self.embedding_method == "cross_attention":
+            cross_proj_in = input_dim  # Full dimension including mutation channel
+            cross_proj_out = in_features
+            self.cross_attn_projection = nn.Linear(cross_proj_in, cross_proj_out, bias=False)
+        
+        # Final MLP
+        layers = []
+        current_dim = in_features
+        for _ in range(num_hidden_layers):
+            layers.append(nn.Linear(current_dim, latent_dim))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(dropout_rate))
+            current_dim = latent_dim
+        layers.append(nn.Linear(current_dim, 1))
+        self.mlp = nn.Sequential(*layers)
+
+    def forward(self, chain1, chain1_mask, chain2, chain2_mask):
+        """
+        chain1, chain2 => [B, L, input_dim] (1153 or 1154 with context)
+        """
+        if "cross_attention" in self.embedding_method:
+            # Process through mutation-specific cross-attention
+            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
+            
+            if self.embedding_method == "cross_attention_swe":
+                # Apply enhanced SWE pooling
+                rep1 = self.swe_pooling(c1_out, chain1_mask)  # [B, num_slices]
+                rep2 = self.swe_pooling(c2_out, chain2_mask)  # [B, num_slices]
+                
+                # Difference aggregator
+                diff = rep1 - rep2
+                if self.normalize_difference:
+                    diff = F.normalize(diff, p=2, dim=1)
+                    
+                # Final prediction
+                preds = self.mlp(diff).squeeze(-1)
+                
+                return preds
+                
+            elif self.embedding_method == "cross_attention":
+                # Use mutation-weighted pooling
+                # Extract mutation channel to guide pooling
+                d_in = c1_out.shape[2]
+                use_context = (d_in > 1153)
+                if use_context:
+                    c1_mut = c1_out[:, :, -2:-1]  # [B, L, 1]
+                    c2_mut = c2_out[:, :, -2:-1]  # [B, L, 1]
+                else:
+                    c1_mut = c1_out[:, :, -1:]  # [B, L, 1]
+                    c2_mut = c2_out[:, :, -1:]  # [B, L, 1]
+                
+                # Weighted pooling - gives higher weight to mutated positions
+                c1_weights = F.softmax(c1_mut * 10, dim=1)  # Sharpen weights
+                c2_weights = F.softmax(c2_mut * 10, dim=1)
+                
+                c1_pool = torch.sum(c1_out * c1_weights, dim=1)  # [B, 1153/1154]
+                c2_pool = torch.sum(c2_out * c2_weights, dim=1)  # [B, 1153/1154]
+                
+                # Create difference representation
+                diff = c1_pool - c2_pool
+                if self.normalize_difference:
+                    diff = F.normalize(diff, p=2, dim=1)
+                
+                # Use pre-defined projection layer instead of creating one dynamically
+                if self.cross_attn_projection is not None:
+                    diff = self.cross_attn_projection(diff)
+                
+                preds = self.mlp(diff).squeeze(-1)
+                
+                return preds
+        
+        elif self.embedding_method == "cosine":
+            # Enhanced SWE => [B, num_slices]
+            rep1 = self.swe_pooling(chain1, chain1_mask)
+            rep2 = self.swe_pooling(chain2, chain2_mask)
+            sim = F.cosine_similarity(rep1, rep2, dim=1).unsqueeze(-1)
+            out = self.mlp(sim).squeeze(-1)
+            
+            return out
+        
+        else:  # "difference"
+            rep1 = self.swe_pooling(chain1, chain1_mask)
+            rep2 = self.swe_pooling(chain2, chain2_mask)
+            diff = rep1 - rep2
+            if self.normalize_difference:
+                diff = F.normalize(diff, p=2, dim=1)
+            out = self.mlp(diff).squeeze(-1)
+            
+            return out
+
+
+class AffinityPredictionModel(pl.LightningModule):
+    """
+    Lightning wrapper for training. Siamese logic in training.py
+    """
+    def __init__(self, predictor: AffinityPredictor, learning_rate=1e-4):
+        super().__init__()
+        self.predictor = predictor
+        self.learning_rate = learning_rate
+
+        self.loss_fn = nn.MSELoss()
+        self.pearson_corr = PearsonCorrCoef()
+        self.spearman_corr = SpearmanCorrCoef()
+        self.r2_score = R2Score()
+        self.mse_metric = MeanSquaredError()
+
+    def forward(self, chain1, chain1_mask, chain2, chain2_mask):
+        return self.predictor(chain1, chain1_mask, chain2, chain2_mask)
+
+    def training_step(self, batch, batch_idx):
+        pass
+
+    def validation_step(self, batch, batch_idx):
+        pass
+
+    def configure_optimizers(self):
+        optimizer = AdamW(self.parameters(), lr=self.learning_rate)
+        steps = self.trainer.estimated_stepping_batches
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer, max_lr=self.learning_rate, total_steps=steps
+        )
+        return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val_loss"}
+
+
+# Enhanced AffinityPredictor with Two-Head Architecture
+# Add this to architectures.py
+
+class DualHeadAffinityPredictor(nn.Module):
+    """
+    Enhanced AffinityPredictor with explicit two-head architecture.
+    Simultaneously processes mutant and wildtype proteins to predict both ΔG and ΔΔG.
+    
+    embedding_method => "difference", "cosine", "cross_attention", or "cross_attention_swe".
+    """
+    def __init__(
+        self,
+        input_dim=1153,  # 1152 (ESM) + 1 (mutation)
+        latent_dim=1024,
+        num_slices=1024,
+        num_ref_points=128,
+        dropout_rate=0.2,
+        freeze_swe=False,
+        embedding_method="difference",
+        normalize_difference=False,
+        num_hidden_layers=2,
+        # cross-attn
+        num_cross_attn_layers=2,
+        num_attention_heads=4,
+        cross_ffn_dim=2048,
+        use_dual_head=True,  # Enable dual-head by default
+        ddg_signal_gain=1.0,  # Initial gain for ddG signal
+        ddg_signal_multiplier=20.0, # FIXED: multiplier for ddG signal (vnew65.0)
+    ):
+        super().__init__()
+        self.embedding_method = embedding_method.lower()
+        self.normalize_difference = normalize_difference
+        self.input_dim = input_dim
+        self.use_dual_head = use_dual_head
+        self._ddg_log_counter = 0
+        
+        # DEBUG: Confirm this version is running
+        print(f"[MODEL INIT] DualHeadAffinityPredictor created: version={ARCH_VERSION}, dual_head={use_dual_head}, method={self.embedding_method}")
+        
+        # ESM dimension (without mutation channel)
+        self.esm_dim = input_dim - 1  # 1152
+        
+        # Define cross-attention stack if needed
+        self.cross_stack = None
+        if "cross_attention" in self.embedding_method:
+            self.cross_stack = MutationSpecificCrossAttentionStack(
+                d_model=self.esm_dim,  # 1152
+                num_heads=num_attention_heads,
+                ffn_dim=cross_ffn_dim,
+                dropout=dropout_rate,
+                num_layers=num_cross_attn_layers
+            )
+        
+        # Enhanced Mutation-Aware SWE Pooling
+        self.swe_pooling = None
+        if self.embedding_method in ["difference", "cosine", "cross_attention_swe"]:
+            # For SWE, we use the full input_dim (1153)
+            self.swe_pooling = MutationAwareSWEPooling(
+                d_in=input_dim,
+                num_slices=num_slices,
+                num_ref_points=num_ref_points,
+                freeze_swe=freeze_swe
+            )
+        
+        # Define aggregator MLP in-dimensions
+        if self.embedding_method == "cosine":
+            in_features = 1
+        else:
+            in_features = num_slices  # difference or cross_attention_swe => [B, num_slices]
+        
+        # Add projection layer for cross_attention to avoid dynamic creation
+        self.cross_attn_projection = None
+        if self.embedding_method == "cross_attention":
+            cross_proj_in = input_dim  # Full dimension including mutation channel
+            cross_proj_out = in_features
+            self.cross_attn_projection = nn.Linear(cross_proj_in, cross_proj_out, bias=False)
+        
+        # Define dG head (main prediction head)
+        layers = []
+        current_dim = in_features
+        for _ in range(num_hidden_layers):
+            layers.append(nn.Linear(current_dim, latent_dim))
+            layers.append(nn.ReLU())
+            layers.append(nn.Dropout(dropout_rate))
+            current_dim = latent_dim
+        layers.append(nn.Linear(current_dim, 1))
+        self.dg_mlp = nn.Sequential(*layers)
+        
+        # Define ΔΔG head for direct prediction
+        # FIX (vnew64.0): Shallow 2-layer MLP + residual skip connection.
+        # DDGACT diagnostics from vnew62-63 showed 7-layer ReLU network causes 
+        # progressive variance collapse (std: 0.1 → 7e-05). Each ReLU zeros ~50%
+        # of activations, so 7 layers → 0.5^7 = 0.8% signal survival.
+        # Solution: (1) 2 layers only, (2) skip connection preserves raw input signal.
+        if self.use_dual_head:
+            # Shallow nonlinear pathway (2 layers)
+            self.ddg_hidden = nn.Sequential(
+                nn.Linear(in_features, latent_dim),
+                nn.GELU(),  # GELU instead of ReLU - no zero-capping, smoother gradients
+                nn.Linear(latent_dim, latent_dim),
+                nn.GELU(),
+            )
+            # Output projection
+            self.ddg_out = nn.Linear(latent_dim, 1)
+            # Skip connection: project input directly to output dimension
+            self.ddg_skip = nn.Linear(in_features, 1)
+        
+        # Source type embedding for conditional inference
+        # User-friendly types for inference:
+        #   0 = "mutant"    - Single mutant predictions (most common)
+        #   1 = "wt_pairs"  - Wildtype pairs with absolute binding affinity
+        #   2 = "antibody"  - Antibody-antigen binding (CDR-focused)
+        self.source_type_embedding = nn.Embedding(3, 32)  # 32-dim embedding
+        self.source_type_projection = nn.Linear(in_features + 32, in_features)  # Project back to in_features
+
+        # FIXED: Restoring the historical 'learnable gain' strategy
+        # This allows the model to amplify the ddG signal early in Stage B
+        self.ddg_signal_gain = nn.Parameter(torch.tensor(float(ddg_signal_gain)))
+        self.ddg_signal_multiplier = float(ddg_signal_multiplier)
+
+    def _extract_features(self, chain1, chain1_mask, chain2, chain2_mask):
+        """
+        Extract feature representation for a protein complex.
+        Returns a vector representation suitable for prediction.
+        """
+        if "cross_attention" in self.embedding_method:
+            # Process through mutation-specific cross-attention
+            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
+            
+            if self.embedding_method == "cross_attention_swe":
+                # Apply enhanced SWE pooling
+                rep1 = self.swe_pooling(c1_out, chain1_mask)  # [B, num_slices]
+                rep2 = self.swe_pooling(c2_out, chain2_mask)  # [B, num_slices]
+                
+                # Difference aggregator
+                diff = rep1 - rep2
+                if self.normalize_difference:
+                    diff = F.normalize(diff, p=2, dim=1)
+                return diff
+                
+            elif self.embedding_method == "cross_attention":
+                # Use mutation-weighted pooling
+                # Extract mutation channel to guide pooling
+                d_in = c1_out.shape[2]
+                use_context = (d_in > 1153)
+                if use_context:
+                    c1_mut = c1_out[:, :, -2:-1]  # [B, L, 1]
+                    c2_mut = c2_out[:, :, -2:-1]  # [B, L, 1]
+                else:
+                    c1_mut = c1_out[:, :, -1:]  # [B, L, 1]
+                    c2_mut = c2_out[:, :, -1:]  # [B, L, 1]
+                
+                # Weighted pooling - gives higher weight to mutated positions
+                c1_weights = F.softmax(c1_mut * 10, dim=1)  # Sharpen weights
+                c2_weights = F.softmax(c2_mut * 10, dim=1)
+                
+                c1_pool = torch.sum(c1_out * c1_weights, dim=1)  # [B, 1153/1154]
+                c2_pool = torch.sum(c2_out * c2_weights, dim=1)  # [B, 1153/1154]
+                
+                # Create difference representation
+                diff = c1_pool - c2_pool
+                if self.normalize_difference:
+                    diff = F.normalize(diff, p=2, dim=1)
+                
+                # Use pre-defined projection layer instead of creating one dynamically
+                if self.cross_attn_projection is not None:
+                    diff = self.cross_attn_projection(diff)
+                
+                return diff
+        
+        elif self.embedding_method == "cosine":
+            # Enhanced SWE => [B, num_slices]
+            rep1 = self.swe_pooling(chain1, chain1_mask)
+            rep2 = self.swe_pooling(chain2, chain2_mask)
+            sim = F.cosine_similarity(rep1, rep2, dim=1).unsqueeze(-1)
+            return sim
+        
+        else:  # "difference"
+            rep1 = self.swe_pooling(chain1, chain1_mask)
+            rep2 = self.swe_pooling(chain2, chain2_mask)
+            diff = rep1 - rep2
+            if self.normalize_difference:
+                diff = F.normalize(diff, p=2, dim=1)
+            return diff
+    
+    def _extract_residue_features(self, chain1, chain1_mask, chain2, chain2_mask):
+        """
+        Extract RESIDUE-LEVEL features (before pooling) for computing differences.
+        Used for ddG to preserve mutation-specific information.
+        
+        Returns:
+            c1_out, c2_out: [B, L1, D] and [B, L2, D] attended residue features
+        """
+        if "cross_attention" in self.embedding_method:
+            c1_out, c2_out = self.cross_stack(chain1, chain2, chain1_mask, chain2_mask)
+            return c1_out, c2_out
+        else:
+            # For non-cross-attention methods, return inputs directly
+            return chain1, chain2
+    
+    def forward(self, mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask, 
+            wt_chain1=None, wt_chain1_mask=None, wt_chain2=None, wt_chain2_mask=None,
+            source_type_ids=None):
+        """
+        Dual-head forward method that can handle both modes:
+        1. Standard mode: Just predict dG for mutant complex
+        2. Dual-head mode: Predict both dG and direct ddG when wildtype is provided
+        
+        For ddG: Uses RESIDUE-LEVEL differences before pooling to preserve mutation info.
+        ddG = ddg_mlp(pool(mut_features - wt_features)) instead of 
+              ddg_mlp(pool(mut_features) - pool(wt_features))
+        
+        Args:
+            source_type_ids: Optional[Tensor] of shape [B], values 0/1/2 for conditioning
+        
+        Returns:
+            If wildtype inputs are None or use_dual_head=False:
+                Returns mutant dG prediction only
+            Else:
+                Returns tuple of (mutant_dG, direct_ddG_prediction)
+        """
+        # ============== OPTIMIZED: Cache residue features ==============
+        # Get mutant RESIDUE-LEVEL features first (used for both dG and ddG)
+        mut_c1_res, mut_c2_res = self._extract_residue_features(
+            mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask)
+        
+        # Pool for dG prediction (reuses cached residue features)
+        if "cross_attention_swe" in self.embedding_method:
+            rep1 = self.swe_pooling(mut_c1_res, mut_chain1_mask)
+            rep2 = self.swe_pooling(mut_c2_res, mut_chain2_mask)
+            mut_features = rep1 - rep2
+            if self.normalize_difference:
+                mut_features = F.normalize(mut_features, p=2, dim=1, eps=1e-8)
+        else:
+            # Fallback pooling
+            mut_features = self._extract_features(mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask)
+        
+        # ============== SOURCE TYPE CONDITIONING ==============
+        # Apply source type conditioning if provided
+        if source_type_ids is not None:
+            # Get source type embedding [B, 32]
+            src_emb = self.source_type_embedding(source_type_ids)
+            # Concatenate with features and project back
+            conditioned_features = torch.cat([mut_features, src_emb], dim=-1)
+            mut_features = self.source_type_projection(conditioned_features)
+        
+        # Predict dG for mutant
+        dg_pred = self.dg_mlp(mut_features).squeeze(-1)
+        
+        # If no wildtype or dual head is disabled, just return mutant dG
+        if not self.use_dual_head or wt_chain1 is None or wt_chain2 is None:
+            return dg_pred
+        
+        # ============== RESIDUE-LEVEL ddG COMPUTATION ==============
+        # Get wildtype RESIDUE-LEVEL features (mutant already cached above)
+        wt_c1_res, wt_c2_res = self._extract_residue_features(
+            wt_chain1, wt_chain1_mask, wt_chain2, wt_chain2_mask)
+        
+        
+        # (Debug logging removed - was polluting training output)
+        
+        # Compute RESIDUE-LEVEL differences BEFORE pooling
+        # This preserves mutation-specific changes at each position
+        # Handle sequence length differences by taking minimum length
+        L_c1 = min(mut_c1_res.shape[1], wt_c1_res.shape[1])
+        L_c2 = min(mut_c2_res.shape[1], wt_c2_res.shape[1])
+        
+        c1_diff = mut_c1_res[:, :L_c1, :] - wt_c1_res[:, :L_c1, :]  # [B, L1, D]
+        c2_diff = mut_c2_res[:, :L_c2, :] - wt_c2_res[:, :L_c2, :]  # [B, L2, D]
+        
+        # Update masks for truncated length
+        c1_diff_mask = mut_chain1_mask[:, :L_c1] if mut_chain1_mask is not None else None
+        c2_diff_mask = mut_chain2_mask[:, :L_c2] if mut_chain2_mask is not None else None
+        
+        # [DIFF CHECK] Diagnostic Logging
+        # Measure cosine similarity at the approximate mutation site to check for embedding collapse.
+        if self._ddg_log_counter % 200 == 1:
+            with torch.no_grad():
+                # Extract center position
+                mid_idx = L_c1 // 2
+                
+                # Vectors at midpoint
+                v_mut = mut_c1_res[:, mid_idx, :]
+                v_wt = wt_c1_res[:, mid_idx, :]
+                
+                # Cosine similarity
+                cos_sim = F.cosine_similarity(v_mut, v_wt, dim=1).mean().item()
+                diff_norm = (v_mut - v_wt).norm(dim=1).mean().item()
+                
+                logger.info(f"[DIFF CHECK] Batch {self._ddg_log_counter}: CosSim at mid={cos_sim:.5f}, DiffNorm={diff_norm:.5f}")
+                
+        # NOW pool the differences
+        if self.swe_pooling is not None:
+            # =================================================================
+            # HYBRID POOLING: Global SWE + Local Mutation-Site-Centric (vnew37.0)
+            # This ensures local mutation signals are not diluted by global pooling
+            # =================================================================
+            
+            # Concatenate chain differences along sequence dimension
+            combined_diff = torch.cat([c1_diff, c2_diff], dim=1)  # [B, L_comb, D=1153]
+            if c1_diff_mask is not None and c2_diff_mask is not None:
+                combined_mask = torch.cat([c1_diff_mask, c2_diff_mask], dim=1)
+            else:
+                combined_mask = None
+
+            # A. Global Component: Standard SWE pooling (capture global stability context)
+            global_diff = self.swe_pooling(combined_diff, combined_mask) # [B, num_slices]
+
+            # B. Local Component: Mutation-Site-Centric Pooling (MSCP)
+            # CRITICAL FIX (v49.0): Extract indicator from RAW INPUT chains, NOT cross-attention output!
+            # The cross-attention stack applies 0.9 convex combination at each layer (5 layers = 0.9^5 = 59% decay).
+            # Using mut_chain1/mut_chain2 (raw inputs) instead of mut_c1_res/mut_c2_res (diluted outputs).
+            
+            # Determine if we have context window (1154-dim) or standard (1153-dim)
+            d_raw = mut_chain1.shape[2]
+            use_context = (d_raw > 1153)
+            
+            # Extract RAW indicator from INPUT chains (before cross-attention!)
+            if use_context:
+                c1_mut_raw = mut_chain1[:, :L_c1, -2:-1]
+                c2_mut_raw = mut_chain2[:, :L_c2, -2:-1]
+            else:
+                c1_mut_raw = mut_chain1[:, :L_c1, -1:]
+                c2_mut_raw = mut_chain2[:, :L_c2, -1:]
+            
+            mut_indicator = torch.cat([c1_mut_raw, c2_mut_raw], dim=1)  # [B, L_comb, 1]
+            
+            # Stability: Clamp indicator to be non-negative and bounded.
+            # In case attention stack produced weird values, we force them back to [0, 2]
+            mut_indicator = mut_indicator.clamp(min=0.0, max=2.0)
+            
+            # Weighted average focusing ONLY on the mutation sites
+            # use 0.001 epsilon to avoid nan for WT samples (mut_sum=0)
+            mut_sum = mut_indicator.sum(dim=1).clamp(min=1e-3)
+            
+            # MSCP Calculation with additional stability guard
+            mscp_esm = (combined_diff[:, :, :1152] * mut_indicator).sum(dim=1) / mut_sum # [B, 1152]
+            mscp_mut = (mut_indicator * mut_indicator).sum(dim=1) / mut_sum # [B, 1] (should be ~1.0)
+            
+            # Project local delta into the same slice space as global features 
+            # using the SHARED theta projection and mut_projection from SWE
+            # This ensures consistent representation between global and local paths
+            local_diff_esm = self.swe_pooling.theta(mscp_esm)
+            local_diff_mut = self.swe_pooling.mut_projection(mscp_mut)
+            local_diff = local_diff_esm + local_diff_mut  # [B, num_slices]
+
+            # C. Combine: Residual-style addition + Gain
+            # FIXED: Apply signal_gain AFTER normalization so it actually has effect
+            # Previously, L2-norm undid the scaling entirely.
+            diff_multiplier = getattr(self, "ddg_signal_multiplier", 20.0)
+            diff_features = (global_diff + local_diff) * diff_multiplier
+            
+            # ===================================================================
+            # SIGNAL FLOW LOGGING: Track local vs global contribution
+            if not hasattr(self, '_ddg_log_counter'):
+                self._ddg_log_counter = 0
+            self._ddg_log_counter += 1
+            
+            should_log = (self._ddg_log_counter % 200 == 1)
+            if should_log:
+                g_mag = global_diff.abs().mean().item()
+                l_mag = local_diff.abs().mean().item()
+                logger.info(f"[DDG SIGNAL] Batch {self._ddg_log_counter}: Global_mag={g_mag:.4f}, Local_mag={l_mag:.4f}")
+                #region agent log
+                try:
+                    # Inspect what the model thinks the mutation indicator is (both tail channels)
+                    d_raw_dbg = int(mut_chain1.shape[2])
+                    # indicator candidate stats on chain1/chain2 for last and second-last channels
+                    def _chan_stats(x):
+                        return {
+                            "min": float(x.min().item()),
+                            "max": float(x.max().item()),
+                            "mean": float(x.float().mean().item()),
+                            "std": float(x.float().std().item()),
+                        }
+                    c1_last = _chan_stats(mut_chain1[:, :L_c1, -1])
+                    c2_last = _chan_stats(mut_chain2[:, :L_c2, -1])
+                    c1_last2 = _chan_stats(mut_chain1[:, :L_c1, -2]) if d_raw_dbg >= 1154 else None
+                    c2_last2 = _chan_stats(mut_chain2[:, :L_c2, -2]) if d_raw_dbg >= 1154 else None
+                    mut_sum_dbg = float(mut_sum.mean().item()) if "mut_sum" in locals() else None
+                    payload = {
+                        "sessionId": "debug-session",
+                        "runId": "pre-fix",
+                        "hypothesisId": "F",
+                        "location": "architectures.py:DualHeadAffinityPredictor:mscp_indicator_debug",
+                        "message": "Indicator channel stats (last vs second-last) to detect double-indicator / wrong channel selection",
+                        "data": {
+                            "ddg_log_counter": int(self._ddg_log_counter),
+                            "d_raw": d_raw_dbg,
+                            "use_context_flag": bool(use_context),
+                            "c1_last": c1_last,
+                            "c2_last": c2_last,
+                            "c1_last2": c1_last2,
+                            "c2_last2": c2_last2,
+                            "mut_sum_mean": mut_sum_dbg,
+                        },
+                        "timestamp": int(time.time() * 1000),
+                    }
+                    with open("/Users/supantha/Documents/code_v2/protein/.cursor/debug.log", "a") as f:
+                        f.write(json.dumps(payload, default=str) + "\n")
+                    logger.info(f"[AGENTLOG MSCP] d_raw={d_raw_dbg} use_context={use_context} c1_last={c1_last} c1_last2={c1_last2} mut_sum_mean={mut_sum_dbg}")
+                except Exception:
+                    pass
+                #endregion
+            
+            if self.normalize_difference:
+                diff_features = F.normalize(diff_features, p=2, dim=1, eps=1e-8)
+            
+            # FIXED: Apply signal_gain AFTER normalization so it actually scales the output
+            diff_features = diff_features * self.ddg_signal_gain
+        else:
+            # Fallback: mean pooling of differences
+            if c1_diff_mask is not None:
+                c1_diff = c1_diff * c1_diff_mask.unsqueeze(-1).float()
+                c1_pool = c1_diff.sum(dim=1) / c1_diff_mask.sum(dim=1, keepdim=True).clamp(min=1)
+            else:
+                c1_pool = c1_diff.mean(dim=1)
+            if c2_diff_mask is not None:
+                c2_diff = c2_diff * c2_diff_mask.unsqueeze(-1).float()
+                c2_pool = c2_diff.sum(dim=1) / c2_diff_mask.sum(dim=1, keepdim=True).clamp(min=1)
+            else:
+                c2_pool = c2_diff.mean(dim=1)
+            diff_features = c1_pool - c2_pool
+            if self.normalize_difference:
+                diff_features = F.normalize(diff_features, p=2, dim=1)
+        
+        # DEBUG: Check for NaNs/Infs in diff_features
+        if torch.isnan(diff_features).any() or torch.isinf(diff_features).any():
+            print(f"[DEBUG MODEL] NaN/Inf in diff_features! Shape: {diff_features.shape}")
+            if c1_diff_mask is not None:
+                print(f"  c1_mask sum: {c1_diff_mask.sum(dim=1).min().item()}")
+            print(f"  c1_diff nan: {torch.isnan(c1_diff).any().item()}")
+            print(f"  c2_diff nan: {torch.isnan(c2_diff).any().item()}")
+
+        # ============== SOURCE TYPE CONDITIONING FOR DDG ==============
+        # Apply same source conditioning to diff_features for ddG prediction
+        # This allows ddG head to learn different behaviors for different data sources
+        if source_type_ids is not None:
+            src_emb = self.source_type_embedding(source_type_ids)
+            conditioned_diff = torch.cat([diff_features, src_emb], dim=-1)
+            diff_features = self.source_type_projection(conditioned_diff)
+
+        # MSCP Hybrid: Skip 20x gain since MSCP provides raw, un-diluted signal
+        # vnew64.0: Shallow 2-layer GELU + skip connection to preserve variance
+        ddg_hidden_out = self.ddg_hidden(diff_features)
+        ddg_pred = (self.ddg_out(ddg_hidden_out) + self.ddg_skip(diff_features)).squeeze(-1)
+
+        #region agent log
+        # Diagnose "train variance but eval constant" which often indicates dropout-only variance or head collapse.
+        try:
+            if should_log:
+                # Diff feature stats across the batch
+                df = diff_features.detach()
+                df_mean = float(df.mean().item()) if df.numel() else None
+                df_std = float(df.std().item()) if df.numel() else None
+                df_abs_mean = float(df.abs().mean().item()) if df.numel() else None
+                # per-sample spread: average std over features (helps detect "all samples identical")
+                df_per_sample_std = float(df.float().std(dim=1).mean().item()) if df.dim() == 2 and df.shape[0] > 0 else None
+
+                # ddg_pred stats (across batch)
+                p = ddg_pred.detach()
+                p_mean = float(p.mean().item()) if p.numel() else None
+                p_std = float(p.std().item()) if p.numel() else None
+
+                # If dropout is active (training), a second forward pass should differ.
+                p2_std = None
+                p_diff_std = None
+                if self.training:
+                    p2 = (self.ddg_out(self.ddg_hidden(diff_features)) + self.ddg_skip(diff_features)).squeeze(-1).detach()
+                    p2_std = float(p2.std().item()) if p2.numel() else None
+                    p_diff_std = float((p2 - p).std().item()) if p2.numel() else None
+
+                # Weight/bias norms (to detect collapse to near-zero weights or bias-only prediction)
+                lin_layers = [m for m in self.ddg_hidden.modules() if isinstance(m, nn.Linear)] if hasattr(self, "ddg_hidden") else []
+                w0 = lin_layers[0] if len(lin_layers) > 0 else None
+                wL = lin_layers[-1] if len(lin_layers) > 0 else None
+                w0_norm = float(w0.weight.detach().norm().item()) if w0 is not None else None
+                wL_norm = float(wL.weight.detach().norm().item()) if wL is not None else None
+                bL_norm = float(wL.bias.detach().norm().item()) if (wL is not None and wL.bias is not None) else None
+
+                payload = {
+                    "sessionId": "debug-session",
+                    "runId": "pre-fix",
+                    "hypothesisId": "I",
+                    "location": "architectures.py:DualHeadAffinityPredictor:ddg_head_eval_vs_train",
+                    "message": "ddG head collapse vs dropout-only variance diagnostics",
+                    "data": {
+                        "ddg_log_counter": int(self._ddg_log_counter),
+                        "model_training": bool(self.training),
+                        "normalize_difference": bool(getattr(self, "normalize_difference", False)),
+                        "ddg_signal_gain": float(self.ddg_signal_gain.detach().item()) if hasattr(self, "ddg_signal_gain") else None,
+                        "diff_features_mean": df_mean,
+                        "diff_features_std": df_std,
+                        "diff_features_abs_mean": df_abs_mean,
+                        "diff_features_per_sample_std_mean": df_per_sample_std,
+                        "ddg_pred_mean": p_mean,
+                        "ddg_pred_std": p_std,
+                        "ddg_pred2_std": p2_std,
+                        "ddg_pred_repeat_diff_std": p_diff_std,
+                        "ddg_w0_norm": w0_norm,
+                        "ddg_wL_norm": wL_norm,
+                        "ddg_bL_norm": bL_norm,
+                    },
+                    "timestamp": int(time.time() * 1000),
+                }
+                # Log first (before file write that may fail on cluster)
+                logger.info(
+                    f"[AGENTLOG DDGHEAD] train={self.training} df_std={df_std:.4f} df_ps_std={df_per_sample_std:.4f} "
+                    f"pred_std={p_std:.4f} pred_mean={p_mean:.4f} rep_diff_std={p_diff_std if p_diff_std is not None else 'NA'} "
+                    f"w0={w0_norm:.2f} wL={wL_norm:.2f} bL={bL_norm if bL_norm is not None else 'NA'}"
+                )
+
+                # Layerwise activation trace to locate where variance collapses (ReLU dead / dropout-only variance)
+                layer_stats = []
+                try:
+                    with torch.no_grad():
+                        x = diff_features.detach()
+                        for li, layer in enumerate(self.ddg_hidden):
+                            x = layer(x)
+                            st = {
+                                "i": int(li),
+                                "t": layer.__class__.__name__,
+                                "mean": float(x.mean().item()) if x.numel() else None,
+                                "std": float(x.std().item()) if x.numel() else None,
+                            }
+                            if isinstance(layer, nn.ReLU):
+                                st["zero_frac"] = float((x == 0).float().mean().item()) if x.numel() else None
+                            layer_stats.append(st)
+                    # Compact summary for logs (first 2 + last 2 layers)
+                    compact = (layer_stats[:2] + (["..."] if len(layer_stats) > 4 else []) + layer_stats[-2:])
+                    logger.info(f"[AGENTLOG DDGACT] train={self.training} layers={compact}")
+                except Exception:
+                    layer_stats = []
+                # File write may fail on cluster - that's OK
+                try:
+                    payload["data"]["ddg_layer_stats"] = layer_stats
+                    with open("/Users/supantha/Documents/code_v2/protein/.cursor/debug.log", "a") as f:
+                        f.write(json.dumps(payload, default=str) + "\n")
+                except Exception:
+                    pass
+        except Exception:
+            pass
+        #endregion
+
+        # DEBUG: Check ddg_pred
+        if torch.isnan(ddg_pred).any() or torch.isinf(ddg_pred).any():
+             print(f"[DEBUG MODEL] NaN/Inf in ddg_pred!")
+        
+        return dg_pred, ddg_pred
+
+
+class DualHeadAffinityPredictionModel(pl.LightningModule):
+    """
+    Lightning wrapper for dual-head training.
+    """
+    def __init__(self, predictor: DualHeadAffinityPredictor, learning_rate=1e-4, ddg_loss_weight=1.0):
+        super().__init__()
+        self.predictor = predictor
+        self.learning_rate = learning_rate
+        self.ddg_loss_weight = ddg_loss_weight
+
+        self.loss_fn = nn.MSELoss()
+        self.pearson_corr = PearsonCorrCoef()
+        self.spearman_corr = SpearmanCorrCoef()
+        self.r2_score = R2Score()
+        self.mse_metric = MeanSquaredError()
+        
+        # Save hyperparameters for checkpointing
+        self.save_hyperparameters(ignore=['predictor'])
+
+    def forward(self, mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask, 
+                wt_chain1=None, wt_chain1_mask=None, wt_chain2=None, wt_chain2_mask=None,
+                source_type_ids=None):
+        return self.predictor(mut_chain1, mut_chain1_mask, mut_chain2, mut_chain2_mask,
+                             wt_chain1, wt_chain1_mask, wt_chain2, wt_chain2_mask,
+                             source_type_ids=source_type_ids)
+
+    def training_step(self, batch, batch_idx):
+        # Mutant data
+        (c1, m1, c2, m2, y_mut) = batch["mutant"]
+        
+        # Wildtype data with valid mask
+        (cw1, w1m, cw2, w2m, y_wt) = batch["wildtype"]
+        has_wt = batch["has_wt"]
+        
+        if self.predictor.use_dual_head and has_wt.sum() > 0:
+            # For samples with wildtype available, use dual-head prediction
+            valid_samples = has_wt.bool()
+            
+            # Get predictions for valid samples
+            dg_pred, ddg_pred = self(
+                c1[valid_samples], m1[valid_samples], 
+                c2[valid_samples], m2[valid_samples],
+                cw1[valid_samples], w1m[valid_samples], 
+                cw2[valid_samples], w2m[valid_samples]
+            )
+            
+            # Calculate losses for valid samples
+            dg_loss = self.loss_fn(dg_pred, y_mut[valid_samples])
+            
+            # Calculate true ddG as difference between mutant and wildtype dG
+            true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
+            ddg_loss = self.loss_fn(ddg_pred, true_ddg)
+            
+            # Combined loss with weighting
+            loss = dg_loss + self.ddg_loss_weight * ddg_loss
+            
+            # Process remaining samples (without wildtype) with standard prediction
+            if (~valid_samples).sum() > 0:
+                standard_dg_pred = self(
+                    c1[~valid_samples], m1[~valid_samples], 
+                    c2[~valid_samples], m2[~valid_samples]
+                )
+                standard_loss = self.loss_fn(standard_dg_pred, y_mut[~valid_samples])
+                
+                # Add to total loss, weighted by proportion of samples
+                n_valid = valid_samples.sum()
+                n_total = len(valid_samples)
+                loss = (n_valid / n_total) * loss + ((n_total - n_valid) / n_total) * standard_loss
+        else:
+            # Standard prediction for all samples
+            dg_pred = self(c1, m1, c2, m2)
+            loss = self.loss_fn(dg_pred, y_mut)
+        
+        # Log metrics
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        # Similar to training_step but with more comprehensive metrics
+        (c1, m1, c2, m2, y_mut) = batch["mutant"]
+        (cw1, w1m, cw2, w2m, y_wt) = batch["wildtype"]
+        has_wt = batch["has_wt"]
+        
+        # Store predictions and targets for all samples
+        all_dg_preds = []
+        all_dg_targets = []
+        all_ddg_preds = []
+        all_ddg_targets = []
+        
+        if self.predictor.use_dual_head and has_wt.sum() > 0:
+            valid_samples = has_wt.bool()
+            
+            # Dual-head prediction for samples with wildtype
+            dg_pred, ddg_pred = self(
+                c1[valid_samples], m1[valid_samples], 
+                c2[valid_samples], m2[valid_samples],
+                cw1[valid_samples], w1m[valid_samples], 
+                cw2[valid_samples], w2m[valid_samples]
+            )
+            
+            # Calculate true ddG
+            true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
+            
+            # Store predictions and targets
+            all_dg_preds.append(dg_pred)
+            all_dg_targets.append(y_mut[valid_samples])
+            all_ddg_preds.append(ddg_pred)
+            all_ddg_targets.append(true_ddg)
+            
+            # Process remaining samples with standard prediction
+            if (~valid_samples).sum() > 0:
+                standard_dg_pred = self(
+                    c1[~valid_samples], m1[~valid_samples], 
+                    c2[~valid_samples], m2[~valid_samples]
+                )
+                all_dg_preds.append(standard_dg_pred)
+                all_dg_targets.append(y_mut[~valid_samples])
+        else:
+            # Standard prediction for all samples
+            dg_pred = self(c1, m1, c2, m2)
+            all_dg_preds.append(dg_pred)
+            all_dg_targets.append(y_mut)
+            
+            # For samples with wildtype, calculate implicit ddG
+            if has_wt.sum() > 0:
+                valid_samples = has_wt.bool()
+                wt_dg_pred = self(cw1[valid_samples], w1m[valid_samples], 
+                                 cw2[valid_samples], w2m[valid_samples])
+                
+                implicit_ddg_pred = dg_pred[valid_samples] - wt_dg_pred
+                true_ddg = y_mut[valid_samples] - y_wt[valid_samples]
+                
+                all_ddg_preds.append(implicit_ddg_pred)
+                all_ddg_targets.append(true_ddg)
+        
+        # Concatenate all predictions and targets
+        if all_dg_preds:
+            all_dg_preds = torch.cat(all_dg_preds)
+            all_dg_targets = torch.cat(all_dg_targets)
+            
+            # Calculate dG metrics
+            dg_mse = self.mse_metric(all_dg_preds, all_dg_targets)
+            dg_pearson = self.pearson_corr(all_dg_preds, all_dg_targets)
+            dg_spearman = self.spearman_corr(all_dg_preds, all_dg_targets)
+            dg_r2 = self.r2_score(all_dg_preds, all_dg_targets)
+            
+            # Log dG metrics
+            self.log('val_dg_mse', dg_mse, on_epoch=True, prog_bar=True)
+            self.log('val_dg_pearson', dg_pearson, on_epoch=True)
+            self.log('val_dg_spearman', dg_spearman, on_epoch=True)
+            self.log('val_dg_r2', dg_r2, on_epoch=True)
+        
+        # Calculate ddG metrics if available
+        if all_ddg_preds:
+            all_ddg_preds = torch.cat(all_ddg_preds)
+            all_ddg_targets = torch.cat(all_ddg_targets)
+            
+            ddg_mse = self.mse_metric(all_ddg_preds, all_ddg_targets)
+            ddg_pearson = self.pearson_corr(all_ddg_preds, all_ddg_targets)
+            ddg_spearman = self.spearman_corr(all_ddg_preds, all_ddg_targets)
+            ddg_r2 = self.r2_score(all_ddg_preds, all_ddg_targets)
+            
+            # Log ddG metrics
+            self.log('val_ddg_mse', ddg_mse, on_epoch=True, prog_bar=True)
+            self.log('val_ddg_pearson', ddg_pearson, on_epoch=True)
+            self.log('val_ddg_spearman', ddg_spearman, on_epoch=True)
+            self.log('val_ddg_r2', ddg_r2, on_epoch=True)
+            
+            # Combined validation metric for early stopping
+            combined_metric = dg_mse + self.ddg_loss_weight * ddg_mse
+            self.log('val_combined_metric', combined_metric, on_epoch=True)
+        
+        return {'val_dg_mse': dg_mse if 'dg_mse' in locals() else None,
+                'val_ddg_mse': ddg_mse if 'ddg_mse' in locals() else None}
+
+    def test_step(self, batch, batch_idx):
+        # Similar to validation_step but returns more detailed metrics
+        return self.validation_step(batch, batch_idx)
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.AdamW(self.parameters(), lr=self.learning_rate)
+        scheduler = torch.optim.lr_scheduler.OneCycleLR(
+            optimizer, max_lr=self.learning_rate, total_steps=self.trainer.estimated_stepping_batches
+        )
+        return {"optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val_combined_metric"}