Title: Rationalized All-Atom Protein Design with Unified Multi-Modal Bayesian Flow
Venue: NeurIPS 2025
Accepted: yes
Keywords: Bayesian Flow Networks

Protein Generation
Abstract: Designing functional proteins is a critical yet challenging problem due to the intricate interplay between backbone structures, sequences, and side-chains. Current approaches often decompose protein design into separate tasks, which can lead to accumulated errors, while recent efforts increasingly focus on all-atom protein design. However, we observe that existing all-atom generation approaches suffering from an information shortcut issue, where models inadvertently infer sequences from side-chain information, compromising their ability to accurately learn sequence distributions. To address this, we introduce a novel rationalized information flow strategy to eliminate the information shortcut. Furthermore, motivated by the advantages of Bayesian flows over differential equation–based methods, we propose the first Bayesian flow formulation for protein backbone orientations by recasting orientation modeling as an equivalent hyperspherical generation problem with antipodal symmetry. To validate, our method delivers consistently exceptional performance in both peptide and antibody design tasks.

Rationalized All-Atom Protein Design with Unified Multi-Modal Bayesian Flow

Abstract

Designing functional proteins is a critical yet challenging problem due to the intricate interplay between backbone structures, sequences, and side-chains. Current approaches often decompose protein design into separate tasks, which can lead to accumulated errors, while recent efforts increasingly focus on all-atom protein design. However, we observe that existing all-atom generation approaches suffering from an information shortcut issue, where models inadvertently infer sequences from side-chain information, compromising their ability to accurately learn sequence distributions. To address this, we introduce a novel rationalized information flow strategy to eliminate the information shortcut. Furthermore, motivated by the advantages of Bayesian flows over differential equation-based methods, we propose the first Bayesian flow formulation for protein backbone orientations by recasting orientation modeling as an equivalent hyperspherical generation problem with antipodal symmetry. To validate, our method delivers consistently exceptional performance in both peptide and antibody design tasks. Our code, checkpoint, and designed PDBs can be found in https://github.com/GenSI-THUAIR/ProBayes .

1 Introduction

Proteins are fundamental to life, carrying out essential functions in biological processes [Clark and Pazdernik, 2012, Whitford, 2013]. Designing functional protein stands as a grand scientific challenge with great potential to modulate life processes Notin et al. [2024]. Prevalent approaches to protein design often decompose this task into a set of more manageable sub-problems, i.e. generating partial protein components ( e.g. backbone), and subsequently designing the remaining components ( e.g. sequence and side-chain) conditioned on those generated sub-units [Watson et al., 2023, Yim et al., 2023a, Bose et al., 2023, Lin and AlQuraishi, 2023].

However, in such a loosely connected pipeline, each compartment can only capture a partial protein component, which may lead to an incongruous whole. Furthermore, the three-dimensional structure of protein is simultaneously determined by both the backbone and the side chain [Schulz and Schirmer, 2013] while each model in pipelines can only capture and design a partial protein component. Moreover, early fixed inaccurate partial components accumulate and propagate errors downstream [Kong et al., 2023]. These issues all call for a unified end-to-end all-atom protein design model.

Understandably, achieving all-atom modeling is inherently challenging, due to dependency between the sequence and the coordinates (namely, the number and arrangement of side-chain atoms vary with the sequence. e.g. Lysine residue has four side-chain rigid groups while the Valine residue only has

one, see Fig. 1). We observe that existing all-atom modeling practices [Li et al., 2024, Chu et al., 2024, Kong et al., 2023, Zhu et al., 2024] suffer from an information shortcut issue: when the network is trained to jointly denoise both the sequence and the side chain, the network can infer sequence from even noised side-chain for hacking the training loss, thereby failing to properly learn the sequence distribution. In this paper, we propose to address this by designing a rationalized information flow, which avoids such problem while still preserving network's ability to capture all-atom geometry.

Furthermore, current works employ stochastic differential equations (SDE) ( e.g. diffusion models [Song et al., 2020, Ho et al., 2020]) or ordinary differential equations (ODE) ( e.g. flow matching [Lipman et al., 2022, Chen and Lipman, 2023]) to generate each protein component. However, these SDE and ODE methods are known to be prone to discretization errors , which can hinder both the quality and efficiency of the generation process [Karras et al., 2022, Tong et al., 2024, Sabour et al., 2024]. Bayesian Flow Networks (BFNs) [Graves et al., 2023] are a new class of generative models that build a generation process through Bayesian inferences over noised data, which effectively circumvents the discretization error without the need to discretize continuous differential equations 1 and has been verified to be effective in molecule generation tasks [Song et al., 2023, Wu et al., 2025].

Motivated by this, we introduce ProBayes, a new all-atom protein design model with Bayesian flows. A key theoretical challenge needs to be tackled: The orientations of the protein backbone frame are SO(3) matrices { R : R ∈ R 3 × 3 , RR T = I, det( R ) = I } , while there is no existing Bayesian flow workable on SO(3). Directly conducting Bayesian inference on SO(3) is impractical, because current SO(3) distributions ( e.g. IGSO(3)) cannot guarantee conjugacy under Bayesian inference, which is pivotal to build Bayesian flow. To overcome this issue, we propose to transform the problem into an equivalent hypersphere one, by generating SO(3) elements on hypersphere with antipodal symmetry. Additionally, ProBayes demonstrates clear advantages over the current prevalent SO(3) flow matching method [Yim et al., 2023a], which relies on a hand-crafted inference annealing technique to adjust the generation flow. Our contributions can be summarized as follows:

Weintroduce ProBayes, a new protein generation approach with Bayesian flow, by transforming SO(3) generation into an equivalent problem of hypersphere generation with antipodal symmetry.

We identify and resolve the issues of information shortcut in all-atom protein generation with a novel rationalized information flow.

Extensive experiments demonstrate that ProBayes consistently achieves significant improvements over existing methods on peptide design and antibody design tasks, e.g. achieving a DockQ score of 0.74 on PepBench.

2 Preliminary

Hypersphere and hypertorus Given an Euclidean space R d , the hypertorus T d is represented as the Cartesian product of d flat toruses T 1 = { ϑ : ϑ ∈ [0 , 2 π ) } . The hypersphere S d -1 located in R d is defined as: S d -1 = { x : x ∈ R d , || x || 2 = 1 } . Superscript d denotes the degree of freedom.

Protein representation A protein P is a biomolecule composed of L amino acid residues, each defined by its type, backbone frame, and side-chain angles Fisher [2001]. All aminoacid types form the protein sequence S = [ s 1 , . . . , s L ] , where s l ∈ { 1 , . . . , 20 } , ∀ l ∈ { 1 , 2 , . . . , L } . The backbone frame of each residue is parameterized by the position vector of the C α atom, p ∈ R 3 , an orientation matrix

𝐂𝐀

𝐍

𝐂

𝐎

𝑹

𝐂𝐁

𝐂𝐆

𝐂𝐃

𝐂𝐄

𝐍𝐙

𝝌𝟏

𝝌𝟑

𝝌𝟒

𝐂𝐀

𝐂

𝐎

𝝍

𝐍

𝐂𝐁

𝐂𝐆𝟏

𝑹

𝝍

𝒮𝑙

=LYS

𝒮𝑙+1

= VAL

𝐂𝐆𝟐

𝝌𝟏

R ∈ SO (3) that determines the relative positions of the N and C atoms, and a torsion angle ψ which governs the position of the O atom [Engh and Huber, 2006, Jumper et al., 2021, Yim et al., 2023b]. The side-chain can be represented by up to four rigid groups with χ denoting the flexible torsion angle between rigid groups. The number of side chain torsions and arrangements depend on the specific residue type as shown in Fig. 1.

Bayesian flow networks As a new class of generative models, Bayesian flow networks (BFNs) [Graves et al., 2023] build a generation process through Bayesian inferences using noised data, which is fundamentally different from diffusion models [Song and Ermon, 2019, Ho et al., 2020, Song et al., 2020] and flow matching [Lipman et al., 2022].

Formally, BFN parameterizes a belief for clean data x over the space with a distribution family, e.g. Gaussian distribution for x ∈ R . The belief starts from an uninformative prior θ 0 , which is updated according to the received noised data y i ∼ p S ( y i | x , α i ) , i ∈ { 1 , 2 , . . . , n } with different levels of signal-to-noise ratio parameter α i . The update transition from θ i -1 to θ i is termed Bayesian update function θ i = h ( θ i -1 , y i , α i ) , which is derived from the Bayesian theorem according to the form of the belief distribution and noise distribution p S . With each θ i -1 as input, the neural network Ψ is trained to transform the belief θ i -1 as receiver distribution p R , to minimize the evidence lower-bound based training loss:

where p F is the Bayesian flow distribution representing the distribution of θ at time step i . Compared to diffusion models and flow matching, BFN enjoys the following advantages: (1) free from discretization error without the need to discretize continuous differential equations; (2) precisely adjusted step sizes according to the entropy parameter α i , allowing fast and accurate generation demonstrated in [Song et al., 2023, Wu et al., 2025]. Motivated by these advantages, we aim to design a new SO (3) Bayesian flow as a more effective and efficient approach to protein design.

3 Bayesian flow for SO(3) generation

In this section, we introduce the first known Bayesian flow for generating SO(3) data, 2 which is derived from constructing a Bayesian Flow on the hypersphere S 3 with antipodal symmetry.

3.1 Theoretical challenge in directly building SO(3) Bayesian flow

Difficulty in guaranteeing Bayesian conjugacy The Bayesian inference approach used in BFN requires the Bayesian conjugacy of the distribution, i.e. the posterior distribution after observing the noised data should be in the same distribution family as the prior belief distribution. However, existing SO(3) distributions ( e.g. the IGSO(3) distribution used in Yim et al. [2023a,b], Bose et al. [2023]) are intractable to derive an analytical Bayesian update to ensure its conjugacy due to its infinite sum. Furthermore, other well-known existing statistical distributions over rotation matrix [Watson, 1966, Downs, 1972, Khatri and Mardia, 1977] also involve intractable probability density function to guarantee their Bayesian conjugacy, e.g. matrix Fisher distribution.

Therefore, we resort to other widely used rotation representations such as Euler angles, axis-angle and unit quaternions, instead of the rotation matrix. As discussed below, we demonstrate that both Euler angles and the axis-angle representation fail to adequately capture the group structure of SO(3), whereas unit quaternions are able to do so effectively, which is used in this paper. Additionally, the unit quaternion's intrinsic manifold, i.e. the hypersphere, is significantly simpler than the structure of SO(3), and there exists a normal distribution over the hypersphere which satisfies Bayesian conjugacy.

Non-covering of Euler angle and axis-angle representation We observe that some rotation representations cannot capture the group structure of SO(3). For example, Euler angles suffer from the singularity problem of gimbal lock [Hoag, 1963]. From the perspective of topology, such problems reveal the fact that the hyper-torus T 3 cannot cover the SO(3). Practically, we validate the ineffectiveness of Euler angles and axis-angle representation via toy data generation experiment (see Appendix C). The definition of the cover map is as follows:

𝑹

𝑝 𝑹 , 𝑹 ∈ 𝐒𝐎(𝟑)

𝒒

-𝒒

𝑝 𝒒

𝑝 𝒒

= 𝑝 -𝒒

⇔𝑝 𝑹 ,𝝓 𝒒

= 𝝓 -𝒒

= 𝑹

𝒒

-𝒒

𝑶

𝐒𝐎(𝟑)

distribution

hypersphere

𝕊 3

distribution

w/o

antipodal symmetry

hypersphere

𝕊 3

distribution

w/

antipodal symmetry

Definition 1 (Covering Map, Hatcher [2005]) . Given topological space X and Y , a covering map is a surjective continuous map f : X → Y such that: every element y ∈ Y has an open neighborhood U y ⊆ Y where f -1 ( U y ) is a disjoint union of open sets, each of which is mapped by f homeomorphically onto U y .

The quaternion representation of rotation can be used to build an SO(3) Bayesian flow because the space of this representation ( i.e. the hypersphere) can cover SO(3). Intuitively, we aim to build a generative model in a larger topological space while preserving the original group structure of SO(3).

3.2 Transforming SO(3) generation into hypersphere S 3 generation

Fig. 2 provides an intuitive illustration of transforming SO(3) generation into hypersphere S 3 generation. Firstly, we introduce the group structure of unit quaternion rotation representation, i.e. SU(2) group and its corresponding manifold hypersphere S 3 . Then, we explain the proposed key constraint of antipodal symmetry to equivalently transform SO(3) generation into hypersphere S 3 generation, and introduce a simple approach to satisfy such constraint.

SU(2) group SU(2) is the special unitary group composed by unitary matrices U ∈ { U : U ∈ C 2 × 2 , UU T = I, det( U ) = 1 } . SU(2) is diffeomorphic to the hypersphere S 3 [Hall, 2013]. Hence, we can represent SU(2) elements using four real numbers, namely the quaternion q = [ a, b, c, d ] T with

Double covering between SO(3) and SU(2) In fact, SU (2) is the double covering group of SO (3) [Hall, 2013] and there exists a surjective homomorphism ϕ : SU (2) → SO (3) such that for every g SU (2) ∈ SU (2) , we have ϕ ( g SU (2) ) = ϕ ( -g SU (2) ) = g SO (3) ∈ SO (3) . Therefore, we can transform the representations of these two groups between one another. Specifically, given a unit quaternion q = a + b i + c j + d k , the corresponding rotation matrix R = ϕ ( q ) = ϕ ( -q ) can be obtained by:

Inversely, rotation matrix R can be transformed into the corresponding double unit quaternion q 1 = a + b i + c j + d k and q 2 = -( a + b i + c j + d k ) from Eq. (2):

As illustrated above, each rotation matrix R corresponds to two unit quaternions q and -q , which are equivalent when representing the three-dimensional rotation operation. This phenomenon is termed as SU(2)'s double-covering of SO(3) [Altmann, 2005], which results in the discrepancy between generating rotation matrices and directly generating unit quaternions. We propose a natural solution to this problem, by adding an antipodal symmetry constraint to the modeled unit quaternion distribution on the hypersphere S 3 :

Definition 2 (Antipodal Symmetry) . For a distribution p ( x ) supported on S d , we say this distribution is antipodal invariant if p ( x ) = p ( -x ) , ∀ x ∈ S d . Correspondingly, we say a function f operating q ∈ S 3 is antipodal equivariant if f ( -q ) = -f ( q ) .

Now, with antipodal symmetry, we can transform the generative modeling problem of SO(3) data into an equivalent one on S 3 with antipodal symmetry. We formulate the proposed transformation with the following preposition:

Proposition 1. Every distribution p SO (3) ( R ) supported on SO(3) can be bijectively mapped to a distribution p S 3 ( q ) supported on S 3 which satisfies p S 3 ( q ) = p S 3 ( -q ) .

In fact, we find that antipodal equivariance can be easily ensured by implementing the Invariant Point Attention (IPA) module [Jumper et al., 2021] based on quaternion representation (instead of the rotation matrix) without extra operations or complexity.

3.3 Bayesian flow on hypersphere S d

Equipped with Proposition 3, we can now derive a Bayesian flow for generating SO(3) on the hypersphere S 3 . In this section, we introduce the first derivation of the Bayesian flow on S d , which can be applied to generate unit quaternions with d = 3 for SO(3) generation.

Von Mises Fisher distribution and Bayesian update function h We choose von Mises Fisher (vMF) distribution [Mardia and Jupp, 2009] to parameterize the belief over S d because of its tractability and Bayesian conjugacy. With the location parameter µ and concentration parameter κ , the probability density function of the von Mises-Fisher distribution is:

where || µ || = || x || = 1 , ν = d/ 2 -1 , κ ≥ 0 , I ν ( κ ) denotes the modified Bessel function of the first kind at order ν . Given the prior belief parameterized by vMF distribution with parameter µ i -1 , κ i -1 , the noised sample y i from sender distribution with unknown clean x , and the accuracy α i describing the entropy of y i , the Bayesian update function h for von Mises Fisher distribution is deducted as:

The proof of the above Eq. (5) can be found in Appendix B. This update function operates in an intuitive manner, analogous to the moment equation in physics, where the accuracy/entropy parameters κ and α represent the mass, and µ and y denote the velocity. We define θ def = κ µ to compactly denote the parameter of the vMF distribution { κ, µ } .

Bayesian flow distribution Based on the Bayesian update function, the Bayesian flow distribution on hypersphere S d over vMF distribution can be obtained by conducting a series of Bayesian inferences. With the notation α 1: i def = { α 1 , α 2 , . . . , α i } , the Bayesian flow distribution is expressed as (details in Appendix B):

Now with Eq. (6) and Proposition 3, we are able to generate the orientation of the backbone frame with a Bayesian flow, which is free from discretization error and enjoys entropy-adjusted dynamic step size. The antipodal symmetry of the modeled distribution is guaranteed as follows:

Proposition 2. With Ψ as an antipodal equivariant function, i.e. Ψ( -θ q ) = -Ψ( θ q ) , and uniform prior p ( θ q 0 ) on S 3 , the marginal distribution defined by multiple Markov transitions p Ψ ( θ q n ) = p ( θ q 0 ) ∫ Π n i =1 p F ( θ q i | Ψ( θ q i -1 ); α 1: i ) d θ q 1: n -1 is antipodal invariant, which is equivalent to a distribution p Ψ ( R ) .

Note that, as the rotation and orientation are ubiquitous in geometric deep learning, such SO(3) Bayesian flow can be used in other applications, e.g. , molecular docking [Corso et al., 2022].

4 Rationalized all-atom protein design

In this section, we present our approach to jointly generating the complete protein structure, including the sequence, backbone and sidechain. We emphasize that such a co-design process is inherently tricky, and naive joint training risks information shortcut between sequence and side-chain.

ProBayes parameterizes the belief over the protein space via each modality's belief parameter θ P = { θ S , θ p , θ R , θ ψ , θ χ } with a predefined number of generation steps n . For training, we first sample the time index i ∼ U { 1 , . . . , n } and θ i -1 ∼ p F for each modality. The neural network Ψ( θ P i -1 , i ) is trained to minimize the training loss according to Eq. (1). After sufficient training, the data distribution p data is well captured by the neural network Ψ . As for sampling, each modality starts from an uninformative prior θ P 0 , which is to be updated according to the network output Ψ( θ P , i ) and the Bayesian flow distribution till i = n .

4.1 Bayesian flow for each protein modality

We explain the definition of the Bayesian flow for each protein modality. For detailed training and sampling algorithm, please see Appendix D.

Bayesian flow for orientation matrix R generation We utilize the proposed Bayesian flow Eq. (6) to generate the orientation of the protein backbone frame R . We begin by uniformly converting the orientation matrix R into its corresponding quaternion representation, yielding one of the two equivalent quaternions { q , -q } with equal probability. The SO(3) Bayesian flow distribution is:

where α q i is the predefined accuracy to make the entropy of Eq. (7) linearly decrease with respect to time step i . Sampling θ q i from p q F , the training objective is derived based on Eq. (1) and the KL divergence between vMF distributions (details in Appendix B):

where dot represents the dot product between two quaternions, and Ψ q represents the quaternion part of the network prediction. We prove in Appendix B that such a loss is proportional to the geodesic distance between the predicted and ground truth orientation matrix up to a constant, which further solidifies the transformation between SO(3) and SU(2).

Bayesian flow for C α position p generation The C α positions of the backbone frame are in Euclidean space p ∈ R 3 . Therefore, we use the Gaussian distribution Bayesian flow [Graves et al., 2023] to generate them. The Bayesian flow distribution of C α position p is:

where σ 1 is the hyperparameter controlling the entropy of the Bayesian flow at t = 1 . The variance of the belief Gaussian distribution ln σ 2 ( t ) = ln(1 -γ ( t )) /t does not need to be modeled because it is deterministic w.r.t. time steps.

Bayesian flow for sequence S generation Given that sequences are discrete variables, we utilize the discrete Bayesian flow [Graves et al., 2023] to generate the sequence. With the notation of projection from the class index j to the length K one-hot vector ( e j ) k def = δ jk , where e j ∈ R K , e S def = ( e s 1 , . . . , e s L ) ∈ R K × L , the Bayesian flow distribution of sequence S is defined as:

where β S ( t ) is the predefined accuracy schedule.

Bayesian flow for angles ψ, χ generation The torsional angles ψ and χ are periodic variables on the hypertorus. Hence, we use the periodic Bayesian flow [Wu et al., 2025] to generate them:

We illustrate our method using χ ; the Bayesian flow for ψ has the same form. While both lie on the same manifold, unlike Li et al. [2024], we treat them differently. Specifically, we distinguish χ from the backbone torsion ψ because χ carries partial sequence information. We expand on this in the next section.

0.0

0.2

0.4

0.6

0.8

1.0

Time Step

t

0.5

1.0

1.5

2.0

2.5

3.0

KL Divergence

DKL

[

||

shortcut

(

t

)]

DKL

[

||

t

]

shortcut benefit

(a)

(b)

(c)

(d)

Train Seq Loss

Val AAR

Train SO(3) Loss

w/oShorcutw/Shortcut

w/oShorcut-w/Shortcut

FMw/oShorcutFMw/Shortcut

3

4

0.7

2.5

0.6

3

2

1.5

0.5

2

1

0.4

0.5

0.3

1

Step

0.2

Step

Step

0

50k100k 150k 200k 250k

0

0

50k100k 150k 200k 250k

4.2 The information shortcut problem

As discussed in Sec. 2 and Fig. 1, the residue type S reflects the types of the side-chain with varying numbers of side-chain torsion angles from 0 to 4 [Jumper et al., 2021, Li et al., 2024]. However, diffusion-like models ( e.g. FM and BFNs) are inflexible in capturing varying numbers of random variables. Prior practices [Li et al., 2024, Zhu et al., 2024] tried to address this by padding the number of variables with a fixed padding value ( e.g. 0 in [Li et al., 2024]), then mixing the side chain information with the sequence and backbone.

Shortcut induced by paddings However, we prove that such padding and mixing operations cause information leakage from the even noised side chain information to the sequence prediction. To illustrate the idea, we design a simple threshold-based shortcut function to "predict" the ground truth sequence for residue Alanine and Glycine, which have no side chain torsion angles i.e. χ = [0 , 0 , 0 , 0] :

where circmean is the circular mean of periodic variables [Mardia and Jupp, 2009], with the first two elements representing the predicted probability of Alanine and Glycine, respectively. We compare the KL divergence between the predicted categorical distribution shortcut ( θ χ t ) and the ground-truth one-hot distribution S , against the KL divergence between the noised sequence θ s t and S , as shown in Fig. 3a. The results demonstrate that the shortcut prediction achieves significantly lower KL divergence than the noised sequence most of the time. This indicates the network can exploit shortcut information from the number of side-chain angles/paddings to minimize sequence training loss, rather than genuinely denoising the sequence by learning the underlying sequence distribution.

Shortcut induced by inherent side-chain distributions Moreover, we emphasize that the information shortcut can still exist even in the absence of padding. This is because the side-chain torsion distributions themselves (i.e., the non-padded values) are inherently dependent on the residue types, as shown in Fig. 4. Consequently, the network can exploit these statistical regularities in the side-chain distributions to reduce the sequence prediction loss, thereby hindering the proper learning of the true sequence distribution.

Rationalized information flow Therefore, we propose a rationalized network information flow to avoid the possibility of shortcuts, as illustrated in Fig. 5. Concretely, we avoid involving side-chain information with the sequence Ψ S = Ψ( θ S , θ p , θ R , θ ψ , ❍ ❍ θ χ ) , while maintaining network's capacity to capture the all-atom geometry by the backbone and side-chain mixing module for denoising other modalities. For side-chain prediction, the loss is computed only for non-padded values to prevent the network from learning harmful patterns. The detailed implementation is provided in Appendix F.

Excluding side-chain information from sequence prediction does not negatively impact the performance ceiling of sequence generation. This stems from the fundamental unfair relationship between sequence and side-chain in the data distribution: The sequence serves as the primary determinant of residue types, while side-chain configurations are dependent on the sequence for both side-chain atom types and conformational specifications.

100

0

100

Angle/°

Density

R

Arginine

=-39.7°

=116.1°

1

100

0

100

Angle/°

Density

=-15.3°

=152.9°

2

100

0

100

Angle/°

Density

=-3.3°

=126.3°

3

100

0

100

Angle/°

Density

=3.0°

=130.4°

4

100

0

100

Angle/°

Density

P

Proline

=4.2°

=25.3°

100

0

100

Angle/°

Density

=-2.6°

=33.5°

Position

Sequence

Orientation

Oxygen Torsion

Sidechain Torsion

Seq.

&Backbone

Mixing

Module

Seq. Head

Backbone

&Sidechain

Mixing

Module

Sidechain

Embedder

෡

𝓢

𝜽 𝒑

𝜽 𝒮

𝜽 𝒒

𝜽 𝝍

𝜽 𝝌

ෝ 𝝌

ෝ 𝒑

ෝ 𝒒

෡ 𝝍

Info. Shortcut

Rational Info.

5 Related work

Protein generation There is a body of meaningful protein generation practices focusing on single protein component, e.g. backbone only Yim et al. [2023b], Bose et al. [2023], Geffner et al. [2025], Watson et al. [2023], sequence only Madani et al. [2020], Dauparas et al. [2022], and side-chain only Zhang et al. [2023]. Recently, motivated by the advantages of joint modeling, an increasing number of sequence-structure co-design approaches have emerged, e.g. , backbone-sequence codesign [Campbell et al., 2024, Lisanza et al., 2024] and all-atom protein generation [Martinkus et al., 2024, Li et al., 2024, Chu et al., 2024, Kong et al., 2024, Zhu et al., 2024]. However, the increased information of all-atom geometry also brings challenges: We observe that all-atom generation is prone to the information shortcut problem due to the dependency between sequence and side-chain, which is highlighted and resolved in this paper.

BFN for molecule generation Inspired by BFN's unified generative modeling for continuous and discrete data, GeoBFN [Song et al., 2023] first incorporates molecular geometric constraints into the BFN framework and achieves exceptional generation performance and efficiency. Recently, Wu et al. [2025] proposed to build BFN on the hypertorus for periodic material generation, demonstrating a superior sampling efficiency. However, the rigidity of the protein backbone requires generative modeling on SO(3), which is inherently difficult due to the intractability of ensuring Bayesian conjugacy with SO(3) distributions. In this paper, we address the problem by transforming the modeled space into a simpler one, thereby allowing Bayesian inferences.

6 Experiments

We verify the effectiveness of ProBayes on two protein design tasks, including peptide design and antibody design. For peptide design task, our evaluation configurations follow Kong et al. [2024], including peptide co-design and peptide binding conformation generation. The description of metrics and implementation details can be found in Appendix F. We also include a toy dataset experiment in Appendix C of SO(3) Bayesian flow compared to SDE / ODE methods [Bose et al., 2023].

6.1 Peptide design

Datasets The datasets used for evaluation include PepBench and PepBDB [Wen et al., 2019]. PepBench is constructed from the Protein Data Bank [Berman et al., 2000], containing 6105 complexes and the LNR dataset Tsaban et al. [2022] is utilized as the test set with 93 complexes. PepBDB is a protein-peptide complex dataset containing 8434, 370, and 190 items for training, validation, and test, respectively.

Baselines We evaluate our approach against the following baselines: PepFlow [Li et al., 2024], a peptide design model utilizing Riemannian flow matching; PepGLAD [Kong et al., 2024], which employs latent diffusion techniques for all-atom peptide modeling; and RFDiffusion , a multi-stage protein design framework [Watson et al., 2023] integrated with ProteinMPNN [Dauparas et al., 2022] for sequence generation.

Results We report the results on PepBench and PepBDB in Table 1 and Table 2 for seq-structure codesign and binding conformation generation task, respectively. Across both PepBench and PepBDB datasets, ProBayes outperforms existing methods in energy, native structural fidelity, and design quality. For peptide binding conformation generation task, our method significantly outperforms all baselines. It achieves the lowest RMSDC α , RMSDatom, and DockQ, indicating ProBayes's superior all-atom structure generation accuracy.

6.2 Antibody design

In this experiment, we focus on designing the third Complementarity-Determining Region of the antibody heavy chain (CDR-H3), as it exhibits the greatest variability and plays a dominant role in antigen binding [MacCallum et al., 1996].

Dataset We utilize the SAbDab database [Dunbar et al., 2014] of antibody-antigen complexes as our training dataset and evaluate model performance using the RAbD bench-

mark [Adolf-Bryfogle et al., 2018]. The data-processing procedures, evaluation configurations, and baseline results are from [Ye et al., 2024]

Baselines We compare our method against three existing approaches: dyMEAN [Kong et al., 2023] employs a full-atom geometric encoder with iterative non-autoregressive generation; DiffAb [Luo et al., 2022] jointly diffuses categorical residue types, C α coordinates, and residue orientations; AbDPO and its variant AbDPO++ [Zhou et al., 2024] are methods tailored for optimizing antibody energy based on direct preference optimization.

Results The results are listed in Tab. 3. With consistently better AAR and RMSD, ProBayes also demonstrates significant improvements in energy-based metrics, achieving state-of-the-art results with a lower order of magnitude, surpassing even the energy-optimization-focused models AbDPO and AbDPO++.

6.3 Ablation study

We validate our design considerations through the following ablation experiments using the antibody design task. Results can be found in Tab. 4: (1) antipodal symmetry : we break the model's antipodal symmetry by removing equivariance in the Markov transition i.e. we replace our quaternion-based implementation with the

original IPA implementation [Jumper et al., 2021], which leads to a significant drop in energy-related metrics, demonstrating the necessity of enforcing the proposed constraints. (2) information shortcut : we reintroduce the information shortcut by incorporating side-chain information during sequence prediction. This experiment results in a ∼ 10% drop in AAR and hurts performance across other metrics, indicating the detrimental impact of such shortcuts on model performance. (3) all-atom capability : we replace model-generated side-chains with those obtained from Rosetta's side-chain packing [Alford et al., 2017]. Our model achieves superior performance on energy-based metrics, demonstrating its capability to generate high-fidelity side-chain structures.

7 Conclusion

Weintroduce ProBayes, a novel all-atom protein design model using Bayesian flows, which effectively resolves critical issues of information shortcut that plagues existing methods. Achieving significant improvements over prior methods on peptide and antibody design tasks, ProBayes demonstrates a more robust pathway towards generating functional proteins.

Acknowledgments

This work is supported by the National Science and Technology Major Project (2022ZD0117502), the National Natural Science Foundation of China (Grant No. 62376133), the Wuxi Research Institute of Applied Technologies, Tsinghua University under Grant 20242001120.

References

David P Clark and Nanette J Pazdernik. Molecular biology . Elsevier, 2012.

David Whitford. Proteins: structure and function . John Wiley & Sons, 2013.

Pascal Notin, Nathan Rollins, Yarin Gal, Chris Sander, and Debora Marks. Machine learning for functional protein design. Nature biotechnology , 42(2):216-228, 2024.

Joseph L Watson, David Juergens, Nathaniel R Bennett, Brian L Trippe, Jason Yim, Helen E Eisenach, Woody Ahern, Andrew J Borst, Robert J Ragotte, Lukas F Milles, et al. De novo design of protein structure and function with rfdiffusion. Nature , 620(7976):1089-1100, 2023.

Jason Yim, Andrew Campbell, Andrew YK Foong, Michael Gastegger, José Jiménez-Luna, Sarah Lewis, Victor Garcia Satorras, Bastiaan S Veeling, Regina Barzilay, Tommi Jaakkola, et al. Fast protein backbone generation with se (3) flow matching. arXiv preprint arXiv:2310.05297 , 2023a.

Avishek Joey Bose, Tara Akhound-Sadegh, Kilian Fatras, Guillaume Huguet, Jarrid Rector-Brooks, Cheng-Hao Liu, Andrei Cristian Nica, Maksym Korablyov, Michael Bronstein, and Alexander Tong. Se (3)-stochastic flow matching for protein backbone generation. arXiv preprint arXiv:2310.02391 , 2023.

Yeqing Lin and Mohammed AlQuraishi. Generating novel, designable, and diverse protein structures by equivariantly diffusing oriented residue clouds. arXiv preprint arXiv:2301.12485 , 2023.

Georg E Schulz and R Heiner Schirmer. Principles of protein structure . Springer Science & Business Media, 2013.

Xiangzhe Kong, Wenbing Huang, and Yang Liu. End-to-end full-atom antibody design. arXiv preprint arXiv:2302.00203 , 2023.

Jiahan Li, Chaoran Cheng, Zuofan Wu, Ruihan Guo, Shitong Luo, Zhizhou Ren, Jian Peng, and Jianzhu Ma. Full-atom peptide design based on multi-modal flow matching. arXiv preprint arXiv:2406.00735 , 2024.

Alexander E Chu, Jinho Kim, Lucy Cheng, Gina El Nesr, Minkai Xu, Richard W Shuai, and Po-Ssu Huang. An all-atom protein generative model. Proceedings of the National Academy of Sciences , 121(27):e2311500121, 2024.

Tian Zhu, Milong Ren, and Haicang Zhang. Antibody design using a score-based diffusion model guided by evolutionary, physical and geometric constraints. In Forty-first International Conference on Machine Learning , 2024.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456 , 2020.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems , 33:6840-6851, 2020.

Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747 , 2022.

Ricky TQ Chen and Yaron Lipman. Riemannian flow matching on general geometries. arXiv preprint arXiv:2302.03660 , 2023.

Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusionbased generative models. Advances in neural information processing systems , 35:26565-26577, 2022.

Vinh Tong, Trung-Dung Hoang, Anji Liu, Guy Van den Broeck, and Mathias Niepert. Learning to discretize denoising diffusion odes. arXiv preprint arXiv:2405.15506 , 2024.

Amirmojtaba Sabour, Sanja Fidler, and Karsten Kreis. Align your steps: Optimizing sampling schedules in diffusion models. arXiv preprint arXiv:2404.14507 , 2024.

Alex Graves, Rupesh Kumar Srivastava, Timothy Atkinson, and Faustino Gomez. Bayesian flow networks. arXiv preprint arXiv:2308.07037 , 2023.

Kaiwen Xue, Yuhao Zhou, Shen Nie, Xu Min, Xiaolu Zhang, Jun Zhou, and Chongxuan Li. Unifying bayesian flow networks and diffusion models through stochastic differential equations. arXiv preprint arXiv:2404.15766 , 2024.

Yuxuan Song, Jingjing Gong, Hao Zhou, Mingyue Zheng, Jingjing Liu, and Wei-Ying Ma. Unified generative modeling of 3d molecules with bayesian flow networks. In The Twelfth International Conference on Learning Representations , 2023.

Hanlin Wu, Yuxuan Song, Jingjing Gong, Ziyao Cao, Yawen Ouyang, Jianbing Zhang, Hao Zhou, Wei-Ying Ma, and Jingjing Liu. A periodic bayesian flow for material generation. In The Thirteenth International Conference on Learning Representations , 2025. URL https://openreview.net/ forum?id=Lz0XW99tE0 .

Matthew Fisher. Lehninger principles of biochemistry, ; by david l. nelson and michael m. cox. The Chemical Educator , 6:69-70, 2001.

RA Engh and R Huber. Structure quality and target parameters. 2006.

John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature , 596(7873):583-589, 2021.

Jason Yim, Brian L Trippe, Valentin De Bortoli, Emile Mathieu, Arnaud Doucet, Regina Barzilay, and Tommi Jaakkola. Se (3) diffusion model with application to protein backbone generation. arXiv preprint arXiv:2302.02277 , 2023b.

Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. Advances in neural information processing systems , 32, 2019.

Geoffrey S Watson. The statistics of orientation data. The Journal of Geology , 74(5, Part 2):786-797, 1966.

Thomas D Downs. Orientation statistics. Biometrika , 59(3):665-676, 1972.

CG Khatri and Kanti V Mardia. The von mises-fisher matrix distribution in orientation statistics. Journal of the Royal Statistical Society Series B: Statistical Methodology , 39(1):95-106, 1977.

David Hoag. Apollo guidance and navigation: Considerations of apollo imu gimbal lock. Canbridge: MIT Instrumentation Laboratory , pages 1-64, 1963.

Allen Hatcher. Algebraic topology . Tsinghua University Press, 2005.

Brian C Hall. Lie groups, lie algebras, and representations. In Quantum Theory for Mathematicians , pages 333-366. Springer, 2013.

Ken Shoemake. Animating rotation with quaternion curves. In Proceedings of the 12th annual conference on Computer graphics and interactive techniques , pages 245-254, 1985.

Simon L Altmann. Rotations, quaternions, and double groups . Courier Corporation, 2005.

Kanti V Mardia and Peter E Jupp. Directional statistics . John Wiley & Sons, 2009.

Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. arXiv preprint arXiv:2210.01776 , 2022.

Tomas Geffner, Kieran Didi, Zuobai Zhang, Danny Reidenbach, Zhonglin Cao, Jason Yim, Mario Geiger, Christian Dallago, Emine Kucukbenli, Arash Vahdat, et al. Proteina: Scaling flow-based protein structure generative models. arXiv preprint arXiv:2503.00710 , 2025.

Ali Madani, Bryan McCann, Nikhil Naik, Nitish Shirish Keskar, Namrata Anand, Raphael R Eguchi, Po-Ssu Huang, and Richard Socher. Progen: Language modeling for protein generation. arXiv preprint arXiv:2004.03497 , 2020.

Justas Dauparas, Ivan Anishchenko, Nathaniel Bennett, Hua Bai, Robert J Ragotte, Lukas F Milles, Basile IM Wicky, Alexis Courbet, Rob J de Haas, Neville Bethel, et al. Robust deep learning-based protein sequence design using proteinmpnn. Science , 378(6615):49-56, 2022.

Yangtian Zhang, Zuobai Zhang, Bozitao Zhong, Sanchit Misra, and Jian Tang. Diffpack: A torsional diffusion model for autoregressive protein side-chain packing. arXiv preprint arXiv:2306.01794 , 2023.

Andrew Campbell, Jason Yim, Regina Barzilay, Tom Rainforth, and Tommi Jaakkola. Generative flows on discrete state-spaces: Enabling multimodal flows with applications to protein co-design. arXiv preprint arXiv:2402.04997 , 2024.

Sidney Lyayuga Lisanza, Jacob Merle Gershon, Samuel WK Tipps, Jeremiah Nelson Sims, Lucas Arnoldt, Samuel J Hendel, Miriam K Simma, Ge Liu, Muna Yase, Hongwei Wu, et al. Multistate and functional protein design using rosettafold sequence space diffusion. Nature biotechnology , pages 1-11, 2024.

Karolis Martinkus, Jan Ludwiczak, Kyunghyun Cho, Wei-Ching Liang, Julien Lafrance-Vanasse, Isidro Hotzel, Arvind Rajpal, Yan Wu, Richard Bonneau, Vladimir Gligorijevic, and Andreas Loukas. Abdiffuser: Full-atom generation of in vitro functioning antibodies, 2024. URL https: //arxiv.org/abs/2308.05027 .

Xiangzhe Kong, Yinjun Jia, Wenbing Huang, and Yang Liu. Full-atom peptide design with geometric latent diffusion. arXiv preprint arXiv:2402.13555 , 2024.

Zeyu Wen, Jiahua He, Huanyu Tao, and Sheng-You Huang. Pepbdb: a comprehensive structural database of biological peptide-protein interactions. Bioinformatics , 35(1):175-177, 2019.

Helen M Berman, John Westbrook, Zukang Feng, Gary Gilliland, Talapady N Bhat, Helge Weissig, Ilya N Shindyalov, and Philip E Bourne. The protein data bank. Nucleic acids research , 28(1): 235-242, 2000.

Tomer Tsaban, Julia K Varga, Orly Avraham, Ziv Ben-Aharon, Alisa Khramushin, and Ora SchuelerFurman. Harnessing protein folding neural networks for peptide-protein docking. Nature communications , 13(1):176, 2022.

Robert M MacCallum, Andrew CR Martin, and Janet M Thornton. Antibody-antigen interactions: contact analysis and binding site topography. Journal of molecular biology , 262(5):732-745, 1996.

James Dunbar, Konrad Krawczyk, Jinwoo Leem, Terry Baker, Angelika Fuchs, Guy Georges, Jiye Shi, and Charlotte M Deane. Sabdab: the structural antibody database. Nucleic acids research , 42 (D1):D1140-D1146, 2014.

Jared Adolf-Bryfogle, Oleks Kalyuzhniy, Michael Kubitz, Brian D Weitzner, Xiaozhen Hu, Yumiko Adachi, William R Schief, and Roland L Dunbrack Jr. Rosettaantibodydesign (rabd): A general framework for computational antibody design. PLoS computational biology , 14(4):e1006112, 2018.

Fei Ye, Zaixiang Zheng, Dongyu Xue, Yuning Shen, Lihao Wang, Yiming Ma, Yan Wang, Xinyou Wang, Xiangxin Zhou, and Quanquan Gu. Proteinbench: A holistic evaluation of protein foundation models. arXiv preprint arXiv:2409.06744 , 2024.

Shitong Luo, Yufeng Su, Xingang Peng, Sheng Wang, Jian Peng, and Jianzhu Ma. Antigen-specific antibody design and optimization with diffusion-based generative models for protein structures. Advances in Neural Information Processing Systems , 35:9754-9767, 2022.

Xiangxin Zhou, Dongyu Xue, Ruizhe Chen, Zaixiang Zheng, Liang Wang, and Quanquan Gu. Antigen-specific antibody design via direct energy-based preference optimization. Advances in Neural Information Processing Systems , 37:120861-120891, 2024.

Rebecca F Alford, Andrew Leaver-Fay, Jeliazko R Jeliazkov, Matthew J O'Meara, Frank P DiMaio, Hahnbeom Park, Maxim V Shapovalov, P Douglas Renfrew, Vikram K Mulligan, Kalli Kappel, et al. The rosetta all-atom energy function for macromolecular modeling and design. Journal of chemical theory and computation , 13(6):3031-3048, 2017.

Daniel Frisch and Uwe D Hanebeck. Deterministic von mises-fisher sampling on the sphere using fibonacci lattices. In 2023 IEEE Symposium Sensor Data Fusion and International Conference on Multisensor Fusion and Integration (SDF-MFI) , pages 1-8. IEEE, 2023.

James A Brofos, Marcus A Brubaker, and Roy R Lederman. Manifold density estimation via generalized dequantization. arXiv preprint arXiv:2102.07143 , 2021.

Michiel Hazewinkel, Nadiya Gubareni, and Vladimir V Kirichenko. Algebras, Rings and Modules: Volume 1 , volume 575. Springer Science & Business Media, 2006.

Sankar Basu and Björn Wallner. Dockq: a quality measure for protein-protein docking models. PloS one , 11(8):e0161879, 2016.

Haitao Lin, Odin Zhang, Huifeng Zhao, Dejun Jiang, Lirong Wu, Zicheng Liu, Yufei Huang, and Stan Z Li. Ppflow: Target-aware peptide design with torsional flow matching. bioRxiv , pages 2024-03, 2024.

Contents

A Proofs of propositions

Proposition 3. Every distribution p SO (3) ( R ) supported on SO(3) can be bijectively mapped to a distribution p S 3 ( q ) supported on S 3 which satisfies p S 3 ( q ) = p S 3 ( -q ) .

Proof. Given each distribution p SO (3) ( R ) support on the SO(3) group, we can build the corresponding S 3 measure p SO (3) ( R ) using the surjective homomorphism ϕ : SU (2) → SO (3) :

Now we prove that p S 3 ( q ) satisfies all the probability distribution axioms and the antipodal symmetry:

Non-negativity: ∀ q ∈ S 3 , p S 3 ( q ) = p SO (3) ( ϕ -1 ( R )) / 2 ≥ 0

For p S 3 ( q ) = p S 3 ( -q ) , we can also build p SO (3) ( R ) := 2 p S 3 ( q ) = 2 p S 3 ( -q ) , ∀ R ∈ SO (3) . We can check that p SO (3) ( R ) ≥ 0 and ∫ SO (3) p SO (3) ( R ) dSO (3) = 1 . Therefore, p SO (3) is a probability density function on SO(3).

Proposition 4. With Ψ as an antipodal equivariant function, i.e. Ψ( -θ q ) = -Ψ( θ q ) , and uniform prior p ( θ q 0 ) on S 3 , the marginal distribution defined by multiple Markov transitions p Ψ ( θ q n ) = p ( θ q 0 ) ∫ Π n i =1 p F ( θ q i | Ψ( θ q i -1 ); α 1: i ) d θ q 1: n -1 is antipodal invariant, which is equivalent to a distribution p Ψ ( R ) .

Proof.

q is antipodal invariant.

Therefore, the marginal distribution p ( θ n )

(a)

κ

= 1

,

µ

= [

1

√

3

,

1

√

3

,

1

√

3

]

T

(b)

κ

= 10

,

µ

= [

1

√

3

,

1

√

3

,

1

√

3

]

T

(c)

κ

= 100

,

µ

= [

1

√

3

,

1

√

3

,

1

√

3

]

T

(d)

κ

= 10

,

µ

= [0

,

0

,

1]

T

B Detailed derivation of SO(3) Bayesian flow

B.1 Hypersphere and von Mises Fisher distribution

Given an Euclidean space R d , the hypersphere S d -1 located in R d is defined as:

The von Mises Fisher distribution [Mardia and Jupp, 2009] is a uni-modal distribution on S d -1 with the location parameter µ and concentration parameter κ . The parameters µ and κ are analogous to the mean µ and variance 1 /σ 2 in the normal distribution: µ represents the central location around which the distribution is concentrated, and κ serves as a measure of concentration. We provide a visualization of vMF distribution with different µ and κ in Fig. 6. The probability density function of vMF distribution is:

where || µ || = || x || = 1 , ν = d/ 2 -1 , κ ≥ 0 , I ν ( κ ) denotes the modified Bessel function of the first kind at order ν . The differential entropy of vMF ( µ , κ ) is:

We use vMF's entropy Eq. (16) to determine the accuracy schedule in the following sections.

B.2 Bayesian update function h

Sender distribution For each time step i , the sender corrupts the clean data x ∼ p data using the sender distribution p S and the signal-to-noise ratio parameter α i , gets the noised information y ∼ p S , and send y to the receiver. Here, we use the vMF distribution to instantiate p S for hypersphere data x ∈ S d -1 :

Input distribution As for the receiver, its belief over the ground truth data x is formulated by a distribution family with parameter θ , e.g. , Gaussian distribution with µ and σ for Euclidean space data x ∈ R 1 , which can be expressed as:

Bayesian update function For each time step i , the receiver updates its prior belief θ i -1 according to the observed noised data y i with known accuracy α i according to the Bayesian theorem. This process is formulated as the Bayesian update function θ i = h ( θ i -1 , y i , α i ) . The Bayesian update

function h for von Mises Fisher distribution is derived as follows:

where κ i µ i = κ i -1 µ i -1 + α i y i . Therefore, the Bayesian update function is:

We define θ def = κ µ to compactly denote the parameter of the vMF distribution { κ, µ } .

B.3 Bayesian flow distribution p F ( θ i | x ; α 1: i )

Given time index i and clean data x , the Bayesian flow distribution for θ i is defined as the marginal distribution after conducting multiple Bayesian updates over noised data y i :

For determining the accuracy schedule α 1: i according to entropy, we first formalize the entropy of the Bayesian flow:

The entropy of the Bayesian flow is dominated by the accuracy schedule α 1: n . Although Eq. (21) is not tractable, we can evaluate its value at each step given the accuracy schedule α 1: n . For synchronizing the entropy across each modality, we find the accuracy schedule that α 1: n such that the entropy of this Bayesian flow linearly decreases with respect to time:

Although Eq. (22) is not analytically tractable, we can solve this equation using binary search to find each α i at each step i .

B.4 Discrete-time loss with n steps L n ( x )

Receiver distribution Given θ i from the Bayesian flow distribution as input, the receiver uses the network Ψ to improve its belief considering interdependency across modalities and dimensions, which is referred to as the output distribution. We choose to parameterize the output distribution as delta distribution to directly predict the ground truth x following [Graves et al., 2023]:

To approximate the sender's distribution, the receiver's distribution is obtained by convolving the output distribution with the sender's distribution:

Discrete-time loss with n steps The discrete time loss is defined as the KL divergence between sender and receiver distribution:

In fact, the term (1 -dot ( x , Ψ( θ , t i ))) in Eq. (25) is proportional to the geodesic distance between the corresponding orientation matrix R = ϕ ( x ) and ˆ R = ϕ = (ˆ x ) = ϕ (Ψ( θ , t i )) .

Proposition 5. The training loss defined in Eq. (8) is proportional to the geodesic distance between the predicted orientation matrix ˆ R = ϕ (Ψ q ( θ P , t )) and the ground truth orientation matrix R up to a constant.

Proof. We first introduce the ( w, v ) representation for unit quaternions. Specifically, for a unit quaternion q = [ a, b, c, d ] T , we define the scalar part as w = a and the vector part as v = b i + c j + d k , so that the quaternion can be written as q = w + v , where i , j , k are the standard basis elements corresponding to the three Cartesian axes. This representation is particularly convenient because the rotation angle ϑ of the rotation matrix R = ϕ -1 ( q ) satisfies the relation

Using the ( w, v ) representation, the geodesic distance between two unit quaternions q = w 1 + v 1 and ˆ q = w 2 + v 2 on S 3 is given by:

where ⟨· , ·⟩ denotes the standard Euclidean inner product.

The geodesic distance between the corresponding rotation matrices R = ϕ ( q ) and ˆ R = ϕ ( ˆ q ) in the rotation group SO(3) is the rotation angle of the relative rotation R r = R ⊤ ˆ R . This angle can be computed via the quaternion representation of R r :

The corresponding rotation angle ϑ of R r is then:

The proof is done.

C SO(3) toy data experiment

In this section, we use the synthetic dataset in Brofos et al. [2021] to verify the effectiveness of the proposed SO(3) Bayesian flow over SDE/ODE based methods following [Bose et al., 2023].

Toy dataset network The antipodal equivariance of the toy dataset network is achieved by the following formula:

where ⊗ denotes the element-wise product between two vectors, and Normalize ( q ) = q / || q || 2 .

Metrics The metrics W1 and W2 refer to the Wasserstein p -distance between the generated distribution and the ground truth test distribution with p = 1 and p = 2 , respectively. These two metrics measure the distributional distance between the generated distribution and the ground truth data distribution.

Results We present the visualized generated results in Fig. 7 and the corresponding evaluation metrics in Tab. 5. As shown in Fig. 7, ProBayes produces higher-quality samples with fewer outliers compared to the SDE and ODE baselines. This observation is consistent with the improved W1 and W2 scores reported in Tab. 5, further demonstrating the effectiveness of the proposed SO(3) Bayesian flow.

We further validate the ineffectiveness of the Euler angle and axis-angle rotation representations in Fig. 7f and Fig. 7g, as discussed in the main text. Specifically, we employ the hypertorus BFN [Wu et al., 2025] to generate the Euler angle representation, and a combination of the hypertorus and hypersphere BFN to produce the axis-angle representation. As shown, both representations fail to capture the underlying SO(3) group structure, resulting in poor SO(3) generative modeling performances.

(a) data dist.

(b) FoldFlow-base

(c) FoldFlow-OT

(d) FoldFlow-SFM

(e) ProBayes

(f) Euler angle

(g) Axis-angle

0.0

0.0

0.5

0.0

0.5

-0.5

LL

D Detailed training and sampling procedure

We describe the training loss functions used for the remaining modalities:

C α position p training loss Sampling µ p i -1 from p p F , the training loss for position p is calculated as the discrete time loss for continuous Euclidean space variable:

Sequence S training loss Sampling θ S i from p S F , the training loss is the n -step discrete-time loss for discrete variable:

where I ∈ R and 1 ∈ R .

Angle χ , ψ training loss Sampling θ χ from p χ F , the training loss is the n -step discrete-time loss for periodic variable:

Note that for χ , we only calculate the loss for those non-padded values. Also, sampling θ ψ from p ψ F , the training loss for ψ is:

Algorithm 1 Training

1: Require: n, K ∈ Z , σ 1 ∈ R + , β 1 ∈ R + , α χ 1: n , α ψ 1: n , α q 1: n ∈ R + , λ q , λ S , λ p , λ χ , λ ψ , λ bb ∈ R +

2: Input: sequence S , position p , orientation R , backbone torsion ψ , sidechain torsion χ , χ mask Mask χ

3: Sample i ∼ U { 1 , n } , t ← ( i -1) n

4: # randomly select one from the two quaternions

6: # sampling from the Bayesian flow distribution of each modality

10: m χ i ∼ p χ F ( m χ i | χ ; α χ 1: i ) , c χ i ∼ p χ F ( c χ i | χ ; α χ 1: i ) ,

13: # feed all protein modalities into the network

14: θ P ←{ θ S , θ p , θ R , θ ψ , θ χ }

16: # calculate losses

Algorithm 2 Sampling

E Limitations and broader impacts

Limitations Despite the promising results of ProBayes, we acknowledge that the binding affinities and complex stability assessed by Rosetta serve as an initial screening step. Real-world properties of the generated candidates still require validation through costly wet-lab experiments. However, this challenge is shared across the research community.

Broader impacts Our work contributes to the advancement of unified, end-to-end protein design, potentially enabling more accurate generation of functional proteins. This has implications for therapeutic discovery, enzyme engineering, and synthetic biology. However, as with any generative biological technology, careful oversight is essential to prevent unintended misuse, such as the design of harmful or dual-use proteins.

F Experiment details

Algorithm 3 Invariant point attention with quaternion implementation (IPAq)

We provide our code in https://github.com/GenSI-THUAIR/ProBayes .

F.1 Rationalized information flow implementation

In fact, the rationalized information flow can be implemented in various ways using different parameter sizes or by mixing modules, with the only constraint being the input and output specifications of each module. Here we illustrate our instantiations for each module. Firstly, we explain the key component, i.e. , IPA with quaternion implementation.

Invariant point attention with quaternion implementation We illustrate the invariant point attention [Jumper et al., 2021] implemented by the unit quaternion in Algorithm 3 using the orange color to denote the modification compared to the original IPA. Specifically, we only replace the original implementation of the SE(3) transformation based on rotation matrix and translation T i = { R , t } with the quaternion-based one T Q i = { q , t } . For a point with homogeneous representation p = [ x, y, z, 0] T and T Q i , the transformation is:

where × is the Hamilton product for two quaternions [Hazewinkel et al., 2006] and the inverse transformation ( T Q i ) -1 = { ( q ) -1 , -t } . Such operations are antipodal invariant because q × p × ( q ) -1 = ( -q ) × p × ( -q ) -1 . Therefore, the IPAq module defined in Algorithm 3 is an antipodal invariant transformation.

Sequence&Backbone Mixing Module implementation The implementation of this module builds upon the graph attention-based architecture FramePred [Yim et al., 2023b], with the following modifications: 1) The initial node embedding is mixed with the noised sequence θ S , the noised backbone torsion angles θ ψ , and the timestep t via an MLP as follows:

where embed S , embed ψ , embed t are embedders for sequence, angles, and time, respectively. 2) The IPA module is modified as stated above and also in Algorithm 3. 3) The prediction of ψ is removed and repositioned at the end of the Backbone & Side-chain Mixing module.

Backbone&Sidechain Mixing Module implementation To capture all-atom geometry during backbone prediction, we introduce an additional multilayer perceptron (MLP) that embeds sidechain features. This sidechain embedding is then passed to the BackboneUpdate module [Jumper et al., 2021], enabling refined updates to the backbone frame. The sidechain embedding module comprises three linear layers with bias terms, followed by a LayerNorm layer. Conversely, to incorporate backbone information into sidechain design, we employ another MLP-based module consisting of linear layers with bias terms-which jointly encodes the node embeddings from the Sequence&Backbone Mixing Module alongside the sidechain features.

F.2 Peptide design

Metrics Consistent with previous research [Kong et al., 2024], we generate 40 candidate structures per receptor for the sequence-structure co-design task and 10 candidates for each receptor-ligand pair in the binding conformation generation task. The evaluation metrics are detailed as follows:

∆G [Kong et al., 2024] The binding energy (in kcal/mol) calculated by Rosetta [Alford et al., 2017] to evaluate the binding affinity of the generated peptide using the function InterfaceAnalyzer .

Success rate [Kong et al., 2024] The proportion of successful designs i.e., those with ∆ G < 0 , out of all generated candidates.

DockQ [Basu and Wallner, 2016] A comprehensive metric that evaluates the all-atom similarity at the interface between a candidate and the reference complex.

RMSD C α [Kong et al., 2024] The root mean square deviation (RMSD) of the C α coordinates between a candidate and the reference structure, measured in Ångströms (Å).

Validity [Lin et al., 2024] refers to the proportion of peptides that are chemically valid, determined by whether the generated atomic bond lengths are within 0.5A above and below the ideal value.

Diversity [Lin et al., 2024] is measured by the average pairwise distance of the generated peptides using TM score and sequence similarity.

Novelty [Lin et al., 2024] A generated peptide is considered novel if both its TM-score and sequence overlap with the reference are both below 0.5.

Note that the diversity and novelty are evaluated combining the validity to avoid false positives.

RMSD atom [Kong et al., 2024] on all atoms to measure the quality of the all-atom geometry for the binding conformation generation task.

Baselines For RFDiffusion, we set the number of cycles to one and disable the empirical force field refinement to ensure a fair comparison with other methods. For PepFlow, we use the official implementations and retrain the models on the same datasets using the default hyperparameters provided in its repository. For PepGLAD, most evaluation metrics are taken directly from the original paper, while the design metrics are computed using the official checkpoint available in its repository.

F.3 Antibody design

Following Ye et al. [2024], we generate 64 candidates for each receptor in the sequence-structure co-design task and report the average performance across these candidates.

Metrics The metrics for antibody design include:

AAR [Ye et al., 2024] The amino acid recovery rate, calculated as the number of residues in the generated CDR-H3 sequences that match the reference antibody;

RMSD [Ye et al., 2024] The root-mean-square deviation, measured between the generated and natural antibodies using the C α coordinates of the CDR-H3 region.

Etotal [Ye et al., 2024] The total energy computed using Rosetta's full-atom score function with the default REF15 weight set, evaluated on the CDR-H3 regions.

∆G [Ye et al., 2024] The binding energy, analogous to that used in peptide design, but here evaluated between the CDR-H3 region and the antigen.

Baselines We directly borrow the baseline results reported in [Ye et al., 2024].

F.4 Other implementation details

Computation resources All experiments in this paper are conducted on a node with 8 NVIDIA A100 80GB. Each training task requires 4 GPUs.

Relaxation For peptide design, we adopt the relaxation procedure described in Kong et al. [2024], which optimizes the full all-atom structure. In contrast, for antibody design, we employ the protocol from Ye et al. [2024], which performs relaxation exclusively on side-chain atoms while keeping the backbone atoms fixed.

Hyperparameters For all experiments, the network is configured with a node embedding size of 128 and an edge embedding size of 64. Node embeddings in the IPAq module have 128 dimensions, and edge embeddings have 64 dimensions. The IPAq attention mechanism comprises 8 heads, with 8 query-key points and 12 value points for geometric attention. Additionally, a sequence transformer is integrated into the encoder, featuring 4 attention heads and 2 layers to enhance contextual representation. The entire encoder architecture consists of 6 stacked IPAq blocks, enabling deep and expressive modeling of the input molecular graph structure. Finally, the network consists of 7.02 million parameters. For optimization, we set the learning rate to 5 × 10 -4 and use a batch size of 40 per distributed node with Adam optimizer.

G Comparison between BFN and SDEs/ODEs

Arecent study by Xue et al. [2024] explores potential connections between continuous-time Bayesian Flow Networks and SDEs. However, we clarify that the BFNs are fundamentally distinct from SDE/ODE with evidence as follows:

The Bayesian flow distribution, as defined in Graves et al. [2023], does not involve differential terms, which are central to the formulation of SDEs and ODEs.

The Bayesian flow formulation employed in this work is discrete-time and does not rely on any discretization of continuous-time dynamics, in contrast to SDEs- and ODEs-based approaches.

In both the hypertorus Bayesian flows [Wu et al., 2025] and the hypersphere Bayesian flows introduced in this paper, the accumulated accuracy parameters κ i are inherently stochastic. These stochastic components cannot be fully represented or captured within a conventional SDE framework.

The connection proposed in Xue et al. [2024] is not mathematically rigorous: it introduces a truncation of the time domain from [0 , 1] to [0 , 1 -η ] , where η is a parameter, thereby deviating from the standard continuous-time SDE formulation.

H More results

H.1 Error bars

We report the error bars for the protein design task, as shown in Tab. 6, based on three independent experiments using different random seeds.

H.2 Visualization

We provide the visualization of designed peptides and antibodies in Fig. 8 and Fig. 9.

I Further discussion of information shortcut

Protpardelle [Chu et al., 2024] tries to remedy the information leakage by noising all 37 unique atom position inputs instead of only the atoms corresponding to the sequence. However, we have proved in the main text that even noisy side-chain information can create an information shortcut for the sequence prediction. Therefore, the information shortcut still exists in Protpardelle.

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Our main experiment and ablation study support our claims.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of this work in Appendix E.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide the proof of the propositions in Appendix A.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental result reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide our implementation details in Appendix F

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have provided our code link.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental setting/details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide the training and test details in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We provide the error bars in Appendix H.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide this in Appendix F.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: This research conforms with the NeurIPS Code of Ethics in every respect.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We provide the discussion of broader impacts in Appendix E.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to

generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [Yes]

Justification: We provide the discussion in Appendix E.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We provide this in Appendix F.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We provide the documentation for our code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and research with human subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional review board (IRB) approvals or equivalent for research with human subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [Yes]

Justification: We provide the discussion of this in Appendix E.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

16. Declaration of LLM usage

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.

Answer: [NA]

Justification: The core method development in this research does not involve LLMs. Guidelines:

The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.

Please refer to our LLM policy ( https://neurips.cc/Conferences/2025/LLM ) for what should or should not be described.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a novel all-atomic Bayesian flow framework for SO(3) generation. It has a rational information flow design to avoid the information shortcut from the side-chain to the sequence in joint generation. The work provides theoretical proofs and achieves good results on multiple benchmarks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1. Introduces Bayesian flow on the rotation group by mapping SO(3) generation to antipodally symmetric hypersphere (S³) sampling, avoiding the numerical discretization errors typical of diffusion or flow-matching approaches.\n\n2. The authors identify and eliminate the “side-chain→sequence” shortcut problem by masking side-chain inputs in the sequence branch. Their ablation demonstrates that this design prevents degenerate learning.\n\n3.  Both the hypersphere and SO(3) bijection (Proposition 1) and the vMF Bayesian update rule (Eq. 5) are stated and proved, giving strong mathematical grounding. \n\nWeaknesses:\n\nThe motivation for applying Bayesian Flow Networks to protein design is important but not adequately explained. The details of the discretization errors and the problems they introduce to protein design should be discussed.\n\nQUESTIONS:\n1. The theory demands $f(-q)=-f(q)$, how the network architecture or training enforces this?\n\n2. Bayesian flows avoid fine-grained ODE/SDE integration and thus should be faster or use fewer steps. Is there any efficiency comparison to diffusion baselines?\n\n3. How sensitive are results to the choices of $\\alpha_i$ for the vMF noise, which controls how “noisy” each update is? \n\n4. Fully excluding side-chain inputs from sequence prediction removes the shortcut. Would this also discards useful structures like sequence correlations? Is there a balance of properly using side-chain hints?\n\nLIMITATIONS:\nThe work has discussed its limitations in the appendix.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents ProteoBayes, an end-to-end framework for all-atom protein design based on a unified multi-modal Bayesian flow. The method jointly generates amino acid sequences, backbone structures, and side-chain conformations while addressing challenges related to modeling inconsistency and information leakage. A key contribution is the formulation of a Bayesian flow over SO(3), achieved by mapping rotations to unit quaternions on the hypersphere $S^3$ and enforcing antipodal symmetry to ensure valid generation. Additionally, the authors propose a rationalized information flow design that prevents side-chain torsion angles from leaking residue-type information during sequence prediction, improving generalization. Experimental results show that ProteoBayes outperforms prior methods in peptide and antibody design tasks.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n\n1. Strong alignment between BFN and protein design modalities\n    The paper systematically applies the BFN framework to all-atom protein design, effectively modeling diverse data types: vMF distributions on $S^3$ for backbone rotations, circular distributions for side-chain torsions, and discrete BFNs for residue types. This highlights BFN’s flexibility and simplicity across continuous, periodic, and discrete modalities in structural biology.\n\n2. Theoretically grounded modeling with principled handling of symmetry and information flow\n    The method extends BFN to non-Euclidean spaces via antipodal symmetry constraints, enabling valid generation over SO(3) through the unit quaternion space $S^3$. \n\nWeaknesses\n\n1. The information leakage issue stems from the authors’ modeling choice rather than the task itself. By representing all residues with fixed-length $\\chi$ vectors padded with zeros, the model implicitly encodes residue-specific torsion patterns into the input. This design choice introduces leakage by construction. The issue could be fundamentally avoided by using dynamic χ lengths or residue-type-based masking.\n\n2. The SO(3) to $S^3$ transformation lacks novelty. Modeling rotations with quaternions and addressing antipodal symmetry are well-established techniques. The paper does not introduce new mathematical insights, but rather applies standard tools within the Bayesian flow framework—an engineering integration rather than a theoretical contribution.\n\nQUESTIONS:\n- Why did the authors choose to represent side-chain torsion angles with a fixed-length $\\chi$ vector padded to four angles, rather than using a residue-type-dependent dynamic length or masking? This design appears to introduce an artificial information leakage issue rather than reflecting an inherent modeling challenge.\n- The paper does not clearly specify the training setup: was a single universal model trained across all design tasks, or were separate task-specific models used? Please clarify the datasets, model checkpoints, and training regimes.\n- In the design setting, what inputs are provided to the model? If receptor structure and sequence are included, the current evaluation does not fully validate the framework in realistic pipeline scenarios—can the authors elaborate on this?\n- Given that the model predicts backbone positions p and orientations R (or quaternions q), why is it necessary to explicitly model the backbone dihedral angle $\\phi$? This variable may be redundant unless justified.\n- Have the authors considered refolding the generated sequences using AlphaFold 3 and computing structural quality metrics (e.g., RMSD, TM-score, binding energy)? Such an experiment could strengthen the validation of both peptide and antibody design results.\n\nLIMITATIONS:\n- The framework primarily integrates existing techniques—such as quaternion-based rotation modeling and antipodal symmetry handling—without introducing substantial theoretical innovation. \n- Experimental details are insufficient, and the evaluation would benefit from folding the designed sequences using AlphaFold3 to verify structural and functional fidelity.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a Bayesian flow method for protein sequence-structure co-generation, including sidechains. Authors propose a new Bayesian flow formulation for SO3 generation, needed for protein backbones. They show, that sequence-structure models that incorporate sidechains, can provide a shortcut for sequence prediction, that negatively impacts model training. They adapt their architecture to avoid this, by sequence part of the model being blind to the current sidechain positions. The method is benchmarked on peptide and antibody (CDR H3) design tasks. Where they show superior performance to baselines.\n\nSTRENGTHS AND WEAKNESSES:\nTo the best of my knowledge this is the first work using Bayesian flows for protein generation. The proposed formulation adapting BFNs to S03 looks sound and is new as far as I know. I think trying out new types of generative models for proteins is quite valuable for the field.\nThe paper is well written and easy to follow.\n\nThe insight made about shortcut learning is quite valuable and new. It also is nicely backed by in-silico KL divergence experiments. \n\nThe experiments are quite standard in the field and reasonable extensive. Including ablations of the proposed method. The antibody experiment (Table 3) experiments could include some newer and harder baselines, like AbX (https://openreview.net/pdf?id=1YsQI04KaN), that now is not cited.\nBut in the end I can't see much wrong or lacking in the paper.\n\nQUESTIONS:\nAre there some other (simpler) solutions for the shortcut issue? For example one could use a training mask and pad with arbitrary (or random) values, that would still be further noised/denoised, just without training loss applied to them during training. At least equation 12 would then stop working. Such non-0 padding strategies have been used, like with 'ghost' sidechain atoms in AbDiffuser (https://arxiv.org/pdf/2308.05027, it also makes sense to include it in the related work). Although never with the aim of improving sequence prediction performance. It would be curious to see what fraction of the issue comes from padding, and which from the simple fact that noised non-padded angles always will carry some information about the sequence, even if it's hard to pull out.\n\nSpeaking about shortcuts, why in Fig 3 b) the training performance with the shortcut is so much worse? One would assume that a properly tuned model should still fit the training set just as well with extra (shortcut) information that makes the task easier.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

comments: SUMMARY:
The paper introduces ProteoBayes, a Bayesian Flow Network-based all-atom protein design model. It takes use of the homomorphism from SU(2) to SO(3) and maps the orientation matrix into a 4-dimensional sphere, which makes Bayesian posterior computation tractable. Further, it resolves the issues of information shortcut in all-atom protein generation by a novel information flow. It also achieves significant improvements in peptide design and the antibody design task.

STRENGTHS AND WEAKNESSES:
Strengths
1. The paper is logically coherent and well written.

2. It gives a rigorous proof of homomorphic mappings and carefully constructs an Antipodal distribution on S^3.

3. The work discovers and proves the existence of an information shortcut and designs an information flow that separates information from sequence and that from sidechain torsion.

Weaknesses
1. This work does not provide the training and sampling time cost and memory cost.

2. It lacks comparison with other methods in the visualization part. 

3. The paper lacks detailed information regarding the calculation of evaluation metrics, such as the specific parameters used.

QUESTIONS:
1. Does information shortcut exist in the next time step? The sequence parameter current time will meet the Sidechain Torsion parameter in the next time step.

3. Most said BFN is equivalent to diffusion in certain conditions. In the last paragraph in the Appendix, it says BFN does not involve differential terms and is discrete-time in this work. Can the differential term cancel out by score matching in the diffusion model? Can a discrete-time BFN be a discrete version of a continuous-time diffusion model?

LIMITATIONS:
yes
TITLE: Official Review
CLARITY: 3
IMPACT: 2
ORIGINALITY: 2
DATE: 2025-07-23