PRISM: Enhancing Protein Inverse Folding through Fine-Grained Retrieval on Structure-Sequence Multimodal Representations

Designing protein sequences that fold into a target 3-D structure, termed as the inverse folding problem, is central to protein engineering. However, it remains challenging due to the vast sequence space and the importance of local structural constraints. Existing deep learning approaches achieve strong recovery rates, however, lack explicit mechanisms to reuse fine-grained structure-sequence patterns conserved across natural proteins. To mitigate this, we present PRISM a multimodal retrieval-augmented generation framework for inverse folding. PRISM retrieves fine-grained representations of potential motifs from known proteins and integrates them with a hybrid self-cross attention decoder. PRISM is formulated as a latent-variable probabilistic model and implemented with an efficient approximation, combining theoretical grounding with practical scalability. Experiments across multiple benchmarks, including CATH-4.2, TS50, TS500, CAMEO 2022, and the PDB date split, demonstrate the fine-grained multimodal retrieval efficacy of PRISM in yielding SoTA perplexity and amino acid recovery, while also improving the foldability metrics (RMSD, TM-score, pLDDT).

💡 Research Summary

The paper tackles the protein inverse folding problem—designing an amino‑acid sequence that adopts a given three‑dimensional backbone. While recent deep learning models (e.g., ProteinMPNN, PiFold, LM‑Design) achieve impressive sequence recovery, they lack an explicit mechanism to reuse fine‑grained structure‑sequence patterns that are conserved across natural proteins. PRISM (Protein Retrieval‑augmented Inverse folding via Structure‑Sequence Multimodal representations) addresses this gap by introducing a retrieval‑augmented generation (RAG) framework operating at the residue level.

Key components:

1. Motif Definitions – A “motif” is a recurring local structural‑sequence pattern; a “potential motif” generalizes this to any residue together with its local 3‑D neighborhood, serving as the basic retrieval unit.
2. Joint Encoder – A multimodal encoder G maps a protein (backbone B and optional sequence S) to per‑residue embeddings E_i that capture both local geometry and global context.
3. Vector Database – All residues from the CATH‑4.2 training set are encoded and stored as (embedding, residue index, protein ID) tuples, forming a prior knowledge memory of potential motifs.
4. Latent‑Variable Formulation – The conditional distribution p(S|B,D) is factorized with latent variables for representation (E), retrieval hypotheses (R), and attribution (Z). The retrieval kernel p(R|E,D) can be stochastic (learnable prior) or deterministic (Top‑K nearest‑neighbor search). Attribution Z is realized via attention weights in a hybrid transformer decoder.
5. Hybrid Decoder – Each decoder block contains self‑attention (global backbone context) and cross‑attention (retrieved motif embeddings). The attention scores constitute Z, which directly influence the final logits Y_i for each residue.
6. Training Objective – With deterministic approximations, training reduces to standard per‑residue cross‑entropy. When the retrieval prior is learnable, a KL‑regularization term aligns the amortized posterior q(R|·) with the prior.

Experimental evaluation spans five benchmarks: CATH‑4.2, TS50, TS500, CAMEO 2022, and a temporally split PDB dataset, plus an orphan‑protein test set. PRISM (using the AIDO.Protein‑IF base estimator) consistently outperforms baselines:

• On CATH‑4.2 test, PPL drops from 7.16 (ProteinMPNN‑CMLM) to 3.74 and AAR rises from 35.42 % to 40.98 %.
• On TS50/TS500, PRISM achieves PPL ≈ 2.4 and AAR ≈ 68‑71 %, surpassing all prior methods.
• Structural metrics (RMSD, sc‑TM, pLDDT) also improve, indicating that retrieved motifs not only boost sequence accuracy but also enhance foldability.
• Ablation studies confirm that (i) removing the retrieval step degrades performance sharply, (ii) deterministic Top‑K retrieval offers a practical speed‑accuracy trade‑off, and (iii) the hybrid self‑cross attention contributes an additional 5‑7 % AAR gain.

In summary, PRISM demonstrates that integrating a fine‑grained multimodal retrieval mechanism into protein inverse folding yields state‑of‑the‑art results across diverse datasets while maintaining computational efficiency. The work opens avenues for dynamic database updates, multi‑chain design, and Bayesian uncertainty quantification in future extensions.