""" AFFINose interaction model — cross-attention with live BERTose encoding Architecture: GLYCAN: WURCS → BPE → BERTose (live, freeze layers 0-3) → [B, Lg, 768] ↓ proj(768→512) PROTEIN: precomputed ESM-C → [B, Lp, 960] ↓ ↓ proj(960→512) ↓ 2× CrossAttentionBlock(d=512, 8 heads, FFN=1024) ↓ ↓ ↓ SHARED mask-aware SWE(d=512, S=512, R=64) ↓ ↓ ↓ [B, 512] [B, 512] ↓ element-wise product + sum [B, 1024] ↓ MLP → binding score This release exposes the manuscript-facing AFFINose architecture: BERTose glycan tokens, per-residue ESM-C protein embeddings, bidirectional cross-attention, pooled fusion and scalar interaction scoring. """ import os import sys import math from pathlib import Path from typing import Dict, Optional, Tuple import torch import torch.nn as nn import torch.nn.functional as F # ============================================================================ # BERTose model imports # ============================================================================ def _default_bertose_root() -> Path: """Resolve the BERTose source root without assuming a specific local path.""" env_root = os.environ.get("BERTOSE_ROOT") or os.environ.get("BERTOSE_REPO_ROOT") if env_root: return Path(env_root).expanduser().resolve() here = Path(__file__).resolve() for parent in here.parents: if (parent / "src").exists() or (parent / "bertose_model.py").exists(): return parent return here.parent BERTOSE_ROOT = _default_bertose_root() def _ensure_bertose_imports(): """Add BERTose source directories to sys.path if not already present.""" source_dir = Path(__file__).resolve().parent roots = [ str(source_dir), str(BERTOSE_ROOT), str(BERTOSE_ROOT / "src"), ] for root in roots: if root not in sys.path: sys.path.insert(0, root) def load_bertose_config(): """Create BERTose config matching the BERTose glycan encoder checkpoint.""" _ensure_bertose_imports() try: from model.bertose_model import MultimodalGlycanBERTConfig except ModuleNotFoundError: from bertose_model import MultimodalGlycanBERTConfig return MultimodalGlycanBERTConfig( seq_vocab_size=2200, use_cnn_frontend=True, ) def load_bertose_encoder( checkpoint_path: str, freeze_layers: int = 4 ): """ Load BERTose sequence encoder with pretrained weights. Args: checkpoint_path: Path to pretrained BERTose checkpoint. freeze_layers: Number of transformer layers to freeze (0-indexed). Returns: Tuple of (bertose_config, seq_embeddings, seq_layers). """ _ensure_bertose_imports() try: from model.bertose_model import ( MultimodalGlycanBERT, MultimodalGlycanBERTConfig, ) except ModuleNotFoundError: from bertose_model import ( MultimodalGlycanBERT, MultimodalGlycanBERTConfig, ) # Load checkpoint ckpt = torch.load(checkpoint_path, map_location="cpu") state_dict = ckpt.get("model_state_dict", ckpt) # Infer vocab size and max position embeddings from checkpoint vocab_size = state_dict["seq_embeddings.token_embeddings.weight"].shape[0] max_pos = state_dict["seq_embeddings.position_embeddings.weight"].shape[0] config = MultimodalGlycanBERTConfig( seq_vocab_size=vocab_size, seq_max_length=max_pos, use_cnn_frontend=True, ) # Instantiate full model, then extract sequence encoder model = MultimodalGlycanBERT(config) missing, unexpected = model.load_state_dict(state_dict, strict=False) loaded = len(state_dict) - len(unexpected) print(f" Loaded {loaded}/{len(state_dict)} pretrained weight tensors") print(f" ({len(missing)} missing in checkpoint, {len(unexpected)} unexpected)") seq_embeddings = model.seq_embeddings seq_layers = model.seq_layers # Freeze embedding layer + first N transformer layers for param in seq_embeddings.parameters(): param.requires_grad = False for i in range(min(freeze_layers, len(seq_layers))): for param in seq_layers[i].parameters(): param.requires_grad = False trainable = sum( p.numel() for p in seq_embeddings.parameters() if p.requires_grad ) trainable += sum( p.numel() for layer in seq_layers for p in layer.parameters() if p.requires_grad ) total = sum(p.numel() for p in seq_embeddings.parameters()) total += sum( p.numel() for layer in seq_layers for p in layer.parameters() ) print( f" BERTose encoder: {total:,} params total, " f"{trainable:,} trainable (frozen layers 0-{freeze_layers - 1})" ) return config, seq_embeddings, seq_layers # ============================================================================ # Differentiable interpolation # ============================================================================ def differentiable_interp1d( x: torch.Tensor, y: torch.Tensor, xnew: torch.Tensor ) -> torch.Tensor: """ Fully differentiable 1D linear interpolation. Gradients flow through y (values) back to theta projection and earlier layers. Args: x: [B, N] sorted input coordinates (detached) y: [B, N] values at x positions (REQUIRES grad flow!) xnew: [B, R] query coordinates Returns: [B, R] interpolated values """ n_pts = x.shape[1] # Find interpolation indices ind = torch.searchsorted( x.contiguous().detach(), xnew.contiguous().detach() ) ind = ind.clamp(1, n_pts - 1) # Gather neighbor values — preserves gradient flow through y x_lo = torch.gather(x, 1, ind - 1) x_hi = torch.gather(x, 1, ind) y_lo = torch.gather(y, 1, ind - 1) y_hi = torch.gather(y, 1, ind) # Linear interpolation weights denom = (x_hi - x_lo).clamp(min=1e-8) alpha = ((xnew - x_lo) / denom).clamp(0, 1) # Interpolated value — fully differentiable w.r.t. y_lo and y_hi return y_lo + alpha * (y_hi - y_lo) # ============================================================================ # SWE pooling # ============================================================================ class SWE_Pooling(nn.Module): """ Sliced-Wasserstein Embedding pooling. Maps variable-length token embeddings [B, L, d_in] => [B, num_slices]. Includes mask-aware sorting and degenerate-sample handling. """ def __init__( self, d_in: int, num_slices: int, num_ref_points: int, freeze_swe: bool = False, ): super().__init__() self.num_slices = num_slices self.num_ref_points = num_ref_points # Learnable reference distribution ref = torch.linspace(-1, 1, num_ref_points).unsqueeze(1).repeat( 1, num_slices ) self.reference = nn.Parameter(ref, requires_grad=not freeze_swe) # Projection directions (weight-normalized) 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) # Weighted aggregation over reference points 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: torch.Tensor, mask: Optional[torch.Tensor] = None ) -> torch.Tensor: """ Args: x: [B, L, d_in] token embeddings mask: [B, L] attention mask (1=valid, 0=padding) Returns: [B, num_slices] fixed-length representation """ batch_size, seq_len, _ = x.shape device = x.device # Degenerate: single token → just project if seq_len == 1: x_slices = self.theta(x) return x_slices.squeeze(1) # Project onto learned directions x_slices = self.theta(x) # [B, L, num_slices] # MASK-AWARE SORTING: set padding → -inf so they sort to bottom if mask is not None: mask_exp = mask.unsqueeze(-1).expand_as(x_slices) x_slices = x_slices.masked_fill(mask_exp == 0, float("-inf")) x_sorted, _ = torch.sort(x_slices, dim=1) # Strip padding from sorted array if mask is not None: valid_counts = mask.sum(dim=1).long() degenerate_mask = valid_counts < 2 safe_counts = valid_counts.clamp(min=2) max_valid = safe_counts.max().item() # Vectorized gather: take last safe_count values starts = (seq_len - safe_counts).unsqueeze(1) offsets = torch.arange(max_valid, device=device).unsqueeze(0) raw_idx = starts + offsets gather_idx = raw_idx.clamp(max=seq_len - 1).unsqueeze(-1).expand( batch_size, max_valid, self.num_slices ) x_sorted = torch.gather(x_sorted, 1, gather_idx) # Replace -inf with 0.0 for degenerate samples x_sorted = x_sorted.masked_fill( x_sorted == float("-inf"), 0.0 ) n_eff = max_valid else: degenerate_mask = None n_eff = seq_len # Interpolate to fixed reference grid x_coord = ( torch.linspace(0, 1, n_eff, device=device) .unsqueeze(0) .expand(batch_size * self.num_slices, -1) ) x_flat = x_sorted.permute(0, 2, 1).reshape( batch_size * self.num_slices, n_eff ) xnew = ( torch.linspace(0, 1, self.num_ref_points, device=device) .unsqueeze(0) .expand(batch_size * self.num_slices, -1) ) y_intp = differentiable_interp1d(x_coord, x_flat, xnew) x_interp = y_intp.view( batch_size, self.num_slices, self.num_ref_points ).permute(0, 2, 1) # Compare with reference distribution r_expanded = self.reference.expand_as(x_interp) embeddings = (r_expanded - x_interp).permute(0, 2, 1) # Weighted aggregation → [B, num_slices] weighted = self.weight(embeddings).sum(dim=-1) # Zero out degenerate samples if degenerate_mask is not None and degenerate_mask.any(): weighted = weighted.masked_fill( degenerate_mask.unsqueeze(-1), 0.0 ) return weighted # ============================================================================ # Cross-attention block # ============================================================================ class CrossAttentionBlock(nn.Module): """ Bidirectional cross-attention between glycan and protein tokens. Glycan tokens attend to protein residues (Q=glycan, KV=protein) Protein residues attend to glycan tokens (Q=protein, KV=glycan) Includes NaN guard for all-masked keys and padding-position zeroing. """ def __init__( self, d_model: int, num_heads: int, ffn_dim: int, dropout: float = 0.1, ): super().__init__() # Glycan → Protein cross-attention self.glycan_cross_attn = nn.MultiheadAttention( d_model, num_heads, dropout=dropout, batch_first=True ) self.glycan_norm1 = nn.LayerNorm(d_model) self.glycan_ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, d_model), nn.Dropout(dropout), ) self.glycan_norm2 = nn.LayerNorm(d_model) # Protein → Glycan cross-attention self.protein_cross_attn = nn.MultiheadAttention( d_model, num_heads, dropout=dropout, batch_first=True ) self.protein_norm1 = nn.LayerNorm(d_model) self.protein_ffn = nn.Sequential( nn.Linear(d_model, ffn_dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(ffn_dim, d_model), nn.Dropout(dropout), ) self.protein_norm2 = nn.LayerNorm(d_model) def forward( self, glycan: torch.Tensor, protein: torch.Tensor, glycan_mask: Optional[torch.Tensor] = None, protein_mask: Optional[torch.Tensor] = None, return_attention: bool = False, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Returns updated (glycan, protein) enriched with cross-modal info. NaN guard: nn.MultiheadAttention produces NaN when ALL keys are masked. We replace NaN→0 so residual preserves query. """ # Convert to key_padding_mask (True = padded) g_key_pad = (~glycan_mask.bool()) if glycan_mask is not None else None p_key_pad = ( (~protein_mask.bool()) if protein_mask is not None else None ) # Glycan attends to protein g_cross, g_attn_weights = self.glycan_cross_attn( query=glycan, key=protein, value=protein, key_padding_mask=p_key_pad, need_weights=return_attention, average_attn_weights=False, ) g_cross = torch.nan_to_num(g_cross, nan=0.0) glycan = self.glycan_norm1(glycan + g_cross) glycan = self.glycan_norm2(glycan + self.glycan_ffn(glycan)) if glycan_mask is not None: glycan = glycan * glycan_mask.unsqueeze(-1) # Protein attends to glycan p_cross, p_attn_weights = self.protein_cross_attn( query=protein, key=glycan, value=glycan, key_padding_mask=g_key_pad, need_weights=return_attention, average_attn_weights=False, ) p_cross = torch.nan_to_num(p_cross, nan=0.0) protein = self.protein_norm1(protein + p_cross) protein = self.protein_norm2(protein + self.protein_ffn(protein)) if protein_mask is not None: protein = protein * protein_mask.unsqueeze(-1) if return_attention: attn_dict = { "glycan_to_protein": g_attn_weights, "protein_to_glycan": p_attn_weights, } return glycan, protein, attn_dict return glycan, protein # ============================================================================ # AFFINose interaction model # ============================================================================ class AffinoseInteractionModel(nn.Module): """ Glycan-protein interaction predictor with cross-attention. Glycan: Live BERTose (partially frozen) → per-token [B, Lg, 768] Protein: Precomputed ESM-C per-residue [B, Lp, 960] Cross-attention: 2 bidirectional layers in shared 512-dim space SWE: Variable-length → fixed [B, 512] for each side Interaction: product + sum → MLP → scalar """ def __init__( self, seq_embeddings: nn.Module, seq_layers: nn.ModuleList, glycan_dim: int = 768, protein_dim: int = 960, shared_dim: int = 512, num_cross_layers: int = 2, num_heads: int = 8, ffn_dim: int = 1024, swe_slices: int = 512, swe_ref_points: int = 64, head_hidden: int = 256, dropout: float = 0.1, separate_swe: bool = False, pooling_mode: str = "swe", interaction_mode: str = "product_sum", use_cross_attention: bool = True, ): """ Args: seq_embeddings: Pretrained BERTose embedding layer. seq_layers: Pretrained BERTose transformer layers. glycan_dim: BERTose output dimension (768). protein_dim: ESM-C per-residue dimension (960). shared_dim: Shared space for cross-attention (512). num_cross_layers: Number of cross-attention blocks. num_heads: Attention heads per block. ffn_dim: FFN hidden dim in cross-attention. swe_slices: Number of SWE projection directions. swe_ref_points: Number of SWE reference distribution points. head_hidden: MLP head hidden dimension. dropout: Dropout rate. separate_swe: If True, use independent SWE modules. pooling_mode: 'swe', 'mean', or 'joint_swe'. interaction_mode: 'product_sum' or 'concat'. use_cross_attention: If False, skip cross-attention. """ super().__init__() self.separate_swe = separate_swe self.pooling_mode = pooling_mode self.interaction_mode = interaction_mode self.use_cross_attention = use_cross_attention print(f" Architecture config:") print(f" cross_attention={use_cross_attention}") print(f" pooling_mode={pooling_mode}") print(f" interaction_mode={interaction_mode}") # === BERTose sequence encoder (partially frozen) === self.seq_embeddings = seq_embeddings self.seq_layers = seq_layers # === Projection to shared space === self.glycan_proj = nn.Sequential( nn.Linear(glycan_dim, shared_dim), nn.LayerNorm(shared_dim), ) self.protein_proj = nn.Sequential( nn.Linear(protein_dim, shared_dim), nn.LayerNorm(shared_dim), ) # === Cross-attention stack (optional) === if use_cross_attention: self.cross_attention = nn.ModuleList([ CrossAttentionBlock( d_model=shared_dim, num_heads=num_heads, ffn_dim=ffn_dim, dropout=dropout, ) for _ in range(num_cross_layers) ]) else: self.cross_attention = nn.ModuleList() # === Pooling === if pooling_mode == "swe": if separate_swe: self.swe_glycan = SWE_Pooling( d_in=shared_dim, num_slices=swe_slices, num_ref_points=swe_ref_points, ) self.swe_protein = SWE_Pooling( d_in=shared_dim, num_slices=swe_slices, num_ref_points=swe_ref_points, ) pool_out_dim = swe_slices else: self.swe = SWE_Pooling( d_in=shared_dim, num_slices=swe_slices, num_ref_points=swe_ref_points, ) pool_out_dim = swe_slices elif pooling_mode == "mean": pool_out_dim = shared_dim elif pooling_mode == "joint_swe": self.swe_joint = SWE_Pooling( d_in=shared_dim, num_slices=swe_slices, num_ref_points=swe_ref_points, ) pool_out_dim = swe_slices # === Regression head === if pooling_mode == "joint_swe": head_input_dim = pool_out_dim else: head_input_dim = 2 * pool_out_dim self.head = nn.Sequential( nn.Linear(head_input_dim, head_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_hidden, head_hidden), nn.GELU(), nn.Dropout(dropout), nn.Linear(head_hidden, 1), ) # Initialize (skip SWE weight-normed params) self.apply(self._init_weights) self._count_params() def _init_weights(self, module: nn.Module) -> None: """Xavier init for Linear, skip weight-normed SWE modules.""" if isinstance(module, nn.Linear): if hasattr(module, "weight_v"): return # Preserve SWE initialization nn.init.xavier_uniform_(module.weight) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) def _count_params(self) -> None: """Log parameter counts.""" total = sum(p.numel() for p in self.parameters()) trainable = sum( p.numel() for p in self.parameters() if p.requires_grad ) print(f"AffinoseInteractionModel: {total:,} total, {trainable:,} trainable") def _masked_mean_pool(self, x, mask): """Masked mean pooling: average valid tokens only.""" mask_expanded = mask.unsqueeze(-1) x_masked = x * mask_expanded summed = x_masked.sum(dim=1) counts = mask.sum(dim=1, keepdim=True).clamp(min=1) return summed / counts def forward( self, token_ids: torch.Tensor, attention_mask: torch.Tensor, branch_depths: torch.Tensor, linkage_types: torch.Tensor, protein_emb: torch.Tensor, protein_mask: torch.Tensor, log_conc: Optional[torch.Tensor] = None, has_conc: Optional[torch.Tensor] = None, return_attention: bool = False, ) -> torch.Tensor: """ Forward pass with cross-attention. Args: token_ids: [B, Lg] BPE token IDs. attention_mask: [B, Lg] glycan attention mask (1=valid, 0=pad). branch_depths: [B, Lg] branch depth per token. linkage_types: [B, Lg] linkage type per token. protein_emb: [B, Lp, protein_dim] per-residue ESM-C embeddings. protein_mask: [B, Lp] protein attention mask (1=valid, 0=pad). Returns: [B] binding score predictions. """ # === 1. BERTose forward: per-token embeddings === x = self.seq_embeddings(token_ids, branch_depths, linkage_types) for layer in self.seq_layers: x = layer(x, attention_mask) # x: [B, Lg, 768] — all tokens (not just CLS!) # Glycan mask: use the attention_mask from BPE tokenizer glycan_mask = attention_mask # [B, Lg], 1=valid, 0=pad # === 2. Project to shared space === glycan = self.glycan_proj(x) # [B, Lg, 512] protein = self.protein_proj(protein_emb) # [B, Lp, 512] # Zero padding positions after projection: # LayerNorm bias produces non-trivial values at padding) glycan = glycan * glycan_mask.unsqueeze(-1) protein = protein * protein_mask.unsqueeze(-1) # === 3. Cross-attention (optional) === all_attention_maps = [] if self.use_cross_attention: for cross_layer in self.cross_attention: if return_attention: glycan, protein, attn_dict = cross_layer( glycan, protein, glycan_mask, protein_mask, return_attention=True, ) all_attention_maps.append(attn_dict) else: glycan, protein = cross_layer( glycan, protein, glycan_mask, protein_mask ) # === 4. Pooling === if self.pooling_mode == "joint_swe": joint_tokens = torch.cat([glycan, protein], dim=1) joint_mask = torch.cat([glycan_mask, protein_mask], dim=1) pooled = self.swe_joint(joint_tokens, joint_mask) return self.head(pooled).squeeze(-1) if self.pooling_mode == "swe": if self.separate_swe: glycan_pooled = self.swe_glycan(glycan, glycan_mask) protein_pooled = self.swe_protein(protein, protein_mask) else: glycan_pooled = self.swe(glycan, glycan_mask) protein_pooled = self.swe(protein, protein_mask) elif self.pooling_mode == "mean": glycan_pooled = self._masked_mean_pool(glycan, glycan_mask) protein_pooled = self._masked_mean_pool(protein, protein_mask) # === 5. Interaction === if self.interaction_mode == "product_sum": interaction = torch.cat([ glycan_pooled * protein_pooled, glycan_pooled + protein_pooled, ], dim=-1) elif self.interaction_mode == "concat": interaction = torch.cat([ glycan_pooled, protein_pooled, ], dim=-1) # === 6. Predict binding score === out = self.head(interaction).squeeze(-1) if return_attention and all_attention_maps: return out, all_attention_maps return out # ============================================================================ # Loss # ============================================================================ class AffinoseInteractionLoss(nn.Module): """MSE loss for regression.""" def __init__(self): super().__init__() self.mse = nn.MSELoss() def forward( self, pred: torch.Tensor, target: torch.Tensor ) -> torch.Tensor: """Compute MSE loss.""" return self.mse(pred, target) # ============================================================================ # Sanity check # ============================================================================ if __name__ == "__main__": print("=" * 60) print("AFFINose interaction model architecture sanity check") print("=" * 60) # Mock BERTose encoder (for testing without cluster) class MockEmbeddings(nn.Module): """Mock BERTose embeddings for local testing.""" def __init__(self, dim: int = 768): super().__init__() self.proj = nn.Linear(64, dim) def forward(self, token_ids, branch_depths, linkage_types): """Return random embeddings.""" batch_size, seq_len = token_ids.shape return torch.randn(batch_size, seq_len, 768) class MockLayer(nn.Module): """Mock transformer layer.""" def forward(self, x, mask): """Identity.""" return x seq_emb = MockEmbeddings() seq_layers = nn.ModuleList([MockLayer() for _ in range(12)]) model = AffinoseInteractionModel( seq_embeddings=seq_emb, seq_layers=seq_layers, glycan_dim=768, protein_dim=960, ) # Simulate batch batch_size = 4 lg = 37 # Glycan: 37 BPE tokens lp = 150 # Protein: 150 residues token_ids = torch.randint(0, 100, (batch_size, lg)) attention_mask = torch.ones(batch_size, lg).float() branch_depths = torch.zeros(batch_size, lg, dtype=torch.long) linkage_types = torch.zeros(batch_size, lg, dtype=torch.long) protein_emb = torch.randn(batch_size, lp, 960) protein_mask = torch.ones(batch_size, lp).float() out = model( token_ids=token_ids, attention_mask=attention_mask, branch_depths=branch_depths, linkage_types=linkage_types, protein_emb=protein_emb, protein_mask=protein_mask, ) print(f"\nInput shapes:") print(f" Glycan tokens: {token_ids.shape}") print(f" Protein emb: {protein_emb.shape}") print(f"\nOutput shape: {out.shape} — values: {out.detach()}") # Test loss loss_fn = AffinoseInteractionLoss() target = torch.rand(batch_size) loss = loss_fn(out, target) print(f"\nLoss: {loss.item():.6f}") # Test backward loss.backward() grad_count = sum( 1 for p in model.parameters() if p.grad is not None and p.requires_grad ) total_trainable = sum( 1 for p in model.parameters() if p.requires_grad ) print(f"Gradients: {grad_count}/{total_trainable} trainable params") print(f"\nAffinose sanity check passed.")