Deep Reinforcement Learning for Modelling Protein Complexes
Authors: Ziqi Gao
Tao Feng
Jiaxuan You
Chenyi Zi
Yan Zhou
Chen Zhang
Jia Li
Venue: ICLR 2024 poster (2024)
Keywords: protein complex structure prediction
docking path prediction
policy network
reinforcement learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: A generative adversarial policy network for fast 3D structure prediction of protein complex.

Abstract: Structure prediction of large protein complexes (a.k.a., protein multimer mod-
elling, PMM) can be achieved through the one-by-one assembly using provided
dimer structures and predicted docking paths. However, existing PMM methods
struggle with vast search spaces and generalization challenges: (1) The assembly
of a N -chain multimer can be depicted using graph structured data, with each
chain represented as a node and assembly actions as edges. Thus the assembly
graph can be arbitrary acyclic undirected connected graph, leading to the com-
binatorial optimization space of N^(N −2) for the PMM problem. (2) Knowledge
transfer in the PMM task is non-trivial. The gradually limited data availability as
the chain number increases necessitates PMM models that can generalize across
multimers of various chains. To address these challenges, we propose GAPN, a
Generative Adversarial Policy Network powered by domain-specific rewards and
adversarial loss through policy gradient for automatic PMM prediction. Specifi-
cally, GAPN learns to efficiently search through the immense assembly space and
optimize the direct docking reward through policy gradient. Importantly, we de-
sign a adversarial reward function to enhance the receptive field of our model. In
this way, GAPN will simultaneously focus on a specific batch of multimers and
the global assembly rules learned from multimers with varying chain numbers.
Empirically, we have achieved both significant accuracy (measured by RMSD
and TM-Score) and efficiency improvements compared to leading complex mod-
eling software. GAPN outperforms the state-of-the-art method (MoLPC) with up
to 27% improvement in TM-Score, with a speed-up of 600×.

Published as a conference paper at ICLR 2024
DEEP REINFORCEMENT LEARNING FOR MODELLING PROTEIN COMPLEXES
Ziqi Gao HKUST
Tao Feng HKUST (GZ)
Jiaxuan You UTUC
Chenyi Zi HKUST (GZ)
Yan Zhou
Chen Zhang
Jia Li
Createlink Technology
Createlink Technology
HKUST (GZ)
ABSTRACT
AlphaFold can be used for both single-chain and multi-chain protein structure pre- diction, while the latter becomes extremely challenging as the number of chains increases. In this work, by taking each chain as node and assembly actions as edges, we show that an acyelic undirected connected graph can used to predict the structure of multi-chain protein complexes (a.k.a., protein complex modelling, PCM). However, there are still two challenges: 1) The huge combinatorial opti- mization space of N (N is the number of chains) for the PCM problem can easily lead to high computational cost. 2) The scales protein complexes exhibit distribution shift due to variance in chain numbers, which calls for the generaliza- tion in modelling complexes of various scales. To address these challenges, we propose GAPN, a Generative Adversarial Policy Network powered by domain- specific rewards and adversarial loss through policy gradient for automatic PCM prediction. Specifically, GAPN learns to efficiently search through the immense assembly space and optimize the direct docking reward through policy gradient. Importantly, we design an adversarial reward function to enhance the receptive field of our model. In this way. GAPN will simultaneously focus on a specific batch complexes and the global assembly rules learned from complexes with varied chain numbers. Empirically, we have achieved both significant accuracy (measured by RMSD and TM-Score) and efficiency improvements compared to leading PCM softwares. GAPN outperforms the state-of-the the-trr method (MoLPC) with up to 27% improvement in TM-Score, with a speed-up of 600x.
INTRODUCTION
Predicting the 3D structures of large protein com- plexes (Jumper et al., 2021; Hickman & Levy, 1988: Benesch et al., 2003) is crucial for under- standing their biological functions and mecha- nisms of action. This knowledge is essential for studying various cellular processes, such sig- nal transduction (Ditzel et al., 1998), gene regu- lation (Will & Lührmann, 2011), and enzymatic reactions (Gao et al., 2006). The groundbreak- ing AlphaFold-Multimer (Evans et al., 2021) has presented great potential in directly predicting 3DD complex structures from their amino acid se- quences. However, it faces difficulties in maintaining high accuracy when dealing with complexes with larger number (> 9) of chains (Bryant et al., 2022a; Burke al., 2023; Bryant et al., 2022b).
Multipi Chains
Complex
ABCD
Assembly
Stp
AUC
C-A
A+D
C-A-D
A-B
Graph
C-
Figure 1: The assembly process applied for the PCM problem. Each incoming chain has the opportunity to dock onto any already docked chain.
Recently, a novel setting for protein complex modelling (PCM) has emerged and achieved signif- icant success. These methods (Bryant et al., 2022b; Aderinwale et al., 2022;
"Equal contribution.
1 "ooresponndecc Jia Li (jial seiiaaeeeest.he

---

Published as a conference paper at ICLR 2024
et al., 2012; Burke et al., 2023) suggest that since numerous advanced protein- prreein docking tech- niques (Ganea et al., 2021; Wang et al., 2023) are available, why not utilize them to sequentially assemble protein chains into complexes? They have demonstrated the ability to perform PCM on very large complexes (with up to 30 chains). However, the combinatorial optimization space (COS) of the PCM problem greatly limits the inference efficiency and effectiveness of these methods. As shown in Fig. I, real-world data indicates that in --chain complex, the interaction of each chain may involve more than two chains. Therefore, the assembly can be determined by an acyclic, undi- rected, and connected (AUC) graph, which leads to the COS of N (Shor, 1995).
Designing a proper deep learning (DL) framework holds great promise in handling the vast COS. However, directly learning from available complexes can be difficult since the data scarcity will increase as the chain number grows. For example, when the chain number exceeds 10, the average number of available complexes from the PDB database (Berman et al., 2000) is less than 100. Due to the inherent variability in the chain numbers of complexes, DL models are likely to face challenges in learning adequate knowledge for assembling large complexes (Aderinwale et al., 2022). Based on the above statements, we summarize two main challenges of the PCM problem as follows: 1) How to efficiently explore the vast combinatorial optimization space to find the optimal assembling action? 2) How to ensure that the DL model can effectively generalize to complex data with limited availability?
In this paper, we propose GAPN, a generative adversarial policy network powered by domain- specific rewards and adversarial loss through policy gradient for automatic PCM prediction. Specif- ically, to address the first challenge, we design GAPN as a deep reinforcement learning (DRL) agent that operates within a complex assembling environment. The DRL agent can actively explore the vast combinatorial optimization space (COS) for assembling large complexes beyond samples through and--oon interaction with the environment (Zhang et al., 2019; Samsuden et al., 2019). Moreover, it incorporates a balance between exploration (trying new actions) and exploitation (stick- ing with known good actions) to ensure that the agent will not get stuck in local optima and continue to explore other potentially better solutions (Ladosz et al., 2022; Gupta et al., 2018a). The above two aspects make the GAPN agent able to efficiently explore the vast combinatorial optimization space to find the optimal assembling action. In response to the second challenge, we introduce an adversarial reward function to guide the learning process and enhance the receptive field of our model. The adversarial reward function incorporates the global assembly rules learned from com- plexes with varied chain numbers to enhance the generalization of GAPN. We empirically validate that the proposed GAPN model has high efficiency in exploring the vast COS in PCM problem, and can effectively generalize to complex data with limited availability. In summary, we make the following contributions:
We propose GAPN, fast, policy-based method for modeling large protein complexes, where GAPN functions as a DRL agent to efficiently explore correct assembly actions.
We design an adversarial reward function that offers a global view of assembly knowledge, ensuring that GAPN can generalize over limited complex data.
Quantitative results show that GAPN achieves superior structure prediction accuracy com- pared to advanced baseline methods, with speed-up of 600x. Moreover, we contribute a dataset containing non-redundant complexes.
2 RELATED WORK
Protein Complex Modelling. The task of protein complex modelling (PCM) aims predict the complex 3D structure given the information of all its chains. As shown in Table 1, existing PCM methods can divided into two categories, namely one-time forming (OTF) and one-by-one as- sembling (OBOA). ObF-based methods can typically produce PCM results directly from the se- quence or chain-level 3D structures. Specifically, the modified AlphaFold system, referred to as AlphaFold-Multimer (Evans et al., 2021). can handle up to 9-chain complexes, but it is very time- consuming. Moreover, Ji et al. (2023) proposed SyNDock, the learnable group synchronization ap- proach for forming complexes in globally consistent manner. SyNDock mainly claims significant speed-up, but the performance and the maximum chain number it can handle are both limited. An- other PCM setting mainly involves two optimization aspects: the dimeric structures and the docking
2

---

Published as a conference paper at ICLR 2024
path. Optimization approaches such as reinforcement learning (Aderinwale al., 2022) and ge- et al., 2012) are proposed for refining the dimeric structures. These methods are typically with low efficiency, and their max chain No. is also not satisfactory. McL.C (Bryant et al., 2022b) assumes that the dimeric structures are given, thus only predicting the correct docking path. In this way, it is the first to demonstrate that the structures of complexes with up to 30 chains can be accurately predicted. However, its application of a tree search-based non-deep learning technique often leads to complete failure due the lack of generalization.
netic algorithm
Methods
Assembly
Opt. Dimer
Opt. Path
Max Chain No.
Speed-Up
CombDock Inber et al. (2005)
6
1.07
Multi-LZen DZeri- Sqquivel-Rodrigueee et al. (2012)
1.33
ZerD Aderinwale et al. (2022)
V
5
1
ESMFold Lin et al. (2023)
X
X
X
4
213.65
AlphaFold-Multimer Evans et al. (2021)
X
4.53
McLPC Bryant al. (2022b)
30
3.02
GAPN
60
1934.27
Table 1: Technical comparison between PCM methods. Assembly: PCM in the assembly fashion; Opt. dimer: Dimer structures will be optimized; Opt. Path: Docking path will be optimized; Max Chain No.: The maximum chain number that the method can handle. Speed-up: Efficiency speed-up compared to the RI-MMIZerD baseline.
Reinforcement Learning. Over the past few years, RL has emerged compelling approach for various optimization tasks. The unique ability of RL. agents to learn optimal policies from interaction with an environment enables them collect and utilize more data, which enhances their ability to find excellent solutions to optimization problems with large search space. Several studies (Vinyals al., 2015; Bello al., 2016; Nazari et al., 2018) utilize RL and attention architecture to solve classic NP-hard problems with large and complex search space, such as TSP (Vinyals et al., 2015) and VRP (Bello et al., 2016; Nazari et al., 2018) problems. Khalil et al. (2017) and Fan et al. (2020) propose to introduce the graph architecture into the RI. framework, so that RL can be used to solve large-scale graph optimization problems, such as finding the most influential nodes on the graph and optimizing the properties of the graph network. However, when applying RL to our problem, it still faces the generalization problem in scenarios with varied chain numbers. Therefore, in our paper, design adversarial reward function guide the learning of the RL agent, which incorporates prior knowledge of multiple optimization scenarios.
Ad Adversarial Training Adversarial training has been widely used to learn the underlying data dis- tribution and derive generator and a discriminator. The generator can generate synthetic data that is similar the real data distribution, while the discriminator outputs the similarity between the input and real data. Adversarial training is widely used in many aspects, such text generation (Zhang et al., 2017), imitation of human behavior (Gupta et al., 2018b), protein generation (Gupta & Zou, 2019). Adversarial training also enhances RL's ability to solve complex optimization problems by incorporating prior knowledge specified by a good oracle (Ho et al., 2016; Liu et al., 2018). Inspired by existing research, in this paper, we design adversarial reward function to enhance the gener- alization ability of the RL. agent based on the framework of adversarial training, which incorporates the global assembly rules learned from complexes with varied chain numbers.
PROPOSED APPROACH
3.1 PROBLEM DEFINITION
We are given set of N chains to form a protein complex. The - th chain consists of n, residues, which is characterized by its amino acid sequence s' and 3D structure in its undocked state X; e 33x.; We focus the bye-by-one assembly setting, where the dimeric 3D structures of all possible pairs of chains are provided.
We define the assembly process follows. Formally, have set of assembly actions A {(A), A2) with each element representing a pair of chain indexes. Then, a transforma-
tion function set T {(R,, (.R.tt)))ar tt) is obtained by R X ao + tx X' where Re, tx
3

---

Published as a conference paper at ICLR 2024
Value Network
Generalive Adversarial Policy Network
Avg
Assembled
Chaliss
Global Embedding
Query
Graph
MLP
Docked Embodding
Softnax
GCN
T ad
Meac Pooling
Uncbooked
Undosied Embedding
&Sigmoid
Key
Action
State dg
Domein-specific reward ?1
State
Assembly Action
Assembled Graph
reward far 'ri
4g
fry
su fra.
dg
Multi-chain Assembly Process
Figure 2: Overview the proposed framework. It contains Generative Adversarial Policy Network (GAPN) (b), which converts the imput state 9 into assembly action aj for the multi-chain assembly process (d). The GAPN is trained with the help of value network (a), which estimates the values come from two aspects: domain-specific rewards obtained through multi-chain assembly process (d) and adversarial rewards from adversarial reward function (c).
represent the rotation and translation transformation for the k - th action, respectively. XA as de- notes structure of chain A extracted from the dimeric structure of the pair of (A). AA) and X A denotes the structure of the docked chain A For the k th action, we dock the chain A onto the docked chain A such that X = R X A3 + tx-
The task to predict the assembly action such that {RX thisen (XihsisN ER SO(3), & R., by taking as inputs the amino acid sequences of all chains (8:) hisgN of an input complex, where X denotes the ground truth of the th chain. In short, our goal is to predict the correct assembly action A such that the assembled complex structure matches the ground truth complex structure.
3.2 SYSTEM OVERVIEW
To model the large protein complexes, propose GAPN powered by domain-specific rewards and adversarial rewards through policy gradient. The structure overview is shown in Figure 2. The system mainly consists of four modules:
Generative Adversarial Policy Network. It converts the input state to corresponding assembly action and interacts with multi-chain assembly process.
Adversarial Reward Function. It incorporates the global assembly rules learned from complexes with varied chain numbers to enhance the generalization of GAPN, which converts the assembled graph into adversarial reward through Graph Convolutional Net- work (GCN).
Value Network. It aims estimate the value of the state and helps train the GAPN.
Multi-chain Assembly Process. It describes the process of multi-chain assembly, state transfer, and assembly action effects, which has been discussed Sec 3.1.
In the following sections, will introduce these modules in detail.
3.3 ASSEMBLY OF PROTEIN COMPLEXES AS MARKOV DECISION PROCESS
We transform the assembly of protein complexes problem into a sequential decision-man making prob- lem, with the goal of predicting the correct assembly action that will bring the assembled complex structure to match the ground truth. Thus, we propose to formulate the problem using Markov Decision Process (MDP) (Puterman, 1990) in RI. setting, which has five parts:

---

Published as a conference paper at ICLR 2024
Agent: We consider the complex M {Mi as the RL agent, who observes set of N chains and n, residues of each chain M4.
State: Since the agent is the complex M, we define state at time interval as (ct, (t, Gr). where e, denotes docked chain feature embeddings, 4n denotes the undocked chain embeddings and G, denotes all chain embeddings in the complex. The chain embed- ding e, of M, is calculated through protein pre-trained model, which will be introduced in Sec 3.4.
Action: The action or at time interval is assemble two chains, which can be denoted by (A), A with two elements representing a pair of chains.
Transition:
This component describes the process that the current state s will transit to the next state St+1 after an action a is taken. Specifically, in our state setting, the stack of all protein embeddings G, will remain unchanged. The stack of assembled protein feature embeddings c++1 = GJle and the stack of unassembled protein embeddings M4+1 = 4 'ey will change due newly docked chains, where ej denotes the embedding of the newly docked chain. With the implementation of assembly action ar, the dimeric structure of chains will be updated, which is achieved with the transformation set defined in Sec 3.1.
This reward function r is to measure how similar the assembled complex struc- ture is compared with the real ones. Here, define the rewards as sum over domain- specific rewards ry and adversarial rewards Tad- The domain-specific rewards are intro- duced to measure the similarity between the assembled complex structure and the corre- sponding ground truth. Specifically, we utilize RMSD to measure this similarity, which is inverted when used as a reward. Since the smaller the RMSD, the higher the similarity, we take negative RMSD values as domain-specific rewards. In order to enhance the general- ization of the agent to limited data, we design an adversarial reward function incorporating the global assembly rules learned from complexes with varied chain numbers (see more discussion in Sec 3.5).
Reward:
3.4 GENERATIVE ADVERSARIAL POLICY NETWORK
Having modeled the assembly of protein complexes as Markov Decision Process, we outline the architecture of GAPN, a policy network learned by the RL agent to act in the multi-chain assembly process. GAPN takes the state 9 input and outputs the action a,, which predicts an assembly action, as described in Section 3.3.
Chain Embedding Calculation. Through ESM (Lin et al., 2023) pre-training models, we obtain residue-level embeddings by inputting the entire sequence of each chain. ESM's pre-training cap- tures the reflection of protein 3D structure within sequence patterns. To obtain the chain embedding take the average of residue embeddings. It's worth noting that refrain from fine-tuning the residue embeddings. This approach allows us to leverage the power of ESM's pre-training while maintaining the stability of the residue-level representations.
Assembly Action Prediction. The assembly action a at time step t is decide which pair of pro- teins assembles together. The traditional RL approach is to perform one-hot encoding on all possible actions. However, these methods are difficult to apply to our problem: (1) In the protein assembly process, the number of possible protein pairs is constantly changing, which makes these methods hard to handle; (2) Since different protein complexes may have varied numbers of proteins, models trained based on these methods are difficult to generalize to different protein complexes. Inspired by the design of attention mechanism for supporting model's inputs with variable lengths (Vaswani al., 2017),
first introduce function fe to convert ct, G,- into dense vector Jt. and then regard the gr and üh as key and query to apply inner product for them and obtain the probability of deciding which pair of proteins assemble together,
gt (a(cce,G).
(8+) ( wx)), or ~ T(x),
(2)
where Sa and fo denote two Multilayer Perceptrons (MLP) and denotes the GAPN.
5

---

Published as a conference paper at ICLR 2024
3.5 ADVERSARIAL REWARD FUNCTION
The vast combinatorial optimization space of large protein complexes makes it hard for RL to search for optimal assembly action sets to bring the assembled complex structure match the ground truth. Moreover, the limited availability of these complex data further aggravates the difficulty of searching. In order to deal with this dilemma, the traditional approach to mix complex data of different chains and utilize a unified RL policy network to learn (Beck et al., 2023). However, the huge differences the features of complex data of different chains lead to poor training stability and performance (Nagabandi et al., 2018). Adversarial training can improve RL's ability to deal with complex problems by learning expert solutions as prior knowledge (Ho et al., 2016; Liu et al., 2018). Inspired by this, we utilize the correct assembly action sets of complex data in different chains as expert solutions, and employ the Generative Adversarial Network (GAN) framework (Goodfellow et al., 2020) to design the adversarial rewards R(TT, Dz) for RL learning:
min max R( D&) = (x- Dog [x(x:)] + Ezwao [log(1 - D*(x)))
where Ty is the policy network parameterized by 0, D& is the discriminator network parameterized by Φ, represents assembled graph whose nodes represent chains and edges depict the connection relationship among the chains determined by the assembly action set. Pdata is the underlying data distribution of the ground truth assembly action sets of complex data with different chains. However, from equation 3. we can observe that only Dø. can be trained through stochastic gradient descent and x as graph is not differentiable in terms of parameters Φ. Therefore, we utilize R(TT.. D₂) as an auxiliary term with domain-specifer rewards, which measures the similarity between the complex structure assembled from the generated assembly action set and the real ones.
The ability of Graph Neural Networks (Xu et al., 2018; Tang et al., 2022) to model the structured data enables it to perform well in many tasks (Gao et al., 2023a; Tang et al., 2023; Li al., 2019b; Gao et al., 2023b). We employ the Graph Convolutional Networks (GCN) (Kipf' & Welling, 2016; Li et al., 2019a; Cheng et al., 2023) to construct the discriminator network. Specifically, the nodes of the input complex graph chain embeddings E et, and the adjacency matrix U of the input graph depicts the docked relationship among the chains determined by assembly action set. the assembly action set contains at (Ak, A2), then U U[A], A U[A Ab 1. We apply message passing with initial node embeddings E = e and adjacency matrix U to compute node embeddings for a total of L layers:
(4) where H' is the node embedding of the 1 ih layer, D is degree matrix of U. Ú U + / and W is a trainable weight matrix of the 1 - th layer. AGG(.) here is employed to denote an aggregation function that could be one of {MEAN, MAX, SUM, CONCAT] (Sun et al., 2023).
H AGG(ReLU DUUURRLLLD--UDD--R AÜD-"H''")),
Based on the output 4 of the last layer of GCN, the discriminator network applies an MLP layer, a SIGMOID function, and a MEAN Pooling function to get the final output (z(x):
(x) = SICMOID((MEAN(((((F1))))
(5)
where fa denotes the MLP layer and the MEAN Pooling function is to average embedding of the chain to obtain the representation of the entire graph.
3.6
MODEL TRAINING
For the optimization of GAPN, we adopt Proximal Policy Optimization (PPO) (Schulman et al., 2017), which is one of the state-of-the-art policy gradient methods and has been widely used for optimizing policy networks. We show the objective function of PPO as follows:
max (Jeup (0)) = EErny [min (r.(0) A ae (s, a), clip(r: -
-,,, e, + A a))l
(6)
To (a(s))
rr (0)
Te (a|s)
where r() represents the ratio of the probability (a(s) of an action under the current policy to the probability no (als) of the same action under the old policy and is clipped the range of

---

Published as a conference paper at ICLR 2024
Algorithm 1 GAPN
Input: ground truth assembly graphs G".
1: Initialize Ta, D& with random weights for policy network and discriminator;
2: for i=0,1,2... do
3: Generate batch of assembly graphs G! ~-r
4: ko 0, 1, 2... do
5: Calculate a batch of rewards based on domain-specitf reward function ry and adversarial function rod for each state-action pair in sequences G.:
6: Sample generated assembly graphs G% from G with batch size B;
7: Sample ground truth G from G with batch size B;
8: Update De based on equation (3) through stochastic gradient descent with the positive
samples G2 and negative samples Ge
9: end for
Update πe by maximizing equation (6) via the PPO method;
11: end for
c, + e| that makes the training of policy network more stable (Schulman et al., 2017), a denotes the estimated advantage function which involves learned value function V-) to reduce the variance estimation. Specially, (-) is designed as an MLP that maps the state 8, into scalar.
We summarize details of the training process of GAPN in Algorithm 1. From the algorithm, can first observe that a batch of generated assembly graphs and ground truth are sampled to train the discriminator (lines 6-8). The generated assembly graphs regarded negative samples and ground truths are regarded as positive samples to train discriminator via an Adam Optimizer (Reddi al., 2019) (line 8). Then, calculate batch of rewards for the generated assembly graphs (line 5). Finally, we train the policy network by maximizing the expectation of reward via PPO algorithm (line 10).
4 EXPERIMENTS
4.1 EXPERIMENTAL SETUP
Dataset. Following MoLPC, conduct experiments protein complexes with the chain number 3 < N < 30. We download all available complexes of their 'first biological assembly' version from the PDB database (Berman et al., 2000). Subsequently, we perform data filtering and sequence clustering using CD-HIT (Huang et al., 2010). We split the data based on the clusters to ensure that the sequence similarity between samples in the training and test will not exceed 50%. Finally, we obtained 7,063 PDBs for the training set and validation set, and 180 for the test set. A more detailed description of the data processing and filtering process can be found in Appendix B.
We apply two types of dimer structures, the ground-truth dimers (GT Dimer) and dimers predicted by AlphaFold-Multimee (AFM Dimer). For GT Dimer, we directly extract the dimer structures of all pairs of chains that have actual physical contacts. For those without physical contacts, apply their dimers predicted with ESMFold (Lin et al., 2023) (due to its inference efficiency).
Baselines. We compare GAPN with recent advanced PCM methods. For non con-assembly methods, we include (1) SahhaFood-Multinn (Evans et al., 2021) (AF-Multimer, an end-to-end deep learning model that takes only protein sequence as input for PCM): (2) ESMFold (Lin et al., 2023)(an end- to-end deep learning-base based protein structure predictor without multiple sequence alignment). For assembly-bate based methods, include (3) Multi-LZerD (Esquivve-Roddíguez et al., 2012) (a PCM predictor with genetic algorithm and a structure refinement procedure); (4) RL-MLZerD (Ader- inwale al., 2022) (a reinforcement learning model for exploring potential dimer structures and docking path); (5) MoLPC (Bryant et al., 2022b) (a non-deep learning method with Monte Carlo tree search and DIDDT information for guidance).
According to the maximum chain numbers that the baselines can handle, we set their available ranges of chain number Table for testing (e.g., Multi-I.ZerD:: 3 < N 10). We test all baselines
7

---

Published as a conference paper at ICLR 2024
Table PCM prediction results. The best performance for each metrie is bold and the second best is underlined.
suggests that the method is not applicable to this range of chain number.
Chain numbers N
Methods
4
5
6-10
11-15
16-20
21-25
26-30
TM-Score (mean) RMSD (mean)
Multi-LZerD
0.54/ 19.17
0.49/21.58
03.05
0.24 56.73
RI.-MI.ZeID
0.62/ 16.56
0.47 / 19.09
0.49/ 28.73
0.31 /49.60
AF-Multimer
0.79 7.87
0.70/ 14.19 0.67 22.25
0.64/ 35.05
0.39 /4.22
ESMFold
0.73 99755
0.52/ / 18.06 0.44 /3.96 0.36
46.12
MoLPC
0.89 / 5.33
0.71/ / 13.49 0.66 / 21.67 0.60 36.46
0.45 0.45/43.00 43.00
0.36/ 54.75
0.34 / 60.58 0.29 / 66.60
GAPN
0.95 1.16
0.81 / 8.52
0.79/ / 16.96
0.65 / 32.37
0.62/ / 37.58
0.48/ / 50.05 0.38 / 55.76 0.36 59.13
1.0
3MSD (22 Dinse))
14
0.
80
0.0
MoLPCIAPM
SLL Dioory
Ose Doser
0est7 Dieer
MLieD
ESMIio
Molioo MA.
Mole-
RLAMEEdD
At Mtttimn Meuti
GAPN
(b)
Figure 3: PCM performance analysis. (a). TM-Score and (b). RMSD distributions all baselines on multimers of < N < 10. (c). TM-Score distribution of MoLPC and GAPN on large-scale multimers of < S 30, with the median TM-Score marked by X.
using 2 GeForce RTX 4090 GPUs. For each sample within the available range, if the testing time exceeds 2 hours or the memory limit is exceeded, its result will be replaced by the optimal value of all baselines for that sample.
Metrics. We evaluate the PCM prediction methods by the following metrics: (1) RMSD: the root-
mean-square deviation of the Cartesian coordinates, measures the distance the ground-truth and predicted multimer structure the residue level; (2) TM-Score: the proportion of residues that similar between the ground-truth and predicted multimer structures with values ranging from 0 to 1.
4.2 PROTEIN COMPLEX MODELLING
Performance. Model performance are summarized in Table 2 and Fig. 3. Our GAPN method achieves state-of-the-art in all ranges of chain numbers. Compared to the second-best baseline with small-scale multimers (3 < N < 10), we achieve improvements of 8.7%, 9.6%. 11.0%, and 12.7% on metrics of mean TM-Score. median TM-Score, mean RMSD and median RMSD, respec- tively. Moreover, our methods consistently outperform the most advanced assembly based MoLPC method, without resource-ing inntrrceintensive exhaustive search or expensive pIDDT information. Fig. 3c shows the comparison MoLPC and our method for handling large-scale multimers. We can find that although relatively imprecise, AFM dimers can still achieve accurate assembly in practice. This further reinforces the practical significance of the assembly-base tantly, method outperforms MoLPC on large-scale multimers in both settings of AFM and GT dimers. For instance, under the setting of GT dimer, we achieve an improvement of about 31% and 53% in terms of the mean and median TM-Score metrics for multimers of 11 & N < 30.
setting that we follow. Impor-
8

---

Published as a conference paper at ICLR 2024
Visualization. In Fig. 6, we show two examples with small (N = 5) and large (N 17) chain number, respectively. These two complexes are both unknown samples, meaning that none of their chains have more than 10% similarity with any chain present in the training set.
Efficiency. We evaluate the two types of efficiency of our GAPN method, computational and exploration efficiency, from both the train- ing and inference perspectives. For inference, since the running time of all methods increases with the number of chains, of running time to chain number (T/N) among different methods. As
Table 3: Inference time.
Methods
T/N(sec)
compare the ratio
Multi-LZZeD
1235.20
RL-MLZZerD
1644.78
shown in Table 3, the inference efficiency of our GAPN is about 600 times that of MoLPC. More importantly, our method is not limited by chain numbers, and can typically predict up to 60-chain complexes in less than 2 minutes. We illustrate the training efficiency (exploration ef- ficiency) with the training curves shown in Fig. 4b. We observe that our GAPN model can converge rapidly and effectively after only ~ 400 episodes. This indicates that the model can efficiently explore vast combination optimization space (N S N S 30), which also demonstrates the high training efficiency of GAPN (requiring hour for complete training on the whole 6,054 samples).
AFM
357.30
ESMFold
7.73
MoLPC
540
GAPN
0.85
(a)
(b)
GAPN
-60
Noralized RMSD
1.0
GAPN
GIPN wiu
-50
0.8
(5)
-40
RMSD
0.6
0.4
-10
62
0.0
11-15
26:30
Chain number
Episode
Figure 4: (a). The RMSD performance of GAPN and the one without Adversarial reward (AR). The bar chart represents the relative difference between the two for multimers of each specific scale. (b). Training curve comparison between our GAPN model and the one without AR. Both models are trained on multimers of 3 < N < 30 with GT dimer structures.
4.3 ABLATION STUDY
We conduct an ablation study on the whole validation set to explore the importance of our core ad- versarial reward (AR). Specifically, we call the process of GAPN assembling a multimer an episode and testing the policy network the validation set after training the policy in each episode. In Fig. 4a, we find that GAPN's performance significantly decreases after removing AR. Importantly, the gain brought by AR will become more significant with the increase of chain number, which indicates that AR can effectively address the problem of knowledge generalization difficulty caused by data scarcity in large-scale multimers. Moreover, also observe in Fig. 4b that AR crucial for model training, as it enables rapid convergence of our model.
CONCLUSION
In this paper, we propose novel generative adversarial policy network (GAPN) to tackle the prob- lem of Protein Complex Modelling (PCM). We follow the assembly-based setting to predict the docking path of chains and finally obtain the complex structures. GAPN has been demonstrated to efficiently and actively explore the vast combinatorial optimization space with a reinforcement learn- ing (RL) agent. Moreover, we design the domain reward and the adversarial reward (AR) so that the RL agent can acquire knowledge from the current and global perspectives. AR can ef- fectively enhance the receptive field of our model, which is beneficial for the model to generalize knowledge across complexes different scales. Extensive experiments on the complex dataset show that GAPN achieves significant accuracy and efficiency improvements.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGEMENTS
This work was supported by NSFC Grant No. 62206067, HKUST(GZZ-Cuunnglnn Graph Data Joint Lab, HKUST-HKUST(GZZ) 20 for 20 Cross-campus Collaborative Research Scheme C019 and Guangzhoo-HKKUSTTGZZ Joint Funding Scheme 2023A0300673.
REFERENCES
Tunde Aderinwale, Charles Christoffer, and Daisuke Kihara. Rl-mizerd: Multimeric protein docking using reinforcement learning. Frontiers Molecular Biosciences, 9:969394, 2022.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shi- mon Whiteson. survey of meta-reinforceement learning. arXiv preprint arXiv:2301.08026 2023.
Irwan Bello, Hieu Pham, Quoc V Le, Mohammad Norouzi, and Samy Bengio. Neural combinatorial optimization with reinforcement learning. arXiv preprint 1r1iv:1611.090 2016.
Justin LP Benesch, Frank Sobott, and Carol V Robinson. Thermal dissociation of multimeric protein complexes by using nanoelectrospray mass spectrometry. Analytical chemistry, 75(10):2208 2214, 2003.
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.
Patrick Bryant, Gabriele Pozzati, and Arne Elofsson. Improved prediction of protein-protein inter- actions using alphafold2. Nature communications, 13(1):1265 2022a.
Patrick Bryant, Gabriele Pozzati, Wensi Zhu, Aditi Shenoy, Petras Kundrotas, and Ame Elofsson. Predicting the structure of large protein complexes using alphafold and monte carlo tree search. Nature communications, 13(1):6028 2022b.
David F Burke, Patrick Bryant, Inigo Hernin-Hrrandez. Danish Memon, Gabricle Pozzati, Aditi Shenoy, Wensi Zhu, Alistair $ Dunham, Pascal Albanese, Andrew Keller, al. Towards struc- turally resolved human protein interaction network. Nature Structural & Molecular Biology, 30 (2):216-225, 2023.
Jiashun Cheng, Man Li, Jia Li, and Fugee Tsung. Wiener graph deconvolutional network improves graph self-supervised learning. In AAAI, pp. 7131-7139, 2023.
Lars Ditzel, Jan Löwe, Daniela Stock, Karl-Otto Stetter, Harald Huber, Robert Huber, and Stefan Steinbacher. Crystal structure of the thermosome, the archaeal chaperonin and homolog of cet. Cell, 93(1):125-133; 1998.
Juan Rodriiguzz Yifeng David Yang, and Daisuke Kihara. Multi-lzerd: multiple protein docking for asymmetric complexes. Proteins: Structure, Function, and Bioinformatics, 80(7): 1818-1833, 2012.
Richard Evans, Michael O'Neill, Alexander Pritzel, Natasha Antropova, Andrew Senior, Tim Green, Augustin Židek, Russ Bates, Sam Blackwell, Jason Yim, et al. Protein complex prediction with Hiphafold-multingg biorxiv, pp. 2021-10, 2021.
Changjun Fan, Li Zeng, Yizhou Sun, and Yang- Yu Liu. Finding key players in complex networks through deep reinforcement learning, Nature machine intelligence, 2(6):317-324, 2020.
Octavian-Eugen Ganea, Xinyuan Huang, Charlotte Bunne, Yatao Bian, Regina Barzilay, Tommi Jaakkola, and Andreas Krause. Independent se (3)-cquivariann models for end-to-end rigid protein 2021.
docking. arXiv preprint arXiv:2/11.07720
Jiali Gao, Shuhua Ma, Dan Major, Kwangho Nam, Jingzhi Pu, and Donald G Truhlar. Mechanisms 2006.
and free energies of enzymatic reactions. Chemical reviews,
10

---

Published as a conference paper at ICLR 2024
Ziqi Gao, Chenran Jiang, Jiawen Zhang, Xiaosen Jiang, Lanqing Li, Peilin Zhao, Huanming Yang, Yong Huang. and Jia Li. Hierarchical graph learning for protein protein-protein interaction. Nature Communications, 14(1):1093, 2023a.
Ziqi Gao, Yifan Niu, Jiashun Cheng, Jianheng Tang, Lanqing Li, Tingyang Xu, Peilin Zhao, Fugee Tsung, and Jia Li. Handling missing data via max-entropy regularized graph autoencoder. In Proceedings of the AAAI Conference Artificial Intelligence, volume 37, pp. 7651-7659 2023b.
Ian Goodfellow, Jean Pouget-Abadie, Mchdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial networks. Communications of the ACM, 63(11):139-11) 2020.
Abhishek Gupta, Russell Mendonea, YuXuan Liu, Pieter Abbeel, and Sergey Levine. Meta- reinforcement learning of structured exploration strategies. Advances in neural information pro- cessing systems, 31, 2018a.
Agrim Gupta, Justin Johnson, Li Fei-Fei, Silvio Savarese, and Alexandre Alahi. Social gan: So- cially acceptable trajectories with generative adversarial networks. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 2255-22644 2018b.
Anvita Gupta and James Zou. Feedback gan for dna optimizes protein functions. Nature Machine Intelligence, (22::05-111: 2019.
RK Hickman and SB Levy. Evidence that tet protein functionn as multimer in the inner membrane of escherichia coli. Journal of bacteriology, 170(4):1155-1720, 1988.
Jonathan Ho, Jayesh Gupta, and Stefano Ermon. Model-free imitation learning with policy opti- mization. In International conference on machine learning, pp. 2760-2769. PMLR, 2016.
Ying Huang, Beifang Niu, Ying Gao, Limin Fu, and Weizhong Li. Cd-hit suite: a web server for clustering and comparing biological sequences. Bioinformatics, 26(5):880-62,, 2010.
Yuval Inbar, Hadar Benyamini, Ruth Nussinov, and Haim Wolfson. Prediction of multimolecular assemblies by multiple docking. Journal of molecular biology, 349(2):435-447 2005.
Yuanfeng Ji, Yatao Bian, Guoji Fu, Peilin Zhao, and Ping Luo. Syndock: N rigid protein docking via learnable group synchronization. arXiv preprint 2305. 15156, 2023.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool Russ Bates, Augustin Židek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. Nature. 596(7873):583-588) 2021.
Elias Khalil, Hanjun Dai, Yuyu Zhang, Bistra Dilkina, and Le Song. Learning combinatorial opti-
mization algorithms over graphs. Advances in neural information processing systems, 30, 2017.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net-
works. arXiv preprint arXiv: (atp::66000200 2016.
Vijay Konda and John N Tsitsiklis. Actor-critic algorithms. In Advances in neural information processing systems, pp. 1008-1014, 2000.
Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforcement learning: A survey. Information Fusion, 85:1-22. 2022.
Jia Li, Zhichao Han, Hong Cheng, Jiao Su, Pengyun Wang, Jianfeng Zhang, and Lujia Pan. Predict- ing path failure in time-evolving graphs. In Proceedings of the 25th ACM SIGKDD international conference on knowledge discovery & data mining, pp. 1279 1289, 2019a.
Jia Li, Yu Rong. Hong Cheng, Helen Meng, Wenbing Huang, and Junzhou Huang. Semi-supervised graph classification: A hierarchical graph perspective. In The World Wide Web Conference, pp- 972-982, 2019b.
Zeming Lin, Halil Akin, Roshan Rao. Brian Hie, Zhongkai Zhu, Wenting Lu, Nikita Smetanin, Robert Verkuil, Ori Kabeli, Yaniv Shmueli, et al. Evolutionary-scale prediction of atomic-level protein structure with language model. Science, 379(6637): 122-1666333:11111133333 123-1130. 2023.

---

Published as a conference paper at ICLR 2024
YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine. Imitation from observation: Learn- ing to imitate behaviors from raw video via context translation. In 2018 JEEE International Con- ference Robotics and Automation (ICRA), pp. 1118-1111 1 125. IEEE, 2018.
Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprini arXiv:1312.5003 2013.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt in dynamic. real-world environments through meta- reinforcement learning. arXiv preprint arXiv:1803. 1347, 2018.
Mohammadreza
Nazari, Afshin Oroojlooy, Lawrence Snyder, and Martin Takáe. Reinforcement learning for solving the vehicle routing problem. Advances in neural information processing systems, 31, 2018.
Martin L Puterman. Markov decision processes. Handbooks in operations research and management science, 2:331-434, 1990.
Sashank Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. arXiv preprint arXiv: 1904:0937,7 2019.
Mohd Azmin Samsuden, Norizan Mat Diah, and Nurazzah Abdul Rahman. A review paper on implementing reinforcement learning technique in optimising games performance. In 2019 IEEE 9th international conference on system engineering and technology (ICSET), pp. 258-263. IEEE, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXivv1707.00637, 2017.
Peter W Shor. A new proof of cayley's ' formula for counting labeled trees. Journal of Combinatorial Theory, Series A, 11(1)1154-158, 1995.
Xiangguo Sun, Hong Cheng, Jia Li, Bo Liu, and Jihong Guan. All in one: Multi-task prompting for graph neural networks. In Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining (KDD 23), pp. 2120-212121111 2130,2111,,201 2023.
Richard S Sutton, David A McAllester, Satinder P Singh, Yishay Mansour, al. Policy gradient methods for reinforcement learning with function approximation. In NIPs, volume 99, pp. 1057- 1063. Citeseer, 1999.
Jianheng Tang, Jiajin Li, Ziqi Gao, and Jia Li. Rethinking graph neural networks for anomaly detection. In International Conference Machine Learning, pp. 21076-21089. PMLR, 2022.
Jianheng Tang, Fengrui Hua, Ziqi Gao, Peilin Zhao, and Jia Li. Gadbench: Revisiting and bench- marking supervised graph anomaly detection. In Thirty-seventh Conference on Neural Informa- tion Processing Systems, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information 2017.
processing systems, pp.
Oriol Vinyals, Meire Fortunato, and Navdeep Jaitly. Pointer networks. Advances in neural informa- tion processing systems, 28, 2015.
Yiqun Wang, Yuning Shen, Shi Chen, Lihao Wang, Fei Ye, and Hao Zhou. Learning harmonic molecular representations on riemannian manifold. arXiv preprint arXiv:2303.15520 2023.
Cindy L Will and Reinhard Lührmann. Spliceosome structure and function. Cold Spring Harbor perspectives in biology, 3(7):0037077, 2011.
Ronald J Williams. Simple statistical gradient algorithms for connectionist reinforcement learning. Machine learning, 8(3-4):22-2556 1992.
12

---

Published as a conference paper at ICLR 2024
Keyulu Xu, Weihua Hu, Jure Leskovee, and Stefanie Jegelka. How powerful graph neural networks? arXiv preprint drXiv:1810.00S26, 2018.
Zhang, Yu Liu, Ke Zhou, Guoliang Li, Zhili Xiao, Bin Cheng, Jiashu Xing, Yangtao Wang, Tianheng Cheng, Li Liu, al. An end-to-end automatic cloud database tuning system using deep reinforcement learning. In Proceedings of the 2019 International Conference on Management of Data, pp. 415-432, 2019.
Yizhe Zhang, Zhe Gan, Kai Fan, Zhi Chen, Ricardo Henao, Dinghan Shen, and Lawrence Carin. Ad- versarial feature matching for text generation. In International conference machine learning. pp. 4006-4015. PMLR, 2017.
13

---

Published as a conference paper at ICLR 2024
A
NOTATIONS
We summarize the commonly used notations of our paper in Table 4.
Table 4: list of commonly used notations.
Notation
Description
M
The complex. The state of RL agent time interval
8t
Cr
Docked chain feature embeddings.
Undocked chain embeddings and G, denotes all chain embeddings. All chain embeddings in the complex.
Gr
The chain embedding of M,-
The action of RL agent at time interval t.
r
The reward function of agent.
r
The domain-specificc rewards at time interval The adversarial rewards.
Tad
fer fa
Multilayer perceptron. The GAPN.
G
Ground-truth assembly graphs.
Tud
The adversarial rewards.
fa
Multilayer perceptron.
Da
Ground-truth assembly graphs.
G
Generated assembly graphs.
Paota
Ground truth assembly action sets of complex data with different chains. Assembled graph.
x
V,, v()
Value estimated by value network.
H'
The node embedding of th layer.
B DATASET
Our goal is to screen high-quality and non-redundant 3D structural data of complexes. We first downloaded the first biological assembly version of all available complexes released in the PDB database before 2022-6-10 (https: hopp://wwwt /www. rcsb. org/ docs / /prp//rrraaa sile-download-ee serviiissrrcesss We have mined a total of 283,346 complexes before the formal data processing. we process and screen multimers according to the following steps:
We only keep PDB files obtained with X-ray or EM.
We remove complexes whose NMR structure resolution below 3.0.
We remove chains with less than residues.
We remove complexes with an average buried surface area below 500 or NMR structure resolution below 3.0.
We cluster all individual chains of remaining complexes on 50% similarity with CD-HIT.
We remove the complexes the respective clusters of its chains are all more than one member.
We randomly select samples according to the required number form the test set.
Please note that different data processing orders may lead to different results. The above process avoids data leakage between the test set and the training set, as well as the occurrence data redundancy within the training set. We only used small-sized test set, which was due to the cine-consuming nature running the involved baselines. However, our dataset allows for random partitioning into any test size, and can always ensure that there is data leakage.
14

---

Published as a conference paper at ICLR 2024
Chain number
Train
Valid
Test
1526
254
20
1041
174
20
869
183
20
6-10
1654
276
40
11-15
267
45
20
16-20
206
34
20
21-25
166
28
20
26-30
98
16
20
Overall
6054
1009
180
Table The dataset statistic after complete data processing and filtering.
C ADDITIONAL EXPERIMENTS
C.I KNOWLEDGE GAP BETWEEN CHAIN NUMBERS.
We validated the existence of knowledge gaps between complexes of different sizes (chain numbers) using a reinforcement learning paradigm on our own GAPN model. As shown in the Fig. 5, the horizontal axis represents the number of sizes of complexes we used train individual models. We tested the well-trained 8 models on a selected dataset containing only 7-chain complexes. The results show that the model trained with the 7-chain complexes is the most powerful to the 7- chain complexes. Both the TM-Score and RMSD curves show a clear V-shape with the peak or valley at
40
0.60
RMSD
0.55
98
0.50
34
0.45
0.40
7
10
3
Chain number
Chain number
Figure Knowledge gap analysis.
C.2 PERFORMANCE OF GAPN ON A LARGER- SIZED TEST SET
We partitioned larger test according to the release time of the complexes test model. However, we are unable to show the results of the baselines due to efficiency limitations. We split the training and sets using the date of 2005-1-1 the threshold, with 5530/1710 samples in training/test sets.
The results Table 6 show that the performance exhibited in large sizes is not significantly different from the performance in Table 2. This indicates that the selected 180 test samples in the main text representative, and our model does exhibit a significant decrease in performance due the smaller training set. Although using AFM as the dimer structure resulted in noticeable decrease in each metric, the results are still competitive. For example, for complexes with chain number of about 10, the average TM-Score can reach 0.55.
15

---

Published as a conference paper at ICLR 2024
Table 6: GAPN performance large-size set.
Chain numbers N
Methods
6-10
11-15
16-20
21-25
26-30
TM-Scome (mean)/ RMSD (mean)
GAPN (AFM Dimer)
0.91 1.55 0.78 6.70 0.69 / 16.63 0.62 38.77
0.52 / 46.56
0.40 /56.17
0.32 62.62
0.30 66.58
GAPN (GT Dimer)
0.99/ 0.85 /.66 0.79/ 15.73
0.68 33.37 0.63 / 38.56 0.46 51.17 0.40 / 52.48
05.9/58.9
C.3 VISUALIZATION
PDB id: 7DHR
TM-Score
TM. Score 0.32
1M Score 0.37
PDB id: 5W66
TM Score 0.91
TM Scoce 0.67
Ground-truth
Ours
MoLPC
AF-Multimer
Ground-truth
MoLPC
Figure 6: Visualization of complexes successfully predicted by GAPN. We show the visual and qualitative results of a small-scale (7DHR, N = 5) and a large-scale (5W66, N = 17) multimer.
D HYPERPARAMETERS
We summarize the hyperparameter choices for GAPN in Table 7.
Hyperparameters
Values
Embedding function dimension (input) GCN layer number
13
2
Dimension MLP in policy network
13, 32,32
Dimension MLP value network
13.32,64,1
Dimeusion of GCN in discriminator
13. 10,10,1
Dropout rate
0.2
Clip ratio
0.2
Gamma RL training
1
Entropy penalty coefficient
0.01
Max grad norm
0.5
Batch size
Initial learning rates policy network, value network, discriminator Optimizer
10 6
x 10-4,33 10
Adam
Table 7: Hyperparameter choices GAPN.
E OVERVIEW OF ACTOR-CRITIC (AC) AND PROXIMAL POLICY OPTIMIZATION (PPO) ALGORITHM
The Actor-Critic (AC) strategy, as proposed by Konda & Tsitsiklis (2000), adeptly integrates strengths from both value-driven Mnih et al. (2013) and policy-driven Sutton al. (1999) ap- proaches. The framework posits two essential components: a critic Vs, and an actor
The critic, VrG. essentially embodies the value-oriented approach. It assesses and predicts the po- tential worth of present state during the learning phase. A core objective is to curtail the Temporal Difference (TD) o, error to ensure a refined valuation of the current state:
6. (r(se, ax) Vn (S1+1) Vr, (s:))
16

---

Published as a conference paper at ICLR 2024
Conversely, the actor, ne. symbolizes the policy-centric method. By engaging actively with the given environment, it generates actions in alignment with the prevailing policy. With the help of the advantage function ATy, the actor, To. enjoys enhanced stability over conventional policy gradient techniques Williams (1992):
Ar (8t, 0r) = r(8,, 0r) + (x1(5)+)) Va (st).
Subsequent this, updates the actor, n0, made via the function [():
VJ(0) = EVV plogTa Amp (st, az)),
(10)
where E signifies the expected value. The PPO algorithm adds policy update constraints on the basis of AC make the training process more stable.
17


Reviews:
[{"review_id": "NQw7yVlYhj", "reviewer": "Reviewer_MBaj", "summary": "This paper proposes a generative adversarial policy network (GAPN) for modeling large protein complexes through sequential assembly of individual protein chains. GAPN uses reinforcement learning to efficiently explore the vast combinatorial search space and find optimal assembly actions. The adversarial reward function in GAPN incorporates global assembly knowledge from complexes of various sizes, enabling it to generalize despite limited training data. Overall, GAPN achieves state-of-the-art accuracy in predicting structures of protein complexes across a wide range of chain numbers with significant speed-up.", "questions": "The authors mention that they apply two types of dimer structures, the ground-truth dimers (GT Dimer) and dimers predicted by AlphaFold-Multimer (AFM Dimer). In Table 6, the GAPN results using either of them are noted and provided. Are the results in Table 2 achieved by using both or one of them?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This work integrates reinforcement learning, graph neural networks, and adversarial training to tackle the challenging problem of protein complex modeling, which is novel. In detail, the authors identify the key issues of huge search space and lack of generalization across protein complexes of different sizes. The proposed GAPN framework well addresses these challenges through policy-based active search and an adversarial reward function that encodes global assembly knowledge. The graph representation of protein chains and complexes is well-motivated and fits naturally with the assembly actions. The experiments comprehensively evaluate performance over a range of complex sizes, convincingly demonstrating GAPN's accuracy and efficiency advantages when compared to existing baselines. The ablation studies validate the benefits of the adversarial reward.", "weaknesses": "It would be more helpful and intuitive to provide the assembly process of GAPN and MoLPC for the examples shown in Figure 4. \n\nFor the efficiency analysis, it would be better to also theoretically analyze the exploration complexity and empirically analyze the relationship between efficiency and chain number N.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ozK8WVXjVX", "reviewer": "Reviewer_Sjm2", "summary": "Paper is interested in protein folding. Given a protein, i.e., a \"graph\" of amino acids where each amino acid is a sequence of \"residues\", paper would like to output the 3D structure of the protein (i.e., the cartesian xyz coordinates of each residue), subject to rotation and translation invariance. The paper achieves that by repeatedly attaching amino acids to one another. Specifically, if a protein consists of $N$ amino acids, then their algorithm runs for $N - 1$ iterations. At each iteration, a pair of amino acids will be paired (the choice of the pair is output by their RL policy, with softmax activation). Their algorithm allows on running on larger proteins (e.g., >20 amino acids)\n\n-------------\n\n# Update\n\nI read the author's response. Accordingly, I am raising my score", "questions": "In addition to the points brought-up in \"Weaknesses\", above, I have the following questions:\n\nQ1. where does $\\theta_{old}$ in Eq.6 and Eq.7 come from?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "# Problem domain\n* The paper is an important problem domain: protein folding. Accurate protein folding could imply faster discovery of drugs, especially given the \"virus era\" erupted by COVID-19.\n* The problem comes with some interesting challenges, specifically, models should be in(/equi)variant to rotation and translation.\n\n# Enabling folding of larger proteins\n\nIn my understanding (per paper text), the earlier methods either take too long to simulate folding for larger proteins (e.g., >9 chains) or give bad accuracy. The intent of the paper is to fill that gap.\n\n# Dataset\n\nThey contribute a dataset containing non-redundant complexes. This is useful to test the generalization of protein modeling.", "weaknesses": "# Missing primer / prelim\nIt would be nice if you give some 4-line summary of terminology (even if brief). The ICLR audience might not be familiar with concepts like \"docking\" and \"dimer\".\n\n# Adversarial reward Eq.3\n\nThe motivation and implementation of Eq.3 are ambiguous. It says that $p_{data}(x)$ is \"the underlying distribution of ground-truth assembly action set\". Can you give details to the support of the distribution? Is it actually \"pairs of amino acid indices\"? Do you even have that information in the dataset? I thought you only have protein chains (not set of pairs of indices). Does every protein become $N-1$ entries in $p_{data}$. Importantly, does the order matter of how you order these edges? Perhaps some amino acid $C$ wouldn't bind to $B$ but would bind (from left) to $B$-$A$.\n\n# GNN\n\nIn Equation 4, the aggregation is done by $D^- U D^-$. Can you please explain function \"AGG\" or remove it, if it was a typo?\n\n# Clarity\n* Second page, around the middle: \"Balance between exploration and exploitation\". Can you please elaborate in the text?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aUi8BuBJd4", "reviewer": "Reviewer_554U", "summary": "The authors proposed a novel deep reinforcement learning model, named GAPN, to efficiently explore the vast combinatorial optimization space in the protein complex modelling (PCM) prediction problem to find the optimal assembling action. GAPN allows active exploration of a vast space of potential solutions to strike a balance between trying new actions and sticking with known good actions, ensuring it avoids local optima. Additionally, an adversarial reward function incorporating global assembly rules enhances the model's learning process and generalization abilities.", "questions": "N.A.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper is well-organized and easy to follow.\n2. Experimental results demonstrates the effectiveness of proposed GAPN in both prediction accuracy and efficacy.", "weaknesses": "This idea is not new, and the models they used are all well established.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BiauLoDKCi", "reviewer": "Reviewer_rNSL", "summary": "The paper introduces GAPN, a model designed to predict docking complexes through the application of reinforcement learning. The agent is tasked with determining the optimal docking path, subsequently utilizing precomputed dimer structures to assemble the entire complex. The paper provides an extensive comparative analysis with established methods such as Multi-LZerD, RL-MLZerD, AF-Multimer, ESMFold, and MoLPC, demonstrating GAPN's superior performance across various protein chain number ranges, as evidenced by metrics like TM-Score and RMSD.", "questions": "1. Could you please clarify the source of the dimer structures used by GAPN in Table 2? If GT-Dimer is utilized, it implies that GAPN has direct access to the ground truth relative positions of different chains, which could potentially skew the results and present a misleading improvement.\n\n2. In Section 4.1, the paper mentions that all possible pairs of chains are precomputed or fetched. Could you provide details on the average time spent in this step and the average number of dimer structures generated?\n\n3. In Table 1, it is indicated that GAPN does not optimize the dimer structure. Could this lead to error accumulation, especially if some precomputed dimer structures significantly deviate from the ground truth dimer structures?", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper is articulate, well-structured, and easy to follow.\n\n2. The authors have provided a detailed description of the RL setup, including the state, action, transition, and reward components, which are crucial for understanding the methodology.", "weaknesses": "1. The model primarily focuses on predicting the docking path, relying on precomputed dimer structures for the completion of the docking task. This reliance on precomputed structures potentially limits the model’s applicability in real-world scenarios, as it is sensitive to the quality of these precomputed dimers.\n\n2. The protein docking process is simplified to a sequential decision-making problem in the proposed method. Each step in this process is determined by a precomputed dimer. While the precomputed dimers (utilizing AFM or ESMFold) effectively capture bi-interactions, they may not adequately represent the complex interactions that can occur across different chains in a protein complex in real-world settings.\n\n3. There appear to be inconsistencies in some of the comparative results presented. GAPN has the advantage of accessing precomputed/prefetched dimer structures, reducing its task to merely finding a docking path. This setup makes the inference time incomparable to other baselines like AFM and ESMFold, which adopt an end-to-end approach for predicting complex structures. Consequently, the inference speedup showcased in the tables may not translate to tangible benefits in practical applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)