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

TL;DR

PRISM tackles the protein inverse folding problem by grounding sequence design in fine-grained, retrieved motifs from a structured memory. It introduces a latent-variable, multimodal retrieval framework with E (representation), R (retrieval), and Z (attribution) latent variables, plus a vector database of potential motifs and a hybrid self-cross attention decoder to emit residues conditioned on retrieved context and backbone information. Across five benchmarks, PRISM achieves state-of-the-art perplexity and amino-acid recovery while improving foldability metrics such as RMSD, TM-score, and pLDDT, all with only a small runtime overhead. The work demonstrates that explicit, residue-level retrieval of conserved local patterns provides a principled and scalable boost for protein sequence design, enabling both higher fidelity and controlled diversity via decoding temperature.

Abstract

Designing protein sequences that fold into a target three-dimensional structure, known as the inverse folding problem, is central to protein engineering but remains challenging due to the vast sequence space and the importance of local structural constraints. Existing deep learning approaches achieve strong recovery rates, yet they lack explicit mechanisms to reuse fine-grained structure-sequence patterns that are conserved across natural proteins. We present PRISM, a multimodal retrieval-augmented generation framework for inverse folding that 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. Across five benchmarks (CATH-4.2, TS50, TS500, CAMEO 2022, and the PDB date split), PRISM establishes new state of the art in both perplexity and amino acid recovery, while also improving foldability metrics (RMSD, TM-score, pLDDT), demonstrating that fine-grained multimodal retrieval is a powerful and efficient paradigm for protein sequence design.

Figures (9)

• Figure 1: Probabilistic graphical model of our proposed approach.
• Figure 2: The overall pipeline of our proposed framework PRISM. ① We start with a joint-embedding model and ② prepare a vector-database by inferring embeddings of known structure--sequence pairs. ③ Our retriever operates on per-token (fine-grained) embeddings, representing the surrounding potential motifs. The color coding shows the retrieved vectors for each corresponding site. ④ A hybrid decoder aggregates the retrieved entities and generates a refined protein sequence, enriched with the 3D structure encoding of the input protein. The legend for elements is provided in parentheses ("[]") at the bottom-right.
• Figure 3: Ablation on $K$ (CATH-4.2 validation split). PPL decreases as $K$ increases, but saturates at $K \geq 35$.
• Figure 4: AAR distribution across protein length bins on CATH-4.2. PRISM consistently outperforms AIDO.Protein, with especially large gains for shorter proteins ($<200$ residues).
• Figure 5: Distribution of lengths of the protein sequences in the benchmark dataset CATH-4.2 cath42.

Theorems & Definitions (13)

• Definition 3.1: Protein Motif
• Definition 3.2: Potential Motif
• Proposition 1
• proof
• Theorem F.1: Variational ELBO
• proof
• Corollary F.1.1: Prior–Jensen bound
• proof
• Proposition 2: Deterministic reduction and tightness
• proof