MSA (evolutionary information)
↓
┌──────────────────────────────┐
│ ROW ATTENTION │ "Which residues interact?"
│ • Find local patterns │
│ • Secondary structure │
│ • Functional motifs │
└──────────────────────────────┘
↓
┌──────────────────────────────┐
│ COLUMN ATTENTION │ "How does each position evolve?"
│ • Conservation │
│ • Functional importance │
│ • Allowed mutations │
└──────────────────────────────┘
↓
┌──────────────────────────────┐
│ OUTER PRODUCT MEAN │ "Which pairs co-evolve?"
│ • Contacts │
│ • Correlated mutations │
│ • MSA → Pair transfer │
└──────────────────────────────┘
↓
Pair (pairwise features)
↓
┌──────────────────────────────┐
│ TRIANGLE ATTENTION │ "Is this geometrically valid?"
│ • 3D consistency │
│ • Triangle inequality │
│ • Error correction │
└──────────────────────────────┘
↓
Refined Pair → Structure PredictionEXample : Enzyme Active site
Catalytic triad (Serine protease):
His 57, Asp 102, Ser 195
ROW ATTENTION:
"These three residues are spatially close"
COLUMN ATTENTION:
His 57: 100% conserved (critical!)
Asp 102: 100% conserved (critical!)
Ser 195: 100% conserved (critical!)
OUTER PRODUCT:
His-Asp: Strong correlation (must maintain charge)
His-Ser: Strong correlation (hydrogen bond)
Asp-Ser: Strong correlation (catalytic mechanism)
TRIANGLE ATTENTION:
His-Asp: 10Å
His-Ser: 8Å
Asp-Ser: 12Å
→ Geometrically consistent! ✓Example 2:Disulfide Bond
Cys 20 - Cys 80 (disulfide bond)
ROW ATTENTION:
"Cys 20 and Cys 80 are in same sequence"
COLUMN ATTENTION:
Cys 20: Highly conserved (structural)
Cys 80: Highly conserved (structural)
OUTER PRODUCT:
Cys 20 - Cys 80: VERY strong correlation
"If one mutates, other must too!"
TRIANGLE ATTENTION:
Cys 20 - Cys 80: ~2Å (disulfide bond length)
Enforces: "Must be exactly this distance!"Example Alpha-helix
Helix: residues 10-20
ROW ATTENTION:
"Residues 10-20 form local pattern"
"i, i+3, i+4 spacing"
COLUMN ATTENTION:
Position 10: Hydrophobic (buried)
Position 14: Hydrophobic (buried)
Position 18: Hydrophobic (buried)
OUTER PRODUCT:
10-14: Close (helix turn)
14-18: Close (helix turn)
TRIANGLE ATTENTION:
10-14: 5.4Å (helix pitch)
14-18: 5.4Å (helix pitch)
10-18: 10.8Å (consistent!)Why all four?
- Evolution (column) + Structure (row) + Contacts (outer product) + Geometry (triangle)
- Each provides unique information
- Combining them = accurate structure prediction!
hwo many types of mask are there
- msa_mask
msa_mask = np.ones((10, 110), dtype=np.float32)
# Example:
# [[1, 1, 1, ..., 1], ← All valid
# [1, 1, 1, ..., 1],
# ...]- seq_mask
seq_mask = np.ones(110, dtype=np.float32)
# Example: [1, 1, 1, 1, 1, ..., 1]- pair_mask
pair_mask = np.ones((110, 110), dtype=np.float32)
# Example:
# [[1, 1, 1, ..., 1],
# [1, 1, 1, ..., 1],
# ...]- bert_mask
bert_mask = (np.random.uniform((10, 110)) < 0.15).astype(np.float32)
# Example:
# [[0, 0, 1, 0, 0, ..., 0], ← Random 15%
# [0, 1, 0, 0, 0, ..., 1],
# ...]- cluster_bias_mask
cluster_bias_mask = np.ones(10, dtype=np.float32)
# Example: [1.0, 1.0, 0.5, 1.0, 0.5, ...]
# ↑ ↑
# Downweighted duplicatesAlphaFold2 Complete Pipeline: From MSA to Structure
Pipeline Overview with Mask Creation and Usage
┌─────────────────────────────────────────────────────────────────────┐
│ PREPROCESSING PHASE │
│ (Happens ONCE per protein) │
└─────────────────────────────────────────────────────────────────────┘
Step 1: Load Raw MSA (A3M format)
═══════════════════════════════════
Input: alignend_10_Seq.a3m (10 sequences, 1739 positions each)
>sp|P01308|INS_HUMAN
---MALWMR----LLPLL-----ALLALWGPDPAAAFVNQHLCG...
>sp|P01315|INS_PIG
---MALWTR----LLPLL-----ALLALWAPAPAQAFVNQHLCG...
...
↓
Step 2: Remove Gaps & Create Deletion Matrix
═════════════════════════════════════════════
Query (no gaps): MALWMRLLPLLALLALWGPDPAAAFVNQHLCG...
Length: 110 residues (after removing gaps)
Deletion Matrix [10, 110]:
┌─────────────────────────────────────────┐
│ Seq 0: [3, 0, 0, 0, 0, 0, 4, 0, 0, ...] │ ← 3 gaps before M, 4 before L
│ Seq 1: [3, 0, 0, 0, 0, 0, 4, 0, 0, ...] │
│ Seq 2: [3, 0, 0, 0, 0, 0, 4, 0, 0, ...] │
│ ... │
└─────────────────────────────────────────┘
↓
Step 3: Integer Encode MSA
═══════════════════════════
MSA [10, 110]:
┌────────────────────────────────────┐
│ [12, 0, 10, 17, 12, 1, 10, 10, ...] │ ← M=12, A=0, L=10, W=17, etc.
│ [12, 0, 10, 17, 16, 1, 10, 10, ...] │
│ ... │
└────────────────────────────────────┘
↓
Step 4: CREATE MASKS (Static - Never Updated!)
═══════════════════════════════════════════════
┌──────────────────────────────────────────────────────────┐
│ MSA Mask [10, 110] │
│ ┌────────────────────────────────────────────────┐ │
│ │ [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1, 1, 1] │ ← All 1s (all valid)
│ │ [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1, 1, 1] │ │
│ │ ... │ │
│ └────────────────────────────────────────────────┘ │
│ Purpose: Mask padding (none in this case) │
└──────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────┐
│ Seq Mask [110] │
│ ┌────────────────────────────────────────────────┐ │
│ │ [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1, 1, 1] │ ← All 1s
│ └────────────────────────────────────────────────┘ │
│ Purpose: Mask invalid residues │
└──────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────┐
│ Pair Mask [110, 110] │
│ ┌────────────────────────────────────────────────┐ │
│ │ [[1, 1, 1, ..., 1], │ │
│ │ [1, 1, 1, ..., 1], │ ← All 1s
│ │ ... │ │
│ │ [1, 1, 1, ..., 1]] │ │
│ └────────────────────────────────────────────────┘ │
│ Purpose: Mask invalid residue pairs │
└──────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────┐
│ BERT Mask [10, 110] (Training Only!) │
│ ┌────────────────────────────────────────────────┐ │
│ │ [0, 0, 1, 0, 0, 0, 0, 1, 0, 0, ..., 0, 1, 0] │ ← 15% are 1s
│ │ [0, 1, 0, 0, 0, 1, 0, 0, 0, 0, ..., 0, 0, 0] │ (randomly masked)
│ │ ... │ │
│ └────────────────────────────────────────────────┘ │
│ Purpose: Masked language modeling loss │
└──────────────────────────────────────────────────────────┘
┌──────────────────────────────────────────────────────────┐
│ Cluster Bias Mask [10] │
│ ┌────────────────────────────────────────────────┐ │
│ │ [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, ...] │ ← Equal weights
│ └────────────────────────────────────────────────┘ │
│ Purpose: Weight sequences (downweight duplicates) │
└──────────────────────────────────────────────────────────┘
↓
Step 5: Initialize Pair Representation
═══════════════════════════════════════
Pair [110, 110, 128]:
┌─────────────────────────────────────────┐
│ Relative position encoding (65 features)│
│ Separation encoding (7 features) │
│ Residue type pairs (40 features) │
└─────────────────────────────────────────┘
↓
Step 6: Save All Features
══════════════════════════
Output: alphafold2_ready.npz
✓ aatype [110]
✓ msa [10, 110]
✓ deletion_matrix [10, 110]
✓ deletion_mean [110]
✓ pair [110, 110, 128]
✓ msa_mask [10, 110] ← CREATED HERE
✓ seq_mask [110] ← CREATED HERE
✓ pair_mask [110, 110] ← CREATED HERE
✓ bert_mask [10, 110] ← CREATED HERE (training)
✓ cluster_bias_mask [10] ← CREATED HERE
┌─────────────────────────────────────────────────────────────────────┐
│ TRAINING/INFERENCE PHASE │
│ (Masks are USED, not modified) │
└─────────────────────────────────────────────────────────────────────┘
Step 7: Load Features
══════════════════════
features = np.load('alphafold2_ready.npz')
msa = features['msa'] # [10, 110]
pair = features['pair'] # [110, 110, 128]
msa_mask = features['msa_mask'] # [10, 110] ← LOADED
pair_mask = features['pair_mask'] # [110, 110] ← LOADED
↓
Step 8: Evoformer (48 blocks)
══════════════════════════════
┌─────────────────────────────────────────────────────────┐
│ Block 1 of 48 │
├─────────────────────────────────────────────────────────┤
│ │
│ ┌────────────────────────────────────────┐ │
│ │ MSA Row Attention │ │
│ │ Input: msa [10, 110, 256] │ │
│ │ Mask: msa_mask [10, 110] ← USED! │ │
│ │ ──────────────────────────── │ │
│ │ attention = softmax(QK^T / √d) │ │
│ │ attention = attention * msa_mask │ ← Masks padding!
│ │ output = attention @ V │ │
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ MSA Column Attention │ │
│ │ Input: msa [10, 110, 256] │ │
│ │ Mask: msa_mask [10, 110] ← USED! │ │
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ Outer Product Mean │ │
│ │ msa → pair update │ │
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ Triangle Attention (Starting) │ │
│ │ Input: pair [110, 110, 128] │ │
│ │ Mask: pair_mask [110, 110] ← USED! │ │
│ │ ──────────────────────────── │ │
│ │ attention = attention * pair_mask │ ← Masks invalid pairs!
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ Triangle Attention (Ending) │ │
│ │ Mask: pair_mask [110, 110] ← USED! │ │
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ Triangle Multiplication (Outgoing) │ │
│ │ Mask: pair_mask [110, 110] ← USED! │ │
│ └────────────────────────────────────────┘ │
│ ↓ │
│ ┌────────────────────────────────────────┐ │
│ │ Triangle Multiplication (Incoming) │ │
│ │ Mask: pair_mask [110, 110] ← USED! │ │
│ └────────────────────────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────┘
↓
(Repeat 47 more times)
↓
Step 9: Structure Module (8 IPA blocks)
════════════════════════════════════════
Input: refined_pair [110, 110, 128]
Mask: seq_mask [110] ← USED!
┌─────────────────────────────────────────┐
│ IPA Block 1 of 8 │
│ Invariant Point Attention │
│ Mask: seq_mask [110] ← USED! │
└─────────────────────────────────────────┘
↓
(Repeat 7 more times)
↓
Step 10: Output 3D Structure
═════════════════════════════
Output: atom_positions [110, 37, 3]
(110 residues, 37 atoms each, xyz coordinates)
┌─────────────────────────────────────────────────────────────────────┐
│ TRAINING LOSS COMPUTATION │
│ (BERT mask used here!) │
└─────────────────────────────────────────────────────────────────────┘
Step 11: Masked MSA Loss (Training Only)
═════════════════════════════════════════
┌────────────────────────────────────────────────────┐
│ Predict masked positions │
│ predicted_msa = model_output['msa'] │
│ true_msa = batch['true_msa'] │
│ bert_mask = batch['bert_mask'] ← USED! │
│ │
│ errors = cross_entropy(predicted, true) │
│ loss = sum(errors * bert_mask) / sum(bert_mask) │ ← Only masked positions!
│ ───────────────────────── │
│ Only compute loss where bert_mask = 1 │
└────────────────────────────────────────────────────┘Example with Your Data
Input MSA (alignend_10_Seq.a3m)
10 sequences × 1739 positions (with gaps)
↓ Remove gaps
10 sequences × 110 residues (no gaps)Deletion Matrix Example
Position: 0 1 2 3 4 5 6 7 8 9 10 11
Sequence 0: [3, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, ...]
↑ ↑
3 gaps before 'M' 4 gaps before 'L'
From: ---MALWMR----LLPLL...
^^^ ^^^^
3 gaps 4 gapsMSA Mask Example
All positions valid (no padding):
[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1], ← Sequence 0
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1], ← Sequence 1
...
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 1]] ← Sequence 9
If we had padding:
[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0], ← Last position padded
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0],
...]BERT Mask Example (Training)
Random 15% of positions masked:
[[0, 0, 1, 0, 0, 0, 0, 1, 0, 0, ..., 0], ← Positions 2 and 7 masked
[0, 1, 0, 0, 0, 1, 0, 0, 0, 0, ..., 1], ← Positions 1, 5, and 109 masked
...]
Model predicts: What amino acid is at masked positions?
Loss computed: Only at positions where bert_mask = 1Key Takeaways
- Masks are CREATED once during preprocessing
- Masks are USED throughout training/inference
- Masks are NEVER UPDATED - they're static binary indicators
- Purpose: Tell the model which data is real vs padding/masked
Mask Usage Summary
| Mask | Created | Used In | Purpose |
|---|---|---|---|
msa_mask | Preprocessing | MSA attention (Evoformer) | Mask padding |
pair_mask | Preprocessing | Triangle attention (Evoformer) | Mask invalid pairs |
seq_mask | Preprocessing | IPA (Structure Module) | Mask padding |
bert_mask | Training data aug | Loss computation | Masked LM loss |
cluster_bias_mask | Preprocessing | MSA weighting | Downweight duplicates |