# DockGame: Cooperative Games for
Multimeric Rigid Protein Docking

Vignesh Ram Somnath
Department of Computer Science, ETH Zürich
&Pier Giuseppe Sessa∗
Department of Computer Science, ETH Zürich
Maria Rodriguez Martinez
IBM Research Zürich
&Andreas Krause
Department of Computer Science, ETH Zürich
Equal Contribution. Correspondence to {vsomnath,
<EMAIL_ADDRESS>
###### Abstract

Protein interactions and assembly formation are fundamental to most biological
processes. Predicting the assembly structure from constituent proteins –
referred to as the protein docking task – is thus a crucial step in protein
design applications. Most traditional and deep learning methods for docking
have focused mainly on binary docking, following either a search-based,
regression-based, or generative modeling paradigm. In this paper, we focus on
the less-studied _multimeric_ (i.e., two or more proteins) docking problem. We
introduce DockGame, a novel game-theoretic framework for docking – we view
protein docking as a cooperative game between proteins, where the final
assembly structure(s) constitute stable equilibria w.r.t. the underlying game
potential. Since we do not have access to the true potential, we consider two
approaches - i) learning a surrogate game potential guided by physics-based
energy functions and computing equilibria by simultaneous gradient updates,
and ii) sampling from the Gibbs distribution of the true potential by learning
a diffusion generative model over the action spaces (rotations and
translations) of all proteins. Empirically, on the Docking Benchmark 5.5
(DB5.5) dataset, DockGame has much faster runtimes than traditional docking
methods, can generate multiple plausible assembly structures, and achieves
comparable performance to existing binary docking baselines, despite solving
the harder task of coordinating multiple protein chains.

##  1 Introduction

Protein function is often dictated by interactions with other proteins,
forming assemblies that regulate most biological processes. Predicting these
assembly structures from constituent proteins – referred to as the
(multimeric) protein docking task – is crucial in protein engineering
applications. Traditional and deep learning methods for protein docking have
largely focused on binary docking (i.e., two proteins). These methods either
use a scoring function to rank millions of candidate assemblies, or model
assembly structure prediction as a regression (Ganea et al., 2022) or
generative modeling problem (Ketata et al., 2023) over the space of rotations
and translations.

Only a handful of approaches exist for the harder _multimeric_ docking task
(i.e. involving more than two proteins). A key challenge here is that the
space of possible actions – rotational, translational and optionally
conformational changes, grows combinatorially with the number of proteins.
Traditional methods for this task therefore first generate pairwise docking
candidates and then combine them using combinatorial optimization algorithms,
making them very inefficient and slow. Deep learning approaches for multimeric
docking, while faster, still utilize some aspect of pairwise decomposition in
their architectures – either in generating multiple sequence alignments (Evans
et al., 2021), or pairwise roto-translation predictions (Ji et al., 2023).
Moreover, they are both limited to predicting a single assembly
structure111AlphaFold-Multimer uses an ensemble of 5 models, each with a
single assembly structure prediction. (while in fact there might be multiple
plausible structures).

In this work, we introduce DockGame, a novel game-theoretic framework for the
rigid multimeric docking problem: we view docking as a game between protein
chains or their relevant subassemblies (smaller assemblies made up of chains),
where the final assembly structure(s) constitute stable equilibria. To the
best of our knowledge, this is the first approach connecting game theory with
protein docking. In particular, we model docking as a _cooperative_ game,
where all proteins have aligned interests (e.g., improved assembly energetics)
described through an underlying potential, and the action space for each
protein corresponds to the continuous space of rotations and translations.
Intuitively, this allows us to simultaneously model interactions between all
proteins, and decouple the combinatorial action space into individual ones.

Figure 1: Overview of DockGame. We introduce a game-theoretic framework for
rigid multimeric docking: we view assembly structures as equilibria w.r.t to
the potential of an underlying cooperative game. In practice, however, we do
not have access to such potential function. We propose two approaches to
tackle this. Left: We learn a surrogate potential using supervision from
traditional physics-based scoring functions, and compute equilibria via a
gradient-based scheme. Right: Viewing assembly structures as samples from the
Gibbs distribution of the underlying potential, we learn a diffusion
generative model over the roto-translation action spaces of all game agents.

In practice however, we do not have access to the true underlying potential.
To tackle this problem, we propose two approaches which we summarize in Figure
1. Our first approach employs supervised learning, learning a differentiable
analogue of traditional (and black-box) physics-based scoring functions. This
allows us to compute equilibria via a gradient-based scheme over the space of
roto-translations w.r.t the learnt potential. While simple, efficient, and
generalizable to various objectives, it crucially relies on the supervision
signal, which may not accurately reflect the game structure underlying
observed assemblies. Our second approach, instead is self-supervised, learning
solely from assembly structures in the data by interpreting them as samples
from the Gibbs distribution of the underlying potential. To this end, we then
define a diffusion generative model (DGM) over the joint roto-translation
spaces of all players (relative to a fixed player), which can then be trained
with standard denoising score-matching objectives.

To summarize, we make the following contributions:

  * •

We formulate (rigid) multimeric protein docking as a cooperative game between
proteins, where the final assembly structure(s) constitute equilibria w.r.t a
underlying potential. To the best of our knowledge, this is the first work
connecting game theory to protein docking. The game-theoretic framework also
allows us to compute multiple equilibria.

  * •

To compute plausible equilibria, we propose two different approaches. Our
first approach utilizes supervision from scoring functions, e.g. physics-
based, to learn a surrogate game potential and computes equilibria via
gradient-based learning updates. Our second approach, instead, is self-
supervised in that only uses observed assemblies data to train a DGM that can
efficiently sample from the game equilibria distribution.

  * •

We evaluate DockGame on the Docking Benchmark 5.5 (DB5.5) dataset. While DB5.5
has been traditionally used for binary docking evaluation, many proteins in
the dataset contain multiple chains, allowing us to evaluate multimeric
docking. Despite the harder task, DockGame generates multiple plausible
assemblies, and achieves comparable (or better) performance on the commonly
used Complex-RMSD (C-RMSD) and TM-score metrics, while exhibiting orders of
magnitude faster runtimes than traditional docking methods.

##  2 Related Work

##### Protein Docking

Traditional protein docking methods (Chen et al., 2003; Mashiach et al., 2010;
Kozakov et al., 2017; de Vries et al., 2015; Yan et al., 2020) have largely
focused on the binary docking setting. These methods typically generate a
large number (possibly millions) of candidate complexes, which are then ranked
using a scoring function. Finally, top-ranked candidates can undergo further
refinement using other energetic models. Deep learning approaches for protein
docking have also mainly focused on the rigid binary docking setting, learning
the best roto-translation of one protein (ligand) relative to the other
protein (receptor) (Ganea et al., 2022), or a distribution thereof (Ketata et
al., 2023). Different from all these works, we focus on the less explored
(rigid) multimeric docking setting, where the goal is to the predict the
assembly structure of $\geq 2$ proteins.

A key challenge in the multimeric docking task, is to efficiently explore the
combinatorial space of possible relative motions between proteins. Traditional
approaches, such as Multi-LZerD (Esquivel-Rodríguez et al., 2012), achieve
this by generating pairwise docking candidates, which are then combined with
combinatorial assembly algorithms, with a similar algorithm in more recent
works (Bryant et al., 2022). The only deep learning approaches for multimeric
docking, AlphaFold-Multimer (Evans et al., 2021), and SynDock (Ji et al.,
2023), also adopt pairwise decompositions, either in generating paired
multiple sequence alignments, or in predicting pairwise roto-translations. In
contrast, we propose a novel game-theoretic framework for multimeric docking
based on cooperative games, and compute assembly structures (equilibria)
efficiently using gradient-based schemes that decouple the search space over
individual proteins. Additionally, game-theory offers a natural way to compute
_multiple_ equilibria, while previous multimer docking methods are limited to
predicting a single assembly structure.

##### Diffusion Generative Models for Proteins and Molecules

Diffusion generative models have become increasingly popular in the molecular
modeling community, used in tasks like conformer generation (Xu et al., 2022;
Jing et al., 2022), molecular and protein docking (Qiao et al., 2022; Corso et
al., 2023; Ketata et al., 2023), protein backbone generation (Watson et al.,
2022; Yim et al., 2023), and generative structure prediction (Jing et al.,
2023). The closest related work to ours are (Corso et al., 2023; Ketata et
al., 2023), of which our diffusion generative model for the cooperative game
can be interpreted as a multi-agent extension. To our best knowledge, this is
the first approach of a ”multi-agent” diffusion generative model in the
context of multimeric docking (and cooperative games). In this regard, our
score network is parameterized such that it can handle a varying number of
proteins during training and inference – a fact of potentially more general
interest.

##  3 Preliminaries and Background

In this section, we introduce relevant preliminaries and background that form
the basis for the remainder of the paper. Throughout, we use the term protein
to denote both individual protein chains, and any subassemblies made up of
protein chains. We also use $[N]$ to denote the set $\\{1,2,\ldots,N\\}$.

##### Problem Definition.

In the multimeric protein docking task, the goal is to predict the protein
assembly structure, given 3D structures of the constituent protein chains. The
3D structure of each protein chain $i\in[N]$ is described by a point cloud
$\mathbf{X}_{i}\in\mathbb{R}^{n_{i}\times 3}$ where $n_{i}$ is the number of
chain’s residues, and the residue position corresponds to the coordinates of
its C$\alpha$ atom.

##### Game-Theoretic Concepts.

Game theory provides a natural framework to model and analyze systems composed
of multiple self-interested agents. Formally, a game can be described by $N$
agents, indexed $i\in[N]$. Each agent $i$ is equipped with an action set
$\mathcal{A}_{i}$, and a cost function
$f_{i}:\prod_{i=1}^{N}\mathcal{A}_{i}\rightarrow\mathbb{R}$, that maps the
actions chosen by all agents to a scalar cost for agent $i$. Solution concepts
to games are defined by equilibria, where no agent has an incentive to deviate
from their equilibrium strategy. In the most general setting (general-sum
games), agents can have differing objectives, i.e. $f_{i}$ are not the same.
In this work, we focus on cooperative games, in which case agents’ interests
are perfectly aligned towards a common goal i.e. $f_{1}=f_{2}\cdots f_{N}=f$.
In this case, $f$ is also referred to as the game’s _potential_ (Monderer &
Shapley, 1996).

##### Diffusion Generative Models.

Diffusion generative models (DGMs) are a class of models that generate data by
defining a diffusion process that transforms the data distribution into a
tractable prior, and learning the time-reversal of this process. More
formally, let the data distribution be $p_{\text{data}}(x)$ that is
transformed into a tractable prior $p_{\text{prior}}(x)$ through a diffusion
process $d\mathbf{x}=f(\mathbf{x},t)dt+g(t)d\mathbf{w}$, where $d\mathbf{w}$
denotes the Wiener process. The time-reversal of this process is described by
$d\mathbf{x}_{t}=\left[f(\mathbf{x}_{t},t)-g(t)^{2}\nabla_{\mathbf{x}_{t}}\log
p(\mathbf{x}_{t})\right]dt+g(t)d\mathbf{w}$, where
$\nabla_{\mathbf{x}_{t}}\log p(\mathbf{x}_{t})$ is the score of the time-
evolving distribution. This quantity is approximated by DGMs using neural
networks $s_{\theta}$ by first sampling from the transition kernel
$p(\mathbf{x}_{t}|\mathbf{x}_{0})$ defined by the forward diffusion process,
computing the score $\nabla_{\mathbf{x}_{t}}\log
p(\mathbf{x}_{t}|\mathbf{x}_{0})$ for $\mathbf{x}_{0}\sim p_{\text{data}}$ and
then regressing $s_{\theta}(\mathbf{x}_{t})$ against
$\nabla_{\mathbf{x}_{t}}\log p(\mathbf{x}_{t}|\mathbf{x}_{0})$.

##  4 Method

In this section, we first formalize our game-theoretic view on multimeric
protein docking. We then present the proposed methods and discuss their
limitations (Sections 4.2-4.3). Finally, we describe the used learning
architectures (Section 4.4).

##### Multimeric Protein Docking as a Game.

In multimeric protein docking, proteins interact with each other through a
combination of relative motions and structural changes to form the assembly
structure. Notably, the space of their relative motions grows combinatorially
with the number of proteins, making it challenging to compute plausible
docking structures. Previous methods tackle this by modeling the relative
motions between pairs of proteins followed by a global synchronization or a
combinatorial assembly step. Instead, we argue that game theory provides a
natural approach to simultaneously model the associated interactions and
differing objectives between proteins.

In this work, we model (rigid) protein docking as a _cooperative_ game, where
agents’ interests are aligned towards a common potential function $f$. In
accordance, we view the constituted assembly structures as the underlying game
equilibria, i.e. stable outcomes from which no agent has an incentive to
deviate. We note that our game-theoretic framework is agnostic to whether the
agents are protein chains, sub-assemblies as long as agents’ incentives are
modeled appropriately.

###  4.1 Overview

Modeling docking as a cooperative game requires three components – i) an
appropriate definition of the action set $\mathcal{A}_{i}$ for each protein
$i$, ii) the game potential function $f$ (either specified or learnt)
capturing relevant properties of the assembly, and iii) a strategy for
computing equilibria. In the rigid docking setting, we can define the action
set to be the continuous space of roto-translations. In principle, this allows
us to _simultanously_ steer the proteins to equilibria configurations and thus
efficiently scale to a large number of proteins.

In practice however, we do not have access to the underlying game potential
$f$. To circumvent this difficulty, we describe two viable approaches in
subsequent sections – i) learning an approximation of $f$ via supervised
learning, using traditional physics-based scoring functions in PyRosetta, and
computing equilibria via simultaneous gradient descent (Section 4.2), and ii)
viewing equilibria as samples from the Gibbs distribution of $f$, and learning
a diffusion generative model over the roto-translation action spaces of all
players (Section 4.3). Intuitively, the former approach corresponds to
learning a differentiable analogue of traditional scoring functions to
facilitate gradient-based equilibrium computation, while the latter approach
can be viewed as a multi-agent extension of the denoising score-matching
paradigm. A key feature of our parameterization is that the learned (reverse)
diffusion can handle varying number of proteins across training and inference.

Before describing our methods, we summarize the following useful notation. We
define the action space for each protein $i$ as
$\mathcal{A}_{i}=SO(3)\times\mathbb{T}(3)$, where $SO(3)$ is the 3D rotation
group, and $\mathbb{T}(3)$ is the 3D translation group. The joint action space
of all proteins $[N]$ is thus the product space
$\mathbf{\mathcal{A}}=(SO(3)\times\mathbb{T}(3))^{N}$. These actions transform
an input point cloud $X$ as $X^{\prime}=R(X-\bar{X})+\bar{X}+r$ for a given
rotation $R$ and translation $r$, and $\bar{X}$ denotes the (unweighted)
center of mass of $X$. This transformation is consistent with the definition
of $\mathcal{A}_{i}$ as a direct product of $SO(3)$ and $\mathbb{T}(3)$, and
is different from the group of 3D rigid motions, $SE(3)$, defined a semi-
direct product of $SO(3)$ and $\mathbb{T}(3)$ (More details in Appendix D). We
use $\mathbf{X}=[X_{1},\ldots,X_{N}]\in\mathbb{R}^{\sum_{i=1}^{N}n_{i}\times
3}$ to denote the protein point clouds, and
$\mathbf{H}=[H_{1},\ldots,H_{N}]\in\mathbb{R}^{\sum_{i=1}^{N}n_{i}\times
n_{H}}$ for residue features.

###  4.2 Gradient-based Learning with Surrogate Potentials

##### Learning the Potential Function

For each assembly structure in the training set, we generate a given number of
_decoys_ by sampling random rotations and translations and applying these to
each protein. We then score each decoy $\mathbf{X}_{j}$ by its potential
$f(\mathbf{X}_{j},\mathbf{H})$ which in the context of this work we assume is
represented by the commonly employed PyRosetta (Chaudhury et al., 2010) energy
function.222The approach is evidently more general and any scoring/ranking
function can be applied. We stick to PyRosetta because it is the one we use in
our experiments. This constitutes our training dataset $\mathcal{D}$ for
learning a surrogate game potential. To back up our choice of using PyRosetta,
we have observed assembly structures in the data to exhibit the lowest
PyRosetta energy compared to the other structures in $\mathcal{D}$ (see Figure
3 in Appendix B). To learn $f$, we employ a parametrized architecture
$f_{\theta}$ which is invariant to permutations, translations and rotations of
the input assembly, as discussed in Section 4.4.

Each protein can have a different size, number and type of atoms, and thus,
the range of energy values produced for different protein assemblies can vary
widely. We observed this fact hinders model generalization across different
proteins when training $f_{\theta}$ using a simple regression (e.g., l2 error
norm) objective. For this reason, we propose to train $f_{\theta}$ solely
based on energy comparisons, employing the widely used ranking loss:

$\mathcal{L}(\theta)=-\mathop{\mathbb{E}}_{\mathbf{X}_{l},\mathbf{X}_{h}\sim\mathcal{D}}\big{[}\log\sigma(f_{\theta}(\mathbf{X}_{h},\mathbf{H})-f_{\theta}(\mathbf{X}_{l},\mathbf{H}))\big{]}$
(1)

where $\mathbf{X}_{h},\mathbf{X}_{l}$ are random pairs of decoys such that
$f(\mathbf{X}_{h},\mathbf{H})>f(\mathbf{X}_{l},\mathbf{H})$, i.e.
$\mathbf{X}_{l}$ has lower (true) energy than $\mathbf{X}_{h}$, and $\sigma$
is the logistic function. This is a widely-used loss in preference-based
learning and corresponds to maximum likelihood under a Bradley-Terry reward
model, see e.g. (Bradley & Terry, 1952; Rafailov et al., 2023).

##### Equilibrium Computation via Gradient-Based Learning

Once $f_{\theta}$ has been trained, we can compute assembly structures by
randomly initializing each protein and updating their roto-translation actions
via simultaneous gradient descent. In a potential game, this ensures
convergence to equilibria (which are local minima of $f_{\theta}$), but in
case of general-sum games more sophisticated update rules could be employed,
see e.g., (Balduzzi et al., 2018). A caveat of PyRosetta scores, though, –
which is also inherited by $f_{\theta}$ – is that assembly structures where
constituent proteins are far away (in 3D space) from each other also have low
energy, which is undesirable as a docking outcome. This was also observed
empirically in our early experiments. To discourage this behavior, we thus add
a distance-based penalty
$d(X_{i},X_{J}):=\text{ReLU}(d_{\text{res}}(X_{i},X_{J})-d_{\text{ths}})$
which computes the minimum Euclidean distance between any pair of residues
belonging to proteins $i$ and $j$ and penalizes the agents only when this
distance is greater than a threshold $d_{\text{ths}}$ (in our experiments,
this was $5$Å). Overall, the action updates for proteins’ reads as:

$(R_{i}^{t+1},r_{i}^{t+1})=(R_{i}^{t},r_{i}^{t})+\eta^{t}\cdot\nabla_{(R_{i},r_{i})}\Big{[}f_{\theta}(\mathbf{X}^{t},\mathbf{H})+\lambda\cdot\sum_{j\neq
i}d(X_{i}^{t},X_{j}^{t})\Big{]},\quad i\in[N],$ (2)

where $\eta^{t}$ is a (decreasing) learning rate, $\nabla_{(R_{i},r_{i})}$ is
the Riemannian gradient (i.e. living in the tangent space of $SO(3)$ and
$\mathbb{T}(3)$) and $\lambda$ is a tunable weight. In practice, the above
updates need not necessarily be simultaneous as one could also update proteins
sequentially in a round-robin fashion.

##### Discussion

On one hand, the above approach is quite general and flexible: The discussed
gradient-based update scheme can be carried with any potential function (also
fully specified, and not necessarily trained) as long as it is differentiable.
Moreover, it can also be easily extended to the case where each protein has
its own objective $f_{i}$ by adding penalty terms that are specific to each
protein. In addition, while not the focus of this work, $f_{\theta}$ can also
be used beyond equilibrium computation, e.g. to characterize the associated
energy landscape (see Jin et al. (2023)). On the other hand, though, it
heavily relies on the supervision signal used to train $f_{\theta}$, which may
not accurately reflect the true underlying game structure. In the next
section, we propose an alternative approach that circumvents the need to
utilize this supervision and learns solely from the observed assemblies.

###  4.3 Diffusion Generative Models over the Proteins’ Action Spaces

Section 4.2 views assembly structures as local minima of $f$ wrt roto-
translation gradients of all proteins. An equivalent characterization of the
assembly structures is interpreting them as sample(s) from the mode(s) of the
Gibbs distribution $p\propto\exp{(-f)}$. Furthermore, it is easy to see that
for any protein $i$, $\nabla_{(R_{i},r_{i})}\log{p}=-\nabla_{(R_{i},r_{i})}f$
where $\nabla_{(R_{i},r_{i})}\log{p}\in T_{R_{i}}SO(3)\oplus
T_{r_{i}}\mathbb{T}(3)$ are the Riemannian gradients in the tangent space of
$SO(3)$ and $\mathbb{T}(3)$ at $R_{i},r_{i}$ respectively. This implies an
equivalence between the roto-translation scores of $p$ and the roto-
translation gradients wrt $f$, and connects equilibrium computation to
sampling for cooperative games. We can thus frame rigid multimeric docking as
a generative modeling problem over the joint roto-translation action space
$\mathcal{A}$ using the framework of diffusion generative models. Intuitively,
this corresponds to learning $f$ implicitly through the (perturbed) scores of
$p$.

##### DGMs over the joint action space $\mathcal{A}$

De Bortoli et al. (2022) showed that the DGM framework for standard Euclidean
spaces (Section 3) also holds for compact Riemannian manifolds with minor
modifications, with the drift function $f(\mathbf{x}_{t},t)$ and score
$\nabla_{\mathbf{x}_{t}}\log p(\mathbf{x}_{t})$ being elements of the tangent
space, and the reverse diffusion corresponding to a geodesic random walk on
the manifold.

As the action space $\mathcal{A}_{i}$ for each protein $i$ is a product space
– following the arguments in Corso et al. (2023); Rodolà et al. (2019) – the
forward diffusion process proceeds independently on each component manifold,
and the tangent space $T_{(R,r)}\mathcal{A}_{i}$ is a direct sum, i.e.,
$T_{(R,r)}\mathcal{A}_{i}=T_{R_{i}}SO(3)\oplus T_{r_{i}}\mathbb{T}(3)$. The
tangent space $T_{(\mathbf{R},\mathbf{r})}\mathcal{A}$ is thus,
$T_{(\mathbf{R},\mathbf{r})}\mathcal{A}=\bigoplus_{i=1}^{N}T_{R_{i}}SO(3)\oplus
T_{r_{i}}\mathbb{T}(3)$ ,where
$\mathbf{R}=\left(R_{1},R_{2},\cdots,R_{N}\right)$ and
$\mathbf{r}=\left(r_{1},r_{2},\cdots,r_{N}\right)$. Note that this is
different from recent $SE(3)$ diffusion models (Yim et al., 2023), where the
semi-direct product (between $SO(3)$ and $\mathbb{T}(3)$ requires defining a
tailored inner product so that the diffusion processes can be decomposed over
each component. Our DGM can also be interpreted as a multi-body extension of
Corso et al. (2023), emerging from the connections between cooperative games
and sampling.

Having established the framework, we can now apply standard denoising score
matching objectives (Song & Ermon, 2019; Song et al., 2020) independently for
each component. Furthermore, to ensure (global) translation invariance, we
keep one protein fixed (with its COM at the origin), and define the diffusion
model only over the joint action space of the remaining $N-1$ proteins.

##### Diffusion Processes over $SO(3)$ and $\mathbb{T}(3)$

For both $SO(3)$ and $\mathbb{T}(3)$, we use the Variance Exploding SDE (VE-
SDE) formulation (Song et al., 2020) to define the forward noising process
$d\mathbf{x}_{t}=\sqrt{\frac{d\left[\sigma^{2}(t)\right]}{dt}}d\mathbf{w}$,
where $\sigma$ is defined as an exponential schedule
$\sigma(t)=\sigma_{\text{min}}^{1-t}\sigma_{\text{max}}^{t}$, with appropriate
values for $SO(3)$ and $\mathbb{T}(3)$. The transition kernel on $SO(3)$ is
the $IGSO(3)$ distribution (Nikolayev et al., 1997; Leach et al., 2022), which
can be sampled in the axis-angle parameterization of $SO(3)$ as described in
(Corso et al., 2023; Yim et al., 2023), by sampling a unit vector
$\hat{\omega}$ uniformly and a $\omega$ according to:

$\displaystyle p(\omega)=\frac{1-\cos\omega}{\pi}f(w)\hskip
5.0pt\text{where}\hskip
5.0ptf(\omega)=\sum_{l=0}^{\infty}(2l+1)\exp(-l(l+1)\sigma^{2})\frac{\sin{((l+1/2)\omega)}}{\sin(\omega/2)}$
(3)

The score of the transition kernel is given by $\left(\frac{d}{d\omega}\log
f(\omega)\right)\hat{\omega}$, which can be precomputed by truncating an
infinite sum. Sampling from the transition kernel can be accomplished by
interpolating the CDF of $p(\omega)$. For $\mathbb{T}(3)\cong\mathbb{R}^{3}$,
the transition kernel is simply the standard Gaussian with variance
$\sigma^{2}(t)$ and the score of the transition kernel is simply
$-r_{t}^{2}/\sigma^{2}(t)$.

##### Computing Equilibria

Under the DGM paradigm, computing equilibria is equivalent to sampling using
reverse diffusion guided by the learnt score. The reverse diffusion, as
mentioned above, corresponds to a discretized geodesic random walk on the
joint action space $\mathcal{A}$.

We note that connections between sampling, SDEs and games (Chen et al., 2021;
Liu et al., 2022) have also been explored in previous works, assuming either
an infinite-player setting or a homogenous player one. In contrast, our
approach is designed for a finite-player setting and places no restriction on
player types. Furthermore, as discussed in Section 4.4, the parameterization
of our score-network allows us to handle a varying number of players across
training and inference.

###  4.4 Architectures

In this section, we summarize the data representations and network
architectures utilized by the proposed approaches. We defer more formal
descriptions and additional details to Appendix A.

All our architectures take as inputs the protein point clouds $\mathbf{X}$ and
residue features $\mathbf{A}$. The point clouds are then used to dynamically
build graphs based on appropriate distance cutoffs between residues of the
same and different proteins. Residue representations are then learnt using
message-passing. The potential network and the score network share similar
architectures for the message-passing layers, with the only differences in the
output layers. The potential network predicts a $SE(3)$-invariant scalar,
while the score-network predicts an output in the tangent space
$T_{(\mathbf{R},\mathbf{r})}\mathcal{A}=\bigoplus_{i=1}^{N}T_{R_{i}}SO(3)\oplus
T_{r_{i}}\mathbb{T}(3)$, where all outputs are $SE(3)$-equivariant vectors.

##### Residue Representations

Residue representations are learnt using message-passing layers based on
tensor products. These message-passing layers are based on $SE(3)$
convolutions using tensor products, as implemented in the e3nn library (Geiger
& Smidt, 2022). We use separate message-passing layers for edges within the
same protein and between different proteins. These messages are then
aggregated to produce scalar and vector representations for each residue.

##### Potential Network

The output of the potential network $f_{\theta}$ is a $SE(3)$-invariant
scalar, which can be interpreted as the energy of the system. We first
generate edge representations by concatenating the scalar components of the
corresponding residue representations, which are then passed to fully
connected layers followed by a mean-pooling step to compute edge energy
contributions. The edge and residue contributions (computed similarly) are
added up to predict the energy.

##### Score Network

For each protein $i$ (except the fixed protein), the score network
$s_{\theta}$ takes the corresponding learnt residue representations as input
and computes the rotational and translational scores (i.e. two
$SE(3)$-equivariant outputs) via a tensor-product convolution with the COM of
protein $i$, as done in Corso et al. (2023). This parameterization allows the
score network to handle varying number of agents across training and
inference, and could be of independent interest.

##  5 Experiments

##### Datasets

We use two datasets in our experiments: Docking Benchmark 5.5 (DB5.5) and the
Database of Interacting Protein Structures (DIPS). DB5.5 is a standard dataset
used in benchmarking docking methods and contains 253 assembly structures.
DIPS is a larger dataset containing 42826 assembly structures, but only
consists of single-chain proteins. While DB5.5 has been traditionally used in
the context of binary protein docking, many examples consist of proteins made
up of multiple chains, allowing us to evaluate DockGame for multimeric
docking. We use the same splits for DIPS and DB5.5 datasets as Ganea et al.
(2022) for our experiments.

##### Experimental Setup

Following Ganea et al. (2022), we first train our models on the DIPS dataset,
and finetune it on the DB5.5 dataset. Unlike Ganea et al. (2022), however, we
train our models at the granularity of protein chains. For the DIPS dataset,
this implies no difference as all examples comprise two single-chain proteins.
However, on the DB5.5 dataset, the examples comprise proteins with 2-8 chains.
To the best of our knowledge, this is the first usage of the DB5.5 dataset for
multimeric docking experiments. More experimental details can be found in
Appendix B, with code available at https://github.com/vsomnath/dockgame/.

##### Baselines

We compare DockGame against traditional binary docking methods Attract,
ClusPro, PatchDock, the traditional multimeric docking method Multi-LZerD, and
recent deep learning methods for binary docking, EquiDock and DiffDock-PP. All
baselines except Multi-LZerd are binary docking methods. Among other
multimeric docking methods, we do not include any comparisons to AlphaFold-
Multimer and SynDock as the DB5.5 dataset is part of AlphaFold-Multimer’s
training set, while SynDock has no open-source implementation available at the
time of writing. More details regarding the baselines can be found in Appendix
B.4.

##### Evaluation

DB5.5 has traditionally been used for benchmarking binary docking methods,
where the task is to predict the best relative orientation (or a set thereof)
between two proteins (referred to as ligand and receptor). For such binary
docking baselines, we utilize the evaluations as provided with code in Ganea
et al. (2022). However, many examples in DB5.5 consist of proteins that are
subassemblies of _multiple_ chains. This fact is neglected by the
aforementioned binary docking, where the relative configuration between all
chains in the subassembly is assumed constant and already specified. Here, we
test DockGame on the significantly harder multimeric docking setting, where
the relative configuration even between chains in the same subassembly needs
to be inferred. Notably, this task has many more degrees of freedom than its
binary counterpart, and, to the best of our knowledge, it has not been
considered in the literature.

Table 1: Assembly Structure Prediction on DB5.5 Dataset. For the first five
baselines, the inputs consist of two proteins, with the goal of predicting
their assembly structure (binary docking). Instead, DockGame and MultiLZerd
take as input constituent protein chains, and are faced with the harder
_multimeric_ docking task. DockGame-E refers to the DockGame model with the
learnt potential function (Section 4.2), while DockGame-SM, refers to the
DockGame model with the learnt score network (Section 4.3). We use both
methods to compute multiple (20 and 40, resp.) assemblies for each complex.
The rule “(X, filtered by TM-score)” implies that, among the X generated
assemblies, we identify the one with the highest TM-score to the true assembly
and consider all other generated assemblies within a 0.05 TM-score absolute
difference. This copes with the fact that, while DockGame (and DiffDock-PP)
can generate multiple plausible equilibria, DB5.5 test set contains only a
single assembly structure per example. The rule “(X, best C-RMSD)” implies
that, among the X generated assemblies, we consider the one with the smallest
C-RMSD.

Method Avg. Runtime C-RMSD $\downarrow$ TM-score $\uparrow$ (in [s]) Mean
Median Std Mean Median Std Binary Docking Attract 570 10.63 12.71 10.05 0.8317
0.8256 0.1668 ClusPro 15507 8.26 3.38 7.92 0.8318 0.8938 0.1535 PatchDock 3290
18.01 18.26 10.12 0.7270 0.7335 0.1237 EquiDock 5 14.72 14.1 5.3 0.7191 0.7107
0.1078 DiffDock-PP (40, filtered by TM-score)∗ 80 17.19 16.29 6.79 0.7086
0.7312 0.1142 (40, best C-RMSD) 80 12.81 11.79 4.61 0.7014 0.6773 0.1125
Multimeric Docking DockGame-E (20, filtered by TM-score) 182 19.28 17.74 7.37
0.6714 0.6820 0.1543 (20, best C-RMSD) 182 14.18 11.54 9.24 0.6182 0.6257
0.1795 DockGame-SM (40, filtered by TM-score) 157 14.12 9.44 4.75 0.7173
0.7773 0.1419 (40, best C-RMSD) 157 8.73 8.72 4.68 0.7246 0.7681 0.1578 Multi-
LZerD 82753 21.23 20.89 6.82 0.6312 0.6266 0.1352

##### Metrics

Docking methods are commonly evaluated by the Complex Root Mean Squared
Deviation (C-RMSD) metric, which measures the deviation between the predicted
and ground truth assembly structure after Kabsch alignment (details in Ganea
et al. (2022)). We also utilize TM-score (Zhang & Skolnick, 2005) as an
alternative metric to quantify structural alignment.

A fundamental evaluation problem for DockGame and DiffDock-PP is that they can
both generate multiple plausible equilibria, while the test set of DB5.5 only
has a single assembly structure per example. We thus compute summary
statistics (as described by Mean, Median and Std. deviation of C-RMSD and TM-
score) on a filtered set of the predicted assemblies, constructed in two ways:

  * •

Identifying the predicted assembly with the highest TM-score (as an effective
measure of structural alignment) to the ground truth, and considering all
predicted assemblies within a TM-score (to ground truth) radius of 0.05 (we
noticed similar results for a threshold of 0.1). We adopted this heuristic to
extract predicted assemblies close to the equilibrium present in the data, and
filter out different (but still potentially plausible) equilibria generated by
DockGame.

  * •

Identifying the predicted assembly with the lowest C-RMSD to the ground truth.

(a)

(b)

Figure 2:  (a) Learned vs. True PyRosetta energy differences, over 50 random
decoy pairs for each complex in DB5.5 test. (b) PyRosetta energetics during
game rounds. Median, $25$-$75^{\text{th}}$ percentiles, over 20 games for each
complex in DB5.5 test.

##### Results and Discussion

Despite the harder multimeric docking task, DockGame-SM, trained only with
assembly structures, achieves comparable performance to binary docking methods
on both the C-RMSD and TM-score metrics (Table 1). In particular, the median
C-RMSD (more robust to outliers) for DockGame-SM across both filtered sets is
better than all baselines except ClusPro (which however requires a
significantly higher runtime). Furthermore, when compared with the traditional
multimeric docking method Multi-LZerD, both DockGame-E and DockGame-SM achieve
significantly better performance (both in terms of C-RMSD and TM-score), with
almost 3 orders of magnitude faster runtimes.

To better assess the performance of DockGame-E, in Figure 2(a) we plot the
energy differences predicted by the trained potential network, versus the
respective differences computed by PyRosetta. We can see that a good
approximation of the PyRosetta energy is achieved from our training set
containing only $10$ pairs of roto-translation perturbed structures per
example. Furthermore, in Figure 2(b) we observe that DockGame-E reduces the
PyRosetta energetics during gradient-based equilbrium computation as desired.
This demonstrates the utility of DockGame-E as a faster, differentiable
alternative for traditional scoring methods. However, the improved performance
of DockGame-SM relative to DockGame-E highlights that PyRosetta might not
offer the best supervision signal if the goal is predict assembly structures
close to the ground truth. DockGame-SM might be the more preferred
alternative, especially in data-rich settings.

Finally, we highlight that while DockGame can naturally predict different
assembly structures (equilibria) for every example, we utilized the
aforementioned filtering schemes to compare these to the only assembly
structure present in the dataset. An exciting avenue for future work would be
to explore ways to assess the plausibility of multiple generated equilibria,
which is currently lacking.

##  6 Conclusion

In this work, we presented DockGame, a novel game-theoretic framework for
rigid multimeric protein docking. We view protein docking as a cooperative
game between protein chains, described by a underlying potential whose
equilibria constitute the assembly structures. To learn the unknown game
potential and compute equilibria, we propose two approaches – i) a supervised
approach using traditional scoring functions to learn a surrogate potential
function, ii) a self-supervised approach based on diffusion generative models,
that learns solely from observed assembly structures. Both approaches can
efficiently decouple the combinatorial docking space into individual protein’s
action spaces. Moreover, our game-theoretic framework allows us to compute
multiple assembly structures as plausible equilibria. We evaluated DockGame on
the harder multimeric docking task (compared to binary docking baselines)
using the DB5.5 dataset, and achieve comparable (or better) performance on
C-RMSD and TM-score metrics, while exhibiting significantly faster runtimes.

While this work focused on cooperative games, the presented ideas are general
and can be extended to general-sum games allowing us to model, e.g., protein’s
_flexibility_. Potential ways to do so would be to add protein structure-
specific penalty terms, or to combine diffusion bridges (Holdijk et al., 2022;
Somnath et al., 2023) with the proposed diffusion generative model. Other
interesting avenues of future work include developing new metrics for
evaluating multiple equilibria.

##  7 Acknowledgements

We thank Ya-Ping Hsieh, Mohammad Reza Karimi and Aditi Shenoy for useful
discussions across different stages of the project. We also thank Scott Sussex
for valuable feedback on multiple drafts of the paper. This publication was
partly supported by the NCCR Catalysis (grant agreement No. 180544), a
National Centre of Competence in Research funded by the Swiss National Science
Foundation, European Union’s Horizon 2020 research and innovation programme
(grant agreement No. 826121), and ELSA (European Lighthouse on Secure and Safe
AI) funded by the European Union (grant agreement No. 101070617).

## References

  * Balduzzi et al. (2018) David Balduzzi, Sebastien Racaniere, James Martens, Jakob Foerster, Karl Tuyls, and Thore Graepel.  The mechanics of n-player differentiable games.  In _International Conference on Machine Learning_ , pp. 354–363, 2018.
  * Bradley & Terry (1952) Ralph Allan Bradley and Milton E Terry.  Rank analysis of incomplete block designs: I. the method of paired comparisons.  _Biometrika_ , 39(3/4):324–345, 1952.
  * Bryant et al. (2022) Patrick Bryant, Gabriele Pozzati, Wensi Zhu, Aditi Shenoy, Petras Kundrotas, and Arne Elofsson.  Predicting the structure of large protein complexes using alphafold and monte carlo tree search.  _Nature communications_ , 13(1):6028, 2022.
  * Chaudhury et al. (2010) Sidhartha Chaudhury, Sergey Lyskov, and Jeffrey J Gray.  Pyrosetta: a script-based interface for implementing molecular modeling algorithms using rosetta.  _Bioinformatics_ , 26(5):689–691, 2010.
  * Chen et al. (2003) Rong Chen, Li Li, and Zhiping Weng.  Zdock: an initial-stage protein-docking algorithm.  _Proteins: Structure, Function, and Bioinformatics_ , 52(1):80–87, 2003.
  * Chen et al. (2021) Tianrong Chen, Ziyi O Wang, Ioannis Exarchos, and Evangelos Theodorou.  Large-scale multi-agent deep fbsdes.  In _International Conference on Machine Learning_ , pp. 1740–1748. PMLR, 2021.
  * Corso et al. (2023) Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola.  Diffdock: Diffusion steps, twists, and turns for molecular docking.  _International Conference on Learning Representations (ICLR)_ , 2023\.
  * De Bortoli et al. (2022) Valentin De Bortoli, Emile Mathieu, Michael Hutchinson, James Thornton, Yee Whye Teh, and Arnaud Doucet.  Riemannian score-based generative modelling.  _Advances in Neural Information Processing Systems_ , 35:2406–2422, 2022.
  * de Vries et al. (2015) Sjoerd J de Vries, Christina EM Schindler, Isaure Chauvot de Beauchêne, and Martin Zacharias.  A web interface for easy flexible protein-protein docking with attract.  _Biophysical journal_ , 108(3):462–465, 2015\.
  * Esquivel-Rodríguez et al. (2012) Juan Esquivel-Rodríguez, Yifeng David Yang, and Daisuke Kihara.  Multi-lzerd: multiple protein docking for asymmetric complexes.  _Proteins: Structure, Function, and Bioinformatics_ , 80(7):1818–1833, 2012.
  * Evans et al. (2021) Richard Evans, Michael O’Neill, Alexander Pritzel, Natasha Antropova, Andrew Senior, Tim Green, Augustin Žídek, Russ Bates, Sam Blackwell, Jason Yim, et al.  Protein complex prediction with alphafold-multimer.  _biorxiv_ , pp. 2021–10, 2021.
  * Ganea et al. (2022) Octavian-Eugen Ganea, Xinyuan Huang, Charlotte Bunne, Yatao Bian, Regina Barzilay, Tommi S. Jaakkola, and Andreas Krause.  Independent SE(3)-equivariant models for end-to-end rigid protein docking.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=GQjaI9mLet.
  * Geiger & Smidt (2022) Mario Geiger and Tess Smidt.  e3nn: Euclidean neural networks.  _arXiv preprint arXiv:2207.09453_ , 2022.
  * Holdijk et al. (2022) Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling.  Path integral stochastic optimal control for sampling transition paths.  _arXiv preprint arXiv:2207.02149_ , 2022.
  * Ji et al. (2023) Yuanfeng Ji, Yatao Bian, Guoji Fu, Peilin Zhao, and Ping Luo.  Syndock: N rigid protein docking via learnable group synchronization, 2023\.
  * Jin et al. (2023) Wengong Jin, Siranush Sarkizova, Xun Chen, Nir Hacohen, and Caroline Uhler.  Unsupervised protein-ligand binding energy prediction via neural euler’s rotation equation.  _arXiv preprint arXiv:2301.10814_ , 2023.
  * Jing et al. (2022) Bowen Jing, Gabriele Corso, Jeffrey Chang, Regina Barzilay, and Tommi Jaakkola.  Torsional diffusion for molecular conformer generation.  _Advances in Neural Information Processing Systems_ , 35:24240–24253, 2022.
  * Jing et al. (2023) Bowen Jing, Ezra Erives, Peter Pao-Huang, Gabriele Corso, Bonnie Berger, and Tommi Jaakkola.  Eigenfold: Generative protein structure prediction with diffusion models.  _arXiv preprint arXiv:2304.02198_ , 2023.
  * Ketata et al. (2023) Mohamed Amine Ketata, Cedrik Laue, Ruslan Mammadov, Hannes Stärk, Menghua Wu, Gabriele Corso, Céline Marquet, Regina Barzilay, and Tommi S Jaakkola.  Diffdock-pp: Rigid protein-protein docking with diffusion models.  _arXiv preprint arXiv:2304.03889_ , 2023.
  * Kingma & Ba (2014) Diederik P Kingma and Jimmy Ba.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.
  * Kozakov et al. (2017) Dima Kozakov, David R Hall, Bing Xia, Kathryn A Porter, Dzmitry Padhorny, Christine Yueh, Dmitri Beglov, and Sandor Vajda.  The cluspro web server for protein–protein docking.  _Nature protocols_ , 12(2):255–278, 2017.
  * Leach et al. (2022) Adam Leach, Sebastian M Schmon, Matteo T Degiacomi, and Chris G Willcocks.  Denoising diffusion probabilistic models on so (3) for rotational alignment.  In _ICLR 2022 Workshop on Geometrical and Topological Representation Learning_ , 2022.
  * Liu et al. (2022) Guan-Horng Liu, Tianrong Chen, Oswin So, and Evangelos Theodorou.  Deep generalized schrödinger bridge.  _Advances in Neural Information Processing Systems_ , 35:9374–9388, 2022.
  * Mashiach et al. (2010) Efrat Mashiach, Dina Schneidman-Duhovny, Aviyah Peri, Yoli Shavit, Ruth Nussinov, and Haim J Wolfson.  An integrated suite of fast docking algorithms.  _Proteins: Structure, Function, and Bioinformatics_ , 78(15):3197–3204, 2010.
  * Monderer & Shapley (1996) Dov Monderer and Lloyd S Shapley.  Potential games.  _Games and economic behavior_ , 14(1):124–143, 1996.
  * Nikolayev et al. (1997) Dmitry I Nikolayev, Tatjana I Savyolov, et al.  Normal distribution on the rotation group so (3).  _Texture, Stress, and Microstructure_ , 29:201–233, 1997\.
  * Qiao et al. (2022) Zhuoran Qiao, Weili Nie, Arash Vahdat, Thomas F Miller III, and Anima Anandkumar.  Dynamic-backbone protein-ligand structure prediction with multiscale generative diffusion models.  _arXiv preprint arXiv:2209.15171_ , 2022.
  * Rafailov et al. (2023) Rafael Rafailov, Archit Sharma, Eric Mitchell, Stefano Ermon, Christopher D Manning, and Chelsea Finn.  Direct preference optimization: Your language model is secretly a reward model.  _arXiv preprint arXiv:2305.18290_ , 2023.
  * Rodolà et al. (2019) Emanuele Rodolà, Zorah Lähner, Alexander M Bronstein, Michael M Bronstein, and Justin Solomon.  Functional maps representation on product manifolds.  In _Computer Graphics Forum_ , volume 38, pp. 678–689. Wiley Online Library, 2019.
  * Somnath et al. (2023) Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne.  Aligned diffusion Schrödinger bridges.  In Robin J. Evans and Ilya Shpitser (eds.), _Proceedings of the Thirty-Ninth Conference on Uncertainty in Artificial Intelligence_ , volume 216 of _Proceedings of Machine Learning Research_ , pp. 1985–1995. PMLR, 31 Jul–04 Aug 2023.
  * Song & Ermon (2019) Yang Song and Stefano Ermon.  Generative modeling by estimating gradients of the data distribution.  _Advances in neural information processing systems_ , 32, 2019.
  * Song et al. (2020) 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.
  * Townshend et al. (2019) Raphael Townshend, Rishi Bedi, Patricia Suriana, and Ron Dror.  End-to-end learning on 3d protein structure for interface prediction.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Vreven et al. (2015) Thom Vreven, Iain H Moal, Anna Vangone, Brian G Pierce, Panagiotis L Kastritis, Mieczyslaw Torchala, Raphael Chaleil, Brian Jiménez-García, Paul A Bates, Juan Fernandez-Recio, et al.  Updates to the integrated protein–protein interaction benchmarks: docking benchmark version 5 and affinity benchmark version 2.  _Journal of molecular biology_ , 427(19):3031–3041, 2015.
  * Watson et al. (2022) 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.  Broadly applicable and accurate protein design by integrating structure prediction networks and diffusion generative models.  _BioRxiv_ , pp. 2022–12, 2022.
  * Xu et al. (2022) Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang.  Geodiff: A geometric diffusion model for molecular conformation generation.  _arXiv preprint arXiv:2203.02923_ , 2022.
  * Yan et al. (2020) Yumeng Yan, Huanyu Tao, Jiahua He, and Sheng-You Huang.  The hdock server for integrated protein–protein docking.  _Nature protocols_ , 15(5):1829–1852, 2020.
  * Yim et al. (2023) 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_ , 2023.
  * Zhang & Skolnick (2005) Yang Zhang and Jeffrey Skolnick.  Tm-align: a protein structure alignment algorithm based on the tm-score.  _Nucleic acids research_ , 33(7):2302–2309, 2005\.

##  Appendix A Architecture Details

For learning the surrogate potential $f_{\theta}$, we initially looked at the
both the E-GNN architectures and architectures based on tensor convolutions,
as in the e3nn library (Geiger & Smidt, 2022), and used in other docking
architectures (Corso et al., 2023; Ketata et al., 2023). In practice, we found
the architectures based on tensor convolutions to learn much faster and
generalize better, and use it for all our experiments.

###  A.1 Graph Representations

Proteins are represented as 3D graphs where nodes correspond to amino acid
residues, located at the position of the C$\alpha$ atom. Edges are added
between residues from the same and different proteins by specifying
appropriate radius cutoffs and a maximum allowed number of neighbors. We
employ the following cutoff constraints to build graphs on the fly:

  1. 1.

For graphs between residues from the same protein, we use a radius cutoff of
$20.0$ with a maximum of $30$ neighbors.

  2. 2.

For graphs between residues from different proteins, we consider two settings.
For the model described in Section 4.2, we use a radius threshold of $80.0$,
with a maximum of $500$ neighbors. For the model described in Section 4.3, we
follow a similar argument as to Corso et al. (2023) and use a dynamic cutoff
$40.0+3\sigma_{tr}$Å, where $\sigma_{tr}$ represents the standard deviation of
the diffusion noise from the translation component.

##### Notations

We use $[N]$ to denote the set $\\{1,2,\ldots,N\\}$. For each protein
$i\in[N]$, let $(\mathcal{V}_{i},\mathcal{E}_{i})$ denote its corresponding
graph representation, where $\mathcal{E}_{i}=(\mathcal{E}_{\text{self,
i}},\mathcal{E}_{\text{cross, i}})$, denoting the set of self (intra-protein)
and cross edges (inter-protein) respectively. Let $\mathcal{V}$,
$\mathcal{E}_{\text{self}}$ and $\mathcal{E}_{\text{cross}}$ denote the set of
nodes, self and cross edges for proteins $[N]$ respectively.

###  A.2 Node and Edge Features

##### Node Features

For each residue, node features include a one-hot encoding of the residue
type, the hydrophobicity value (as measured by the Kyle-DooLittle scale),
volume, charge, the polarity of the residue, and whether the residue is a
donor or acceptor (or both). The hydrophobicity, volume and charge features
were expanded into a radial basis (add values later for ranges). Our initial
experiments also used language model features from ESM2, but a significant
performance gain was not noticed. The residue features were concatenated with
sinusoidal embeddings of the diffusion time $t$, with an embedding scale of
10000 and an embedding size of 32. The total number of residue features
(including time embeddings) was $177(145+32)$.

##### Edge Features

Edges were encoded using a radial basis expansion upto a fixed cutoff. For
edges between residues of the same protein, we used a cutoff of $20.0$, while
for edges between residues from different proteins, we used a cutoff of
$80.0$. We did not notice any significant improvement from using sinusoidal
embeddings of sequence positions in the edges for the same protein. The total
number of edge features used for all graphs was $32$.

###  A.3 Convolutional Layers

Our architecture retains the same message-passing components as Corso et al.
(2023), with variations in the output layer (Appendix A.4) depending on
whether $f$ is learnt directly using supervision from PyRosetta or implicitly
through the score network. For completeness, we include the details of the
message-passing architecture here, but interested readers can also refer to
Corso et al. (2023) for the same. We use $\mathbf{h}^{i,(l)}$ to denote the
node representations of protein $i$ at layer $l$, and $\mathbf{h}_{k}^{i,(l)}$
to denote the representations of node $k\in[n_{i}]$ of protein $i$.

At every layer, the node features form a concatenated tensor of scalars,
vectors and other parity-driven combinations, each of which transform
differently under the action of rotations and translations. The message
passing between nodes across all graphs uses tensor products on these features
with the spherical harmonic representation of the (normalized) edge vector.
The learnable set of weights associated with the tensor products are informed
by the edge features (from Appendix A.2), and the scalar feature components in
the current node features. These messages are then aggregated (addition) at
every node, and used to update the current node features.

$\displaystyle\mathbf{h}_{k}^{i,(l+1)}$
$\displaystyle=\mathbf{h}_{k}^{i,(l)}\oplus\mathbf{h}_{k,\text{self}}^{i,(l)}\oplus\mathbf{h}_{k,\text{cross}}^{i,(l)}$
(4) $\displaystyle\mathbf{h}_{k,\text{self}}^{i,(l)}$
$\displaystyle=\frac{1}{|\mathcal{N}_{k,\text{self}}^{i}|}\sum_{j\in\mathcal{N}_{k,\text{self}}^{i}}Y(\hat{r}_{kj})\otimes_{\psi_{kj,i}^{\text{self}}}\mathbf{h}_{j}^{i,(l)}$
(5) $\displaystyle\mathbf{h}_{k,\text{cross}}^{i,(l)}$
$\displaystyle=\frac{1}{|\mathcal{N}_{k,\text{cross}}^{i}|}\sum_{j\in\mathcal{N}_{k,\text{cross}}^{i}}Y(\hat{r}_{kj})\otimes_{\psi_{(k,i),(j,i_{j})}^{\text{cross}}}\mathbf{h}_{j}^{i_{j},(l)}$
(6) $\displaystyle\psi_{kj,i}^{\text{self}}$
$\displaystyle=\Psi(e^{i}_{kj},\mathbf{h}_{k}^{i,(l,\text{scalar})},\mathbf{h}_{j}^{i,(l,\text{scalar})})$
(7) $\displaystyle\psi_{(k,i),(j,i_{j})}^{\text{cross}}$
$\displaystyle=\Psi(e_{kj}^{(i,i_{j})},\mathbf{h}_{k}^{i,(l,\text{scalar})},\mathbf{h}_{j}^{i_{j},(l,\text{scalar})})$
(8)

where $\mathcal{N}_{k,\text{self}}^{i}$ and $\mathcal{N}_{k,\text{cross}}^{i}$
denote intra-protein and inter-protein neighbors for node $k$ in protein $i$.
$\Phi^{\text{self}},\Psi^{\text{cross}}$ correspond to MLPs with learnable
weights for the graphs between the same protein and different proteins
respectively, and $Y$ refers to the spherical harmonics upto $l=2$. The
convolutional layers use these weights to generate scalar & vector
representations for each node.

###  A.4 Output Layers

####  A.4.1 Learning the Potential function $f_{\theta}$

We use the node features $\\{\mathbf{h}^{i,(L)}\\}_{i=1}^{N}$, after $L$
layers to predict a scalar output. Intuitively, we can think of this scalar
capturing the energetics/and or other desirable properties of the set of
proteins The scalar output is predicted as follows:

$\displaystyle s$
$\displaystyle=s_{\text{self}}+s_{\text{cross}}+s_{\mathcal{V}}$ (9)
$\displaystyle s_{\text{self}}$
$\displaystyle=\frac{1}{|\mathcal{E}_{\text{self}}|}\sum_{(k,j)\in\mathcal{E}_{\text{self}}}\text{MLP}_{s}(\mathbf{h}_{k}^{i_{k},(L,\text{scalar})},\mathbf{h}_{j}^{i_{k},(L,\text{scalar})})$
(10) $\displaystyle s_{\text{cross}}$
$\displaystyle=\frac{1}{|\mathcal{E}_{\text{cross}}|}\sum_{(k,j)\in\mathcal{E}_{\text{cross}}}\text{MLP}_{c}(\mathbf{h}_{k}^{i_{k},(L)},\mathbf{h}_{j}^{i_{j},(L)})$
(11) $\displaystyle s_{\mathcal{V}}$
$\displaystyle=\frac{1}{|\mathcal{V}|}\sum_{k\in\mathcal{V}}\text{MLP}_{v}(\mathbf{h}_{k}^{i_{k},(L,\text{scalar})})$
(12)

where $\text{MLP}_{s},\text{MLP}_{v},\text{MLP}_{v}$ are fully-connected
networks with learnable weights for energetic contributions from intra-protein
edges, inter-protein edges and residues respectively.

####  A.4.2 Learning the Score Network $s_{\theta}$

We use the node features $\\{\mathbf{h}^{i,(L)}\\}_{i=1}^{N}$ after $L$ layers
to predict the translational and rotational scores for all non-stationary
proteins relative to the stationary protein. We follow the same intuition as
Corso et al. (2023) in that the translational and rotational scores represent
the linear acceleration of the center of mass, and the angular acceleration of
the rest of the protein around the center. Similar to them, we apply a
convolution of $h^{(L)}$ about the center of mass of the system, with the
natural adaptation to the multi-agent setting. For completeness, this
convolution is described by the following equation, for the center of mass
$c$,

$\displaystyle\mathbf{v}^{i}$
$\displaystyle=\frac{1}{|\mathcal{V}_{i}|}\sum_{k\in\mathcal{V}_{i}}Y(\hat{r}_{ck,i})\otimes_{\psi_{ck,i}}\mathbf{h}_{k}^{i,(L)}$
(13) $\displaystyle\psi_{ck,i}$
$\displaystyle=\Psi(e_{ck}^{i},\mathbf{h}_{k}^{i,(L)})$ (14)

The output of $\mathbf{v}^{i}$ is restricted to a single even and single odd
vector, for each score, which are then summed up, owing to the coarse-grained
representation of the protein (in terms of C$\alpha$ atom graphs). The
magnitudes of the two scores are adjusted with an MLP and the final outputs
are normalized by scaling the translational score with
$\frac{1}{\sigma_{\text{tr}}}$ and the rotational score by the expected
magnitude of the $SO(3)$ score for $\sigma_{\text{rot}}$.

##  Appendix B Experimental Details

###  B.1 Datasets

In our experiments, we use the Dataset of Interacting Protein Structures
(DIPS) (Townshend et al., 2019) and the Docking Benchmark 5.5 (DB5.5) (Vreven
et al., 2015). We utilize the same protein-family splits as Ganea et al.
(2022) for DIPS, and evaluate on the DB5.5 test set. DIPS is a large dataset
with 42826 examples, well suited for the rigid binary docking task, with each
example containing single-chain proteins. DB5.5 is a gold standard dataset
used for evaluating docking methods, but only contains 253 examples. While
DB5.5 has been traditionally used in the context of binary protein docking,
many examples in DB5.5 consist of proteins made up of multiple chains. This
allows us to utilize DB5.5 to evaluate multimeric docking. For DB5, we
download the data and protein structures from https://github.com/octavian-
ganea/equidock_public. For DIPS, we download the processed datasets from the
link provided under https://github.com/ketatam/DiffDock-PP.

###  B.2 Training and Equilibrium Computation

####  B.2.1 Training

##### Potential Network

Training the potential network required an additional step in generating
decoys – these are perturbed assemblies generated by applying a randomly
sampled rotation and translation to the original assembly structure. For the
DIPS dataset, we only had 5000 examples for which PyRosetta was able to
generate decoys in reasonable time, and we used the same examples for training
both the score and potential network. For each training example, we generated
$10$ pairs of decoys which were scored with PyRosetta, and utilized these to
train our model with the ranking loss described in Equation 1. The potential
network was trained for 10 epochs using the Adam (Kingma & Ba, 2014) optimizer
with no weight decay using 4 32GB V100 GPUs. The best validation model
(according to validation loss computed on decoys generated from the DB5.5
validation set) was used for equilibrium computation.

##### Score Network

The score network was first trained on DIPS, and then finetuned on DB5.5.
Across both DIPS and DB5.5, the score network was trained at the granularity
of protein chains (2 chains for DIPS, and ranging between 2-8 for DB5.5). We
use Adam (Kingma & Ba, 2014) as our optimizer with no weight decay. The score
network was trained for 100 epochs on 4 32GB V100 GPUs. As described above, we
trained our model only on 5000 examples from DIPS with, each example repeated
twice. We used the model with the best validation loss as the finetuning model
for DB5.5. The score network was then finetuned on DB5.5 for 100 epochs, also
using 4 32GB V100 GPUs. We used a constant batch size of 4 across both
training and finetuning. When finetuning on DB5.5, we ran inference on 10
complexes in the validation set every 10 epochs, using the C-RMSD percentage
$<$ 5Åto save the best inference model observed so far. The best inference
model after training was used for equilbrium computation.

####  B.2.2 Equilibrium Computation

Equilibrium computation was carried out on a single GPU wherever applicable.

##### Gradient-Based Learning Approach

We generate 20 equilibria for each complex, using the gradient-based updates
of Equation 2 to update the roto-translational actions of all the protein
chains (except one receptor chain which we keep fixed). We set $\lambda=0.5$,
learning rate $\eta^{t}=t^{-0.5}$, and perform 60 gradient steps.

##### Diffusion Generative Model Approach

Here, equilibrium computation corresponds to reverse diffusion by solving the
(discretized) reverse SDE with the score network finetuned on DB5.5. As our
goal was to to evaluate multimeric docking, we ran inference at the
granularity of protein chains, different from the traditional evaluation
setting of binary dockin, which other baselines adopt. We generate 40
equilibria during inference, with the reverse diffusion running for 50 steps.

###  B.3 Hyperparameters

##### Potential Network

For the potential network, we used $16$ scalar features and $4$ vector
features, with $4$ tensor product convolution layers, and a dropout rate of
$0.1$. We use learning rate of 0.001.

##### Score Network

For the score network, we used $16$ scalar features and $4$ vector features,
with $4$ tensor product convolution layers, and a dropout rate of $0.1$,
similar to the potential network. The total number of parameters in the model
was 0.243M. Our hyperparameter tuning was limited to the learning rate for the
score network trained on the DIPS dataset (0.001, 0.0005, 0.0003). For the
translation variance schedule, we used a $\sigma_{\text{max, tr}}$ of $25.0$
and a $\sigma_{\text{min, tr}}$ of $0.01$. For the rotation variance, we used
a $\sigma_{\text{max, rot}}$ of $1.65$ and a $\sigma_{\text{min, rot}}$ of
$0.01$.

###  B.4 Baselines

#####  ClusPro, EquiDock, Attract, PatchDock

. For these baselines, we used the provided prediction files from
https://github.com/octavian-ganea/equidock_public to compute additional
metrics such as TM-Score.

##### DiffDock-PP

We used the code associated with the paper (Ketata et al., 2023) at
https://github.com/ketatam/DiffDock-PP. We use the db5_esm_inference
configuration to run inference, after modifying the parameters $\text{ns}=32$
and $\text{nv}=6$. Inference was run using the large_model_dips without any
additional finetuning on the DB5.5 dataset.

##### Multi-LZerd

We downloaded the relevant binaries from
https://kiharalab.org/proteindocking/multilzerd.php to run Multi-LZerd.
Empirically, we found Multi-LZerd to be extremely slow, and had to submit
parallel jobs for the DB5.5 test complexes, and also modified the parameters
generations=50 and population=50.

#####  AlphaFold-Multimer, SynDock

Both and SynDock are relevant baselines for multimer docking. We however do
not compare to AlphaFold-Multimer as the DB5.5 test set is part of their
training set, leading to potentially over-optimistic numbers. For SynDock,
there was no open-source implementation available at the time of writing this
paper. We will add a relevant comparison once their implementation is open-
sourced.

##### Runtimes

For Cluspro, EquiDock, Attract, PatchDock, we used the runtimes are reported
in Ganea et al. (2022). For DiffDock-PP, we use the runtime as reported by the
inference script. For Multi-LZerD, since we submitted multiple jobs in
parallel (to speed up inference), we used the average of the wall-clock times
of all jobs.

###  B.5 Code

All our code used in this work is available at
https://github.com/vsomnath/dockgame/.

##  Appendix C Additional Analyses

###  C.1 PyRosetta as a ground-truth potential

Figure 3: PyRosetta energy of DB5.5 bound complexes and of 50 random
perturbations (i.e. applying random roto-traslational actions to the protein
chains) for each of them. Bound complexes in the data are found to have the
lowest PyRosetta energy.

Traditional physics-based energy functions are commonplace, which are tuned
such that the available crystal structures have low energy. For this work, we
use PyRosetta, a commonly used protein modelling suite with several energy
functions. In particular, we use the interchain_cen function tailored for
docking. In Figure 3 we plot the PyRosetta energy for each bound complex in
the DB5.5 test dataset, as well as energies computed for 50 random complex
perturbations. This backs up our underlying assumption that ground-truth
assemblies have low PyRosetta energy.

##  Appendix D $SE(3)$ vs $SO(3)\times\mathbb{T}(3)$

The group of orientation-preserving 3D rigid motions is denoted by $SE(3)$.
Each member of the group $T\in SE(3)$ can be decomposed into two components
$T=(R,r)$ where $R\in SO(3)$ is a $3\times 3$ rotation matrix, while
$r\in\mathbb{T}(3)\cong\mathbb{R}^{3}$ denotes a translation. The group action
$A_{SE(3)}(\cdot,x):SE(3)\times\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$ can be
defined as $Rx+r$. An equivalent characterization thus of $SE(3)$ is through
the semi-direct product $SE(3)\cong SO(3)\rtimes(\mathbb{R}^{3},+)$.
Intuitively, this is a semi-direct product because the effect of rotation is
not completely disentangled from that of the translation, as the rotation $R$
also updates the position of the center of mass (COM) $\bar{x}$ of $x$.

In contrast, for the purpose of this work we consider the action space
$\mathcal{A}_{i}=SO(3)\times\mathbb{T}(3)$. To ensure $\mathcal{A}_{i}$ can be
considered a direct-product group (under the $+$ action of $\mathbb{T}(3)$),
it is sufficient to define the group action $A_{\mathcal{A}_{i}}$ that retains
the properties of the group, and disentangles the effect of a rotation $R\in
SO(3)$ is from that of a translation $r\in\mathbb{T}(3)\cong\mathbb{R}^{3}$.
To this end, we follow Corso et al. (2023) and define the group action $A$ as
$A(\cdot,x):\mathcal{A}_{i}\times\mathbb{R}^{3}\rightarrow\mathbb{R}^{3}$ as
$A((R,r),x)=R(x-\bar{x})+\bar{x}+r$. It is easy to see that the rotation does
not update the COM position
($A((R,0),\bar{x})=R(\bar{x}-\bar{x})+\bar{x}=\bar{x}$), thus disentangling
the effect of rotations from translations.