RFD4-Proteina: Designing Proteins with Reward-Guided Inference-Time Search

Protein Design Is Essentially a Search Problem

Designing a protein sequence with just eight amino acids already yields more than 25 billion possible combinations. In a technical article, the NVIDIA developer blog frames protein design directly as a "search problem" — the sequence space explodes exponentially, exhaustive enumeration is infeasible, and any generative model can only sample a limited number of candidates from this space.

This is precisely the context in which RFD4 and Proteina are combined: generative models propose candidates, a reward model scores them, and search capability is layered onto the inference stage rather than relying on a single forward generation pass.

What Roles Do RFD4 and Proteina Play?

RFD (RoseTTAFold Diffusion) is a diffusion-model approach to protein structure generation, and RFD4 is a newer version along this line, responsible for "denoising" foldable protein backbones from noise. Proteina, meanwhile, is a generative model for protein sequences and structures.

The engineering implication of combining the two is: RFD4 provides a structure-level generative prior, Proteina provides sequence-level modeling capability, and the reward model screens and guides candidates during inference. The authors call this entire pipeline reward-guided inference-time search.

The core mechanism of diffusion models in protein structure generation is "noising-denoising": during training, noise is gradually added to real structures and the model learns to reverse the denoising; during inference, it starts from pure noise and reconstructs the structure through multiple iterative steps. The RFD series (RoseTTAFold Diffusion) combines this mechanism with RoseTTAFold's structure prediction capability, making generated results more physically foldable. RFD4, as an iterative version, typically improves sampling efficiency, structural diversity, and complexity. Proteina, as a joint sequence-structure model, often adopts architectures similar to language models or graph neural networks to learn the mapping between sequences and 3D structures. The significance of combining the two is: the structure prior constrains "which backbones are reasonable," while the sequence model determines "which amino acid combinations can fold into that backbone" — the two are complementary and reduce wasted search.

Inference-Time Search: Spending Compute on Screening Rather Than Retraining

The traditional approach is to train a larger model or add preferences to the training objective. Inference-time search takes a different path: model weights stay fixed, and a reward signal is introduced at each sampling step or on each candidate to guide the search direction.

The advantage of this paradigm is modularity. The reward function can be swapped — to optimize folding stability, swap in a stability scorer; to optimize binding affinity, swap in an affinity model — without retraining the generative model. The cost is increased inference compute, since multiple candidates must be generated and evaluated. The article's core argument is precisely this: for problems with enormous search spaces like protein design, spending more inference compute in exchange for higher-quality candidates is worthwhile.

Best-of-N sampling is the simplest form of inference-time search: independently generate N candidates, score them with the reward model, and take the highest-scoring one. More refined strategies include step-by-step guidance (e.g., injecting reward gradients into each denoising step of a diffusion model) or particle-filter-style search. Unlike RLHF in the language model domain, inference-time search does not update the generative model's parameters, so it does not introduce training instability and makes it easy to switch objectives quickly. The theoretical basis is: when the generative model covers enough reasonable candidates, an accurate reward model can pick high-quality solutions from the candidates, thereby trading inference compute for quality improvement. But this also requires the reward model to be highly aligned with the true objective; otherwise, the search will amplify the reward model's biases.

How Reward Guidance Affects Sampling

The key to reward guidance lies in "where to intervene." Intervening too late means merely ranking already-generated results, with limited gains; intervening too early or too strongly compresses sampling toward the reward model's biased regions, sacrificing diversity.

The approach given in the article is to use the reward as a guiding signal during the search process, evaluating while generating candidates, gradually pushing the search toward high-reward regions. This is conceptually akin to best-of-N sampling in the language model domain and reward-model-based decoding strategies, except that the scoring target shifts from text to protein structures and sequences.

In the sampling process of diffusion models, reward guidance is typically implemented by adjusting the denoising direction: at each step, compute the gradient of the reward function with respect to the current sample on the predicted noise or denoised result, and add it to the denoising update with a certain weight. This is similar to classifier guidance or classifier-free guidance, but the guiding signal comes from an external reward model rather than a conditional label. Guidance strength (i.e., the weight of the reward gradient) is a key hyperparameter: too strong leads to reduced sample diversity and mode collapse, while too weak makes the guidance ineffective. In practice, annealing strategies are often used — weak guidance early in sampling to preserve diversity, strong guidance later to improve quality.

Practical Implications for Protein Design Workflows

For teams doing protein engineering, the practical value of this approach lies in decoupling "generation" and "evaluation." The generative model can remain general-purpose, while the evaluation stage is customized per project — simply adjust the reward function for a specific target, a certain stability requirement, or a particular expression system.

This also means the bottleneck of the workflow shifts from training to the quality of the reward function's definition. If the reward model is misaligned with the true experimental objective, the harder the search, the further it deviates. This is worth verifying first when using such methods.

Open Questions Worth Watching

How the compute cost of inference-time search scales with the number of candidates — the article does not provide a complete cost curve. Distribution shift between the reward model and the generative model, and how to set reward weights under multi-objective optimization (simultaneously constraining stability, affinity, and expressibility), are engineering details that will be encountered in practical deployment.

Protein design is moving from "generating a structure" to "targeted search in a vast sequence space," methodologically closer to recommender systems and combinatorial optimization than to pure generative modeling. The RFD4-Proteina combination is a concrete example of this shift.