Spaces:
Sleeping
Sleeping
File size: 7,194 Bytes
5b3b0dc 8f66d04 5b3b0dc 8f66d04 5b3b0dc 8f66d04 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """Cross-attention module for model-native atom--residue attribution.
The module returns an attention matrix ``M`` over drug atoms and protein
residues. Its values are learned model weights that help inspect a prediction;
they are not molecular contacts, a binding pocket, or a structural mechanism.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class CrossAttentionInteraction(nn.Module):
"""
Multi-head cross-attention between drug atoms and protein residues.
Returns:
- Fused representation for prediction
- Attention weights for model-native atom--residue attribution
"""
def __init__(self,
hidden_dim: int = 256,
num_heads: int = 8,
dropout: float = 0.1):
super().__init__()
assert hidden_dim % num_heads == 0
self.hidden_dim = hidden_dim
self.num_heads = num_heads
self.head_dim = hidden_dim // num_heads
self.scale = math.sqrt(self.head_dim)
# drug atoms attend to protein residues
self.W_q_drug = nn.Linear(hidden_dim, hidden_dim)
self.W_k_prot = nn.Linear(hidden_dim, hidden_dim)
self.W_v_prot = nn.Linear(hidden_dim, hidden_dim)
# protein residues attend to drug atoms (bidirectional)
self.W_q_prot = nn.Linear(hidden_dim, hidden_dim)
self.W_k_drug = nn.Linear(hidden_dim, hidden_dim)
self.W_v_drug = nn.Linear(hidden_dim, hidden_dim)
# output projections
self.out_proj_drug = nn.Linear(hidden_dim, hidden_dim)
self.out_proj_prot = nn.Linear(hidden_dim, hidden_dim)
# layer norms
self.ln_drug = nn.LayerNorm(hidden_dim)
self.ln_prot = nn.LayerNorm(hidden_dim)
self.dropout = nn.Dropout(dropout)
def _attention(self, Q, K, V, mask=None):
"""
Standard scaled dot-product attention.
Returns:
output: attended values
attn_weights: softmax attention weights (for interpretability)
"""
# Q: (B, H, Lq, d), K: (B, H, Lk, d), V: (B, H, Lk, d)
scores = torch.matmul(Q, K.transpose(-2, -1)) / self.scale # (B,H,Lq,Lk)
if mask is not None:
# mask shape: (B, 1, 1, Lk) — broadcast over heads and queries
scores = scores.masked_fill(~mask, float('-inf'))
attn_weights = F.softmax(scores, dim=-1)
attn_weights = self.dropout(attn_weights)
output = torch.matmul(attn_weights, V) # (B, H, Lq, d)
return output, attn_weights
def _reshape_to_heads(self, x, batch_size):
"""(B, L, D) → (B, H, L, d)"""
return x.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
def forward(self,
drug_repr: torch.Tensor,
drug_mask: torch.Tensor,
prot_repr: torch.Tensor,
prot_mask: torch.Tensor):
"""
Bidirectional cross-attention between drug atoms and protein residues.
Args:
drug_repr: (B, N_atoms, D) — per-atom drug representations
drug_mask: (B, N_atoms) — boolean mask for drug atoms
prot_repr: (B, L_residues, D) — per-residue protein representations
prot_mask: (B, L_residues) — boolean mask for protein residues
Returns:
drug_updated: (B, N_atoms, D) — drug repr enriched by protein context
prot_updated: (B, L_residues, D) — protein repr enriched by drug context
interaction_map: (B, N_atoms, L_residues) — attention-based interaction matrix
this is the key output for interpretability analysis
"""
B = drug_repr.size(0)
# --- Drug → Protein attention ---
Q_d = self._reshape_to_heads(self.W_q_drug(drug_repr), B)
K_p = self._reshape_to_heads(self.W_k_prot(prot_repr), B)
V_p = self._reshape_to_heads(self.W_v_prot(prot_repr), B)
# mask: (B, L_residues) → (B, 1, 1, L_residues)
prot_attn_mask = prot_mask.unsqueeze(1).unsqueeze(2) if prot_mask is not None else None
drug_attended, drug_to_prot_attn = self._attention(Q_d, K_p, V_p, prot_attn_mask)
# drug_to_prot_attn: (B, H, N_atoms, L_residues)
drug_attended = drug_attended.transpose(1, 2).contiguous().view(B, -1, self.hidden_dim)
drug_attended = self.out_proj_drug(drug_attended)
drug_updated = self.ln_drug(drug_repr + drug_attended)
# --- Protein → Drug attention ---
Q_p = self._reshape_to_heads(self.W_q_prot(prot_repr), B)
K_d = self._reshape_to_heads(self.W_k_drug(drug_repr), B)
V_d = self._reshape_to_heads(self.W_v_drug(drug_repr), B)
drug_attn_mask = drug_mask.unsqueeze(1).unsqueeze(2) if drug_mask is not None else None
prot_attended, _ = self._attention(Q_p, K_d, V_d, drug_attn_mask)
prot_attended = prot_attended.transpose(1, 2).contiguous().view(B, -1, self.hidden_dim)
prot_attended = self.out_proj_prot(prot_attended)
prot_updated = self.ln_prot(prot_repr + prot_attended)
# --- Interaction map (averaged over heads) ---
# This is what we visualise and validate against binding sites
interaction_map = drug_to_prot_attn.mean(dim=1) # (B, N_atoms, L_residues)
return drug_updated, prot_updated, interaction_map
class GatedPooling(nn.Module):
"""
Gated pooling for aggregating atom/residue-level features into
a fixed-size vector for prediction.
Instead of mean/max pooling, we learn which atoms and residues
are most important for the final prediction. The gate weights
are themselves interpretable signals.
"""
def __init__(self, hidden_dim: int):
super().__init__()
self.gate = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
nn.Linear(hidden_dim, 1),
)
self.transform = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.Tanh(),
)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
"""
Args:
x: (B, L, D) — sequence of feature vectors
mask: (B, L) — boolean mask
Returns:
pooled: (B, D) — single feature vector per sample
"""
gate_scores = self.gate(x).squeeze(-1) # (B, L)
if mask is not None:
gate_scores = gate_scores.masked_fill(~mask, float('-inf'))
gate_weights = F.softmax(gate_scores, dim=-1) # (B, L)
transformed = self.transform(x) # (B, L, D)
pooled = torch.bmm(gate_weights.unsqueeze(1), transformed).squeeze(1) # (B, D)
return pooled
|