PRISM: Enhancing PRotein Inverse Folding through Fine- Grained Retrieval on Structure-Sequence Multimodal Representations
Authors: Sazan Mahbub
Souvik Kundu
Eric P. Xing
Venue: ICLR 2026 Poster (2026)
Keywords: Retrieval Augmented Generation
Protein Language Modeling
Protein Inverse Folding
Protein Sequence Design
Multimodal Representation
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: We present PRISM, a multimodal retrieval-augmented generation framework that enhances protein inverse folding by dynamically integrating fine-grained structure-sequence multimodal representations from a larger protein database.

Abstract: 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).

Published as a conference paper at ICLR 2026
PRISM: ENHANCING PROTEIN INVERSE FOLDING THROUGH FINE-GRAINED RETRIEVAL ON STRUCTURE- SEQUENCE MULTIMODAL REPRESENTATIONS
Sazan Mahbub", Souvik Kundu', Eric Xing
CCarnegie Mellon University, "'Mohamed bin Zayed Univeristy of AI, 9GenBio AI, 'Intel
ABSTRACT
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, how- ever, lack explicit mechanisms to reuse fine-grained structure-sequence patterns conserved across natural proteins. To mitigate this, present PRISM, mul- timodal retrieval-auge -aggeeend generation framework for inverse folding. PRISM retrieves fine- gree-grrined representations of potential motif's from known proteins and integrates them with hybrid self-cross attention decoder. PRISM is formulated as a latent-variable probabilistic model and implemented with an efficient approxi- mation, 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 effi- cacy of PRISM in yielding SoTA perplexity and amino acid recovery, while also improving the foldability metries (RMSD, TM-score, pLDDT).
INTRODUCTION
Designing protein sequences that fold into prescribed three-dimensional structure-the inverse fold- ing problem-is a long-standing challenge in computational biology with far-reaching implications in biophysics, enzyme engineering. and drug discovery. Unlike structure prediction, where methods such as AlphaFold2 (John et al., 2021) have achieved transformative success, inverse folding must contend with vast combinatorial search space: many distinct amino acid sequences can realize the same structural fold, and subtle local variations often determine stability and function. This underdetermined nature has made inverse folding scientifically important and challenging.
Recent deep learning approaches have made significant progress (Tian et al., 2026), Autoregressive sequence generators such as ProteinMPNN (Dauparas al., 2022) demonstrated strong sequence recovery and practical utility across monomers, oligomers, and designed nanoparticles. PiFold (Gao et al., 2022) combined expressive encoders with efficient decoders, offering substantial speedups while maintaining competitive accuracy. More recent works have exploited pretrained protein language models. LM-Design, DPI.M, and DPLM-2 (Zheng et al., 2023; Wang et al., 2024a;b) leverage large-scale sequence modeling and diffusion-bases generation, while AIDO.Protein (Sun al., 2024) scaled it to billions of parameters using mixture-offexpert training. Despite these advances, current architectures remain limited: they lack explicit mechanisms reuse fine- mnneggriied structure-sequence patterns (e.g., recurring motifs) that are evolutionarily conserved and central to protein function.
Our key insight that inverse folding can benefit from an explicit retrieval mechanism that grounds predictions in the rich diversity of known proteins at men-igrinedd level. By treating local structure- sequence neighborhoods as reusable building blocks, one can supplement end-to-end generative modeling with memory-based context. This motivates PRISM, multimodal sarrival--augmener generation (RAG) framework that reframes inverse folding through explicit representation, retrieval, and attribution. Instead of relying solely on monolithic encoder, PRISM retrieves embeddings of
Contact: smahhub6e & cs. cmu.edu, kaviikkkknndu @ intel.com, spxing@cs.cmu.cdu

---

Published as a conference paper at ICLR 2026
potential motif's from vector database of proteins, and aggregates them with hybrid transformer decoder to refine sequence emission. This introduces an explicit inductive bias: each residue prediction is guided by retrieved local fragments, while the hybrid decoder integrates these fragment- level priors with global backhone context.
Our major contributions are:
petrieval-augmenn frameewee framework. We propose PRISM, the first RAG framework for protein inverse folding that operates at residue-level granularity, retrieving fine- fine-grained multimodal representations for potential motifs and reusing conserved local patterns during sequence design.
A theoretically grounded formulation. We derive latent-variable model that factorizes representa- tion, retrieval, attribution, and emission, and provide an efficient approximation for implementation, ensuring both theoretical soundness and computational efficiency.
Extensive empirical validation. Through comprehensive experiments across five benchmarks and multiple evaluation metrics, we establish new state of the art in both sequence recovery and structural fidelity, while incurring only negligible runtime overhead. Detailed ablations validate the role of each design choice in our framework.
2
PRELIMINARIES
The protein inverse folding problem aims to design an amino acid sequence that compatible with given three-dimensional protein backbone. Formally, let backbone structure be specified by atomic coordinates
(Pi,...,PP, Pr), where each p; E IR denotes the position of the i-th backbone atom. The goal is to predict sequence S (S1.... Sz), where each residue S, is drawn from the standard amino acid vocabulary V. A inverse folding model thus learns the following distribution
P(S B)
II P(S; B),
which assigns probabilities to candidate sequences consistent with the target backbone. To represent protein structures, modern approaches often construct a residue-level graph G = (V, E),, where nodes U, V correspond residues and edges eij E E' capture spatial or physicochemical interac- tions (Dauparas et al., 2022; Alam et al., 2024; Gao et al., 2022; Mahbub & Bayzid, 2022). A model then encodes G' and outputs a distribution over residues for each position, either autoregressively (predicting residues sequentially) or non autoregressively (predicting all positions in parallel). The designed sequence is obtained by sampling or decoding from this distribution. Detailed discussion on related work has been provided in Appendix A.
PRISM: A MULTIMODAL RAG FRAMEWORK FOR INVERSE FOLDING
We introduce PRISM, a multimodal retrieval-augmenned generation (RAG) framework for protein inverse folding that operates at residue-level granularity. We first formalize fine-grained structural- sequential regularities via motifs and potential motifs, then derive a latent-variable model that factors retrieval, attribution, and emission. We conclude with conerete instantiations of the representation, vector database, retrieval kernel, and the training objective.
3.1 MOTIFS AND POTENTIAL MOTIFS
Definition 3.1 (Protein Motif). It is recurring local structurul-sequeniee pattern of residues that is evolutionarily conserved and often functionally significant. Formally, it can be described as short stretch of amino acids together with its surrounding 3-D conformation, capturing local folding rules and biochemical properties independent of the global protein context.
Definition 3.2 (Potential Motif). We generalize motifs by treating each residue together with its local 3-D neighborhood as potential motif. potential motif may or may not align with a canonical structural motif, but serves as a fine-grained motif-like unit that encodes transferable stracture-seuuence information. These representations are the building blocks for retrieval and sequence emission in our RAG framework.
3.2 LATENT- VARIABLE FORMULATION
Modeling Objective. Given target backbone 3-D structure B and fixed residue-level vector database D (whose entries represent potential motifs in local neighborhoods), our goal is to model
2

---

Published as a conference paper at ICLR 2026
the conditional distribution over amino-acid sequences p(S B, D). Di- rectly parameterizing this conditional challenging due to combinatorial sequence space and long-range dependencies. We therefore introduce latent variables that capture retrieval of locally similar neighbors and their attribution each site before emitting the final sequence.
R
3
Z
S
Latents for representation, retrieval, and attribution. Let 3 {E denote latent variables for the
Figure 1: Probabilistic
representation, and
graphical model of our pro- posed approach.
R = (=: denote latent retrieval hypothesis, where R, are neighbors retrieved from D for the (potential) motif in residue i's locality (Fig. 2, Point () Sec. 3.4). We define the retrieval kernel as p(R]
B, D). Let
denote attribution variables with conditional p R. &, B, D) that specifies how retrieved neighbors contribute emissions S {St) with S p Z, R, E, B, D).
Basic generative factorization.
The joint distribution factorizes as
PS, R. Z D) =
pR 3, B, D) (Z R,E, B, pS Z, R, E,
(I)
representation retrieval kemel
attribution
sequenee emission
Using the conditional independences D B, R.LB E, and IL R. we obtain
E,,,, B,D) = E(E B
(R E
p & B) (S Z, R,
representation retrieval kernel
attribution
sequence emission
Marginalizing over the latents yields
P(S B,
(EII)) B
Z, R, B)
(3)
The corresponding probabilistic graphical model is shown in Fig. 1. This formulation induces family of valid objectives arising under different approximations or parameterizations of the latents E, R, and Z. We describe the model components and their efficient deterministic instantiations below. Fig. 2 depicts the overall pipeline of our proposed framework.
3.3 STRUCTURE-SEQUENCC MULTIMODAL REPRESENTATION OF POTENTIAL MOTIFS
We represent residues way that captures both structural and sequential context of any poten- tial motif around the residue, so that each residue embedding itself summarizes the local motif.
Joint encoder. Let & be joint coder of 3-D structure and ID sequence (Fig. Point ())
Vector-Delabsse
True
68 druuhuuee
E E ((B Blxd, (4)
estimation
Cinarated sequence wi improved residues
where E = -E and d is the embedding dimension.
sequanoe
Aggropation and Ginneration (with M+SCA)
+ sy
Potential-motif
representation.
Resteval Keme!
embedding
Each vector & Rd contextualizes
Trainable soructuee
residue i E [I. by its local 3-D neighborhood and its placement in the global protein P, and is used for both retrieval and emission.
embedding
Quary embedding
Retrieved
Retrieval
embeddings
direction
Figure 2: The overall pipeline of PRISM. ! We start with joint- Query proteins with unknown se- embedding model and 2 prepare vector-database by inferring quence. At inference we observe embeddings of known diructure-squeeet only a query backbone B², We can operates on per-token (fine (fine-graine grained) embeddings, representing the sample an initial sequence estimate 5y from an off tff-thh-shelf inverse folding
pairs. ④ retriever
surrounding potential motifs. The color coding shows the retrieved vectors for each corresponding site. + hybrid decoder aggre- gates the retrieved entities and generates a refined protein sequence, enriched with the 3-D structure encoding of the input protein.
model (Sun et al., 2024; Dauparas et al.,
2022; Wang et al., 2024b), which we refer to as the base estimator in this article. With this, we can form a crude query embedding Eч = G(B9, 5*), which we treat sample from the marginal, i.e., Ex p(E B 139..
3

---

Published as a conference paper at ICLR 2026
3.4 VECTOR DATABASE OF POTENTIAL MOTIFS
We treat the vector database D as prior-knowledge memory of potential motif representations which retrieval is performed. Given M proteins with structures and sequences = (B",5" )P) [M]] we encode each PP via Eq. 4 to obtain E {E")p= p==1 The database is
D D= (E r; p) : e [M], r e [|P*]].
Each residue embedding E. summarizes the locality around residue r in protein p. Let Φ(.) map residue neighborhood to motif representation in a metric space (M,d). Retrieval by similarity of & to {dd) effectively searches for nearby motifs in M. Implementation note: Our vector-DB is created with the training split of the CATH-442 dataset (Orengo et al., 1997), and any search in it runs entirely on GPU, substantially reducing search time (Sec. 4.2.3).
3.5 RETRIEVAL KERNEL
The retrieval kernel p(R &, D) may be instantiated in multiple ways consistent with the latent- variable formulation. In Appendices D and E.2 we outline representative, non-exhaustive choices for stochastic and deterministic variants.
A trainable stochastic kernel constitutes the most faithful realization of the graphical model, defining proper distribution over residue-level neighbors and enabling learned retrieval priors (Appendix E.2). A deterministic Topk operator offers an efficient approximation (A (Appendix D, Eq. 20). Formally,
L
PRR E, D) -I (R,, & D), p(R, E D) a(R, ( D)),
1=1
with Dirac distribution 8(.). Both variants arise from the same underlying model class, differing only in how the latent retrieval variable instantiated.
While the trainable variant provides direct path toward more expressive retrieval mechanism that potentially lead to improved performance at the expense additional computing resource (Tab. 9), the deterministic approximation yields efficient yet powerful solution (Tab. 1-5 in Sec. 4.2).
3.6 ATTRIBUTION MARGINAL
Retrieval provides candidates but not how they are used. We realize attribution via attention weights computed by T hybrid transformer blocks in our aggregation- and and--eneration module Fez (Sec. F), (f.h)
parameterized by 0z. Each head h E f1...., H in block t E {1...., ,.. computes Oik
softmax: (q: (t.h) 1i (t..))) /dd... with query vector qi (t.h) for residue and key (t.h) for neighbor
Rik- Thus a deterministic function E,
(t.h) {a J., = A(R, B),
p(Z R,E, oZ A(R, B)).
(6)
Fig. 7 details the hybrid htff//croos-attettonn design. Sec. 4.3.5 we provide ablation to demonstrate the effectiveness of this design.
3.7 SEQUENCE EMISSION
Given B and R. the module Fez forms retrieval-aware residue representations through Z and outputs per-residue logits Y B, R) Faz( Fe,, (B), E, R) e Lx20 where Føn structure encoder (Appendix. F). The emission distribution factorizes:
I II
pSS B,
Cat(Sn Softmax(Y(E,,
Remark. Although we write p(S ε, R,R,R Z), the logits Y B, R) already incorporate the deterministic attribution Z computed by Fez.
3.8 TRAINING OBJECTIVE
Under our deterministic reduction (Appendix E.7), the objective collapses to maximizing the stan- dard log- likeelold via per-residue cross-entropy, = arg max Ep log pe (S -)], with learnable
4

---

Published as a conference paper at ICLR 2026
Table I: Comparison of protein inverse folding methods on the CATH-4.2 test split (Short, Single-chain. and All). Best and second-bess scores are shown bold and italic. Foldability metrics are computed using ESMFold (Lin et al., 2023), and the subscript aido denotes P.Dt..rotein-I was used as the base estimator.
Method
Uses
Uses
Short
Single
All
PLM
RAG
PPL
AAR
5c-TM
PLDDT
StructTrans
X
X
8.39
28.14
8.83
28.46
6.63
35,82
GVP
X
7.23
7.84
28.95
39,47
ProteisMPNN
X
6.21
36.35
6.68
34.43
4.61
45,96
ProMMPNN-CMEL
X
X
7.16
35.42
7.25
35.71
5.03
48.62
1.61
0.910
0.812
PRISM (str. enc.)
4.26
35.29
3.40
48.97
3.39
49.17
PiFald
X
6.04
39.84
6.31
38.53
4.55
51,66
1.64
0.911
0.816
LM-Design
7.01
DPLM
X
54.54
IIDO.PrtteinIIF
58.87
1.58
0.912
0.821
PRISM
3.74
40.98
2.68
60.89
2.71
60.43
1.56
0.913
0,835
parameters 0 = {0z, 0g). When the retrieval prior is trainable (Appendix E.2), the objective is augmented with a KL-regularization term defined with respect to an amortized posterior q(R -). i.e. @ arg max Eggl log (S :)] Eg ( 1-) pg( -))] with @ {0z, 0B, 0g). where e 8g represents the learnable parameters of the joint encoder G. We leave the exploration of jointly learning amortized posteriors for & and Z via their KL terms to future work.
4 EXPERIMENTS AND RESULTS
Table 2: Comparison TS50 and TS500.
TS50
TS500
Models
4.1 EXPERIMENTAL SETUP
PPL
%
AAR
GVP
4.71
44.14
4,20
49.14
Datasets. We evaluate PRISM on several widely used benchmarks: CATH-4.2, TS50, TS500, CAMEO 2022, PDB date split, and orphan proteins dataset. We use the CATH-4.2 training split both for train- ing and to create the vector database, preventing any form of data leakage by construction (Orengo al., 1997). TS50 and TS500 are used only for evaluation to test cross-dataset generalization (Li et al., 2014), while CAMEO 2022 and the PDB date split assess robustness on proteins outside the CATH classification and under temporally disjoint conditions (Campbell et al., 2024). To fur- ther test PRISM's generalizability to entirely novel backbones, we evaluate it the orphan proteins dataset by Jing et al. (2024). Full dataset descriptions and statistics are provided in Appendix B.
ProtciaMPNN
ProMMPNN-
3.46
53.68
3.35
56.45
(str. enc.)
3.10
PiPald
3.86
$8.72
3.44
60.42
LM- Design
AIDO Protcin-IF
2.68
66. 19
2.42
69.66
PRISMande
Evaluation Metrics. We report two sequence-level Table 3: Comparison on CAMEO 2022 and PDB metrics and three structure-level metrics. Sequence accuracy is assessed by amino acid recovery (AAR) and perplexity (PPL), while foldability is assessed with RMSD, TcM and pLDDT. Together these capture both sequence correctness and structural realizability. Detailed formulations of all metries are provided in Appendix B
date split. Multiflow, ESM3, and DPLM-2 results taken from (Wang et al., 2024a).
CAMEO 2022
PDB date split
Models
PPL
AAR
PHUMPNN-CMMMM
3.62
50.14
3.42
$2.98
PRISM (str. enc.)
MulaFlow
33.58
37.59
ESM3
46.24
DPLM2-3B
53.73
57.91
Baselines. We compare against a comprehen- sive suite of statc-of-the-art inverse folding meth- ods, including StruetTrans, GVP, ProteinMPNN, ProteiMMPNN-CMLAN PiFold, LM-Design, DPLM, MultiFlow, ESM-3, DPLM-2, and the large- scale AIDO.Protein, which we also adopt as our joint encoder. Appendix C provides details of these baselines.
Protei-F
2.68
PRISMando
2.53
64.63
2.35
67.47
4.2 RESULTS AND DISCUSSION
4.2.1 RESULTS ON CATH 4.2, TS50, TS500, CAMEO 2022, AND PDB DATE SPLIT
CATH-4.2 (Tab. 1): shows that PRISM consistently improves over strong baselines across all three CATH-4.2 settings. Compared to L.DO.Protein-IF. PRISM reduces perplexity from 4.09 to 3.74 on short-chains, from 2.91 to 2.68 single-chains, and from 2.94 to 2.71 the full test set. These PPL gains also translate to higher recovery. Specifically, AAR increases by 2.52 (40.98 vs. 38.46) on short, 2.02 (60.89 vs. 58.87) single-chain, and 1.83 (60.43 vs. 58.60) all. Notably, even PRISM
5

---

Published as a conference paper at ICLR 2026
Table 4:
Foldability comparison using AF2 protein folding model. The median and the mean are provided outside and inside the parenthesis, respectively.
Models
CAMBO 2022
PDB date splir
RMSD
T- T
PLDDT1 1
RMSD
sc-TM1
PLDDT 1
DPLM2-38
1.67 (1.833)
0,926 (0.846)
0.923 (0.898)
1.21 (1.399)
0.954 (0.918)
0.94 (0.919)
AIDO, Protein-IF
(1.66)
0.942 (0.862)
0.922 (0.916)
1.1 (1.231)
0.963 (0.936)
0,953 (0:937)
1.49 (1.621)
0.948 (0.867) 0,934 (0.916)
1.04 (1.2)
0.964 (0.938)
0.953 (0.938)
Table Runtime analysis seconds per protein) across different benchmarks. We decompose runtime into the base estimator SAIDO.Protein (F)e retrieval, and decoding. The total time is the sum of all components.
Model
TS50
TS500
CAMBO2022
PDB date split
CATH 4.2 test
CATH 1.2 val
Average
estimator (AIDO.Protcin
0.57
0.89
0.92
RRetrieal
3.1e
I.e-3
1.3-->
6.0c
5.0e-
6.0e
1.2e
+Decoding
0.08
0.17
0.17
0.12
0.10
0.11
0.13
Total
0.91
1.20
1.17
1,03
0.97
1.00
1O5
(str. enc.), the structure-only variant, already outperforms its corresponding baseline, ProteinMPNN- CMLM. while our full framework. with PiDo..rotein-IF IF as the base estimator and joint encoder, yields the best overall trade-off in sequence-level and structure-leve metrics. All scores are obtained with deterministic decoding, where we use the deterministic approximation of our retriever and chose a aremax sampling with the final logits. Additional results on this dataset, including threshold-based designability analysis, are provided in Appendix I. and Tab. 17.
TS50 and TS500 (Tab. 2): On TS50, PRISM sets new SoTA on both metrics, with PPL 2.43 vs. 2.68. and AAR of 67.92 vs. 66.19 19 consistently improving over its base estimator AIDO. Protein-IF. On TS500, PRISM achieves the best AAR (70.53) and a strong PPL. of 2.27, while LM-Design reports a lower PPL TS500, its AAR is substantially lower (64.50), indicating that PRISM's conditioning yields sequences that align better with native residues.
CAMEO 2022 and PDB date split (Tab. 3): PRISM improves both confidence and recovery on distribution shifts. On CAMEO 2022, PPL decreases from 2.68
IF) to 2.53, while AAR rises improves by 1.11. On the PDB date split, PPL drops from 2.49 to 2.35 and AAR improves from 66.27 to 67.47. These trends on these four stand-alone test sets underscore that our proposed approach contributes stable gains even when test distributions diverge from CATH-4.2.
4.2.2 FOLDABILITY ANALYSIS
Tab. 4 and Appendix Tab. 14 show end-to-end foldability of designed sequences via AlphaFold2 (John et al., 2021). PRISM consistently improves structural fidelity over AIDO. Protein-IF across datasets:
on TS50, RMSD drops from 1.075 to 0.985, sc-TM rises from 0.956 to 0.964, and PLDDT slightly improves 11.999-0.9500 on TS500, RMSD improves (1.18~11.12) with sc-TM also higher (0.964). On CAMEO 2022 and the PDB date split, PRISM attains the best RMSD and sc-TM alongside competitive or best PLDDT. These consistent gains indicate that PRISM's s higher AAR is not merely superficial residue matching, rather translates to sequences that fold closer to the target backbones with stronger global topology (sc- TM) and comparable or better local accuracy (pLDDT).
4.2.3 RUNTIME ANALYSIS
A key advantage of PRISM Table 6: Inverse folding performance on the orphan proteins dataset (Jing et al., 2024). Here ESMFold (Lin et al., 2023) was used as the protein folding model.
is that its substantial accu-
racy gains come at negligi- ble runtime cost. As shown in Tab. 5, the base estimator AADDO.Protein--. on average 0.92 seconds per protein, while our full frame- work adds only lightweight
Model
AAR (%)
sc-TM +
RMSD
PLDDT
Native Sequence
100.00
0,360 (0.449)
3.61 (3.654)
0.457 (0.535)
requires
Proteini
28.05 (30.22)
0.369 (0.420)
3.78 (3.615)
0.538 (0.535)
PRISM.nds
28.33 (30.85)
0.391 (0.437)
3.92 (3.846)
0.556 (0.531)
ProtinMPNN-CMLM
29.27 (31.24)
0,359 (0.428)
3.60 (3.712)
0.569 (0.525)
PRISM
27.12 (30.76)
0.404 (0.437)
3.60 (3.587)
0.533 (0.538)
retrieval (~ 1.2 x 10 1003s) and decoding (0.13s), resulting in total runtime of 1.05s. This corre- sponds to a relative overhead of merely 14.3% compared to the base estimator.
In contrast, the improvements in accuracy are much larger. Averaged across benchmarks (Appendix Tab. 18), PRISM reduces perplexity from 2.68 to 2.43 (9.3% improvement) and boosts AAR from 63.0% to 66.9% (+3.9 absolute points). In other words, PRISM delivers significant and consistent accuracy gains across all test sets while incurring only a negligible runtime overhead. This balance demonstrates the efficiency of memory- -ased retrieval: it enriches the model's representations without

---

Published as a conference paper at ICLR 2026
sacrificing throughput, making PRISM a practically viable and scientifically impactful extension the base estimator.
4.2.4 ORPHAN PROTEIN DESIGN
Because natural-sequee sequence recovery does not fully capture generalization novel backbones, we additionally evaluate PRISM on the orphan proteins dataset Jing et al. (2024), a set of eleven proteins with detectable sequence homologs in major databases. As shown in Tab. 6, PRISM consistently improves over its respective base estimators and achieves foldability scores close to the native sequences, despite the lack of homologous entries in its retrieval database. This indicates that PRISM is not merely memorizing natural sequences, but can leverage retrieved structural context to enhance design even for highly novel backbones.
4.3 ABLATION STUDIES AND ADDITIONAL ANALYSES
To better understand the contributions of individual components and design choices in PRISM, we conduct series of ablation studies and supplementary analyses. These highlight the effectiveness, efficiency, and robustness our framework.
4.3.1 ABLATION THE NUMBER OF RETRIEVED ENTRIES
We conducted ablation study analyze the effect of the number of retrieved vectors K on model performance. As shown in Fig. 3, increasing K consistently reduces perplexity (PPL) on the CATH-4.2 validation split. The improvement is sharp for small K (e.g., from 2.788 at K to 2.709 at K 5). but gradually saturates as K increases further. Beyond K 2 35, PPL stabilizes around 2.681, showing no further significant gains. Therefore, we choose K = 35 as the optimal setting, striking balance between efficiency and accuracy.
2.78
Perplciity (PR)
2.76
2.74
2.72
2.70
2.68
15 25 30 40
(number retrieved vectors)
Figure 3: Ablation on K (CATH-4.2 validation split).
4.3.2 EFFECT OF PROTEIN SIZE ON RECOVERY
Fig. 4 shows the distribution of amino-acid recovery rates (AAR) across protein length bins on the CATH-4.2 test set. PRISM con- sistently outperforms Prttenrotei across all lengths, with particularly notable gains for shorter proteins (< 200 residues) where inverse folding more challenging. These results confirm that PRISM's improvements are robust across varying sequence lengths, rather than confined to a narrow subset of proteins.
1.0
AAR 196)
AIDO
0.2
[0.
1100, 1200, 1300. (4.9. Protein Length Bins
4.3.3 CONTRIBUTION OF RETRIEVAL
Figure 4: AAR distribution across A central question in study is whether retrieval itself con- protein length bins on CATH-4.2. tributes meaningfully to inverse folding, beyond what large pre-
PRISM consistently outperforms trained models structural encoders already achieve. Across AIDO.Protens with especially large all benchmarks, PRISM with retrieval consistently outper- residues). forms AUtt.Protein-IF in both perplexity and recovery metrics (Tab. 1, 2, 3, 4). For instance, on CATH-4.2 PRISM improves
gains for shorter proteins (< 200
AAR by nearly two percentage points over AIDO. Proteint,,, while TS50 and TS500 it reduces perplexity and boosts recovery sinultaneously-demooee across datasets of varying scale and diversity.
that retrieval provides tangible benefits
To isolate the effect of retrieval from that of multimodal representation. we introduce controlled unimodal ablation variant, PRISM (str. enc.), which replaces the AIDO.Protein-IF joint encoder with a purely structure-based encoder ProteinMPNN-CMMLM structural embeddings. Remarkably, even in this restricted setting, our retrieval mechanism delivers consistent gains over the baseline ProteinMPNN-CMLA across all datasets (Tab. 1, 2, 3). This result isolates retrieval as an independent driver of performance even without sequence-leve priors, fine- iine-grained retrieval improves recovery by supplying complementary local context that single encoder cannot capture. Together, these findings establish retrieval to not be an auxiliary feature, but core contributor in PRISM.
and performs retrieval solely over
7

---

Published as a conference paper at ICLR 2026
Table 7: Ablation myyrid-attentton cross-attention-only
Models
TS30
TS500
CATH-4.2 test split
CATH-422 val split
PPI. 1
AAR 56
P
AAR %
PPI.
AAR S
PPI.
AAR %
PRISM (wle MHSA)
2.56
64.23 (65. 12/64 98)
2.36
S4 (68 43/70044)
2,82
59.26 44/60.41)
2.79
59.51 (58.4261.11)
PRISM (fall)
2,43
67.92 (66.98666.0) 30)
(69.57770.97)
2.71
60.43
(56614)1411
(59.2861.89)
Table 8: Ablation on the number of MHSCA blocks.
of blocks
TS50
TS500
CATH-4.2 test split
CATH-4.2 val split
PPL 4
AAR %
J
AAR %
PPL
AAR %
PPL
AAR %
N/A (base est.)
2,68
66. (64, 661.9/64.666 .6)
2.42
69 .6 (68 04/69 60)
2.94
58.60 (57.27/60. 13)
2.90
58.73 (58, 000000666062 62)
66.90 (66. 84/66.88)
2.26
70.93 (69.58/70.97)
2.72
(58.53/61.5)
2.69
(59.20/6188)
2.43
67.92 (66.98/66.00 70)
2.27
70.53 (69. (69.57/10.97)
2.71
60.43
(58.55661 41)
2.68
60.26
(59.28/6.89)
2.44
67.71 (66.75/66.48)
2.26
70.59 (69.58/70.99)
2.71
60.35 (58.59/61442)
2.68
60.24 (59.3061.91)
4.3.4 EFFECT OF EXTENDING RETRIEVAL DATABASE
A natural question is whether enlarging the retrieval memory at inference time further improves performance. Our theoretical analysis (Appendix G) establishes that once the vector database achieves near-complete E-coverage of the motif space, additional entries predominantly duplicate existing motifs and thus provide diminishing returns. Empirically, we confirm this saturation effect: augmenting the database with new PDB entries yields almost identical results across all benchmarks (Appendix G, Tab. 13), with differences well within retrieval noise. For instance, for CAMEO 2022 the AAR remains ~44.6% whether using only the CATH-4.2 memory, the PDB extension, their combination. This finding highlights that PRISM's fixed vector database already captures the relevant structural landscape, making post-hoc memory growth unnecessary. Crucially, validates our design choice of treating the vector database as prior knowledge store rather than an expaedinnd index, achieving saate-of-the-ee recovery while avoiding uncontrolled growth in memory size.
4.3.5 CONTRIBUTION OF HYBRID DECODER WITH MHSCA
We next ablate the design of the aggregation module by comparing our hybrid multihead self cross attention (MHSCA) decoder with a simplified variant that relies only on multihead rross-attention (MHCA). As shown in Tab. 7 and Appendix Tab. 15, removing the self-attention component degrades performance across all benchmarks. While the cross- catss-tttetiioo--n oonll variant already improves over the base estimator by attending to retrieved vectors, it lacks the ability to contextualize and refine these fragments jointly. Incorporating MHSA within the block allows the model to propagate information among retrieved neighbors before align- ing them with the query, yielding consistent gains. For example, 64.2% to 67.9%, and perplexity drops from 2.56 to 2.43; on the CATH-4.2 test split, AAR rises from
Table 9: Effect of finetuning the joint encoder G (denoted "ft. G"") and a simple inference-time iter- ative refinement approuch (denoted "iter."). Here "orig. g" denotes the original joint encoder without finetuning.
TS50, the AAR increases from
59.3% to 60.4% with a corresponding reduction in PPL (2.82 + 2.71). Similar improvements are ob- served on TS500 and the PDB date split, with relative gains of +..8%-1.2% AAR. Importantly, these gains are consistent across both in-distribution (CATH-4.2) and out-of-distribution (CAMEO 2022, PDB date split) settings, highlighting that the hybrid MHSCA architecture provides more expressive aggregation by jointly leveraging self- and cross-attention. This vali- dates the adoption of MHSCA as the default decoding module in PRISM.
Model
PPL. 4
AAR (%)
PiFold
4.55
$1.66
PRISMpifala (orig. 9)
2.72
60.09
PRISM Paraa (orig. G. iter:)
2.73
61.51
PRISM (ft. G)
2.63
62.01
PRISMaraat (ft. G. iter.)
2.67
63.09
AIDO Proteinnt
2.94
58.60
PRISM (orig. ))
2.71
60,43
PRISMaiso (orig. G. iter.)
2.73
61.87
PRISMadoo (ft.
2.59
62.07
PRISMadoo (ft. iter.)
2.65
63.33
4.3.6 EFFECT OF AGGREGATION DEPTH (MHSCA LAYERS)
We next study how the number of multihead self-cross attention (MHSCA) blocks in the aggregation module affects performance (Tab. and Appendix Tab. 16). Adding even a single block over the base
Table 10: Ablation on the base estimator. Here the subscripts pmpnn, pr fold, and aido respectively denote whether ProteinMPNN-CME
PiFold, W.DDO..oteinnIl IF were used as the base estimator for PRISM. All results of PRISM here use the original (non- finetuned) joint encoder
Model
4.1.
CATH AAR (%)
TS50 PPI.]
T850 AAR (9)
TS%00 PPI.
TS500 (5)
ProteiMPPNNCCM (base est.)
5n3
48.62
3.46
53.68
3.35
56.45
PRISM
2.82
58.02
2.52
65.03
2.35
69.0
PiFold (base
4.55
$1.66
3.86
58.72
3.44
60.42
PRISM
60.09
66.56
71.16
T.tteirotein IP (base est.)
2.94
58.60
2.68
66.19
2.42
89.00
PRISMasi
2.71
60.43
2.43
67.92
2.27
70.53
8

---

Published as a conference paper at ICLR 2026
Table 11: PRISM's AAR improvements across proteins grouped by PiFold initial prediction quality. Here lower bins correspond to lower-quality initial sequences.
PRISM (original ()) Meun AAR (%)
PRISMpefaie (ft. G. iter.)
PiFold
PiFold Mean AAR (5)
A (original ()
AAR Bin (%)
Proteins
(ft. G, iter.)
Meun AAR (%)
0-20
12
15.89
18.48
2.59
25.93
10:03
20-30
74
25.97
29.52
3.55
33.79
7.82
30-40
131
35.51
42.63
7.12
48.44
12.94
40-50
272
45.95
55.22
9.27
59.21
13.26
50 60
464
54.88
64.60
9.72
67.53
12.66
60-70
157
62.79
71-47
8.68
73.42
10.63
70
73.57
80.03
6.46
78.76
yields large gain: on the CATH-4.2 test split, AAR improves from 58.6% to over 60.2%, and perplexity drops from 2.94 to 2.72. Increasing to two blocks provides the best overall trade-off, achieving the strongest tied-best results across nearly all benchmarks (c.g., CAMEO 2022 with PPL 2.53 and AAR 64.6%, CATH-4.2 validation with PPL 2.68 and AAR 60.3%). Using three blocks maintains similar accuracy but shows consistent benefit, with small oscillations likely due to noise. These results indicate that the aggregation mechanism quickly saturates, and two MHSCA layers suffice capture the additional context from retrieved fragments while avoiding redundancy or overfitting.
estimator
4.3.7 EFFECT OF FINETUNING JOINT ENCODER AND ITERATIVE REFINEMENT Our latent-variable formulation naturally allows finetuning the joint encoder G via a trainable retrieval kernel (Appendix E.2). To assess its impact, we evaluate PRISM with and without finetuning G. As shown Table 9, finetuning 9 yields consistent improvements across two base estimators of different strengths: PiFold and result indicates that adapting the joint encoder further enhances the retrieval, and subsequently, the generation quality.
respectively denoted with subscripts old and aido. This
We additionally ablate a simple inference-time refinement procedure, inspired by recent ap- proaches (Zheng et al., 2023; Wang et al., 2024a). Here, the output sequence is iteratively fed back into the PRISM framework to improve the generation. This refinement tends to further increase AAR, with only marginal change in PPI., suggesting that PRISM can effectively leverage its own predictions to refine sequences without any extra training. more systematic study of different refinement paradigms is left future work.
4.3.8 IMPACT OF INITIAL SEQUENCE QUALITY
To investigate the sensitivity of PRISM's final performance to the quality of initial sequence estimate, we first compare PRISM across three different base estimators of varying strengths: ProteinMPNN- CMLM, PiFold, and improves each of these estimators by a large margin. Remarkably, PRISMpi and PRISM pmppn (both using single-iteration refinement) already achieve performance comparable to the stronger Proteintt IF baseline, which itself uses iterative refinement. These results clearly indicates that PRISM effectiveness is not tied to AIDO. Protein-IF or any specific initializer. Because PRISM is base estimator agnostic, it can in principle be paired with surface-based generators such as SurfPro (Song et al., 2024), SurfDesign (Wu et al., 2024), or biochemistry-awann generators such as BC-Design (Tang et al., 2025). Such combinations would allow PRISM to retrieve surface- or chemistry-aligned exemplars, extending its benefits beyond
The results, summarized in Tab. 10. show that PRISM
conditiooe inverse folding.
To further assess PRISM's robustness to, the quality of the initial sequence provided by the base estimator, we stratify proteins by their PiFold AAR and evaluate PRISM's performance within each bin. This analysis reveals how much improvement PRISM achieves when starting from weak, moderate, or strong initial estimates. As shown in Table 11, PRISM consistently boosts AAR across all bins, with particularly large gains in the low-AAR regimes.
5 CONCLUSIONS
We present PRISM, a multimodal retrieval-augmentt framework for protein inverse folding that integrates fine-grained retrieval of potential motif embeddings with a hybrid self-cross attention decoder. PRISM achieves new state of the art across multiple benchmarks in sequence recovery and foldability, while adding only negligible runtime overhead. Our vatent-variable formulation provides theoretical grounding, and ablations confirm the central role of different design choices, including retrieval, hybrid attention, and aggregation mechanism. These results establish fine-grained retrieval principled and scalable approach for advancing protein sequence design.

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their valuable insights and recommendations, which have gready improved our work. This research has been graciously funded by the National Science Foundation (NSF) awards CNS2414087 and IIS2123952; by the Defense Advanced Research Projects Agency (DARPA) award HR001 12390063; and, the Semiconductor Research Corporation (SRC) AIHW award 2024AH3210. Any opinions, findings, and conclusions recommendations expressed in this publication are those of the author(s) and do not necessarily reflect the views of the National Science Foundation, the Defense Advanced Research Projects Agency, and the Semiconductor Research Corporation.
REFERENCES
Ramisa Alam, Sazan Mahbub, and Md Shamsuzzoha Bayzid. Pair-egret: enhancing the prediction of protein-protein interaction sites through graph attention networks and protein language models. Bioinformatics, 40( 10).bbtte5888 2024.
Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications protein co-design. arXiv preprint arXiv:2402.04997, 2024.
Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob de Haas, Neville Bethel, et Robust deep learning- based protein sequence design using proteinmpnn. Science, 378(661 378(6615):49-56, 2022.
Zhangyang Gao, Cheng Tan, Pablo Chacón, and Stan Z Li. Pifold: Toward effective and efficient protein inverse folding. arXiv preprint arXiv:2209. 2022.
Zhangyang Gao, Cheng Tan, Yijie Zhang, Xingran Chen, Lirong Wu, and Stan Li. Proteininybench: Benchmarking protein inverse folding on diverse tasks, models, and metrics. Advances in Neural Information Processing Systems, 1166207--6220, 2023.
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. arXiv preprint arXiv: 2019.
Thomas Hayes, Roshan Rao, Halil Akin, Nicholas Sofroniew, Deniz Oktay, Zeming Lin, Robert Verkuil, Vincent Tran, Jonathan Deaton, Marius Wiggent, al. Simulating 500 million years of evolution with a language model. Science, 38766736:850-8888 2025.
Neil Houlsby, Andrei Giurgiu. Stanislaw Jastrzebski, Bruna Morrone, Quentin De Laroussilhe, Andrea Gesmundo, Mona Attariyan, and Sylvain Gelly. Parameter-efficient transfer learning for nlp. In International conference on machine learning, pp. 2790-2799 PMLR, 2019.
Chloe Hsu, Robert Verkuil, Jason Liu, Zeming Lin, Brian Hie, Tom Sercu, Adam Lerer, and Alexander Rives. Learning inverse folding from millions predicted structures. ICML, 2022. doi: 10. 10.1101/2022.04 10.487779. URL ht attps: www.. www.:wworxiv.oommmmooo org/ wwt///nnteettcccec eeaee/eeaally 04. 10. 487779.
2022 / /04 11////1//222/2 2022.
John Ingraham, Vikas Garg, Regina Barzilay, and Tommi Jaakkola. Generative models for graph- based protein design. Advances neural information processing systems, 32, 2019.
Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael John Lamarre Townshend, and Ron Dror. Learning from protein structure with geometric vector perceptrons. In International Conference on Learning Representations, 2020.
Xiaoyang Jing, Fandi Wu, Xiao Luo, and Jinbo Xu. Single-sequence protein structure prediction by integrating protein language models. Proceedings of the National Academy of Sciences, 121(13): e2308788121, 2024.
Jumper John, Evans Richard, Pritzel Alexander, Green Tim, Figurnov Michael, Ronneberger Olaf, Tunyasuvunakool Kathryn, Bates Russ, Židek Augustin, Potapenko Anna, et al. Highly accurate protein structure prediction with alphafold. Nature, 2021.
10

---

Published as a conference paper at ICLR 2026
Zhixiu Li, Yuedong Yang, Eshel Faraggi, Jian Zhan, and Yaoqi Zhou. Direct prediction of profiles of sequences compatible with a protein structure by neural networks with fragment- based local and energy- nonlocal profiles. Proteins: Structure, Function, and Bioinformaties, 82(10):
2565-2573, 2014.
Zeming Lin, Halil Akin, Roshan Rao, Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, Allan dos Santos Costa, Maryam Fazel-Zarandi, Tom Sercu, Salvatore Candido, and Alexander Rives. Evolutionary-scale prediction of atomic-level protein structure with a language model. Science, 379(6637):1123-
2023.
Sazan Mahbub and Md Shamsuzzoha Bayzid. Egret: edge aggregated graph attention networks and transfer learning improve protein protein interaction site prediction. Briefings in Bioinformatics 23(2):bbab578, 2022.
Sazan Mahbub, Christoph Feinauer, Caleb N Ellington, Song. and Eric P Xing. Uncertainty-aware discrete diffusion improves protein design. bioRxiv, pp. 2025-06, 2025.
Weian Mao, Muzhi Zhu, Zheng Sun, Shuaike Shen, Lin Yuanbo Wu, Hao Chen, and Chunhua Shen. De novo protein design using geometric vector field networks. arXiv preprint arXiv:2310.11802, 2023.
Joshua Meier, Roshan Rao, Robert Verkuil, Jason Liu, Tom Sercu, and Alexander Rives. Language models enable zero-shot prediction of the effects of mutations on protein function. bioRxiv, 2021. doi: 00.10000001.00.09945555 1101/2021.0 07 .7..0.45064.4.111
URL http::://///// htttps://ww.... http:://wwwww.wwx........ biwrxiv org/ content/10.
Brandes Nadav, Goldman Grant, Wang Charlotte, H., Ye Chun, Jimmie, and Ntranos Vasilis. Genome- wide prediction of disease variant effects with deep protein language model. Nature Genetics, 2023.
Christine A Orengo, Alex D Michie, Susan Jones, David T Jones, Mark B Swindells, and Janet M Thortoo.. Cath-a hierarchic classification of protein domain structures. Structure, 10931199,,, 1997.
Zhenqiao Song, Tinglin Huang, Lei Li, and Wengong Jin. Surfpro: Functional protein design based on continuous surface. arXiv preprint urXiv:2405. 06693, 2024.
Ning Sun, Shuxian Zou, Tianhua Tao, Sazan Mahbub, Dian Li, Yonghao Zhuang, Hongyi Wang, Xingyi Cheng. Le Song, and Eric P Xing. Mixture of experts enable efficient and effective protein understanding and design. bioRxiv, pp. 2024-11 11, 2024.
Xiangru Tang, Xinwu Ye, Fang Wu, Daniel Shao, Dong Xu, and Mark Gerstein. Bc-design: A biochemistry-aware framework for highly accurate inverse protein folding. In ICML 2025 Genera- tive Al and Biology (GenBio) Workshop, 2025.
Jiayi Tian, Seyedarmin Azizi, Yequan Zhao, Erfan Baghaei Potraghloo, Sean McPherson, Sharath Nit- Sridhar, Zhengyang Wang. Zheng Zhang, Massoud Pedram, and Souvik Kundu. Skipkv: Selective skipping of kv generation and storage for efficient inference with large reasoning models. MLSys, 2026.
Vaswani. Attention all you need. Advances Neural Information Processing Systems, 2017.
Xinyou Wang. Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Diffusion language models are versatile protein learners. arXiv preprini arXiv:2402. 18567, 2024a.
Xinyou Wang, Zaixiang Zheng, Fei Ye, Dongyu Xue, Shujian Huang, and Quanquan Gu. Dplm-2: A multimodal diffusion protein language model. arXiv preprint vrXiv:2410.13782. 2024b.
Fang Wu, Shuting Jin, Jianmin Wang, Zerui Xu, Jinbo Xu, Brian Hie, et al. Surfdesign: Effective protein design molecular surfaces. 2024.
Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu. Amaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint a 230 2023.
11

---

Published as a conference paper at ICLR 2026
Zaixiang Zheng, Yifan Deng, Dongyu Xue, Yi Zhou, Fei Ye, and Quanquan Gu. Structure- Struture-informedd informed language models are protein designers. In International conference on machine learning, pp. 42317-42338. PMLR, 2023.
Shuxian Zou, Tianhua Tao, Sazan Mahbub, Caleb N Ellington, Robin Algayres, Dian Li, Yonghao Zhuang, Hongyi Wang, Le Song, and Eric P Xing. A large-scale foundation model for rna function and structure prediction. bioRxiv, pp. 2024-11, 2024.
A RELATED WORKS
Protein inverse folding, the process of designing amino acid sequences that fold into specific three- dimensional structures, has been focal point of computational biology research (Gao et al., 2022; Mahbub et al., 2025; Sun et al., 2024; Hsu et al., 2022; Zheng et al., 2023). In 2022, Dauparas et al. (2022) proposed ProteinMPNN, widely popular autoregressive method for designing protein sequences that fold into desired structures. It achieved an impressive sequence recovery rate on native backbones, outperforming traditional methods, showing versatility extending designing monomers, cyclic oligomers, nanoparticles, and target-binding proteins, Gao et al. (2022) introduced PiFold, a method that effectively combines expressive features with an autoregressive sequence decoder to enhance both the accuracy and efficiency of protein design. PiFold achieved high recovery rate on the benchmark dataset and demonstrated a speed advantage, being 70 times faster than some autoregressive counterparts. That same year, Hsu et al. (2022) proposed sequence-to- sequence transformer model trained using predictions by AlphaFold2. a sate-offthe-art structure prediction method (John et al., 2021). By leveraging putative structures of millions of proteins, their approach achieved a notable improvement in the field. Zheng et al. (2023) introduced the usage of protein language models (Nadav et al., 2023; Meier et al., 2021) for cruuctrr-conditioned protein sequence design, or in other words, inverse folding. Another work by Wang et al. (2024a) extended this by incorporating diffusion language modeling for effective sequence generation. Sun et al. (2024) pretrained a 16 billion parameter protein language model with a mixture-of-expert architecture, which they further adapted for prediction and sequence generation tasks, and surpassing the previous methods. To address the need for standardized evaluation, Gao et al. (2023) also proposed ProteinInv Bench, a comprehensive benchmark for protein design. This framework includes extended design tasks, integrated models, and diverse evaluation metrics, facilitating more rigorous comparisons across different methods. Mao et (2023) proposed VFN-IF and and its ESM-equipped variant VEN-IFF, advancing inverse folding by using Vector Field Networks to perform learnable coordinate- based vector operations that jointly model residue frames and atomic geometry. Recent work has also explored surface-based and biochemistry-awware protein design, such SuríPro (Song et al., 2024) and SurfDesign (Wu et al., 2024), which operate directly on molecular surfaces, and BC-Design (Tang et al., 2025), which incorporates explicit biochemical constraints into the generation process. These approaches differ from inverse folding models in their use of geometric or physicochemical signals rather than backbone-sequence likelihood, and represent an orthogonal line of structure-conditioned design.
B EXPERIMENTAL SETUP
B. DATASETS
We evaluate our framework on six widely used benchmarks: CATH-4.2 (Orengo et al., 1997), TS50 (Li et al., 2014), TS500 (Li et al., 2014), CAMEO 2022 (Campbell et al., 2024), PDB date split (Campbell et al., 2024), and orphan proteins dataset (Jing et al., 2024).
CATH-4.2 is standard benchmark containing proteins with fewer than 500 residues, and is widely adopted for training, validation, and testing of inverse folding models (Zheng et al., 2023; Wang et al., 2024a). Following prior work, we further analyze three subsets of the CATH-4.2 test set: short chains (length < 100, 1165%) single chains (~92.86%) and the full test split. Appendix Fig. 5 shows the sequence length distribution.
TS50 is a compact benchmark of 50 proteins (maximum length 173), while TS500 provides greater variability, ranging from very short chains (43 residues) long proteins (>1600 1600 residues). Following
12

---

Published as a conference paper at ICLR 2026
Table 12: Statistics of CATH-4,2, TS50, TS500, CAMEO 2022, PDB date split, and Orphan proteins benchmark datasets. Here "seq.". "res.". "len.", and "St. Dev." represent "sequence", "residue", "length". and "standard deviation", respectively.
Data split
of seq.
of res.
Mean Len.
Median Len.
St. Dev. Len.
CATH-4.2 Train 18,024
3,941,775
218.7
204.0
109.93
CATH-4.2 Validation
608
105,926
174.22
146.0
92.44
CATH-4.2 Test
1,120
181,693
162.23
138.0
82.22
CAT'-4.22 Combined
19,752
4,229,394
214.12
196.0
109.06
TS50
50
6,861
137.22
145.0
25.96
TS500
500
130,960
261.92
225.0
167.30
CAMEO 2022
183
44,539
243.38
228.0
144.86
PDB date split
449
86,698
193.09
178.0
81.06
Orphan proteins
11
1,310
119.09
124.0
55.04
convention (Zheng et al., 2023; Gao al., 2022), we use these only for evaluation after training on the CATH-4.2 training split, thereby testing cross-dataset generalization.
CAMEO 2022 comprises 183 recently released structures (average length 243 residues), providing an evaluation on proteins outside the CATH classification and closer to real-world modeling tar- gets (Campbell al., 2024). The PDB date split (449 proteins; mean length 193) follows the protocol of previous studies such as Campbell et al. (2024) and Wang et al. (2024b), where training and evaluation proteins are separated strictly by deposition date in the Protein Data Bank. This ensures robustness against temporal leakage and simulates forward-lon sowward-lookig generalization.
The recovery rate of natural proteins does not fully assess the ability of inverse folding models to generalize to novel or orphan backbones. We therefore evaluate PRISM the orphan proteins dataset prepared by Jing al. (2024), which contains 11 challenging proteins with no detectable sequence homologs in major databases. Since retrieval index is built from the CATH-4.2 training split (PDB-derived), these orphan proteins are guaranteed to be non-overlapping with both our training set and retrieval index.
B.2 EVALUATION METRICS
We report two sequence-level metrics and three structure-level metrics.
Sequence-level metrics. Amino Acid Recovery (AAR): Median sequence recovery is the most widely used metric for inverse folding (Zheng al., 2023; Wang et al., 2024a; Sun et al., 2024). It measures the percentage of positions where the predicted amino acid matches the native sequence:
AAR median
Σ
1(S S,) 100%
(8)
where T. is the protein length and 1 is the indicator function.
Perplexity (PPL): Perplexity evaluates how confidently a model predicts the native sequence (Wang al., 2024b; Tian et al., 2026; Zou et al., 2024). For autoregressive models:
PPLAR exp
log P(S, Sci, B)
(9)
4
For our non-autoregressive setting:
M L
PPL = cxp
Σ
Σ
log P(S. S,
(10)
j=1
where $ is noisy initialization of the native sequence.
13

---

Published as a conference paper at ICLR 2026
Structure-level metrics. To evaluate whether generated sequences foldable into the target backbone, we use three complementary metrics, following recent studies (Dauparas et al., 2022; Zheng et al., 2023; Wang et al., 2024a;b):
Root-Mean-Square Deviation (RMSD): measures the average distance (in Å) between backbone alpha carbon atoms of the predicted and native structures after optimal alignment. Lower RMSD indicates higher structural fidelity.
RMSD(X, (X,)
Σ
d
(11)
where d; is the distance between the alpha carbons of i-th pair of residues in two structures X and Y.
TM-score (sc-TM):
similarity metric that robust to protein size and is widely used assess global fold correctness; values 0.5 typically indicate correct fold topology.
Feligned
sc-TM(X, = max alignments Ltargot 2i
(12)
1
d, do do(Tuaa
where do(L)= 1.24) 15)) 1.8
Predicted Local Distance Difference Test (pLDDT): a per-residue confidence score from Al- phaFold2 (John al., 2021), used here to assess the stability and reliability of folded structures generated from designed sequences.
PLDDT I Σ conf(i)
(13)
Together, these metrics evaluate both sequence accuracy and structural realizability, which is critical in physical sciences applications protein design.
B.3 HARDWARE SPECIFICATIONS
These experiments were conducted with single AMD Milan processor (64-core 2.55 GHz), 256 GB of RAM, and four NVIDIA A100 GPUs (each with 40 GB VRAM). The vector-database takes about 15.1 GB of hard disk space.
B.4 IMPLEMENTATION DETAILS OF VECTOR DATABASE
Our vector database is created with the CATH-4.2 training split, where each residue corresponds to database entry as defined in Sect. 3.4 (a total of 3,941,775 residues). Since we leverage this dataset to train our model as well, we additionally mask the query sequence itself to prevent trivial self-retrieval during training.
C BASELINES
StructTrans (Ingraham et al., 2019) proposed conditional generative model for protein sequences given 3D structures based on graph representations. GVP (Jing et al., 2020) introduced geometric vector perceptrons, which extend standard dense layers to operate collections of Euclidean vectors. ProteinMPNN (Dauparas al., 2022) proposes an autoregressive protein sequence generation approach conditioned on structure. ProteinMPNN-CMN CMMM (Zheng et al., 2023), a non-autoregressive variant of the original ProteinMPNN, has been trained with the conditional masked language modeling (CMLM) objective (Ghazvininejad et al., 2019) and achieves higher score than the original version. LM-Design (Zheng et al., 2023) is another non -ooreuttreeressivee model trained with CMLM that leverages pretrained protein language models for inverse folding. DPLM (Wang et al., 2024a) extends this work by using discrete diffusion language modeling objective to enhance sequence generation capabilities of languange models. Multiflow (Campbell et al., 2024), ESM3 (Hayes al., 2025),
14

---

Published as a conference paper at ICLR 2026
Distribution of lengths
100
Frequency
60
20
100
200
300
400
500
Length
Figure 5: Distribution of lengths of the protein sequences in the benchmark dataset CATH-4.2 (Orengo al., 1997).
and DPLM-2 Wang al. (2024b) also take a generative approach, with flow-based and diffusion language modeling. AIDO. (Sun et al., 2024) is 16 billion parameter pretrained protein language model that has been further adapted for inverse folding with conditional discrete diffusion language modeling objective. VFN-IF and VFN-IFE extend the Vector Field Network to inverse folding by using learnable vector computations over virtual atom coordinates (Mao et al., 2023).
D RETRIEVAL KERNEL
We model R (R as latent retrieval hypothesis. The kernel p(R. D) admits both stochastic definitions and deterministic approximations.
Stochastic retrieval. For residue i, let the cosine similarity between query embedding & and entity d E D (with embedding (d)) be
(E,, E(d)
(.(d) [ [[E()]]
(14)
Convert nonnegative weights using temperature and normalize:
w:(d)
wi (d) cxp(ai(d)/7). (c(d) EVeD w,(')
(15)
Sample K distinct entities R, C D without replacement under a Plackett-L kernel. For an ordered K- K-tuple n, (da,...,, dik) with distinct elements,
we we(de)))
Pr(m: &,, D) = II
(16)
Σ
wa (i)d)
dED' (d.........
For the unordered set R,,
K
We wi(ckk)
pRR, D)=
Σ
II
(17)
wr(d)
K.EPerm (R))=) k=1
dED dED (d
15

---

Published as a conference paper at ICLR 2026
PDB 4W66, Chain: (embeddings stored in Vector-DB)
"MNEPIILRY... (true sequence)
PDB ID: 3ZU, Chain: (from CATH-4.2 sest
"MKRIEYTQQ.." (initial estimation)
Structurally similar
Joint embedding of full protein
Joint embedding of full protein
[8]
E
El
El 235
[E4, Ed
E9 .... 8476
2>
58
High cosine similarity (0.98)
Fine-grained (per-residue) joint embeddings, representing their localities
Retrieve the Vector-DB entry
corresponding to E
Figure 6: An example of how our vector DB search works. Here leverage embedding E (representing the potential motif in residue j's locality protein i) to search representations of similar 3D localitics in other proteins in the vector database to enrich context for later generation step.
The kernel factorizes across residues:
L
p E = &g,
R = (R
(18)
We leverage this stochastic process (together with the full probabilistic model) when sampling diverse sequences (Appendix M).
Deterministic approximation. A 0, Eq. 15 concentrates maximizers of (,(d) and Eq. 16 sequentially selects the K largest scores (ties broken arbitrarily). Thus p(R, | & D) Eq. 17 collapses to point mass the top-K K set:
(Ep arg max JCD IJEK dEJ
Σa (aild).
(19)
Formally,
I. p D) Il (R, E),), p(R, E,, D) o(R, - (E)
(20)
with Dirac distribution 8(:); see Appendix D.I for proof.
16

---

Published as a conference paper at ICLR 2026
D.1 CONVERGENCE OF STOCHASTIC RETRIEVAL TO DETERMINISTIC TOP-K
Proposition 1. Let ag(d) denote the cosine similarity score between query embedding & and entity d E D. Define weights and sofimax probabilities
wx(d)
wx(d) exp(ax(d)/T),
P((d)
EvED wi(d')
Consider the Plackett-Liee sampler that draws ordered K K-suple (dit...,,dx) without replacement:
K II
w.((ssk)
Pr(m; & D)
Eden (ia (i)d)
1=1
As + O, the distribution over unordered retrieval sets R; = {dii.... dek) } converges to a point
mass on the deterministic Top-K set of scores TopK(E;; D), up to uniform randomness among exact fies.
Proof. Step 1: Single-choice limit. Fix and write a(d) == a;(d). Let M max, a(d') and
T = (d ad) M) be the argmax set. Then
cadd//
c()(d)-M)/r
Pa(T)
Ex ca(d')/ Ex e
As + the numerator converges lif T and otherwise. Hence
[ d E d d
lim pa(T) T-0
0,
If [T] = 1, the maximizer d' is selected with probability 1.
Step Sequential without replacement. Plackett-Luce draws items by repeating the softmax on the remaining set. If [T'] 1, then d, = d w.p. 1. Removing d", the argument applies inductively the reduced set, so each step the current maximum is selected. Thus the ordered tuple म; the Top-K scores in descending order.
there are ties, the probability mass split uniformly among tied maxima; once one is chosen, the argument recurses on the remaining set.
Conclusion. Therefore, for the unordered set R.,
[1, 0, otherwise,
R, = Top D),
lim (RR & D) -
up to uniform randomness under exact ties. This proves the claim.
E FROM LATENT MODEL TO TRAINING OBJECTIVE: ELBO, PRIOR-JENTE BOUND, AND DETERMINISTIC REDUCTION
Model recap. Given backbone B and database D, the latent variables are the embedding &, the retrieval R (R.) and the attribution Z. The emission factorizes across residues with logits Y R)):
Cat softmax BY pS Z, R, = II
i=1
The joint and marginal are
Z, R
B, D) = P(E (R &, D) R.E, () 2 Z,,,E,
(21)
SS D) E p(E BR E, p( R.E. B) pS. S.E. B)
(22)
Z.E.R
(23)
For notational simplicity, the summation symbol Σ denote marginalization over all latent variables, encompassing both summation (for discrete variables) and integration (for continuous embeddings).
17

---

Published as a conference paper at ICLR 2026
E. VARIATIONAL ELBO
Theorem E.1 (Variational ELBO). For any density q(E, R, Z S, with support contained in that of pEE, R,Z Z B, D),
log D) 2 (ε)BO(a), where the ELBO can be written in either of the equivalent forms
(24)
CFLBO( = E, log Z, E, B)] - KI(G(E, R Z S, B, R, B, D)), (25)
Eyy logg((S Z, R. E, B) + Poglog B) log R E D)
+ B1 p Z R, E, B)] llogog g)
(26)
Moreover,
log D) LELBO(9) + KL((()) || P(E, R, S, B,), so that - lo B, D) S - LELBO(9) (the negative ELBO upper- bounds the true NLL).
(27)
Proof. Start from 23 and multiply and divide by q(E, R, Z B): (()) ( E, R. Z D) QEE, R. Z B)
p(S E, R. Z B) 19(E, R, Z S,B,
log B, D) log Σ
log Eg
Z,E,R
By Jensen's inequality (concavity of log):
p E, R. Z B, D)
( R. B,
log p(S B, D) = log Eg
2 Eg log
[q(E, R. Z S, B, D)
(E R, 2 B)
= Ey [log peepp(SS,.. R, B, B, Eg (log g).
Expanding the joint via the model factorization gives 26, and grouping terms yields 25. For the decomposition with the posterior, observe by Bayes:
P R, S. B, E 8, R. Z B, D)
(( B.
Hence
KI |p(- S,B = E, log p(. S, B, D)
E.y Dog ES, B, D)]
log p(S B)
i.c.
log PS D) Egg llg plogg((s,,.)) 3, R, Z B, Eg log +KI(g |p( S B,
Since KL 2 0, the inequality follows.
E.2 TRAINABLE RETRIEVAL KERNEL
Under the factorizations of the amortized posterior
9(E, R, Z S. B, D) = q(E | () D) [, R, B),
and the prior
pτE, R, Z D) = P(E (R P) R,E,
the exact ELBO from Eq. 25 becomes
LELBO (q)
E,llogp(S Z, R. R.E
KK(G(E S, B) B))
[R S,E, (R E, D))] (KL S,,,, B) (Z R,E, B))]
With qE SS p(E) B) and 9( S, R, p(Z R)E,B), LELBO (g) F. log r(S Z, R. E, B)] F. [T |., D) ( E)D) which enables training the retrieval prior p(R E, D) through minimization of the KT.-divergence.
18

---

Published as a conference paper at ICLR 2026
E.3 PRIOR-JENSEN LOWER BOUND
Corollary E.1.1 (Prior-Jensen bound). Let plE, R, Z D) be the latent prior induced by the model. Then
logp(S) | D) 2 (.(E.R.ZIB.D) ( 2, R, 8, B)],
(28)
equivalently
- -logp(S) D) v [Ep (S 2, &,
Proof. Take q(E, R. Z S, = R. B, D) in Theorem E. Then (L(glpp- B, D)) = in 25, and the ELBO reduces Epllog p(S Z, R, &, B)], which therefore a lower bound ((g (S B, D). Alternativel,, apply Jensen directly to
logp(S B, D) log (.. (p() Z, E, B)] 2 Exp [og -)]-
E.4 DETERMINISTIC REDUCTION TIGHTNESS
Proposition 2 (Deterministic reduction and tightness). Assume
P(E B) -E-EE)) E.) p(R E) II a(R, TopK D)).
p(Z R. B) 8Z -A(E,, B", B))R)). and Z = A(,,,B,R*) Then
Let R =
p(S B) = log p(S Z" Z²,R", E" B'))) and the right hand side equals the standard per-residue log- log-likelihood
(29)
log (ogp(S) Z' R' Z' "R, 89,
Y (B,
(30)
Consequently,
1
- logp(S B, D) =
Σ Dogsoflmaxx(e)) (Y(E9, B4, R R*);
Proof. Under the stated Dirae measures, the summations 23 collapse the single configuration (E,))*Z**):
p(S B, D) = p(S Z" R², 8,,
By the emission factorization, this conditional equals IT Taking logs yields the stated sum of per-residue log-softmax terms.
(Y(É,, B1, R ))).
E.5 CONSEQUENCES FOR THE TRAINING OBJECTIVE
Combining Corollary E.1.1 and Proposition 2:
log i( D) 2 Ep [log p(S E, B)] and logp(S) B, B) = logg Z' R² RT,E,,88 E9,
so the plug-in negative log-likelihood
LNLL
log soft soft (Y (9, B9, R*):)
(i) exactly the NLL of the deterministic latent model, and (ii) an upper bound on the true NLL of the stochastic latent model:
- lo B, D) < LNLL
Equality holds when the stochastic latents degenerate Dirac measures (our current design), when variational posterior collapses point mass concentrated at (E9, R Z*).
19

---

Published as a conference paper at ICLR 2026
E.6 DERIVED PROBABILISTIC MODEL
Substituting the kernels into Eq. gives
wy(dit)
(S B)), Σ
pe, p(e By
Σ
II
Σ
wi(d)
ERR.Z
T,EPerm (R,)
dED (daamaamdd
(2 R. &, B) II Catl softmmxx R))). (31)
Under the deterministic approximations,
pS
B, D)
Σ R.Z
II
aR. TopK D)) 11--A(87) B,R))
i=1
II
(a(oftmaax((())))).)a (32)
t=
E.7 TRAINING OBJECTIVE UNDER DETERMINISTIC REDUCTION
We target the true marginal (S B, D) and optimize its prior-Jensen lower bound (formal proof in App. E):
log D) 2 Ep logp(S :) Ep
Σ log Cat(Se; Y, B, R:)
(33)
Equivalently, we minimize corresponding Jensen negative ELBO (NELBO) to learn parameter set
0 = arg min Ey(., log
Y,(E, B, R,
(34)
Under deterministic reduction for any query protein Ba (App. E, Prop. 2) with & E9, R TopK(E* (É);D), D). Z A(Ea, B3, R*), the prior-Jensen lower bound in Eq. 28 is light and the objective collapses to standard per-residue cross-entropy:
i arg min
Σ log softmax softmax(Y (",
(35)
with gradients flowing through the learnable parameters 0 = {0z,0p) 0p] the retrieval Topk is treated as fixed and non-differentiable in this deterministic setting.
F AGGREGATION AND GENERATION WITH OUR DECODER
We aggregate the retrieved entities R* to generate new sequence 59, We do this with a series of T' consecutive learnable blocks, each consisting of one multihead self-attention layer (MHSA), one multihead cross-attention layer (MHCA), and two bottleneck multilayer perceptrons (Houlsby et al., 2019) (T is hyperparameter). In the rest of this article, we refer to this hybrid block multihead self-cross attention block (MHSCA).
As shown Fig. (Point ) and Fig. 7, for E [1, [S]] we first extract the embedding vectors {(E)), 4+ l E [1, KT} from our retrieved entities. We then merge them with the query vector E/ and linearly project the output to create a matrix HC E (R(+1)xxd as,
H = concat( ((E)), )x 1 e [1, KJ U (E7)) WH.
(36)
where concat(.) performs a concatenation operation on the vectors in the union set, WH E dd learnalle parameter, and d' is the output embedding dimension. Then for the whole query
20

---

Published as a conference paper at ICLR 2026
protein Pt get tensor 719 = [M4. N2..... ....... H2 Suj
[))* which is used query,
key, and value of MHSA (see Vaswani (2017) for definitions). To ensure that the generator can effectively leverage any residual 3D structural information, we also encode the input structure 3 separately using structural encoder. where No sequence information is provided. Similar to Sun al. (2024), we leverage ProteinMPNN-CMLMM (Zheng al., 2023) for structure encoding, which is variant of the original ProteinMPNN method (Dauparas et al., 2022) trained with conditional masked language modeling objective (Ghazvininejad et al., 2019). This generates structural encoding This encoding is then linearly transformed and merged with linear projection of
pe E | query encoding 8", creating new representation matrix 02 [gl"%)xd'' .where each element 04 concat( (.. )) E Rd", with two learnable parameters W, E dRd" and We E B 4 For our MHCA blocks, we use 8 as the query, and 7 as both the key and value. The motivation behind such design of MHSCA is, while MHSA layers can help jointly attend to multiple parts of the input protein as well as their corresponding retrieved embeddings, MHCA can help extract any kind of residual structural information needed to better decode the sequence. Moreover, since the MHCA here preserves the same dimension 09, the output representation has |$9 vectors which can directly pass through another linear layer to generate the output logits E R15"|xd" Sampling with Y provides us with a newly generated sequence $
G POST-HOC MEMORY-GROWE ANALYSIS
In our design, each residue embedding EL summarizes the local motif (or potential motif) around residue r protein p. Let (6-) map residue neighborhood to a motif representation space (M, d), and let the database (memory) be D = {d (E7, r.)).. At inference, for each query residue i we retrieve neighbors D by similarity of & to E {d)) which effectively searches for nearby motifs in M.
metric
Definition G.1 s-cover--oo For tolerance s > and query distribution Q over motifs, the coverage of D is CD) = Prm minaeD d,m, ((d) S E
Proposition (Coverage saturation). Assume i.i.d. sampling of database motifs from the same distribution Q test queries, and that the motif space admits finite e-cover number Ne < co. Then (Dx) +1 as [r.] + xc, and the expected marginal coverage gain from adding a batch of k new entries satisfies E [Cov:(Dn+t) - Cove [D.]] = =O( A )
Intuition. Each database item covers an e-ball in M. Under i.i.d. sampling, uncovered mass shrinks geometrically with until most query motifs lie within of at least memory item. Beyond this point, additional samples mostly fall into already covered regions, yielding negligible retrieval improvements and thus minimal downstream gains.
Implication. Because each residue embedding encodes its local motif, sufficiently large memory D achieves high e-coverage of the motif space. Once coverage saturates, top-K neighbors (or their stochastic variants) are already oear-optimal, post-hoc accretion of similar residues contributes little to retrieval quality or sequence recovery, absent targeted diversification or retraining.
Gonontraate
Tatrieeed vectore
Buillrock MUP
Bonsenack NLP.
6401))
Quayy vectirn
Liner
Concatraan
Unte
get
XT
Figure 7: Our aggregation and generation module. It uses T consccutive blocks of Multi-Head Self-Cross Astention (MHSCA). Here the super-script "[t]" corresponds to the index of the current MHSCA block (and also its components and their inputs). See Sec. for details.
21

---

Published as a conference paper at ICLR 2026
Post Post-hoc Memory Growth Diagnostics
Coverage proxy (before after)
Redundancy
1.0
Before
0.8
0.8
0.4
0.4
0.035
0.019
Top-1 (s-0.0476)
Mean Top 35 dist
Prop S 1s-0.0476)
Mean NN dist.
Figure 8: Post-hoc memory growth diagnostics. Left: coverage proxy before/after adding fractintion of queries with Top-1 cosine distance < e (e = 0.0476) and mean Top-35 distance. Coverage rises only slightly (0.951 + 0.956) and mcan Top-35 distance improves marginally (0.0355 0.0340), indicating ncar-saturation. Right: redundancy of the added memory D+ measured by nearest-neighhor distance the original memory D. The vast majority (96.4%) new entries already lie within of an existing item (mean NN distance 0.0192), explaining the negligible coverage gain.
G.1 DIAGNOSTICS
Does enlarging the memory post-hoc help? To test the hypothesis that our memory already E-covers most query motifs, we augment the fixed memory D with small disjoint batch D from newer PDB entries (no parameter updates) and re-index. We report two diagnostics:
Coverage proxy. For each query residue embedding &,, we compute the cosine distance of its nearest neighbor in the memory, d'dd (i; D), and the mean of its top- K distances, d(k) (i; D). We summarize by the fraction of queries with dled (i; D) S and by E, [deos (i; D)J, where is fixed to the q-th percentile of (2) (i; D) on the original memory (e.g.. 9955..
Redundancy of D+ For each new memory item d E D+, we compute its nearest-neighhor cosine distance to the original memory D, dNN (d; D). We report the proportion with d&N (d; D) S s and the mean D+ NN (d: D)].
Finding. We tested our hypothesis on the TS50 test set. The results are depicted in Fig. 8. The fraction of queries with a nearest neighbor within & = 0.0476 cosine distance increased only marginally from 95.1% to 95.6%, indicating that coverage cosine distance improved slightly (0.0355 + 0.0340), a negligible gain given the scale of retrieval noise. By contrast, the added entries themselves were highly redundant: 96.4% of D+ items had nearest neighbor within in the original memory, with mean nearest-neighbor distance of 0.0192. These results confirm that post-hoc memory growth mostly contributes redundant motifs and provides no meaningful benefit for retrieval or sequence recovery. This supports our design choice to treat the vector database as fixed mnoo-knoweedge memory.
already near saturation. Similarly, the mean Top-35
Detailed ablation. For completeness, we ablate the impact database size by comparing three PRISM configurations: (i) D constructed from the CATH-4.2 training set, (ii) D built from new PDB entries, and (ii) their union. The results in Tab. 13 show that all variants achieve virtually indistinguishable perplexity and recovery scores, with differences well within statistical noise. For example, on CAMEO 2022 the AAR stabilizes at ~-6..6% across all database choices, and on the CATH-4.2 and validation splits the gap remains below 0.2%. These findings suggest that the CATH-4.2-based database already provides near-complete motif coverage, and that additional PDB entries primarily add redundant fragments rather than new information. This empirical evidence supports our design decision to fix the database as a compact, prior-knowledge memory: it ensures efficiency while preserving accuracy.
all samples form the PBD split here: ftttp:://zendd..rre/ httpp://eenddd.rr///eee////////
22

---

Published as a conference paper at ICLR 2026
Table 13: Ablation on the size of database. Results rounded two decimal points, hence very small changes are not reflected.
CAMEO 2022
CATH-4.2 test split
CATH-4,2 val split
Models
PPL L
AAR $
PPL.
AAR %:
PPL
AAR
estimator Proteino))))))F)
2.68
63.52 (60.56/64. 17)
2.94
56.60 (57 23).0.13)
2.90
58.73 (58, 00/00 (58.00065 62)
PRISM (VDB: CATH 4.2 train)
2.53
64.63 (61, (014481))))
2.71
60.43 (58.55661 .41)
2.68
60.26 (59. 28/61 88)
PRISM (VDB: PDB now)
2.53
64.67 611.25/64.88 81)
2.71
60.23
(58 388.5666144) (1)
2.68
60.26 (59.29/61 91)
PRISM (VDB: CATH train + new)
(61.27/64.2))
2.71
60.43 (58.56661.4)
(59. (59.29961.99)) (0))
H LENGTH VS. RECOVERY
To further examine how model performance scales with protein length, we stratified the CATH-4.2 test set into length bins and compared amino-acid recovery rates (AAR) between AIDO.Protein and PRISM. As shown in Fig. 4, both models exhibit the expected trend of improved recovery with increasing sequence length, reflecting richer structural context longer backbones. Crucially, across all bins, the distribution of PRISM's 's recovery rates consistently shifts upward relative to AIDO. Prote.Potein-If, indicating that the gains are not restricted to a narrow subset of proteins but hold robustly across varying sequence lengths. Notably, in the shorter length regimes (< 200 residues), where inverse folding is traditionally more challenging, PRISM delivers marked improvements in both median and interquartile range, suggesting that fine- eee-iga retrieval particularly benefits proteins with limited contextual information. At larger lengths (>> 300 residues), the advantage remains evident, with PRISM maintaining higher medians and tighter variability, underscoring its scalability. This distributional analysis complements the average recovery metrics and highlights that PRISM achieves consistent, robust gains across protein lengths, reinforcing its generality beyond aggregate statistics.
I FOLDABILITY ANALYSIS
In Appendix Tab. 14, we provide additional foldability analysis results on TS50 an TS500 datasets.
Table 14: Foldability comparison using AF2 protein folding model. The median and the mean are provided outside and inside the parenthesis, respectively.
Models
TS50
TS500
RMSD 4
sc-TM f
PLDDT
RMSD
Sc-TM f
PLDDT t
Prote.rrt
1.075 (1.2)
0.956 (0.938)
0.949 (0.937)
1.18 (1.372)
0.96 (0.904)
0.951 (0.931)
PRISM (ours)
0.985 (1.13)
0.964 (0.943)
0.95 (0.939)
1.125 (1.351)
0.964 (0.905)
0.952 (0.929)
J CONTRIBUTION OF HYBRID DECODER WITH MHSCA
In Appendix Tab. 15 we provide further ablation on hybrid-attention vs cross-attn-only on test sets CAMEO 2022 and PDB date split.
Table 15: Ablation on hybrid-attention vs cross-attn-only (CAMEO 2022 and PDB date split).
CAMEO 2022
PDB date split
Models
PPL
AAR %
PPL.
AAR % t
PRISM (w/o MHSA)
2.60
64.63 (60.39/64.07)
2.43
66.67 (66.56/67.77) 67.47 (67.37/68.51)
PRISM (full)
2.53
64.63 (61.30/64.81)
2.35
K CONTRIBUTION OF AGGREGATION DEPTH
In Appendix Tab. 16 we provide additional ablation results on the contribution of aggregation depth, i.e. the impact of the number of MHSCA blocks the test CAMEO 2022 and PDB date split.
23

---

Published as a conference paper at ICLR 2026
Table 16: Ablation on the number of MHSCA blocks (CAMEO 2022 and PDB date split).
CAMEO 2022
PDB date split
# of blocks
PPL
AAR
PPL
AAR % 1
N/A (base est.)
2.68
63.52 (60.556/44.1) 17)
2.49
66.27 (66.37/67.64 (64))
1
2.54
64.67 (61.31/64.81)
2.36
67.20 (67. (67.33/68,49
2
2.53
64.63 (61.30/64.81)
2.35
67.47 (67 (67.37/68.5
3
2.54
64.61 (61.27/64.80)
2.36
67.41 (67.33/68.47
L SEQUENCE RECOVERY AND DESIGNABILITY ANALYSIS
Appendix Tab. 17 reports additional results on sequence recovery and designability for the CATH-4.2 test split. We follow the standard evaluation protocol used in Mao et al. (2023) and Yim et al. (2023), computing the proportions of generated sequences whose ESMFold (Lin et al., 2023) structures (%) Here report scores for two PRISM variants: (i) using the original frozen joint encoder G (denoted "orig. 9"), and (ii) using fine-tuned G and iterative refinement (denoted "ft. G, iter."). The results show that PRISM substantially improves both sequence recovery and structural designability, achieving state-of-the-art performance across all metrics.
satisfy the designability criteria: scTM > 0.5 (%) and scRMSD
Table 17: Performance comparison the CATH-4.2 test split for sequence recovery and designability. All baseline reported from Mao al. (2023). Best and second-best scores shown in bold and italic, respectively. Here subscript aido denotes that PrttoProtein-II was used as the base estimator.
Model
PPL
AAR (%)
scTM 0.5 (96)
scRMSD < (%)
PiFold
4.55
51.66
90.98
60.35
LM-Design
4.41
54.41
89.42
58.41
VEN-IF
4.17
54.74
92.37
62.89
VEN-IFE
3.36
62.67
93.29
64.16
PRISM.ndo (orig. 9)
2.71
60.43
94.82
69.73
PRISMaido (ft. G, iter.)
2.65
63.33
95.36
69.55
M RECOVERY-DIVERSITY TRADE-OFF VIA TEMPERATURE SAMPLING
An important question for inverse folding is whether improvements in recovery come at the cost of reduced sequence diversity, since practical design framework must balance both fidelity the native sequence and exploration of alternative solutions. To probe this trade-off, we performed controlled sampling experiments with PRISM by varying the decoding temperature while holding all other factors fixed. For each backbone, we generated 100 candidate sequences at temperatures ranging from 0. to 1.3 and evaluated (i) Recovery Rate, measured as mean sequence identity to the native, and (ii) Diversity, measured 1- average pairwise identity (PID) among the sampled sequences.
-ivverry Frontier: Fig. (left) illustrates the recovery recovery-diversity frontier achieved by PRISM. At very low temperatures (e.g., T 0.1), the model collapses to near-deterministie decoding, yielding high recovery (~0.58) but very limited diversity (~~0.66). As the temperature increases, diversity rises monotonically, reaching 0.63 at T 1.3. Importantly, this comes with only gradual reduction in recovery, which remains above 0.40 even at the highest temperatures. This smooth frontier indicates that PRISM does degenerate into trivial random sampling; instead, maintains meaningful distributional match to the native even under exploratory sampling.
Recovery vs. Temperature. The middle panel confirms that recovery declines as temperature increases, consistent with expectations that flatter distributions produce more varied but less native- like sequences (Fig. 9 (middle)). However, the slope this decline shallow: from T = 0.1 to T = 1.3, recovery drops by only ~0.16 absolute. This robustness suggests that the retrieval- augmented architecture sharpens conditional probabilities enough to preserve signal even under stochastic decoding.
24

---

Published as a conference paper at ICLR 2026
lal Bivessiy-licooeer Fronker
[b] Nicevery vs limprrotree
Id Diversity va Temperature
0.7
66
66-
05
04-
03
02
Daunway
Figure Recovery-diversity trade-off via temperature sampling. (a) Diversity-Recovvery frontier. PRISM main- tains high recovery while offering controllable diversity as temperature increases. (b) Recovery vs. temperature: recovery decreases gradually under stochastic sampling. demonstrating robustness. (c) Diversity vs. temperature: diversity increases nearly linearly, enabling rich alternative designs. Together, these results highlight PRISM's ability support tunable accuracy accuracy-diversity trade-off without collapse.
Table Comparison Base Estimator AJDDO.Proteine and PRISM across multiple benchmarks. We report perplexity (PPL., lower better) and amino acid recovery (AAR, higher better), along with absolute and percentage improvements of PRISM over the base model.
T5500
CAMEO
P9T.1
PM.
AAR
PV1
AAR
P95.
AAR
191
AAR
Basc osiiaate Peacce--))
2n8
M6.19
2,42
69.66
2.68
63.52
2:49
56.27
2.94
58.60
2.90
58.73
2,855
63.83
PRISM (hal))
243
62.92
2.27
70.33
2.55
64.83
2.35
67.47
2.11
04.43
2.08
60.26
2.495
08.87
Relativ 5%
Diversity vs. Temperature. Conversely, Fig. (right) shows that diversity scales nearly linearly with temperature, highlighting PRISM's ability generate rich alternative sequences when encour- aged to explore. Notably, diversity gains are not confined to "noise": even at moderate temperatures (c.g.. 0.7), diversity is doubled relative T 0.1 while recovery remains 0.50.
Takeaway. These results demonstrate that PRISM supports a controllable accuracy siuuraccy-dvvrsity trade- off without collapsing at either extreme. By adjusting a single temperature parameter, users can shift scamlessly between high- sighffideiity recovery (for benchmarking) and diverse sequence generation (for design). This flexibility is rarely observed in prior inverse folding systems, which often either maximize recovery the cost of trivial diversity sacrifice fidelity under high-temperature sampling. The ablation therefore underscores PRISM's strength as not only an accurate but also a versatile framework for conditional protein design.
N
USAGE OF LARGE LANGUAGE MODELS (LLMS) IN PAPER WRITING
LLMs used polish the writing. It not used for retrieval, discovery, research ideation.
25


Reviews:
[{"review_id": "rix9i8cGTa", "reviewer": "Reviewer_5aXk", "summary": "The paper introduces PRISM, a novel framework for protein inverse folding.\nThe authors identify a key limitation in existing deep learning models: they lack an explicit mechanism to reuse conserved, fine-grained structure-sequence patterns (e.g., motifs) found in known proteins. To address this, PRISM is proposed as a multimodal retrieval-augmented generation (RAG) framework.\nThe core idea is to supplement end-to-end generation with a memory-based approach. The method works in several steps:\n(1) construct a vector database by encoding known protein structures and sequences into contextualized representations.\n(2) given a target 3D structure, leverage a base model to generate an initial sequence estimate, which is used to retrieve the vector database to obtain the most similar potential motifs.\n(3) a novel hybrid self-cross attention (MHSCA) decoder then integrates the backbone information with the retrieved motifs to generate the sequence.\nThe method is theoretically grounded as a latent-variable probabilistic model. Empirically, PRISM improves upon all baselines in both sequence accuracy (AAR) and structural foldability (RMSD and TM-score) across five major benchmarks (CATH-4.2, TS50, TS500, CAMEO 2022, and PDB date split). Extensive and thorough ablation and analysis demonstrate the effectiveness of the proposed modules.", "questions": "See above weaknesses.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The core idea of explicitly retrieving fine-grained, residue-level motifs is an intuitive and novel solution to reuse the conserved local patterns which are central to protein function.\n2. The method demonstrates clear and consistent improvement over existing models across all five evaluation benchmarks. This improvement is shown not only in sequence recovery metrics but also in more practical foldability scores.\n3. The authors show that the significant improvement brought by retrieval-based method are achieved with a negligible runtime overhead (only a 14.3% compared to the base estimator), making PRISM a practical and scalable solution.\n4. The authors provide extensive ablation studies to justify their design choices. They successfully prove the contribution of the retrieval mechanism and improvement over varying sequence lengths, optimize the number of retrieved entries, validate the hybrid MHSCA decoder design, and analyze the effect of decoder depth.", "weaknesses": "1. The main results in Table 1 should include structure-related metrics (e.g., scTM and RMSD) instead of relying solely on AAR and PPL for evaluation. Although scTM and RMSD are assessed in Table 4, this comparison is limited to DPLM2-3B and AIDO and lacks a broader set of baselines.\n2. Several critical details are missing from the main body text:\n- What protein library was the vector database built upon?\n- In Figure 2, do the structure embedding and joint embedding originate from the same model? Are they kept fixed during training?\n- A new term, \"base estimator,\" appears in line 315. I speculate it is used to estimate the initial sequence during inference. If so, the authors should explicitly clarify it.\n3. The sampling process relies on the \"base estimator\" to generate an initial sequence, which is then used for subsequent retrieval. This raises a key question: does the quality of this initial sequence significantly impact the final generated sequence? The authors have not provided an analysis of this potential dependency.\n4. The retrieval-based generation method may cause the generated protein sequences to be overly similar to natural proteins (as suggested by the high AAR of ~70% in Table 2), which will limit the novel protein design. Furthermore, it's uncertain if the method's performance will drop significantly when used on a novel backbone that not similar to any entries in the database.\n5. Minor issue: there is a formatting error where Table 3 overlaps with the main text.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KgjEbia7Ev", "reviewer": "Reviewer_3CLu", "summary": "This paper proposes PRISM, a multimodal retrieval-augmented generation (RAG) framework for protein inverse folding. The key idea is to augment a generative inverse-folding model with fine-grained structure–sequence retrieval from a large protein database. Instead of relying solely on end-to-end transformer inference, PRISM retrieves localized structural motifs or fragments and conditions decoding through a hybrid self-cross-attention decoder that integrates both query (target structure) and retrieved (reference) contexts.", "questions": "- Some of the figures have overlaps with the text, and some tables are too small to read.\n- Section 4.2 is empty.\n- The figures look too busy to read. You can either extend it to full size or remove some details.", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The problem setup and latent-variable formulation provide theoretical grounding for the proposed method.\n- The ablation studies and runtime analysis provide clear evidence of the effectiveness of the proposed method.", "weaknesses": "- The author did not provide the database construction process. What structures have been used in the retrieval? Even without exact overlap, similar folds from the same CATH family may leak local sequence priors, effectively making the task easier.\n- Equation (13): The $p(Z \\mid B, \\mathcal{R} )$ should be $p(Z \\mid B, \\mathcal{R} , \\mathcal{E})$, where $\\mathcal{E}$ is the encoding of the retrieved fragments.\n- The PGM derivation seems to prove something trivial. The actual training process is just a dense retrieval with some next-token prediction loss.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hlJjsgF02M", "reviewer": "Reviewer_1DYG", "summary": "The paper proposes a RAG system for the protein inverse folding task. It uses AIDO.Protein-IF as a joint protein sequence-structure encoder to generate residue-level embeddings for every protein in the knowledge base, which are then used to construct a RAG system. During inference, an initial guess of the protein sequence is first generated from a base estimator, then used to index and retrieve top-k candidates from the RAG. The paper claims SOTA performance of their approach on various benchmarks compared to existing baselines.", "questions": "Please address my two major concerns in the “Weaknesses” section first. I will reassess after the rebuttal. Two other miscellaneous questions are as follows: \n* The RAG system likely limits the generalization ability of the model on unseen protein backbones. Have you tested the performance of the model on novel or orphan proteins that are dissimilar to any known protein in the RAG? \n* How much disk space do we need to store residue-level embeddings for every protein in the knowledge base? And what size of RAM did you use for experiments in section 4.6?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper addresses the protein inverse folding problem, which is a core challenge in proteomics and biology. \n* The residue-level RAG system allows flexible reuse of local structure and sequence patterns. \n* The performance of the model is impressive on different benchmarks for its SOTA amino-acid recovery rate, perplexity and foldability.", "weaknesses": "* The model relies heavily on the AIDO.Protein-IF model in both RAG system construction and inference. The not-so-significant improvement in performance of the proposed method compared to AIDO.Protein-IF makes the entire work more like a second-stage refinement rather than a completely novel end-to-end framework. \n* I suspect that the performance of the model might rely heavily on the quality of the initial guess from the base estimator AIDO.Protein-IF. This may be a particular problem when the retrieval system is at residue-level. The paper provides little further interpretation or ablation study on this matter.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rqXfxaI5eC", "reviewer": "Reviewer_CkVi", "summary": "The paper addresses the protein inverse folding problem: designing a sequence for a target 3D structure. The authors identify that existing methods lack mechanisms to reuse conserved, fine-grained structure-sequence patterns. They propose PRISM, a retrieval-augmented generation (RAG) framework that retrieves fine-grained \"potential motif\" representations from a vector database. These retrieved embeddings are integrated with global backbone context using a hybrid self-cross attention (MHSCA) decoder. The method is formulated as a latent-variable probabilistic model. Experiments across five benchmarks show that PRISM achieves new state-of-the-art results in sequence recovery and perplexity, while also improving structural foldability metrics.", "questions": "Let’s denote the training set as A, the RAG database as B. The paper compares models trained on A and augmented with B, with models trained on A. However, since the information in B is also used, could you compare models trained on A and augmented with B, with models trained on A+B?\n\nCould you explain the independences for Eq. 3?\n\nIs the joint encoder $\\mathcal{G}$ (AIDO.Protein-IF) kept frozen during training? If so, have the authors considered methods to make the retrieval step differentiable (e.g., via relaxation) to enable end-to-end fine-tuning of the query encoder?\n\nAt inference, an \"initial sequence estimate $\\tilde{S}^q$\" is required to generate the query embedding $\\hat{\\mathcal{E}}^q$. How sensitive is the model's final performance to the quality of this initial sequence? For instance, what is the drop in AAR if $\\tilde{S}^q$ is generated by a weaker baseline like ProteinMPNN instead of AIDO.Protein-IF?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper demonstrates strong empirical results, consistently outperforming existing methods, including its own base model (AIDO.Protein-IF), across five challenging benchmarks (CATH-4.2, TS50, TS500, CAMEO 2022, PDB date split).\n\nThe application of a retrieval-augmented generation (RAG) framework to inverse folding at a fine-grained, residue-level (\"potential motif\") is a novel contribution. This provides an explicit mechanism to reuse conserved local patterns, which is a limitation in prior work\n\nThe paper includes an extensive set of ablations that validate the core design choices, including the impact of the number of retrieved entries, the contribution of the hybrid decoder, the effect of aggregation depth, and the saturation of the vector database. The analysis of the recovery-diversity trade-off is also a valuable addition.", "weaknesses": "Presentation needs improvement. Some tables and figures overlap with text body, e.g., tab 3 and fig 3. Sec. 4.2 is missing. ‘Fig.’ and ‘Figure’ are both used in text body.\n\nSome baselines, like SPDesign and VFN-IFE, that report higher metrics (on CATH), are not compared.\n\nThe training objective (Eq. 20) optimizes the parameters of the aggregation module ($\\theta_Z$) and a structure encoder ($\\theta_B$). However, the joint encoder $\\mathcal{G}$ (identified as AIDO.Protein-IF) that produces the query embeddings $\\hat{\\mathcal{E}}^q$ and populates the database appears to be frozen, as the TopK retrieval step is non-differentiable. This prevents end-to-end training and means the model cannot learn to improve its representations for better retrieval; it can only learn to use the results from a fixed, and potentially suboptimal, retrieval set.\n\nThe \"PRISM (str. enc.)\" ablation (Tables 1, 2, 3) is vaguely described. The paper states this variant replaces the joint encoder with a \"purely structure-based encoder (ProteinMPNN-CMLM)\" and retrieval operates \"only over structural embeddings\". This seems to contradict the framework's core \"multimodal representation\" $\\mathcal{E} = \\mathcal{G}(B,S)$ (Eq. 5) and the definition of the database $D$. It is unclear how the multimodal vector database is constructed or queried in this \"structure-only\" setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)